diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py new file mode 100644 index 000000000000..6aa51d66bfdd --- /dev/null +++ b/.buildkite/check-torch-abi.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Audit vLLM compiled libraries for PyTorch stable ABI compliance.""" + +import fnmatch +import sys +from pathlib import Path + +from torch_abi_audit import inspect_package +from torch_abi_audit.report import ExtensionReport, PackageReport + +# Temporary allowlist of extensions not yet on the stable ABI. +# Shrink and remove over time. +ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( + "_flashkda_C.abi3.so", + "vllm_flash_attn/_vllm_fa3_C.abi3.so", + "third_party/deep_gemm/_C*.so", +) + + +def _relative_path(lib: ExtensionReport, package_root: Path) -> str: + try: + return lib.path.relative_to(package_root).as_posix() + except ValueError: + return lib.path.name + + +def _is_torch_unstable(lib: ExtensionReport) -> bool: + return lib.error is None and lib.torch.uses_torch and not lib.torch.stable + + +def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns) + + +def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]: + return (*report.extensions, *report.bundled_libs) + + +def _collect_unstable(report: PackageReport) -> list[str]: + return sorted( + _relative_path(lib, report.root) + for lib in _iter_libs(report) + if _is_torch_unstable(lib) + ) + + +def _find_stale_allowlist_entries( + report: PackageReport, patterns: tuple[str, ...] +) -> list[str]: + """Allowlist patterns that match a built library which is no longer unstable.""" + stale: list[str] = [] + for pattern in patterns: + for lib in _iter_libs(report): + if lib.error is not None: + continue + if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern): + continue + if not _is_torch_unstable(lib): + stale.append(pattern) + break + return stale + + +def check_torch_abi( + package: str = "vllm", + patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES, +) -> int: + report = inspect_package(package) + if report.error: + print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr) + return 2 + + unstable = _collect_unstable(report) + unexpected = [ + rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns) + ] + stale = _find_stale_allowlist_entries(report, patterns) + + if unexpected or stale: + if unexpected: + print( + "Not allowed: torch-unstable libraries outside " + f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}", + file=sys.stderr, + ) + if stale: + print( + "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: " + f"{', '.join(stale)}", + file=sys.stderr, + ) + return 1 + + print("Torch stable ABI check passed.") + return 0 + + +if __name__ == "__main__": + print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance") + sys.exit(check_torch_abi()) diff --git a/.buildkite/ci_config.yaml b/.buildkite/ci_config.yaml index 21ffa1b9b8d7..9e1e46db67eb 100644 --- a/.buildkite/ci_config.yaml +++ b/.buildkite/ci_config.yaml @@ -14,6 +14,7 @@ run_all_patterns: - "setup.py" - "csrc/" - "cmake/" + - ".buildkite/check-torch-abi.py" run_all_exclude_patterns: - "docker/Dockerfile." - "csrc/cpu/" diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index ebfd1c7524ad..fd2d45be5a06 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -18,6 +18,8 @@ steps: - 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 + - tests/kernels/mamba/test_causal_conv1d.py + - tests/kernels/mamba/test_mamba_ssm.py commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " @@ -28,7 +30,9 @@ steps: 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 - 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/mamba/test_causal_conv1d.py + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" # Note: SDE can't be downloaded from CI host because of AWS WAF # - label: CPU-Compatibility Tests diff --git a/.buildkite/intel_jobs/benchmarks_intel.yaml b/.buildkite/intel_jobs/benchmarks_intel.yaml new file mode 100644 index 000000000000..9ce373e4f42e --- /dev/null +++ b/.buildkite/intel_jobs/benchmarks_intel.yaml @@ -0,0 +1,26 @@ +group: Benchmarks +depends_on: + - image-build-xpu +steps: +- label: Benchmarks CLI Test + key: benchmarks-cli-test + timeout_in_minutes: 40 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/benchmarks/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s benchmarks/' diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index d1dc95b1d401..d0aabd6d3299 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -2,6 +2,44 @@ group: Engine Intel depends_on: - image-build-xpu steps: +- label: Engine + key: engine + timeout_in_minutes: 40 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - vllm/config/ + - vllm/engine/ + - vllm/entrypoints/logger.py + - vllm/envs.py + - vllm/logger.py + - vllm/logging_utils/ + - vllm/platforms/ + - vllm/sequence.py + - vllm/triton_utils/ + - vllm/utils/ + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port + - tests/jit_monitor/test_hooks.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py' + - label: Engine (1 GPU) timeout_in_minutes: 30 device: intel_gpu @@ -18,8 +56,50 @@ steps: source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ + - tests/test_config/ commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py' + pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py && + VLLM_XPU_ENABLE_XPU_GRAPH=1 pytest -v -s test_config.py' + +- label: V1 e2e (2 GPUs) + timeout_in_minutes: 30 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/logger.py + - vllm/logging_utils/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/e2e/spec_decode/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism' diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index 1407b02055b6..e914c10d8bfe 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -22,4 +22,5 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && + pytest -v -s ir && pytest -v -s kernels/ir' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 047b00c49c76..679e626b3ed6 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -125,13 +125,13 @@ steps: pytest -v -s v1/kv_offload && pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' -- label: NixlConnector PD accuracy (2 GPUs) +- label: NixlConnector PD accuracy (4 GPUs) timeout_in_minutes: 60 - num_devices: 2 + num_devices: 4 device: intel_gpu agent_tags: label: production - gpu: 2+ + gpu: 4+ mem: 16+ no_plugin: true working_dir: "." @@ -140,7 +140,7 @@ steps: REPO: "vllm-ci-test-repo" VLLM_TEST_DEVICE: "xpu" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/ - vllm/v1/worker/kv_connector_model_runner_mixin.py - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/xpu.py @@ -148,7 +148,10 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' + bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' - label: Regression key: regression @@ -259,3 +262,25 @@ steps: pytest -v -s detokenizer && pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' + +- label: Fusion Unit Tests + timeout_in_minutes: 30 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - tests/compile/passes/test_qk_norm_rope_fusion.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s compile/passes/test_qk_norm_rope_fusion.py' \ No newline at end of file diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml new file mode 100644 index 000000000000..14a853544eca --- /dev/null +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -0,0 +1,33 @@ +group: Model Executor Intel +depends_on: + - image-build-xpu +steps: +- label: Model Executor (Intel) + key: model-executor-intel + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/engine/arg_utils.py + - vllm/config/model.py + - vllm/model_executor + - tests/model_executor + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'apt-get update && apt-get install -y curl libsodium23 && + pip3 install tensorizer==2.10.1 && + pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + export PYTHONFAULTHANDLER=1 && + cd tests && + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml index 0311b5dffb72..85a23eaef09e 100644 --- a/.buildkite/intel_jobs/model_runner_v2_intel.yaml +++ b/.buildkite/intel_jobs/model_runner_v2_intel.yaml @@ -8,7 +8,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -28,7 +28,9 @@ steps: 'export VLLM_USE_V2_MODEL_RUNNER=1 && cd tests && pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" && + pytest -v -s v1/e2e/general/test_context_length.py && ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" && + pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" && pytest -v -s v1/e2e/general/test_min_tokens.py' - label: Model Runner V2 Examples (Intel) @@ -60,3 +62,54 @@ steps: python3 basic/offline_inference/generate.py --model facebook/opt-125m && python3 generate/multimodal/vision_language_offline.py --seed 0 && python3 features/automatic_prefix_caching/prefix_caching_offline.py' + +- label: Model Runner V2 Distributed (2 GPUs) + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - tests/basic_correctness/test_basic_correctness.py + - tests/v1/distributed/test_async_llm_dp.py + - tests/v1/distributed/test_eagle_dp.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd tests && + TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"' + +- label: Model Runner V2 Spec Decode + timeout_in_minutes: 50 + 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/v1/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - tests/v1/spec_decode/test_max_len.py + - tests/v1/spec_decode/test_rejection_sampler_utils.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd tests && + pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py' diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml new file mode 100644 index 000000000000..acd94803202c --- /dev/null +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -0,0 +1,29 @@ +group: Samplers Intel +depends_on: + - image-build-xpu +steps: +- label: Samplers Test (FlashInfer) + key: samplers-test-flashinfer-intel + timeout_in_minutes: 40 + 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/model_executor/layers + - vllm/sampling_metadata.py + - tests/samplers + - tests/conftest.py + - vllm/entrypoints/generate/beam_search + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 5c717c9201c8..a7e6fa87f540 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -99,7 +99,7 @@ steps: pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && - pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]" && + pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py && pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py' - label: "XPU server test" @@ -145,7 +145,32 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s quantization/test_auto_round.py' + pytest -v -s quantization/test_auto_round.py && + pytest -v -s quantization/test_online.py' + - label: "XPU GPQA Eval (GPT-OSS)" + depends_on: + - image-build-xpu + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/evals/gpt_oss/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'pip install "gpt-oss[eval]==0.0.5" && + cd tests && + pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' - label: "XPU compressed tensors FP8 test" depends_on: - image-build-xpu @@ -168,4 +193,4 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8' \ No newline at end of file + pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8' diff --git a/.buildkite/performance-benchmarks/tests/serving-tests.json b/.buildkite/performance-benchmarks/tests/serving-tests.json index 2cbd472295e7..ac0eb213b9c8 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests.json @@ -28,11 +28,6 @@ "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json" } }, - { - "dataset_name": "sharegpt", - "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json" - } - }, { "test_name": "serving_llama8B_tp1_random_128_128", "server_parameters": { @@ -81,6 +76,7 @@ "test_name": "serving_llama70B_tp4_random_128_128", "server_parameters": { "model": "meta-llama/Llama-3.3-70B-Instruct", + "tensor_parallel_size": 4, "async_scheduling": "", "no_enable_prefix_caching": "", "max_num_batched_tokens": 8192 diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index fd33933e4c38..5d291d1b9b84 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -978,14 +978,22 @@ steps: - step: build-cpu-release-image-arm64 allow_failure: true - - label: "Publish release images to DockerHub" + - label: "Publish {{matrix}} release images to DockerHub" depends_on: - block-publish-release-images key: publish-release-images-dockerhub agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/publish-release-images.sh" + - "bash .buildkite/scripts/publish-release-images.sh {{matrix}}" + matrix: + - "cuda-13-0" + - "cuda-12-9" + - "cuda-13-0-ubuntu-24-04" + - "cuda-12-9-ubuntu-24-04" + - "rocm" + - "xpu" + - "cpu" plugins: - docker-login#v3.0.0: username: vllmbot diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 3b9a9a1d7136..31d56bdc2996 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="1" IMAGE_EXISTED_BEFORE_BUILD=0 @@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() { metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")" metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")" metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")" - metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")" - metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}" + metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")" + metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}" metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")" metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}" metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")" @@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() { metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" - metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}" + metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}" metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}" metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" @@ -1686,7 +1686,7 @@ extract_dependency_pins() { return 0 fi - for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do + for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do if [[ -n "${!var:-}" ]]; then echo "Using provided ${var}: ${!var}" continue @@ -1706,30 +1706,30 @@ extract_dependency_pins() { compute_dependency_cache_keys() { local bake_dir="" local dockerfile_rocm="" - local rixl_branch="" + local nixl_branch="" local ucx_branch="" local rocshmem_branch="" local deepep_branch="" - local rixl_material="" + local nixl_material="" local rocshmem_material="" local deepep_material="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") dockerfile_rocm="${bake_dir}/Dockerfile.rocm" - rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH") + nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH") ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH") - if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then - rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl") - RIXL_CACHE_KEY=$( + if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then + nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl") + NIXL_CACHE_KEY=$( compose_dependency_cache_key \ - "${rixl_branch}-ucx-${ucx_branch}" \ - "${rixl_material}" + "${nixl_branch}-ucx-${ucx_branch}" \ + "${nixl_material}" ) - export RIXL_CACHE_KEY - echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}" + export NIXL_CACHE_KEY + echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}" fi if [[ -n "${rocshmem_branch}" ]]; then @@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() { local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" case "${target}" in - rixl-rocm-ci) - if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then - printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}" - elif [[ -n "${RIXL_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}" + nixl-rocm-ci) + if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then + printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}" + elif [[ -n "${NIXL_BRANCH:-}" ]]; then + printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}" fi ;; rocshmem-rocm-ci) @@ -1815,7 +1815,7 @@ add_dependency_cache_target() { resolve_ci_base_dependency_targets() { local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}" - local rixl_ref="" + local nixl_ref="" local rocshmem_ref="" local deepep_ref="" @@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() { case "${mode}" in always) echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially" - for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do + for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then add_dependency_cache_target "${target}" fi @@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() { ;; esac - if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then - rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci") - if dependency_cache_ref_exists "${rixl_ref}"; then - echo "RIXL dependency cache exists: ${rixl_ref}" + if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then + nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci") + if dependency_cache_ref_exists "${nixl_ref}"; then + echo "NIXL dependency cache exists: ${nixl_ref}" else - echo "RIXL dependency cache missing; will seed: ${rixl_ref}" - add_dependency_cache_target "rixl-rocm-ci" + echo "NIXL dependency cache missing; will seed: ${nixl_ref}" + add_dependency_cache_target "nixl-rocm-ci" fi fi diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 4e8735c841a7..8b5afe5a9b06 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -35,7 +35,7 @@ set -o pipefail : "${PY_COLORS:=1}" : "${ROCM_DOCKER_TTY:=1}" : "${PYTHONFAULTHANDLER:=1}" -: "${PYTEST_TIMEOUT:=2100}" +: "${PYTEST_TIMEOUT:=2400}" if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes" fi @@ -45,9 +45,9 @@ fi if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0" fi -# Dump stacks after 15 minutes, then stop an individual test after 35 minutes. +# Dump stacks after 25 minutes, then stop an individual test after 40 minutes. if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900" + PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500" fi if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* && " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then @@ -387,6 +387,7 @@ initialize_native_environment() { local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}" local job_id_suffix="" local native_root="" + local hf_fstype="" local hf_mount="" if [[ "$(id -u)" -ne 0 ]]; then @@ -400,16 +401,19 @@ initialize_native_environment() { native_root="/tmp/vllm-native-${job_id}" TMPDIR="/tmp/vllm-${job_id_suffix}/tmp" VLLM_RPC_BASE_PATH="/tmp" - : "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}" - : "${TRITON_CACHE_DIR:=${native_root}/cache/triton}" - : "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}" - : "${XDG_CACHE_HOME:=${native_root}/cache/xdg}" + TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor" + TRITON_CACHE_DIR="${native_root}/cache/triton" + VLLM_CACHE_ROOT="${native_root}/cache/vllm" + XDG_CACHE_HOME="${native_root}/cache/xdg" : "${HF_HOME:=/home/buildkite-agent/huggingface}" + # datasets uses POSIX locks that are unsupported by the shared HF NFS cache. + # Keep processed datasets job-local while retaining the persistent Hub cache. + HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets" : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" : "${HF_HUB_ETAG_TIMEOUT:=60}" export TMPDIR VLLM_RPC_BASE_PATH export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME - export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT + export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT export PYTORCH_ROCM_ARCH="" mkdir -p "${TMPDIR}" \ @@ -417,7 +421,10 @@ initialize_native_environment() { "${TRITON_CACHE_DIR}" \ "${VLLM_CACHE_ROOT}" \ "${XDG_CACHE_HOME}" \ - "${HF_HOME}" || return 1 + "${HF_HOME}" \ + "${HF_DATASETS_CACHE}" || return 1 + + echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}" if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then if ! command -v findmnt >/dev/null 2>&1; then @@ -430,6 +437,18 @@ initialize_native_environment() { return 1 fi fi + + if command -v findmnt >/dev/null 2>&1; then + hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true) + fi + if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then + # Keep hf-xet state local and avoid vectored writes on shared NFS. + export HF_XET_CACHE="${native_root}/cache/hf-xet" + export HF_XET_HIGH_PERFORMANCE=0 + export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0 + mkdir -p "${HF_XET_CACHE}" || return 1 + echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}" + fi } run_native_preflight() { diff --git a/.buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh index 232673f01a0b..69557258a5b8 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh @@ -1,10 +1,11 @@ #!/bin/bash set -euox pipefail -export VLLM_CPU_KVCACHE_SPACE=1 +export VLLM_CPU_KVCACHE_SPACE=1 export VLLM_CPU_CI_ENV=1 -# Reduce sub-processes for acceleration -export TORCH_COMPILE_DISABLE=1 +# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of +# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op +# when callers specify fullgraph=True. export VLLM_ENABLE_V1_MULTIPROCESSING=0 SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz" @@ -49,15 +50,15 @@ wait_for_pid_and_check_log() { } # Test Sky Lake (AVX512F) -./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 & +./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 & PID_TEST_0=$! # Test Cascade Lake (AVX512F + VNNI) -./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 & +./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 & PID_TEST_1=$! # Test Cooper Lake (AVX512F + VNNI + BF16) -./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 & +./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 & PID_TEST_2=$! wait_for_pid_and_check_log $PID_TEST_0 test_0.log diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 2d11dd477eac..7c8ee86eedc3 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -40,7 +40,9 @@ function cpu_tests() { 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/moe/test_cpu_int4_moe.py - pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py + pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits @@ -54,10 +56,11 @@ function cpu_tests() { set -e python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m" - # Run model tests + # Test encoder-decoder and encoder-only models docker exec cpu-test bash -c " set -e - pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model" + pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model + pytest -x -v -s 'tests/models/language/pooling/test_embedding.py::test_models[sentence-transformers/all-MiniLM-L12-v2]'" # Run quantized model tests docker exec cpu-test bash -c " @@ -97,3 +100,4 @@ function cpu_tests() { # All of CPU tests are expected to be finished less than 40 mins. export -f cpu_tests timeout 2h bash -c cpu_tests + diff --git a/.buildkite/scripts/hardware_ci/run-gh200-test.sh b/.buildkite/scripts/hardware_ci/run-gh200-test.sh index 06e0f7af87ca..c0d1081f0573 100644 --- a/.buildkite/scripts/hardware_ci/run-gh200-test.sh +++ b/.buildkite/scripts/hardware_ci/run-gh200-test.sh @@ -15,7 +15,6 @@ DOCKER_BUILDKIT=1 docker build . \ -t gh200-test \ --build-arg max_jobs=66 \ --build-arg nvcc_threads=2 \ - --build-arg RUN_WHEEL_CHECK=false \ --build-arg torch_cuda_arch_list="9.0+PTX" # Setup cleanup diff --git a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh index 963a9dc4639c..2f9522b6f122 100644 --- a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh @@ -35,7 +35,7 @@ case "${test_suite}" in pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py pytest -v -s v1/structured_output pytest -v -s v1/test_serial_utils.py - pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]" + pytest -v -s v1/e2e/general/test_correctness_sliding_window.py pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py ;; diff --git a/.buildkite/scripts/hardware_ci/run-intel-test.sh b/.buildkite/scripts/hardware_ci/run-intel-test.sh index 83cde9ad16cc..edbb771623e5 100755 --- a/.buildkite/scripts/hardware_ci/run-intel-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-test.sh @@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK -e CMDS \ --name "${container_name}" \ "${IMAGE}" \ - bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \ + bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \ >/dev/null } 9>/tmp/docker-pull.lock diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index 91b5c3ace1be..dd54b1717e10 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -2,12 +2,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Publish release Docker images from ECR to DockerHub. +# Publish a release Docker image family from ECR to DockerHub. # Pulls per-arch images, tags with latest and versioned tags, pushes them, # then creates and pushes multi-arch manifests. set -euo pipefail +TARGET="${1:-all}" +case "${TARGET}" in + cuda-13-0 | cuda-12-9 | cuda-13-0-ubuntu-24-04 | \ + cuda-12-9-ubuntu-24-04 | rocm | xpu | cpu | all) ;; + *) + echo "Usage: $0 {cuda-13-0|cuda-12-9|cuda-13-0-ubuntu-24-04|cuda-12-9-ubuntu-24-04|rocm|xpu|cpu|all}" + exit 2 + ;; +esac + +target_enabled() { + [ "${TARGET}" = "all" ] || [ "${TARGET}" = "$1" ] +} + RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') if [ -z "${RELEASE_VERSION}" ]; then echo "ERROR: release-version metadata not set" @@ -15,12 +29,10 @@ if [ -z "${RELEASE_VERSION}" ]; then fi COMMIT="$BUILDKITE_COMMIT" -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) echo "========================================" -echo "Publishing release images v${RELEASE_VERSION}" +echo "Publishing ${TARGET} release images v${RELEASE_VERSION}" echo " Commit: ${COMMIT}" -echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" echo "========================================" # Login to ECR to pull staging images @@ -29,122 +41,137 @@ aws ecr-public get-login-password --region us-east-1 | \ # ---- CUDA (default: 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai:latest-x86_64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker push vllm/vllm-openai:latest-aarch64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 - -docker manifest rm vllm/vllm-openai:latest || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker manifest push vllm/vllm-openai:latest -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +if target_enabled cuda-13-0; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai:latest-x86_64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker push vllm/vllm-openai:latest-aarch64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + + docker manifest rm vllm/vllm-openai:latest || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker manifest push vllm/vllm-openai:latest + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +fi # ---- CUDA 12.9 ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 -docker push vllm/vllm-openai:latest-x86_64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker push vllm/vllm-openai:latest-aarch64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 - -docker manifest rm vllm/vllm-openai:latest-cu129 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true -docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker manifest push vllm/vllm-openai:latest-cu129 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +if target_enabled cuda-12-9; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + docker push vllm/vllm-openai:latest-x86_64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker push vllm/vllm-openai:latest-aarch64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + + docker manifest rm vllm/vllm-openai:latest-cu129 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true + docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker manifest push vllm/vllm-openai:latest-cu129 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +fi # ---- Ubuntu 24.04 (CUDA 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +if target_enabled cuda-13-0-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +fi # ---- Ubuntu 24.04 (CUDA 12.9) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +if target_enabled cuda-12-9-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +fi # ---- ROCm ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base +if target_enabled rocm; then + ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + echo "ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" + + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker push vllm/vllm-openai-rocm:latest -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker push vllm/vllm-openai-rocm:latest + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base -docker push vllm/vllm-openai-rocm:latest-base -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker push vllm/vllm-openai-rocm:latest-base + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base +fi # ---- XPU ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu +if target_enabled xpu; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai-xpu:latest-x86_64 -docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-xpu:latest-x86_64 + docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker manifest rm vllm/vllm-openai-xpu:latest || true -docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend -docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend -docker manifest push vllm/vllm-openai-xpu:latest -docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + docker manifest rm vllm/vllm-openai-xpu:latest || true + docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend + docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend + docker manifest push vllm/vllm-openai-xpu:latest + docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} +fi # ---- CPU ---- # CPU images are behind separate block steps and may not have been built. @@ -153,44 +180,46 @@ docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} # arch would leave `:latest-x86_64` pointing at the new release while the # `:latest` multi-arch manifest still resolves to the previous release. -CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} -CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} - -CPU_X86_AVAILABLE=false -CPU_ARM_AVAILABLE=false -docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true -docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true - -if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then - docker pull "${CPU_X86_TAG}" - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - docker push vllm/vllm-openai-cpu:latest-x86_64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - - docker pull "${CPU_ARM_TAG}" - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker push vllm/vllm-openai-cpu:latest-arm64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - - docker manifest rm vllm/vllm-openai-cpu:latest || true - docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true - docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 - docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker manifest push vllm/vllm-openai-cpu:latest - docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} -elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then - echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" -else - # Partial state: one arch built, the other did not. Fail loudly rather than - # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) - # disagree on which release they point at. - echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." - echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." - echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." - exit 1 +if target_enabled cpu; then + CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} + CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} + + CPU_X86_AVAILABLE=false + CPU_ARM_AVAILABLE=false + docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true + docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true + + if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then + docker pull "${CPU_X86_TAG}" + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-cpu:latest-x86_64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + + docker pull "${CPU_ARM_TAG}" + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker push vllm/vllm-openai-cpu:latest-arm64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + + docker manifest rm vllm/vllm-openai-cpu:latest || true + docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 + docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker manifest push vllm/vllm-openai-cpu:latest + docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} + elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then + echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" + else + # Partial state: one arch built, the other did not. Fail loudly rather than + # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) + # disagree on which release they point at. + echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." + echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." + echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." + exit 1 + fi fi echo "" -echo "Successfully published release images for v${RELEASE_VERSION}" +echo "Successfully published ${TARGET} release images for v${RELEASE_VERSION}" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index bbde36e970b2..5845844dd393 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -40,7 +40,7 @@ ##################################################################################################################################### # # # IMPORTANT: # -# * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are # +# * Currently AMD CI has MI250 agents, MI300 agents, and MI355 agents. All upcoming feature improvements are # # tracked in: https://github.com/vllm-project/vllm/issues/34994 # # # #-----------------------------------------------------------------------------------------------------------------------------------# @@ -81,10 +81,8 @@ # the above test.) Also run if model initialization test file is modified. # # * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model # # source is modified, or when specified test files are modified. # -# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to # -# run plamo2 model in vLLM. # -# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) # -# and to run plamo2 model in vLLM. # +# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). # +# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). # # * [Multi-Modal Models (Standard) 1-4]: # # - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. # # * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. # @@ -168,23 +166,9 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion==1.1.0 + - pip install helion==1.4.0 - pytest -v -s kernels/helion/ -- label: Kernels Mamba Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/mamba/ - - tests/kernels/mamba - - vllm/model_executor/layers/mamba/ops - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/mamba - #------------------------------------------------------ mi250 · models / basic -------------------------------------------------------# - label: Basic Models Test (Other CPU) # TBD @@ -207,6 +191,7 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -240,6 +225,20 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model +- label: Multi-Modal Processor (CPU) %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + parallelism: 6 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + - tests/models/registry.py + commands: + - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# - label: Batch Invariance (H100-MI250) # TBD @@ -331,7 +330,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + - pytest -v -s v1/e2e/spec_decode/draft_model/ - label: Spec Decode Speculators + MTP # TBD timeout_in_minutes: 180 @@ -349,23 +348,24 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ -- label: V1 attention (H100-MI250) # TBD +- label: V1 others (CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py + - vllm/ + - tests/v1 commands: - - pytest -v -s v1/attention + - pytest -v -s -m 'cpu_test' v1/core + - pytest -v -s v1/structured_output + - pytest -v -s v1/test_serial_utils.py + - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/metrics #------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# @@ -408,6 +408,19 @@ steps: - pytest -v -s transformers_utils - pytest -v -s config +- label: Python-only Installation # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - tests/standalone_tests/python_only_compile.sh + - setup.py + - vllm/platforms/rocm.py + commands: + - bash standalone_tests/python_only_compile.sh + #------------------------------------------------------------ mi250 · rust -----------------------------------------------------------# - label: Rust Frontend Cargo Style + Clippy # TBD @@ -445,6 +458,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 no_gpu: true + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - .buildkite/scripts/docker-build-metadata-args.sh @@ -507,7 +521,7 @@ steps: - tests/models/ 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)' + - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py @@ -657,6 +671,30 @@ steps: - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions +- label: Distributed Compile + RPC Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/compilation/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/compile/fullgraph/test_basic_correctness.py + - tests/compile/test_wrapper.py + - tests/entrypoints/llm/test_collective_rpc.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s entrypoints/llm/test_collective_rpc.py + - pytest -v -s ./compile/fullgraph/test_basic_correctness.py + - pytest -v -s ./compile/test_wrapper.py + #----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------# - label: Platform Tests # TBD @@ -679,6 +717,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -870,6 +909,71 @@ steps: commands: - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep +- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + - vllm/platforms/rocm.py + commands: + - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown + - pytest -v -s v1/worker/test_worker_memory_snapshot.py + +- label: Distributed Compile + Comm (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - tests/distributed/test_pynccl + - tests/distributed/test_events + - tests/compile/fullgraph/test_basic_correctness.py + - tests/distributed/test_symm_mem_allreduce.py + - tests/distributed/test_multiproc_executor.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s compile/fullgraph/test_basic_correctness.py + - pytest -v -s distributed/test_pynccl.py + - pytest -v -s distributed/test_events.py + - pytest -v -s distributed/test_symm_mem_allreduce.py + - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node + +#---------------------------------------------------------- mi300 · engine -----------------------------------------------------------# + +- label: Engine # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port + commands: + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py + #-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------# - label: Entrypoints Unit Tests # TBD @@ -878,6 +982,7 @@ steps: dind: false agent_pool: mi300_1 fast_check: true + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/entrypoints @@ -982,6 +1087,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + optional: true fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -1010,6 +1116,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + optional: true fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -1024,6 +1131,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + optional: true fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -1345,6 +1453,27 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 +- label: LM Eval Large Models (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_DEEP_GEMM=0 + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 + #--------------------------------------------------------- mi300 · examples ----------------------------------------------------------# - label: Examples # TBD @@ -1419,11 +1548,12 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels Core Operation Test # TBD +- label: Kernels Core Operation Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + parallelism: 3 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/ @@ -1434,7 +1564,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - 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 + - 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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: Kernels KDA Test # TBD timeout_in_minutes: 180 @@ -1447,10 +1577,29 @@ steps: - vllm/third_party/flash_linear_attention/ops/kda.py - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py - vllm/third_party/flash_linear_attention/ops/l2norm.py - - tests/kernels/test_kda.py + - vllm/models/kimi_k3/nvidia/kda.py + - vllm/models/kimi_k3/nvidia/kda_metadata.py + - vllm/models/kimi_k3/nvidia/ops/third_party/kda/ + - tests/models/kimi_k3/test_kda.py + - tests/models/kimi_k3/test_kda_metadata.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/test_kda.py + - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py + +- label: Kernels Mamba Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/mamba/ + - tests/kernels/mamba + - vllm/model_executor/layers/mamba/ops + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/mamba - label: Kernels MoE Test %N # TBD timeout_in_minutes: 180 @@ -1593,7 +1742,7 @@ steps: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" - - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" + - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" - pytest -v -s v1/e2e/general/test_context_length.py - pytest -v -s v1/e2e/general/test_min_tokens.py - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" @@ -1688,7 +1837,7 @@ steps: - tests/v1/spec_decode/test_max_len.py - tests/v1/spec_decode/test_rejection_sampler_utils.py - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - tests/v1/e2e/spec_decode/test_spec_decode.py + - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - set -x @@ -1696,7 +1845,9 @@ steps: - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp" - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp" + - pytest -v -s v1/e2e/spec_decode/eagle/ + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ #------------------------------------------------------ mi300 · models / basic -------------------------------------------------------# @@ -1797,6 +1948,37 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +- label: Language Models Test (Extended Generation) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation + commands: + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' + +- label: Language Models Tests (Hybrid) %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + parallelism: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation + commands: + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + #---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------# - label: Multi-Modal Models (Extended Generation 1) # TBD @@ -1899,20 +2081,32 @@ steps: commands: - pytest -v -s models/multimodal/processing/test_tensor_schema.py -- label: Multi-Modal Processor (CPU) %N # TBD +- label: Multi-Modal Models (Extended Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/pooling + commands: + - pytest -v -s models/multimodal/pooling -m 'not core_model' + +- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 - parallelism: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - tests/models/multimodal - - tests/models/registry.py commands: - - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" + - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model #----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------# @@ -1933,12 +2127,12 @@ steps: #-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------# -- label: Transformers Nightly Models (Shardable) %N # TBD +- label: Transformers Nightly Models (Initialization) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 - parallelism: 4 + parallelism: 6 optional: true working_dir: "/vllm-workspace/" source_file_dependencies: @@ -1954,6 +2148,27 @@ steps: commands: - pip install --upgrade git+https://github.com/huggingface/transformers - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +- label: Transformers Nightly Models (Processing) %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + parallelism: 8 + optional: true + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/multimodal/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/models/ + commands: + - pip install --upgrade git+https://github.com/huggingface/transformers - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: Transformers Nightly Models (Single) # TBD @@ -2081,7 +2296,7 @@ steps: - 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/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming" - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" @@ -2339,7 +2554,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + - pytest -v -s v1/e2e/spec_decode/draft_model/ - label: Spec Decode Eagle # TBD timeout_in_minutes: 180 @@ -2357,7 +2572,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + - pytest -v -s v1/e2e/spec_decode/eagle/ - label: Spec Decode Ngram + Suffix # TBD timeout_in_minutes: 180 @@ -2375,7 +2590,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + - pytest -v -s v1/e2e/spec_decode/ngram_suffix/ - label: Spec Decode Speculators + MTP # TBD timeout_in_minutes: 180 @@ -2394,7 +2609,8 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ - label: Speculators Correctness # TBD timeout_in_minutes: 180 @@ -2472,11 +2688,12 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration -- label: V1 attention (H100-MI300) # TBD +- label: V1 attention (H100-MI300) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + parallelism: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2488,7 +2705,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: V1 Core + KV + Metrics # TBD timeout_in_minutes: 180 @@ -2517,23 +2734,6 @@ steps: # - export HSA_NO_SCRATCH_RECLAIM=1 - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 others (CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1 - commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics - - label: V1 Sample + Logits # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2611,7 +2811,10 @@ steps: - vllm/ - tests/v1/e2e commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" + - >- + pytest -v -s + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD timeout_in_minutes: 180 @@ -2773,20 +2976,6 @@ steps: - 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] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/e2e - commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" - - label: V1 e2e (4xH100-4xMI300) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2830,195 +3019,6 @@ steps: commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt -######################################################################################################################################### -# # -# MI325 (gfx942) tests # -# # -######################################################################################################################################### - -#---------------------------------------------------------- mi325 · compile ----------------------------------------------------------# - -- label: Distributed Compile + RPC Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/compilation/ - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/compile/fullgraph/test_basic_correctness.py - - tests/compile/test_wrapper.py - - tests/entrypoints/llm/test_collective_rpc.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s entrypoints/llm/test_collective_rpc.py - - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - - pytest -v -s ./compile/test_wrapper.py - -#-------------------------------------------------------- mi325 · distributed --------------------------------------------------------# - -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/distributed/ - - tests/v1/shutdown - - tests/v1/worker/test_worker_memory_snapshot.py - - vllm/platforms/rocm.py - commands: - - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - - pytest -v -s v1/worker/test_worker_memory_snapshot.py - -- label: Distributed Compile + Comm (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - tests/distributed/test_pynccl - - tests/distributed/test_events - - tests/compile/fullgraph/test_basic_correctness.py - - tests/distributed/test_symm_mem_allreduce.py - - tests/distributed/test_multiproc_executor.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s compile/fullgraph/test_basic_correctness.py - - pytest -v -s distributed/test_pynccl.py - - pytest -v -s distributed/test_events.py - - pytest -v -s distributed/test_symm_mem_allreduce.py - - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node - -#---------------------------------------------------------- mi325 · engine -----------------------------------------------------------# - -- label: Engine # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/engine - - tests/test_sequence - - tests/test_config - - tests/test_logger - - tests/test_vllm_port - commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py - -#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------# - -- label: LM Eval Large Models (4xH100-4xMI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_USE_DEEP_GEMM=0 - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 - -#----------------------------------------------------- mi325 · models / language -----------------------------------------------------# - -- label: Language Models Test (Extended Generation) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/generation - commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - -- label: Language Models Tests (Hybrid) %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - parallelism: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/generation - commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - -#---------------------------------------------------- mi325 · models / multimodal ----------------------------------------------------# - -- label: Multi-Modal Models (Extended Pooling) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling - commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - -#----------------------------------------------------------- mi325 · misc ------------------------------------------------------------# - -- label: Python-only Installation # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - tests/standalone_tests/python_only_compile.sh - - setup.py - - vllm/platforms/rocm.py - commands: - - bash standalone_tests/python_only_compile.sh - ######################################################################################################################################### # # # MI355 (gfx950) tests # @@ -3030,6 +3030,7 @@ steps: - label: Attention Benchmarks Smoke Test (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/" @@ -3046,6 +3047,7 @@ steps: - label: Distributed Tests (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3090,6 +3092,7 @@ steps: - label: Entrypoints Integration (API Server) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true fast_check: true @@ -3107,6 +3110,7 @@ steps: - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3122,6 +3126,7 @@ steps: - label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3138,6 +3143,7 @@ steps: - label: Entrypoints Integration (API Server Generate) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3158,6 +3164,7 @@ steps: - label: Entrypoints Integration (Speech to Text) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3171,6 +3178,7 @@ steps: - label: Entrypoints Integration (Multimodal) timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3184,6 +3192,7 @@ steps: - label: Entrypoints Integration (Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3199,6 +3208,7 @@ steps: - label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3221,6 +3231,7 @@ steps: - label: LM Eval Qwen3-5 Models (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3243,6 +3254,7 @@ steps: - label: LM Eval Small Models (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3262,6 +3274,7 @@ steps: - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace" @@ -3282,6 +3295,7 @@ steps: - label: LM Eval Large Models (4xH100-4xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_4 num_gpus: 4 optional: true @@ -3304,6 +3318,7 @@ steps: - label: Examples # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/examples" source_file_dependencies: @@ -3339,6 +3354,7 @@ steps: - label: Kernels (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/" source_file_dependencies: @@ -3364,6 +3380,7 @@ steps: - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 2 working_dir: "/vllm-workspace/tests" @@ -3381,6 +3398,7 @@ steps: - label: Kernels MoE Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 5 working_dir: "/vllm-workspace/tests" @@ -3401,6 +3419,7 @@ steps: - label: Kernels Quantization Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 2 working_dir: "/vllm-workspace/tests" @@ -3418,6 +3437,7 @@ steps: - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" @@ -3437,6 +3457,7 @@ steps: - label: Language Models Test (Extended Generation) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3450,6 +3471,7 @@ steps: - label: Language Models Test (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3462,7 +3484,9 @@ steps: - label: Language Models Test (PPL) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/model_executor/models/qwen3_5.py @@ -3489,6 +3513,7 @@ steps: - label: Language Models Tests (Standard) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3503,6 +3528,7 @@ steps: - label: Multi-Modal Models (Extended Generation 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3517,6 +3543,7 @@ steps: - label: Multi-Modal Models (Extended Generation 3) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3529,6 +3556,7 @@ steps: - label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3541,6 +3569,7 @@ steps: - label: "Multi-Modal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3554,6 +3583,7 @@ steps: - label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3570,6 +3600,7 @@ steps: - label: Quantized Models Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3586,6 +3617,7 @@ steps: - label: Quantization # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3602,6 +3634,7 @@ steps: # - label: Quantized MoE Test (B200-MI355) # TBD # timeout_in_minutes: 180 # mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false # agent_pool: mi355_1 # working_dir: "/vllm-workspace/" # source_file_dependencies: @@ -3627,10 +3660,12 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -- label: V1 attention (B200-MI355) # TBD +- label: V1 attention (B200-MI355) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 + parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/config/attention.py @@ -3641,11 +3676,12 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: V1 Core + KV + Metrics # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3672,6 +3708,7 @@ steps: - label: V1 Sample + Logits # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3692,6 +3729,7 @@ steps: - label: V1 Spec Decode # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3705,6 +3743,7 @@ steps: - label: Weight Loading Multiple GPU # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" @@ -3717,6 +3756,7 @@ steps: - label: Weight Loading Multiple GPU - Large Models # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 working_dir: "/vllm-workspace/tests" num_gpus: 2 @@ -3732,6 +3772,7 @@ steps: - label: Regression # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index aadea2908c68..71d2975a5f41 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -16,8 +16,9 @@ steps: parallelism: 2 mirror: amd: - device: mi325_1 - timeout_in_minutes: 95 + dind: false + device: mi300_1 + timeout_in_minutes: 125 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index 88f3e0548250..6f0c22cc8b54 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -18,7 +18,8 @@ steps: - pytest -v -s basic_correctness/test_cpu_offload.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 70 + dind: false + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 8aac6545cfa5..034fb5884004 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -15,6 +15,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 40 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 2076eb27fc8e..431ce07af4d1 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -16,6 +16,7 @@ steps: commands: - pytest -v -s cuda/test_cuda_context.py - pytest -v -s cuda/test_platform_no_cuda_init.py + - pytest -v -s cuda/test_cuda_compatibility_path.py - label: Cudagraph device: h200_35gb @@ -26,7 +27,10 @@ steps: - vllm/v1/cudagraph_dispatcher.py - vllm/config/compilation.py - vllm/compilation + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py commands: - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - pytest -v -s v1/cudagraph/test_cudagraph_mode.py - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py + - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 4012f1ebd539..f1a89b396821 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -17,7 +17,7 @@ steps: amd: dind: false device: mi300_4 - timeout_in_minutes: 85 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -68,7 +68,7 @@ steps: amd: dind: false device: mi300_4 - timeout_in_minutes: 60 + timeout_in_minutes: 40 depends_on: - image-build-amd source_file_dependencies: @@ -94,7 +94,7 @@ steps: amd: dind: false device: mi300_4 - timeout_in_minutes: 85 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -120,7 +120,7 @@ steps: amd: dind: false device: mi300_4 - timeout_in_minutes: 80 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: @@ -131,6 +131,22 @@ 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: NixlConnector PD edge case test (2 GPUs) + key: nixlconnector-pd-edge-cases-2-gpus + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - tests/v1/kv_connector/nixl_integration/ + env: + PREFILL_GPU_ID: "0" + DECODE_GPU_ID: "1" + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh + - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus timeout_in_minutes: 25 @@ -177,7 +193,7 @@ steps: amd: dind: false device: mi300_2 - timeout_in_minutes: 70 + timeout_in_minutes: 45 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index d9fcff19f056..258a60451ed8 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -41,6 +41,7 @@ steps: amd: dind: false device: mi300_2 + timeout_in_minutes: 45 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 6e88451af861..cb311ab6a1d3 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -23,13 +23,15 @@ steps: - tests/test_config - tests/test_logger - tests/test_vllm_port - - tests/test_jit_monitor.py + - tests/jit_monitor/test_hooks.py + - tests/jit_monitor/test_hooks_gpu.py commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 50 + dind: false + device: mi300_1 + timeout_in_minutes: 40 depends_on: - image-build-amd @@ -39,13 +41,15 @@ steps: source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ + - tests/v1/test_tensor_ipc_queue.py commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/test_tensor_ipc_queue.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 55 + device: mi250_1 + timeout_in_minutes: 45 depends_on: - image-build-amd @@ -60,8 +64,8 @@ steps: - pytest -v -s v1/e2e/general/test_async_scheduling.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 70 + device: mi250_1 + timeout_in_minutes: 55 depends_on: - image-build-amd @@ -76,8 +80,8 @@ steps: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 60 + device: mi250_1 + timeout_in_minutes: 50 depends_on: - image-build-amd source_file_dependencies: @@ -108,45 +112,21 @@ steps: - vllm/triton_utils/ - vllm/utils/ - vllm/v1/ - - tests/v1/e2e/spec_decode + - tests/v1/e2e/spec_decode/ commands: # Only run tests that need exactly 2 GPUs - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" + - >- + pytest -v -s + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism + v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism mirror: amd: dind: false device: mi300_2 + timeout_in_minutes: 30 depends_on: - image-build-amd -- label: V1 e2e (4 GPUs) - key: v1-e2e-4-gpus - timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability - optional: true - num_devices: 4 - source_file_dependencies: - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/envs.py - - vllm/forward_context.py - - vllm/inputs/ - - vllm/logger.py - - vllm/logging_utils/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/triton_utils/ - - vllm/utils/ - - vllm/v1/ - - tests/v1/e2e/spec_decode - commands: - # Only run tests that need 4 GPUs - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" - - label: V1 e2e (4xH100) key: v1-e2e-4xh100 timeout_in_minutes: 35 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 9d571ff0af53..47c0db109c90 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -30,9 +30,9 @@ steps: - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests mirror: amd: - device: mi325_1 - # TODO(akaratza): Test after Torch >= 2.12 bump - soft_fail: true + dind: false + device: mi300_1 + timeout_in_minutes: 55 depends_on: - image-build-amd @@ -52,7 +52,9 @@ steps: - pytest -v -s entrypoints/scale_out mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 65 depends_on: - image-build-amd @@ -70,7 +72,8 @@ steps: - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd @@ -90,8 +93,9 @@ steps: - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 80 + dind: false + device: mi300_1 + timeout_in_minutes: 70 depends_on: - image-build-amd @@ -113,7 +117,8 @@ steps: - pytest -v -s entrypoints/anthropic mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd @@ -176,7 +181,9 @@ steps: - pytest -s entrypoints/openai/correctness/ mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index d2278408facc..1d1609d46b89 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -18,6 +18,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: @@ -51,4 +52,5 @@ steps: - vllm/compilation/ - tests/distributed/ commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - pytest -v -s distributed/test_elastic_ep.py diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml new file mode 100644 index 000000000000..e2f700a8bd86 --- /dev/null +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -0,0 +1,26 @@ +group: Fault Tolerance +depends_on: + - image-build +steps: +- label: Fault Tolerance E2E (2xH100) + key: fault-tolerance-e2e-2xh100 + timeout_in_minutes: 35 + device: h100 + num_devices: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/fault_tolerance/ + - vllm/v1/worker/sentinel/ + - vllm/entrypoints/serve/fault_tolerance/ + - vllm/distributed/elastic_ep/ + - vllm/distributed/device_communicators/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/v1/fault_tolerance/ + - tests/v1/distributed/test_external_lb_dp.py + commands: + # Base image has no nixl; install it or has_nixl_ep() skips the tests. + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + # https://github.com/NVIDIA/nccl/issues/1838 + - export NCCL_CUMEM_HOST_ENABLE=0 + - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml new file mode 100644 index 000000000000..7119c489c016 --- /dev/null +++ b/.buildkite/test_areas/jit_monitor.yaml @@ -0,0 +1,22 @@ +group: JIT Monitor +depends_on: + - image-build +steps: +- label: No Runtime JITs e2e tests + key: jit-monitor-no-runtime-jit + device: h200_35gb + timeout_in_minutes: 45 + source_file_dependencies: + - vllm/utils/jit_monitor.py + - vllm/v1/worker/gpu_worker.py + - vllm/model_executor/warmup/ + - vllm/config/observability.py + - tests/jit_monitor/test_no_runtime_jit.py + - tests/models/registry.py + commands: + # Boot a curated JIT-heavy model set with the JIT monitor in "error" mode + # and run generation; any post-warmup JIT compilation fails the test. + # Per-test watchdog so a wedged engine/CUDA init fails with a traceback + # instead of running until the build timeout. + - export PYTHONFAULTHANDLER=1 + - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 3618cb6e9696..6a99106205c5 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -61,9 +61,45 @@ steps: source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ + - vllm/models/deepseek_v4/nvidia/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/models/test_deepseek_v4_mega_moe.py commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s models/test_deepseek_v4_mega_moe.py + +# Catch-all for test files at the tests/kernels root. This job collects +# the whole root so new files are wired by default. +# Files with dedicated jobs elsewhere in this file are excluded via --ignore +# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in +# their own jobs / Kernels (B200)). +- label: Kernels Root Misc Test (B200) + key: kernels-root-misc-test-b200 + timeout_in_minutes: 45 + device: b200-k8s + source_file_dependencies: + - csrc/ + - vllm/ + - tests/kernels/ + commands: + - pytest -v -s kernels/ + --ignore=kernels/attention + --ignore=kernels/core + --ignore=kernels/helion + --ignore=kernels/ir + --ignore=kernels/mamba + --ignore=kernels/moe + --ignore=kernels/quantization + --ignore=kernels/test_concat_mla_q.py + --ignore=kernels/test_fused_qk_norm_rope_gate.py + --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + --ignore=kernels/test_top_k_per_row.py + --ignore=kernels/test_kda.py + --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py + --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_shuffle_rows.py + # BROKEN on main, pending kernel fixes (B200): + # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) - label: Kernels Attention Test %N key: kernels-attention-test @@ -80,7 +116,8 @@ steps: parallelism: 2 mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 timeout_in_minutes: 90 depends_on: - image-build-amd @@ -106,6 +143,24 @@ steps: commands: - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py +- label: Kernels FlashMLA Test (H100) + key: kernels-flashmla-test-h100 + timeout_in_minutes: 25 + device: h100 + num_devices: 1 + source_file_dependencies: + - cmake/external_projects/flashmla.cmake + - vllm/v1/attention/ops/flashmla.py + - vllm/v1/attention/backends/mla/flashmla.py + - vllm/v1/attention/backends/mla/flashmla_sparse.py + - tests/kernels/attention/test_flashmla.py + - tests/kernels/attention/test_flashmla_sparse.py + - tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py + commands: + - pytest -v -s kernels/attention/test_flashmla.py + - pytest -v -s kernels/attention/test_flashmla_sparse.py + - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py + - label: Kernels Quantization Test %N key: kernels-quantization-test timeout_in_minutes: 60 @@ -118,7 +173,9 @@ steps: parallelism: 2 mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 120 source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization @@ -148,8 +205,9 @@ steps: parallelism: 5 mirror: amd: - device: mi325_1 - timeout_in_minutes: 65 + dind: false + device: mi300_1 + timeout_in_minutes: 55 source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ @@ -174,17 +232,6 @@ steps: commands: - pytest -v -s kernels/mamba -- label: Kernels KDA Test - timeout_in_minutes: 25 - device: h200_18gb - source_file_dependencies: - - vllm/third_party/flash_linear_attention/ops/kda.py - - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py - - vllm/third_party/flash_linear_attention/ops/l2norm.py - - tests/kernels/test_kda.py - commands: - - pytest -v -s kernels/test_kda.py - - label: Kernels DeepGEMM Test (H100) key: kernels-deepgemm-test-h100 timeout_in_minutes: 35 @@ -240,8 +287,20 @@ steps: - vllm/cute_utils/ - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/ - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py + - cmake/external_projects/fmha_sm100.cmake + - vllm/models/minimax_m3/common/sparse_attention.py + - vllm/models/minimax_m3/nvidia/ + - vllm/config/attention.py + - vllm/v1/attention/backends/registry.py + - vllm/_custom_ops.py + - csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu + - csrc/libtorch_stable/ops.h + - csrc/libtorch_stable/torch_bindings.cpp + - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py + - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py + - tests/kernels/attention/test_minimax_m3.py - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py commands: @@ -254,6 +313,8 @@ steps: - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py + - pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py + - pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - pytest -v -s tests/kernels/test_top_k_per_row.py # Quantization - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8' @@ -273,6 +334,7 @@ steps: - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py + - pytest -v -s tests/kernels/attention/test_minimax_m3.py - pytest -v -s tests/kernels/test_ll_bf16_gemm.py # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py @@ -285,7 +347,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==1.1.0 + - pip install helion==1.4.0 - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 95a8bafdd589..5507fa6023c1 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -14,8 +14,9 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt mirror: amd: - device: mi325_1 - timeout_in_minutes: 55 + dind: false + device: mi300_1 + timeout_in_minutes: 45 depends_on: - image-build-amd source_file_dependencies: @@ -141,7 +142,7 @@ steps: amd: dind: false device: mi300_8 - timeout_in_minutes: 60 + timeout_in_minutes: 40 depends_on: - image-build-amd commands: @@ -232,7 +233,7 @@ steps: - 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 + # - 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 @@ -336,7 +337,7 @@ steps: - label: LM Eval KV-Offload (2xH100) key: kv-offload-medium - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h100 num_devices: 2 source_file_dependencies: @@ -346,7 +347,7 @@ steps: - vllm/v1/simple_kv_offload/ - tests/evals/gsm8k/test_gsm8k_offloading.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b" + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" - label: LM Eval KV-Offload (4xH100) key: kv-offload-large diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 46a3710ea808..a79196ffbd8d 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -14,9 +14,10 @@ steps: parallelism: 4 mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 65 + timeout_in_minutes: 85 source_file_dependencies: - vllm/lora - tests/lora @@ -46,4 +47,4 @@ steps: - pytest -v -s -x lora/test_qwen3_with_multi_loras.py - pytest -v -s -x lora/test_olmoe_tp.py - pytest -v -s -x lora/test_gptoss_tp.py - - pytest -v -s -x lora/test_qwen35_densemodel_lora.py \ No newline at end of file + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 11d469cfb383..f89abadbe052 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -2,33 +2,6 @@ group: Miscellaneous depends_on: - image-build steps: -- label: V1 Spec Decode - device: h200_35gb - key: v1-spec-decode - timeout_in_minutes: 40 - source_file_dependencies: - - vllm/config/ - - vllm/distributed/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - - tests/v1/spec_decode - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - # TODO: create another `optional` test group for slow tests - - pytest -v -s -m 'not slow_test' v1/spec_decode - mirror: - amd: - dind: false - device: mi300_1 - timeout_in_minutes: 75 - depends_on: - - image-build-amd - - label: V1 Sample + Logits key: v1-sample-logits timeout_in_minutes: 83 @@ -59,7 +32,9 @@ steps: - pytest -v -s v1/test_outputs.py mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 70 depends_on: - image-build-amd @@ -90,6 +65,7 @@ steps: - tests/v1/kv_offload - tests/v1/simple_kv_offload - tests/v1/worker + - tests/v1/streaming_input - tests/v1/kv_connector/unit - tests/v1/ec_connector/unit - tests/v1/metrics @@ -103,6 +79,7 @@ steps: - pytest -v -s v1/kv_offload - pytest -v -s v1/simple_kv_offload - pytest -v -s v1/worker + - pytest -v -s v1/streaming_input - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics @@ -111,8 +88,9 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine mirror: amd: - device: mi325_1 - timeout_in_minutes: 75 + dind: false + device: mi300_1 + timeout_in_minutes: 65 depends_on: - image-build-amd @@ -143,6 +121,8 @@ steps: - pytest -v -s -m 'cpu_test' v1/core - pytest -v -s v1/structured_output - pytest -v -s v1/test_serial_utils.py + - pytest -v -s v1/test_kv_cache_spec_registry.py + - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics @@ -206,7 +186,7 @@ steps: - vllm/multimodal - examples/ commands: - - pip install tensorizer # for tensorizer test + - pip install --no-deps tensorizer # for tensorizer test # for basic - python3 basic/offline_inference/chat.py - python3 basic/offline_inference/generate.py --model facebook/opt-125m @@ -230,7 +210,9 @@ steps: - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 75 source_file_dependencies: - vllm/entrypoints - vllm/multimodal @@ -257,6 +239,7 @@ steps: - vllm/utils/ - vllm/v1/ - tests/v1/tracing + - tests/tracing/ commands: - "pip install \ 'opentelemetry-sdk>=1.26.0' \ @@ -264,12 +247,14 @@ steps: 'opentelemetry-exporter-otlp>=1.26.0' \ 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing + - pytest -v -s tracing mirror: amd: - device: mi325_2 + dind: false + device: mi300_2 + timeout_in_minutes: 30 depends_on: - image-build-amd - optional: true - label: Python-only Installation key: python-only-installation @@ -284,8 +269,9 @@ steps: - bash standalone_tests/python_only_compile.sh mirror: amd: - device: mi325_1 - timeout_in_minutes: 45 + device: mi250_1 + timeout_in_minutes: 55 + soft_fail: true depends_on: - image-build-amd source_file_dependencies: @@ -385,7 +371,7 @@ steps: - label: Batch Invariance (A100) key: batch-invariance-a100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: a100 source_file_dependencies: - vllm/v1/attention @@ -395,11 +381,11 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: Batch Invariance (H100) key: batch-invariance-h100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: h100 source_file_dependencies: - vllm/v1/attention @@ -410,12 +396,12 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - label: Batch Invariance (B200) key: batch-invariance-b200 - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: b200-k8s source_file_dependencies: - vllm/v1/attention @@ -426,11 +412,14 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py - + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: Acceptance Length Test (Large Models) # optional device: h200_35gb key: acceptance-length-test-large-models diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index 503bc1654939..07d4aea871c0 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -10,7 +10,9 @@ steps: - vllm/engine/arg_utils.py - vllm/config/model.py - vllm/model_executor + - vllm/model_executor/warmup - tests/model_executor + - tests/model_executor/test_jit_warmup.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py commands: - apt-get update && apt-get install -y curl libsodium23 @@ -28,13 +30,16 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: - vllm/engine/arg_utils.py - vllm/config/model.py - vllm/model_executor + - vllm/model_executor/warmup - tests/model_executor + - tests/model_executor/test_jit_warmup.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 3601aeee1172..53b35ce82935 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -41,7 +41,7 @@ steps: commands: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - - pip install tensorizer # for tensorizer test + - pip install --no-deps tensorizer # for tensorizer test - python3 basic/offline_inference/chat.py # for basic - python3 basic/offline_inference/generate.py --model facebook/opt-125m #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO @@ -103,18 +103,19 @@ steps: - label: Model Runner V2 Spec Decode device: h200_35gb key: model-runner-v2-spec-decode - timeout_in_minutes: 30 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/v1/worker/gpu/ - vllm/v1/worker/gpu_worker.py - tests/v1/spec_decode/test_max_len.py - tests/v1/spec_decode/test_rejection_sampler_utils.py - - tests/v1/e2e/spec_decode/test_spec_decode.py + - tests/v1/e2e/spec_decode/ commands: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp" - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp" + - pytest -v -s v1/e2e/spec_decode/eagle/ + - pytest -v -s v1/e2e/spec_decode/speculators/ diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 95827a894588..e231e6d90f73 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -42,10 +42,39 @@ steps: - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 50 depends_on: - image-build-amd +- label: Inkling Unit Tests (B200) + key: inkling-unit-tests-b200 + timeout_in_minutes: 40 + device: b200-k8s + source_file_dependencies: + - vllm/models/inkling/ + - vllm/cute_utils/ + - cmake/external_projects/tml_fa4.cmake + - tests/models/inkling/ + commands: + # FA4 kernel tests require SM100; the suite skips them elsewhere. + - pytest -v -s models/inkling + +- label: Kimi K3 Unit Tests (B200) + key: kimi-k3-unit-tests-b200 + timeout_in_minutes: 40 + device: b200-k8s + source_file_dependencies: + - vllm/models/kimi_k3/ + - csrc/libtorch_stable/kimi_k3/ + - tests/models/kimi_k3/ + - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py + - tests/kernels/test_bf16_skinny_gemm.py + commands: + # The native NVIDIA Kimi K3 kernels require the SM100 family. + - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py + - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu depends_on: @@ -55,7 +84,8 @@ steps: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/test_adapters.py - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 2aea721e56c5..2d3f14d1715c 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -17,6 +17,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 45 depends_on: - image-build-amd @@ -39,6 +40,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 40 depends_on: - image-build-amd source_file_dependencies: @@ -61,7 +63,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' # Shard the hybrid language model tests that are numerically stable on Hopper. @@ -69,8 +70,9 @@ steps: parallelism: 2 mirror: amd: - device: mi325_1 - timeout_in_minutes: 70 + dind: false + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd commands: @@ -102,7 +104,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' @@ -130,8 +131,9 @@ steps: - pytest -v -s models/language/pooling -m 'not core_model' mirror: amd: - device: mi325_1 - timeout_in_minutes: 120 + dind: false + device: mi300_1 + timeout_in_minutes: 95 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 57f559c59fb6..095af12d2850 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -14,7 +14,9 @@ steps: - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 65 depends_on: - image-build-amd @@ -31,7 +33,9 @@ steps: - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 55 depends_on: - image-build-amd @@ -47,7 +51,8 @@ steps: - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model mirror: amd: - device: mi325_1 + device: mi250_1 + timeout_in_minutes: 55 depends_on: - image-build-amd @@ -65,7 +70,9 @@ steps: - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 50 depends_on: - image-build-amd @@ -109,6 +116,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 35 depends_on: - image-build-amd source_file_dependencies: @@ -131,7 +139,9 @@ steps: - pytest -v -s models/multimodal/test_mapping.py mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 + timeout_in_minutes: 90 depends_on: - image-build-amd @@ -166,8 +176,9 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' mirror: amd: - device: mi325_1 - timeout_in_minutes: 75 + dind: false + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 59a3632d42a1..f53f399946b3 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -107,13 +107,6 @@ steps: - tests/compile/passes commands: - pytest -s -v compile/passes --ignore compile/passes/distributed - mirror: - amd: - dind: false - device: mi300_1 - timeout_in_minutes: 65 - depends_on: - - image-build-amd - label: PyTorch Fullgraph Smoke Test device: h200_35gb @@ -236,6 +229,7 @@ steps: amd: dind: false device: mi300_1 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 16782f727f0e..0a79b7112129 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -24,8 +24,7 @@ steps: - uv pip install --system conch-triton-kernels # The SM90-only checkpoint currently contains a removed weight_chan_scale # parameter. It was not exercised by the previous L4 job. - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 8 + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' - label: Quantized Fusions device: h200_35gb @@ -68,5 +67,4 @@ steps: - vllm/model_executor/layers/quantization - tests/models/quantization commands: - - pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 3 + - pytest -v -s models/quantization diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index c1599c1c6ea2..47cb37f36d12 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -26,7 +26,7 @@ steps: - 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/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming" - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds" diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 2e7cd4a623e6..929cbec2aeb5 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -19,8 +19,18 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers mirror: amd: - device: mi325_1 + device: mi250_1 + timeout_in_minutes: 40 depends_on: - image-build-amd + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - vllm/v1/sample/ + - vllm/entrypoints/generate/beam_search/ + - tests/samplers + - tests/conftest.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - pytest -v -s samplers diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 096c324bb8e2..464c4c021dbd 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -2,20 +2,48 @@ group: Spec Decode depends_on: - image-build steps: +- label: V1 Spec Decode + device: h200_35gb + key: v1-spec-decode + timeout_in_minutes: 40 + source_file_dependencies: + - vllm/config/ + - vllm/distributed/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/spec_decode + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # TODO: create another `optional` test group for slow tests + - pytest -v -s -m 'not slow_test' v1/spec_decode + mirror: + amd: + dind: false + device: mi300_1 + timeout_in_minutes: 50 + depends_on: + - image-build-amd + - label: Spec Decode Eagle key: spec-decode-eagle timeout_in_minutes: 25 - device: h200_18gb + device: h200_35gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + - pytest -v -s v1/e2e/spec_decode/eagle/ mirror: amd: - device: mi325_1 - timeout_in_minutes: 60 + dind: false + device: mi300_1 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: @@ -27,7 +55,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Eagle Nightly B200 +- label: Spec Decode Eagle Nightly (B200) key: spec-decode-eagle-nightly-b200 timeout_in_minutes: 25 device: b200-k8s @@ -37,12 +65,12 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + - pytest -v -s v1/e2e/spec_decode/eagle/ - label: Spec Decode Speculators + MTP key: spec-decode-speculators-mtp - timeout_in_minutes: 20 - device: h200_18gb + timeout_in_minutes: 50 + device: h200_35gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ @@ -50,11 +78,13 @@ steps: - vllm/transformers_utils/configs/speculators/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ mirror: amd: - device: mi325_1 - timeout_in_minutes: 65 + dind: false + device: mi300_1 + timeout_in_minutes: 75 depends_on: - image-build-amd source_file_dependencies: @@ -67,7 +97,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Speculators + MTP Nightly B200 +- label: Spec Decode Speculators + MTP Nightly (B200) key: spec-decode-speculators-mtp-nightly-b200 timeout_in_minutes: 30 device: b200-k8s @@ -78,24 +108,26 @@ steps: - vllm/transformers_utils/configs/speculators/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" - + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ + - label: Spec Decode Ngram + Suffix key: spec-decode-ngram-suffix timeout_in_minutes: 20 - device: h200_18gb + device: h200_35gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ + - tests/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + - pytest -v -s v1/e2e/spec_decode/ngram_suffix/ + - python3 spec_decode/test_custom_proposer.py mirror: amd: - device: mi325_1 - timeout_in_minutes: 55 - # TODO(akaratza): Test after Torch >= 2.12 bump - soft_fail: true + dind: false + device: mi300_1 + timeout_in_minutes: 35 depends_on: - image-build-amd source_file_dependencies: @@ -116,10 +148,11 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + - pytest -v -s v1/e2e/spec_decode/draft_model/ mirror: amd: - device: mi325_1 + dind: false + device: mi300_1 timeout_in_minutes: 55 depends_on: - image-build-amd @@ -132,7 +165,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Draft Model Nightly B200 +- label: Spec Decode Draft Model Nightly (B200) key: spec-decode-draft-model-nightly-b200 timeout_in_minutes: 40 device: b200-k8s @@ -142,9 +175,9 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: Speculators Correctness +- label: Speculators Correctness Nightly (H100) key: speculators-correctness timeout_in_minutes: 30 device: h100 @@ -152,21 +185,66 @@ steps: num_devices: 1 source_file_dependencies: - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/dflash/ - vllm/model_executor/models/qwen3_dflash.py - tests/v1/spec_decode/test_speculators_correctness.py commands: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test -- label: Spec Decode MTP hybrid (B200) - timeout_in_minutes: 20 +- label: Spec Decode DeepSeek MTP Parallel Load (B200) + key: spec-decode-deepseek-mtp-parallel-load-b200 + timeout_in_minutes: 30 device: b200-k8s optional: true + num_devices: 2 + source_file_dependencies: + - vllm/v1/spec_decode/llm_base_proposer.py + - vllm/v1/spec_decode/eagle.py + - vllm/v1/worker/gpu/spec_decode/eagle/ + - vllm/model_executor/models/deepseek_mtp.py + - vllm/model_executor/models/deepseek_v2.py + - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py + commands: + - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py + +- label: Spec Decode AL DFlash Nightly + key: spec-decode-dflash-nightly + timeout_in_minutes: 90 + device: h200_35gb + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/laguna_dflash.py + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ + +- label: Spec Decode AL DSpark Nightly + key: spec-decode-dspark-nightly + timeout_in_minutes: 60 + device: h200_35gb + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/gemma4_dspark.py + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ + +- label: Spec Decode AL MTP + Other Acceptance Nightly + key: spec-decode-mtp-other-acceptance-nightly + timeout_in_minutes: 60 + device: h200_35gb + optional: true source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/models/qwen3_5.py - - vllm/model_executor/models/qwen3_5_mtp.py + - tests/v1/e2e/spec_decode/conftest.py - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid" + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml new file mode 100644 index 000000000000..eaef3551664b --- /dev/null +++ b/.buildkite/test_areas/torch_abi.yaml @@ -0,0 +1,14 @@ +group: Torch ABI +depends_on: + - image-build +steps: +- label: Torch Stable ABI Audit + key: torch-stable-abi-audit + timeout_in_minutes: 5 + source_file_dependencies: + - .buildkite/check-torch-abi.py + - csrc/ + - cmake/ + - setup.py + commands: + - python3 /vllm-workspace/.buildkite/check-torch-abi.py diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 46a874178b5a..2e05e5b5bab8 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -17,6 +17,7 @@ steps: amd: dind: false device: mi300_2 + timeout_in_minutes: 35 depends_on: - image-build-amd commands: diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4f63a75f045b..fa5f276759c1 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,10 +12,11 @@ /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn /vllm/model_executor/layers/batch_invariant.py @yewentao256 +/vllm/model_executor/custom_op.py @shen-shanshan /vllm/ir @ProExpertProg /vllm/kernels/ @ProExpertProg @tjtanaa /vllm/kernels/helion @ProExpertProg @zou3519 -/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa +/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa @shen-shanshan /vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni /CMakeLists.txt @tlrmchlsmth @LucasWilkinson @Harry-Chen /cmake @tlrmchlsmth @LucasWilkinson @Harry-Chen @@ -47,6 +48,7 @@ # Rust Frontend /rust/ @BugenZhao @njhill +/rust/src/bench @esmeetu /build_rust.sh @BugenZhao @njhill /rust-toolchain.toml @BugenZhao @njhill /.buildkite/test_areas/rust* @BugenZhao @njhill @@ -78,6 +80,8 @@ /vllm/v1/executor @njhill /vllm/v1/worker @njhill /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium +/vllm/v1/worker/encoder_cudagraph.py @shen-shanshan +/vllm/v1/worker/encoder_cudagraph_defs.py @shen-shanshan # Model runner V2 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256 @@ -99,7 +103,7 @@ /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa /tests/models @DarkLight1337 @ywang96 @AndreasKaratzas -/tests/multimodal @DarkLight1337 @ywang96 @NickLucche +/tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas /tests/test_inputs.py @DarkLight1337 @ywang96 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm @@ -175,22 +179,26 @@ mkdocs.yaml @hmellor /vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy # ROCm related: specify owner with write access to notify AMD folks for careful code review -/vllm/**/*rocm* @tjtanaa @dllehr-amd +/vllm/**/*rocm* @tjtanaa @dllehr-amd @shen-shanshan /docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas /vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd /vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd /vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd /vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd -/csrc/rocm @tjtanaa @dllehr-amd +/csrc/rocm @tjtanaa @dllehr-amd @hongxiayang /requirements/*rocm* @tjtanaa @AndreasKaratzas /tests/**/*rocm* @tjtanaa @AndreasKaratzas -/docs/**/*rocm* @tjtanaa +/docs/**/*rocm* @tjtanaa @hongxiayang /vllm/**/*quark* @tjtanaa /tests/**/*quark* @tjtanaa @AndreasKaratzas -/docs/**/*quark* @tjtanaa +/docs/**/*quark* @tjtanaa @hongxiayang /vllm/**/*aiter* @tjtanaa @AndreasKaratzas /tests/**/*aiter* @tjtanaa @AndreasKaratzas +# AMD per-model amd/ subtrees +/vllm/models/*/amd @tjtanaa @dllehr-amd @hongxiayang +/vllm/models/deepseek_v4/amd @tjtanaa @dllehr-amd @hongxiayang @zyongye + # TPU /vllm/v1/worker/tpu* @NickLucche /vllm/platforms/tpu.py @NickLucche diff --git a/.github/mergify.yml b/.github/mergify.yml index 4333c6e646da..a7e25ad92a68 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -19,6 +19,7 @@ pull_request_rules: description: Comment on PR when pre-commit check fails conditions: - check-failure=pre-commit + - -check-cancelled=pre-commit - -closed - -draft - or: @@ -181,6 +182,18 @@ pull_request_rules: add: - performance +- name: label-quantization + description: Automatically apply quantization label + conditions: + - label != stale + - or: + - files~=^vllm/model_executor/layers/quantization/ + - title~=(?i)quant + actions: + label: + add: + - quantization + - name: label-qwen description: Automatically apply qwen label conditions: @@ -220,6 +233,31 @@ pull_request_rules: add: - gpt-oss +- name: label-kimi + description: Automatically apply kimi label + conditions: + - label != stale + - or: + - files~=(?i)kimi + - files~=(?i)moonshot + - title~=(?i)(?:kimi|moonshot) + actions: + label: + add: + - kimi + +- name: label-k3 + description: Automatically apply k3 label (launch triage; retire after ramp-down) + conditions: + - label != stale + - or: + - files~=(?i)kimi[-_]?k3 + - title~=(?i)(?:kimi[-\s]?k3|\bk3\b) + actions: + label: + add: + - k3 + - name: label-nvidia description: Automatically apply nvidia label conditions: @@ -337,17 +375,19 @@ pull_request_rules: add: - speculative-decoding -- name: label-v1 - description: Automatically apply v1 label +- name: label-mrv2 + description: Automatically apply mrv2 label conditions: - label != stale - or: - - files~=^vllm/v1/ - - files~=^tests/v1/ + # Model Runner V2 subtree. The trailing slash is required: the MRv1 + # files (gpu_model_runner.py, gpu_worker.py, ...) sit one level up. + - files~=^vllm/v1/worker/gpu/ + - title~=(?i)(?:\bmrv2\b|model[-\s]?runner[-\s]?v2) actions: label: add: - - v1 + - mrv2 - name: label-tpu description: Automatically apply tpu label diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index e5d1accf477c..5e07d228b7f6 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -130,6 +130,66 @@ jobs: }, ], }, + kimi: { + keywords: [ + { term: "Kimi", searchIn: "both" }, + { term: "Moonshot", searchIn: "both" }, + ], + substrings: [ + { term: "moonshotai/", searchIn: "both" }, + { term: "kimi", searchIn: "title" }, + ], + }, + k3: { + keywords: [ + { term: "Kimi K3", searchIn: "both" }, + { term: "K3", searchIn: "title" }, + ], + substrings: [ + { term: "moonshotai/kimi-k3", searchIn: "both" }, + ], + }, + quantization: { + keywords: [ + { + term: "quantization", + searchIn: "both" + }, + { + term: "quantized", + searchIn: "both" + }, + ], + }, + "intel-gpu": { + // Keyword search - matches whole words only (with word boundaries) + keywords: [ + { + term: "B50", + searchIn: "both" + }, + { + term: "B60", + searchIn: "both" + }, + { + term: "B70", + searchIn: "both" + }, + { + term: "intel gpu", + searchIn: "both" + }, + { + term: "Arc GPU", + searchIn: "both" + }, + { + term: "BMG", + searchIn: "both" + }, + ], + }, // Add more label configurations here as needed // example: { // keywords: [...], @@ -491,4 +551,4 @@ jobs: issue_number: context.issue.number, body: message, }); - core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`); \ No newline at end of file + core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`); diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index 4124583d96d0..a2f09eb8a45e 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -80,9 +80,9 @@ jobs: '', '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.', '', - 'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.', + 'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.', '', - 'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.', + 'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.', '', 'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.', '', diff --git a/.github/workflows/notify-ci-authorized.yml b/.github/workflows/notify-ci-authorized.yml new file mode 100644 index 000000000000..30ae33c376d2 --- /dev/null +++ b/.github/workflows/notify-ci-authorized.yml @@ -0,0 +1,49 @@ +name: Notify CI authorization + +on: + pull_request_target: + types: [labeled] + workflow_run: + workflows: [Record CI approval] + types: [completed] + +concurrency: + group: >- + notify-ci-authorized-${{ + github.event.pull_request.number || + github.event.workflow_run.pull_requests[0].number + }} + cancel-in-progress: false + +permissions: + contents: read + issues: write + pull-requests: read + +jobs: + notify: + if: >- + (github.event_name == 'pull_request_target' && + github.event.action == 'labeled' && + (github.event.label.name == 'ready' || + github.event.label.name == 'ready-run-all-tests')) || + (github.event_name == 'workflow_run' && + github.event.workflow_run.event == 'pull_request_review' && + github.event.workflow_run.conclusion == 'success' && + github.event.workflow_run.pull_requests[0]) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: ${{ github.event.repository.default_branch }} + persist-credentials: false + - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + python-version: "3.12" + - name: Notify PR author that CI is available + run: >- + uv run --no-project --python 3.12 + .github/workflows/scripts/run_ci_command.py + env: + CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} + GH_TOKEN: ${{ github.token }} diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 143fc427a49e..aa1f437ab6a0 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -41,7 +41,7 @@ jobs: if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) { core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`); } else { - core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`); + core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`); } pre-commit: diff --git a/.github/workflows/record-ci-approval.yml b/.github/workflows/record-ci-approval.yml new file mode 100644 index 000000000000..f668b2302966 --- /dev/null +++ b/.github/workflows/record-ci-approval.yml @@ -0,0 +1,14 @@ +name: Record CI approval + +on: + pull_request_review: + types: [submitted] + +permissions: {} + +jobs: + approved: + if: github.event.review.state == 'approved' + runs-on: ubuntu-latest + steps: + - run: echo "Approval recorded" diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml new file mode 100644 index 000000000000..2c9f36812fef --- /dev/null +++ b/.github/workflows/run-ci-command.yml @@ -0,0 +1,40 @@ +name: Run CI from PR comment + +on: + issue_comment: + types: [created] + +concurrency: + group: run-ci-comment-${{ github.event.issue.number }} + cancel-in-progress: false + +permissions: + contents: read + issues: write + pull-requests: write + +jobs: + run-ci-command: + if: >- + github.event.issue.pull_request && + (github.event.comment.body == '/ci run' || + github.event.comment.body == '/ci retry') + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: ${{ github.event.repository.default_branch }} + persist-credentials: false + - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + python-version: "3.12" + - name: Authorize and run CI command + run: >- + uv run --no-project --python 3.12 + .github/workflows/scripts/run_ci_command.py + env: + BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} + BUILDKITE_ORGANIZATION: vllm + BUILDKITE_PIPELINE: ci + CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} + GH_TOKEN: ${{ github.token }} diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py new file mode 100644 index 000000000000..905dacc6b1b1 --- /dev/null +++ b/.github/workflows/scripts/run_ci_command.py @@ -0,0 +1,911 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +import random +import sys +import time +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +COMMAND_RUN_CI = "/ci run" +COMMAND_RETRY_FAILED = "/ci retry" +CI_AUTHORIZED_COMMENT_MARKER = "" +READY_LABELS = {"ready", "ready-run-all-tests"} +TRUSTED_PERMISSIONS = {"admin", "maintain", "write"} +ACTIVE_BUILD_STATES = { + "blocked", + "creating", + "scheduled", + "running", + "failing", + "canceling", + "waiting", + "waiting_failed", +} +RETRY_STATES = "failed,timed_out,expired" +SETUP_STEP_KEYS = { + "ensure-ci-base-amd", + "pre-commit", + "refresh-rocm-base-amd", +} + + +class ApiError(RuntimeError): + def __init__(self, status: int | None, message: str) -> None: + super().__init__(message) + self.status = status + + +def rate_limit_jitter() -> float: + return random.uniform(1, 5) + + +class HttpTransport: + def __init__( + self, + *, + max_retries: int = 3, + jitter: Callable[[], float] = rate_limit_jitter, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + self.max_retries = max_retries + self.jitter = jitter + self.sleep = sleep + + def request( + self, + url: str, + *, + body: Mapping[str, Any] | None = None, + headers: Mapping[str, str] | None = None, + method: str = "GET", + ) -> Any: + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + url, + data=data, + headers=dict(headers or {}), + method=method, + ) + for attempt in range(self.max_retries + 1): + try: + with urllib.request.urlopen(request, timeout=30) as response: + response_body = response.read().decode() + break + except urllib.error.HTTPError as error: + response_body = error.read().decode() + if error.code == 429 and attempt < self.max_retries: + delay = self._rate_limit_delay(error, response_body) + print( + "API rate limit reached; " + f"retry {attempt + 1}/{self.max_retries} " + f"in {delay:g} seconds.", + file=sys.stderr, + ) + self.sleep(delay) + continue + message = self._error_message(response_body, error.reason) + raise ApiError( + error.code, + f"API returned {error.code}: {message}", + ) from error + except urllib.error.URLError as error: + raise ApiError( + None, + f"API request failed: {error.reason}", + ) from error + + if not response_body: + return None + try: + return json.loads(response_body) + except json.JSONDecodeError as error: + raise ApiError(None, "API returned a non-JSON response.") from error + + @staticmethod + def _error_message(response_body: str, fallback: str) -> str: + try: + parsed = json.loads(response_body) + except json.JSONDecodeError: + return fallback + return str(parsed.get("message", fallback)) + + def _rate_limit_delay( + self, + error: urllib.error.HTTPError, + response_body: str, + ) -> float: + try: + parsed = json.loads(response_body) + except json.JSONDecodeError: + parsed = {} + + scope = parsed.get("scope") + reset_header = ( + "RateLimit-User-Reset" if scope == "rest_user" else "RateLimit-Reset" + ) + candidates = [ + error.headers.get(reset_header), + parsed.get("reset"), + ] + for candidate in candidates: + try: + delay = float(candidate) + except (TypeError, ValueError): + continue + if delay >= 0: + return delay + self.jitter() + return 60 + self.jitter() + + +class GitHubClient: + def __init__( + self, + token: str, + repository: str, + transport: HttpTransport | None = None, + ) -> None: + if not token: + raise RuntimeError("GH_TOKEN is not set.") + self.owner, self.repo = repository.split("/", maxsplit=1) + self.transport = transport or HttpTransport() + self.headers = { + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": "vllm-ci-command", + "X-GitHub-Api-Version": "2022-11-28", + } + + def _request( + self, + path: str, + *, + body: Mapping[str, Any] | None = None, + method: str = "GET", + ) -> Any: + return self.transport.request( + f"https://api.github.com{path}", + body=body, + headers=self.headers, + method=method, + ) + + def _repo_path(self, suffix: str) -> str: + owner = urllib.parse.quote(self.owner, safe="") + repo = urllib.parse.quote(self.repo, safe="") + return f"/repos/{owner}/{repo}{suffix}" + + def _paginate(self, path: str) -> list[dict[str, Any]]: + results: list[dict[str, Any]] = [] + separator = "&" if "?" in path else "?" + for page in range(1, 101): + response = self._request(f"{path}{separator}per_page=100&page={page}") + if not isinstance(response, list): + raise ApiError(None, "GitHub API returned an invalid list response.") + results.extend(response) + if len(response) < 100: + return results + raise ApiError(None, "GitHub API pagination exceeded 10,000 results.") + + def get_pr(self, number: int) -> dict[str, Any]: + return self._request(self._repo_path(f"/pulls/{number}")) + + def get_permission(self, actor: str) -> str: + username = urllib.parse.quote(actor, safe="") + try: + response = self._request( + self._repo_path(f"/collaborators/{username}/permission") + ) + except ApiError as error: + if error.status == 404: + return "none" + raise + return str(response["permission"]) + + def get_review_decision(self, number: int) -> str | None: + query = """ + query($owner: String!, $repo: String!, $number: Int!) { + repository(owner: $owner, name: $repo) { + pullRequest(number: $number) { + reviewDecision + } + } + } + """ + response = self._request( + "/graphql", + body={ + "query": query, + "variables": { + "number": number, + "owner": self.owner, + "repo": self.repo, + }, + }, + method="POST", + ) + return response["data"]["repository"]["pullRequest"]["reviewDecision"] + + def list_reviews(self, number: int) -> list[dict[str, Any]]: + return self._paginate(self._repo_path(f"/pulls/{number}/reviews")) + + def list_issue_comments(self, number: int) -> list[dict[str, Any]]: + return self._paginate(self._repo_path(f"/issues/{number}/comments")) + + def list_reactions(self, comment_id: int) -> list[dict[str, Any]]: + return self._paginate( + self._repo_path(f"/issues/comments/{comment_id}/reactions") + ) + + def add_reaction(self, comment_id: int, content: str) -> None: + self._request( + self._repo_path(f"/issues/comments/{comment_id}/reactions"), + body={"content": content}, + method="POST", + ) + + def add_comment(self, issue_number: int, body: str) -> None: + self._request( + self._repo_path(f"/issues/{issue_number}/comments"), + body={"body": body}, + method="POST", + ) + + +class BuildkiteClient: + def __init__( + self, + token: str, + organization: str, + pipeline: str, + transport: HttpTransport | None = None, + ) -> None: + self.token = token + self.transport = transport or HttpTransport() + organization = urllib.parse.quote(organization, safe="") + pipeline = urllib.parse.quote(pipeline, safe="") + self.base_url = ( + "https://api.buildkite.com/v2/organizations/" + f"{organization}/pipelines/{pipeline}/builds" + ) + + def _headers(self) -> dict[str, str]: + if not self.token: + raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.") + return { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json", + "User-Agent": "vllm-ci-command", + } + + def _request_url( + self, + url: str, + *, + body: Mapping[str, Any] | None = None, + method: str = "GET", + ) -> Any: + if ( + url != self.base_url + and not url.startswith(f"{self.base_url}?") + and not url.startswith(f"{self.base_url}/") + ): + raise ApiError(None, "Buildkite API returned an invalid pagination URL.") + return self.transport.request( + url, + body=body, + headers=self._headers(), + method=method, + ) + + def _request( + self, + *, + body: Mapping[str, Any] | None = None, + method: str = "GET", + path: str = "", + query: Sequence[tuple[str, str]] = (), + ) -> Any: + url = f"{self.base_url}{path}" + if query: + url = f"{url}?{urllib.parse.urlencode(query)}" + return self._request_url(url, body=body, method=method) + + def list_builds( + self, + commit: str | None, + *, + metadata: tuple[str, str] | None = None, + ) -> list[dict[str, Any]]: + query = [ + ("exclude_jobs", "true"), + ("exclude_pipeline", "true"), + ("per_page", "100"), + ] + if commit: + query.append(("commit", commit)) + if metadata: + key, value = metadata + query.append((f"meta_data[{key}]", value)) + response = self._request(query=query) + if not isinstance(response, list): + raise ApiError(None, "Buildkite API returned an invalid build list.") + return response + + def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]: + return self._request(body=body, method="POST") + + def retry_failed_jobs( + self, + build_number: int, + states: str, + ) -> dict[str, Any]: + number = urllib.parse.quote(str(build_number), safe="") + return self._request( + body={"states": states}, + method="PUT", + path=f"/{number}/retry_failed_jobs", + ) + + def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: + number = urllib.parse.quote(str(build_number), safe="") + query = [ + ("state[]", "failed"), + ("state[]", "timed_out"), + ("state[]", "expired"), + ("include_retried_jobs", "false"), + ("per_page", "100"), + ] + url = f"{self.base_url}/{number}/jobs?{urllib.parse.urlencode(query)}" + jobs: list[dict[str, Any]] = [] + while url: + response = self._request_url(url) + if not isinstance(response, Mapping): + raise ApiError(None, "Buildkite API returned an invalid job list.") + items = response.get("items") + links = response.get("links") + if not isinstance(items, list) or not isinstance(links, Mapping): + raise ApiError(None, "Buildkite API returned an invalid job list.") + jobs.extend(items) + next_url = links.get("next") + if next_url is not None and not isinstance(next_url, str): + raise ApiError( + None, "Buildkite API returned an invalid pagination URL." + ) + url = next_url + return jobs + + +def parse_command(body: str) -> str | None: + if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}: + return body + return None + + +def parse_trusted_users(value: str = "") -> set[str]: + return { + user.casefold() for item in value.split(",") for user in item.split() if user + } + + +def has_ready_label(pr: Mapping[str, Any]) -> bool: + return any(label["name"] in READY_LABELS for label in pr["labels"]) + + +def is_trusted_permission(permission: str) -> bool: + return permission in TRUSTED_PERMISSIONS + + +def authorize( + *, + actor: str, + permission: str, + pr: Mapping[str, Any], + trusted_approval: bool = False, + trusted_users: set[str] | None = None, +) -> tuple[bool, str]: + trusted_users = trusted_users or set() + if is_trusted_permission(permission): + return True, f"repository {permission} permission" + if actor.casefold() in trusted_users: + return True, "configured trusted contributor" + if actor.casefold() != pr["user"]["login"].casefold(): + return ( + False, + "Only reviewers with write access can run CI before it is " + "delegated to the PR author.", + ) + if pr["draft"]: + return False, "PR authors cannot run CI while the PR is a draft." + if has_ready_label(pr): + return True, "ready label" + if trusted_approval: + return True, "approval from a trusted reviewer" + return ( + False, + "A reviewer with write access must run `/ci run`, approve the PR, " + "or add the `ready` label first.", + ) + + +def has_trusted_approval( + github: GitHubClient, + number: int, + trusted_users: set[str], +) -> bool: + if github.get_review_decision(number) != "APPROVED": + return False + + latest_review_states: dict[str, tuple[str, str]] = {} + for review in github.list_reviews(number): + user = review.get("user") or {} + login = user.get("login") + state = review.get("state") + if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}: + latest_review_states[login.casefold()] = (login, state) + + for login, state in latest_review_states.values(): + if state != "APPROVED": + continue + if login.casefold() in trusted_users: + return True + if is_trusted_permission(github.get_permission(login)): + return True + return False + + +def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool: + pull_request = build.get("pull_request") + if isinstance(pull_request, Mapping): + build_pr_number = pull_request.get("id", pull_request.get("number")) + if build_pr_number is not None: + return str(build_pr_number) == str(pr_number) + metadata = build.get("meta_data") or {} + return str(metadata.get("github-pr-number")) == str(pr_number) + + +def is_active_build(build: Mapping[str, Any]) -> bool: + return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES + + +def select_latest_build( + builds: Sequence[dict[str, Any]], + pr_number: int, +) -> dict[str, Any] | None: + matching = [build for build in builds if is_build_for_pr(build, pr_number)] + return max(matching, key=lambda build: build.get("created_at", ""), default=None) + + +def create_build_payload( + *, + actor: str, + comment_id: int, + pr: Mapping[str, Any], +) -> dict[str, Any]: + return { + "commit": pr["head"]["sha"], + "branch": pr["head"]["ref"], + "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}", + "pull_request_id": pr["number"], + "pull_request_base_branch": pr["base"]["ref"], + "pull_request_repository": pr["head"]["repo"]["clone_url"], + "pull_request_labels": [label["name"] for label in pr["labels"]], + "ignore_pipeline_branch_filters": True, + "env": { + "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id), + "VLLM_CI_TRIGGERED_BY": actor, + }, + "meta_data": { + "github-comment-id": str(comment_id), + "github-pr-number": str(pr["number"]), + "github-triggered-by": actor, + }, + } + + +def create_retry_build_payload( + *, + actor: str, + comment_id: int, + pr: Mapping[str, Any], + source_build: Mapping[str, Any], + step_keys: Sequence[str], +) -> dict[str, Any]: + payload = create_build_payload( + actor=actor, + comment_id=comment_id, + pr=pr, + ) + source_number = str(source_build["number"]) + payload["message"] = f"PR #{pr['number']} {COMMAND_RETRY_FAILED} by @{actor}" + payload["env"]["VLLM_CI_ONLY_STEP_KEYS"] = json.dumps( + step_keys, separators=(",", ":") + ) + payload["meta_data"].update( + { + "github-retry-source-build": source_number, + "github-retry-source-commit": str(source_build.get("commit", "")), + } + ) + return payload + + +def add_reaction_safely( + github: GitHubClient, + comment_id: int, + content: str, +) -> None: + try: + github.add_reaction(comment_id, content) + except Exception as error: + print(f"Could not add {content} reaction: {error}", file=sys.stderr) + + +def command_comment_marker(comment_id: int) -> str: + return f"" + + +def has_bot_comment_marker( + github: GitHubClient, + issue_number: int, + marker: str, +) -> bool: + return any( + marker in str(comment.get("body", "")) + and (comment.get("user") or {}).get("login") == "github-actions[bot]" + for comment in github.list_issue_comments(issue_number) + ) + + +def is_already_handled( + github: GitHubClient, + issue_number: int, + comment_id: int, +) -> bool: + terminal_reaction = any( + reaction.get("content") in {"rocket", "-1"} + and (reaction.get("user") or {}).get("login") == "github-actions[bot]" + for reaction in github.list_reactions(comment_id) + ) + if terminal_reaction: + return True + return has_bot_comment_marker( + github, + issue_number, + command_comment_marker(comment_id), + ) + + +def notify_authorized( + event: Mapping[str, Any], + github: GitHubClient, + trusted_users_value: str = "", +) -> None: + pr = event["pull_request"] + if pr["state"] != "open" or pr["draft"]: + return + + trusted_users = parse_trusted_users(trusted_users_value) + author = pr["user"]["login"] + author_permission = github.get_permission(author) + if ( + is_trusted_permission(author_permission) + or author.casefold() in trusted_users + or has_bot_comment_marker( + github, + pr["number"], + CI_AUTHORIZED_COMMENT_MARKER, + ) + ): + return + + if "label" in event: + if ( + event.get("action") != "labeled" + or event["label"]["name"] not in READY_LABELS + or has_trusted_approval(github, pr["number"], trusted_users) + ): + return + elif "review" in event: + if ( + event.get("action") != "submitted" + or str(event["review"].get("state", "")).casefold() != "approved" + or has_ready_label(pr) + or not has_trusted_approval(github, pr["number"], trusted_users) + ): + return + else: + return + + github.add_comment( + pr["number"], + ( + f"✅ @{author}, CI is now available for this PR. Comment `/ci run` " + "to run full CI or `/ci retry` to retry failed jobs.\n\n" + f"{CI_AUTHORIZED_COMMENT_MARKER}" + ), + ) + + +def handle_run_ci( + *, + actor: str, + buildkite: BuildkiteClient, + comment_id: int, + github: GitHubClient, + pr: Mapping[str, Any], +) -> str: + duplicate_builds = buildkite.list_builds( + pr["head"]["sha"], + metadata=("github-comment-id", str(comment_id)), + ) + duplicate = select_latest_build(duplicate_builds, pr["number"]) + if duplicate: + return f"CI was already requested by this comment: {duplicate['web_url']}" + + current_builds = buildkite.list_builds(pr["head"]["sha"]) + active_build = next( + ( + build + for build in current_builds + if is_build_for_pr(build, pr["number"]) and is_active_build(build) + ), + None, + ) + if active_build: + return f"CI is already running for this commit: {active_build['web_url']}" + + current_pr = github.get_pr(pr["number"]) + if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: + return ( + "The PR head changed while processing the command. Comment `/ci run` again." + ) + + build = buildkite.create_build( + create_build_payload( + actor=actor, + comment_id=comment_id, + pr=current_pr, + ) + ) + return ( + f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) " + f"for commit `{current_pr['head']['sha'][:12]}`." + ) + + +def handle_retry_failed( + *, + actor: str, + buildkite: BuildkiteClient, + comment_id: int, + github: GitHubClient, + pr: Mapping[str, Any], +) -> str: + builds = buildkite.list_builds(pr["head"]["sha"]) + build = select_latest_build(builds, pr["number"]) + if build: + metadata = build.get("meta_data") or {} + if str(metadata.get("github-comment-id")) == str(comment_id): + return f"CI was already requested by this comment: {build['web_url']}" + + retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES) + if retried["retried_jobs_count"] == 0: + return ( + "No failed, timed-out, or expired jobs need retrying: " + f"{build['web_url']}" + ) + return ( + f"Queued {retried['retried_jobs_count']} failed job(s) for retry in " + f"[Buildkite CI #{build['number']}]({build['web_url']})." + ) + + previous_builds = buildkite.list_builds( + None, + metadata=("github-pr-number", str(pr["number"])), + ) + previous_builds = [ + candidate + for candidate in previous_builds + if candidate.get("commit") != pr["head"]["sha"] + ] + source_build = select_latest_build(previous_builds, pr["number"]) + if not source_build: + return "No earlier CI build exists for this PR. Use `/ci run` first." + if not source_build.get("finished_at") or is_active_build(source_build): + return f"The previous CI build is still running: {source_build['web_url']}" + + failed_jobs = buildkite.list_failed_jobs(source_build["number"]) + failed_script_jobs = [job for job in failed_jobs if job.get("type") == "script"] + missing_step_keys = [job for job in failed_script_jobs if not job.get("step_key")] + if missing_step_keys: + return ( + f"[Buildkite CI #{source_build['number']}]" + f"({source_build['web_url']}) has failed jobs without stable step " + "keys, so they cannot be retried on a new commit. Use `/ci run`." + ) + + failed_step_keys = {str(job["step_key"]) for job in failed_script_jobs} + setup_failures = sorted( + step_key + for step_key in failed_step_keys + if step_key.startswith("image-build") or step_key in SETUP_STEP_KEYS + ) + if setup_failures: + return ( + f"[Buildkite CI #{source_build['number']}]" + f"({source_build['web_url']}) failed during CI setup, so its test " + "failure set is incomplete. Use `/ci run` for the new commit." + ) + + step_keys = sorted(failed_step_keys) + if not step_keys: + return ( + "No failed, timed-out, or expired jobs need retrying in " + f"[Buildkite CI #{source_build['number']}]" + f"({source_build['web_url']})." + ) + + current_pr = github.get_pr(pr["number"]) + if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: + return ( + "The PR head changed while processing the command. " + "Comment `/ci retry` again." + ) + + retry_build = buildkite.create_build( + create_retry_build_payload( + actor=actor, + comment_id=comment_id, + pr=current_pr, + source_build=source_build, + step_keys=step_keys, + ) + ) + return ( + f"Triggered [Buildkite CI #{retry_build['number']}]" + f"({retry_build['web_url']}) for commit " + f"`{current_pr['head']['sha'][:12]}`, running {len(step_keys)} failed " + f"step(s) from [Buildkite CI #{source_build['number']}]" + f"({source_build['web_url']})." + ) + + +def run( + event: Mapping[str, Any], + github: GitHubClient, + buildkite: BuildkiteClient, + trusted_users_value: str = "", +) -> None: + command = parse_command(event["comment"]["body"]) + if not command or "pull_request" not in event["issue"]: + return + + issue_number = event["issue"]["number"] + comment_id = event["comment"]["id"] + actor = event["comment"]["user"]["login"] + + if is_already_handled(github, issue_number, comment_id): + print(f"Comment {comment_id} was already handled.") + return + add_reaction_safely(github, comment_id, "eyes") + + try: + pr = github.get_pr(issue_number) + permission = github.get_permission(actor) + if pr["state"] != "open": + github.add_comment( + issue_number, + "❌ CI commands require an open PR.\n\n" + f"{command_comment_marker(comment_id)}", + ) + return + + trusted_users = parse_trusted_users(trusted_users_value) + should_check_approval = ( + not is_trusted_permission(permission) + and actor.casefold() not in trusted_users + and actor.casefold() == pr["user"]["login"].casefold() + and not pr["draft"] + and not has_ready_label(pr) + ) + trusted_approval = should_check_approval and has_trusted_approval( + github, + issue_number, + trusted_users, + ) + allowed, reason = authorize( + actor=actor, + permission=permission, + pr=pr, + trusted_approval=trusted_approval, + trusted_users=trusted_users, + ) + if not allowed: + github.add_comment( + issue_number, + f"❌ @{actor}, {reason}\n\n{command_comment_marker(comment_id)}", + ) + return + + print(f"Authorized @{actor}: {reason}") + if command == COMMAND_RUN_CI: + message = handle_run_ci( + actor=actor, + buildkite=buildkite, + comment_id=comment_id, + github=github, + pr=pr, + ) + else: + message = handle_retry_failed( + actor=actor, + buildkite=buildkite, + comment_id=comment_id, + github=github, + pr=pr, + ) + add_reaction_safely(github, comment_id, "rocket") + github.add_comment(issue_number, f"✅ {message}") + except Exception: + add_reaction_safely(github, comment_id, "confused") + raise + + +def main() -> None: + event_path = os.environ["GITHUB_EVENT_PATH"] + with open(event_path, encoding="utf-8") as event_file: + event = json.load(event_file) + + github = GitHubClient( + os.environ.get("GH_TOKEN", ""), + os.environ["GITHUB_REPOSITORY"], + ) + event_name = os.environ.get("GITHUB_EVENT_NAME", "issue_comment") + if event_name == "pull_request_target": + notify_authorized( + event, + github, + os.environ.get("CI_TRUSTED_USERS", ""), + ) + return + if event_name == "workflow_run": + pull_requests = event["workflow_run"].get("pull_requests") or [] + if not pull_requests: + return + pr = github.get_pr(int(pull_requests[0]["number"])) + notify_authorized( + { + "action": "submitted", + "pull_request": pr, + "review": {"state": "approved"}, + }, + github, + os.environ.get("CI_TRUSTED_USERS", ""), + ) + return + + if not parse_command(event["comment"]["body"]): + return + + buildkite = BuildkiteClient( + os.environ.get("BUILDKITE_API_TOKEN", ""), + os.environ.get("BUILDKITE_ORGANIZATION", "vllm"), + os.environ.get("BUILDKITE_PIPELINE", "ci"), + ) + run( + event, + github, + buildkite, + os.environ.get("CI_TRUSTED_USERS", ""), + ) + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py new file mode 100644 index 000000000000..7d583d470fdb --- /dev/null +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -0,0 +1,831 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import io +import json +import unittest +import urllib.error +from typing import Any +from unittest.mock import patch + +from run_ci_command import ( + CI_AUTHORIZED_COMMENT_MARKER, + COMMAND_RETRY_FAILED, + COMMAND_RUN_CI, + RETRY_STATES, + BuildkiteClient, + HttpTransport, + authorize, + create_build_payload, + has_trusted_approval, + is_active_build, + is_build_for_pr, + notify_authorized, + parse_command, + parse_trusted_users, + run, + select_latest_build, +) + + +def make_pr(**overrides: Any) -> dict[str, Any]: + pr = { + "base": {"ref": "main"}, + "draft": False, + "head": { + "ref": "feature", + "repo": {"clone_url": "https://github.com/contributor/vllm.git"}, + "sha": "0123456789abcdef", + }, + "labels": [], + "number": 42, + "state": "open", + "user": {"login": "author"}, + } + pr.update(overrides) + return pr + + +def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]: + return { + "comment": { + "body": command, + "id": 99, + "user": {"login": actor}, + }, + "issue": { + "number": 42, + "pull_request": {}, + }, + } + + +class FakeGitHub: + def __init__( + self, + *, + permission: str = "write", + permissions: dict[str, str] | None = None, + pr: dict[str, Any] | None = None, + review_decision: str = "REVIEW_REQUIRED", + reviews: list[dict[str, Any]] | None = None, + comments: list[str] | None = None, + ) -> None: + self.comments = comments or [] + self.permission = permission + self.permissions = permissions or {} + self.pr = pr or make_pr() + self.reactions: list[str] = [] + self.review_decision = review_decision + self.reviews = reviews or [] + + def get_pr(self, number: int) -> dict[str, Any]: + return self.pr + + def get_permission(self, actor: str) -> str: + return self.permissions.get(actor, self.permission) + + def get_review_decision(self, number: int) -> str: + return self.review_decision + + def list_reviews(self, number: int) -> list[dict[str, Any]]: + return self.reviews + + def list_issue_comments(self, number: int) -> list[dict[str, Any]]: + return [ + { + "body": body, + "user": {"login": "github-actions[bot]"}, + } + for body in self.comments + ] + + def list_reactions(self, comment_id: int) -> list[dict[str, Any]]: + return [] + + def add_reaction(self, comment_id: int, content: str) -> None: + self.reactions.append(content) + + def add_comment(self, issue_number: int, body: str) -> None: + self.comments.append(body) + + +class FakeBuildkite: + def __init__( + self, + build_lists: list[list[dict[str, Any]]] | None = None, + failed_job_lists: list[list[dict[str, Any]]] | None = None, + ) -> None: + self.build_lists = build_lists or [] + self.created_builds: list[dict[str, Any]] = [] + self.failed_job_lists = failed_job_lists or [] + self.job_list_calls: list[int] = [] + self.list_calls: list[tuple[str | None, tuple[str, str] | None]] = [] + self.retry_calls: list[tuple[int, str]] = [] + + def list_builds( + self, + commit: str | None, + *, + metadata: tuple[str, str] | None = None, + ) -> list[dict[str, Any]]: + self.list_calls.append((commit, metadata)) + return self.build_lists.pop(0) + + def create_build(self, body: dict[str, Any]) -> dict[str, Any]: + self.created_builds.append(body) + return { + "number": 123, + "web_url": "https://buildkite.example/builds/123", + } + + def retry_failed_jobs( + self, + build_number: int, + states: str, + ) -> dict[str, Any]: + self.retry_calls.append((build_number, states)) + return {"retried_jobs_count": 3} + + def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: + self.job_list_calls.append(build_number) + return self.failed_job_lists.pop(0) + + +class FakeTransport: + def __init__(self, response: Any) -> None: + self.calls: list[dict[str, Any]] = [] + self.response = response + + def request(self, url: str, **kwargs: Any) -> Any: + self.calls.append({"url": url, **kwargs}) + if isinstance(self.response, tuple): + return self.response[len(self.calls) - 1] + return self.response + + +class FakeHttpResponse: + def __init__(self, response: Any) -> None: + self.response = response + + def __enter__(self) -> "FakeHttpResponse": + return self + + def __exit__(self, *args: Any) -> None: + return None + + def read(self) -> bytes: + return json.dumps(self.response).encode() + + +class RunCiCommandTest(unittest.TestCase): + @patch("run_ci_command.urllib.request.urlopen") + def test_http_transport_retries_buildkite_rate_limit(self, urlopen: Any) -> None: + body = { + "message": "Please wait 9 seconds before making more requests.", + "reset": 9, + "scope": "rest", + } + rate_limit_error = urllib.error.HTTPError( + "https://api.buildkite.com/v2/builds", + 429, + "Too Many Requests", + { + "RateLimit-Limit": "400", + "RateLimit-Remaining": "0", + "RateLimit-Reset": "9", + }, + io.BytesIO(json.dumps(body).encode()), + ) + urlopen.side_effect = [rate_limit_error, FakeHttpResponse({"ok": True})] + delays: list[float] = [] + + response = HttpTransport( + jitter=lambda: 2.5, + sleep=delays.append, + ).request("https://api.buildkite.com/v2/builds") + + self.assertEqual(response, {"ok": True}) + self.assertEqual(delays, [11.5]) + self.assertEqual(urlopen.call_count, 2) + + @patch("run_ci_command.urllib.request.urlopen") + def test_http_transport_retries_rate_limit_three_times(self, urlopen: Any) -> None: + def rate_limit_error() -> urllib.error.HTTPError: + body = { + "message": "Please wait 9 seconds before making more requests.", + "reset": 9, + "scope": "rest", + } + return urllib.error.HTTPError( + "https://api.buildkite.com/v2/builds", + 429, + "Too Many Requests", + { + "RateLimit-Limit": "400", + "RateLimit-Remaining": "0", + "RateLimit-Reset": "9", + }, + io.BytesIO(json.dumps(body).encode()), + ) + + urlopen.side_effect = [rate_limit_error() for _ in range(4)] + delays: list[float] = [] + + with self.assertRaisesRegex(RuntimeError, "API returned 429"): + HttpTransport( + jitter=lambda: 2, + sleep=delays.append, + ).request("https://api.buildkite.com/v2/builds") + + self.assertEqual(delays, [11, 11, 11]) + self.assertEqual(urlopen.call_count, 4) + + @patch("run_ci_command.urllib.request.urlopen") + def test_http_transport_does_not_retry_permission_error(self, urlopen: Any) -> None: + permission_error = urllib.error.HTTPError( + "https://api.github.com/repos/vllm-project/vllm/issues/1/comments", + 403, + "Forbidden", + {}, + io.BytesIO(b'{"message":"Resource not accessible by integration"}'), + ) + urlopen.side_effect = permission_error + delays: list[float] = [] + + with self.assertRaisesRegex( + RuntimeError, + "Resource not accessible by integration", + ): + HttpTransport( + jitter=lambda: 2, + sleep=delays.append, + ).request( + "https://api.github.com/repos/vllm-project/vllm/issues/1/comments" + ) + + self.assertEqual(delays, []) + self.assertEqual(urlopen.call_count, 1) + + def test_only_exact_ci_commands_are_accepted(self) -> None: + self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI) + self.assertEqual( + parse_command(COMMAND_RETRY_FAILED), + COMMAND_RETRY_FAILED, + ) + self.assertIsNone(parse_command("/ci run please")) + self.assertIsNone(parse_command(" /ci run")) + + def test_write_access_authorizes_reviewers_and_authors(self) -> None: + allowed, _ = authorize( + actor="reviewer", + permission="write", + pr=make_pr(), + ) + self.assertTrue(allowed) + + def test_configured_trusted_contributors_can_run_ci(self) -> None: + trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO") + allowed, _ = authorize( + actor="trusted-two", + permission="read", + pr=make_pr(), + trusted_users=trusted_users, + ) + self.assertTrue(allowed) + + def test_authors_need_an_approval_or_ready_label(self) -> None: + pending, _ = authorize( + actor="author", + permission="read", + pr=make_pr(), + ) + approved, _ = authorize( + actor="author", + permission="read", + pr=make_pr(), + trusted_approval=True, + ) + ready, _ = authorize( + actor="author", + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + self.assertFalse(pending) + self.assertTrue(approved) + self.assertTrue(ready) + + def test_non_author_contributors_without_write_are_denied(self) -> None: + allowed, _ = authorize( + actor="contributor", + permission="read", + pr=make_pr(), + trusted_approval=True, + ) + self.assertFalse(allowed) + + def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None: + allowed, _ = authorize( + actor="author", + permission="read", + pr=make_pr(draft=True, labels=[{"name": "ready"}]), + trusted_approval=True, + ) + self.assertFalse(allowed) + + def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None: + approved_review = { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + trusted = FakeGitHub( + permission="read", + permissions={"reviewer": "write"}, + review_decision="APPROVED", + reviews=[approved_review], + ) + untrusted = FakeGitHub( + permission="read", + review_decision="APPROVED", + reviews=[approved_review], + ) + self.assertTrue(has_trusted_approval(trusted, 42, set())) + self.assertFalse(has_trusted_approval(untrusted, 42, set())) + + def test_build_matching_is_scoped_to_the_pr(self) -> None: + self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42)) + self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42)) + self.assertTrue( + is_build_for_pr( + {"meta_data": {"github-pr-number": "42"}}, + 42, + ) + ) + + def test_latest_build_selection_ignores_other_prs(self) -> None: + latest = select_latest_build( + [ + { + "created_at": "2026-07-28T02:00:00Z", + "number": 3, + "pull_request": {"id": 43}, + }, + { + "created_at": "2026-07-28T01:00:00Z", + "number": 2, + "pull_request": {"id": 42}, + }, + { + "created_at": "2026-07-28T00:00:00Z", + "number": 1, + "pull_request": {"id": 42}, + }, + ], + 42, + ) + self.assertEqual(latest["number"], 2) + + def test_active_build_states_prevent_duplicate_runs(self) -> None: + self.assertTrue(is_active_build({"state": "scheduled"})) + self.assertTrue(is_active_build({"state": "running"})) + self.assertTrue(is_active_build({"state": "waiting"})) + self.assertTrue(is_active_build({"blocked": True, "state": "passed"})) + self.assertFalse(is_active_build({"state": "failed"})) + + def test_build_payload_preserves_pr_context(self) -> None: + payload = create_build_payload( + actor="reviewer", + comment_id=99, + pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]), + ) + self.assertEqual( + payload, + { + "commit": "0123456789abcdef", + "branch": "feature", + "message": "PR #42 /ci run by @reviewer", + "pull_request_id": 42, + "pull_request_base_branch": "main", + "pull_request_repository": ("https://github.com/contributor/vllm.git"), + "pull_request_labels": ["ready", "v1"], + "ignore_pipeline_branch_filters": True, + "env": { + "VLLM_CI_GITHUB_COMMENT_ID": "99", + "VLLM_CI_TRIGGERED_BY": "reviewer", + }, + "meta_data": { + "github-comment-id": "99", + "github-pr-number": "42", + "github-triggered-by": "reviewer", + }, + }, + ) + + def test_write_reviewer_runs_ci_without_delegation(self) -> None: + github = FakeGitHub() + buildkite = FakeBuildkite([[], []]) + run(make_event(COMMAND_RUN_CI), github, buildkite) + + self.assertEqual(len(buildkite.created_builds), 1) + self.assertEqual( + buildkite.created_builds[0]["message"], + "PR #42 /ci run by @reviewer", + ) + self.assertEqual(github.reactions, ["eyes", "rocket"]) + self.assertTrue(github.comments[0].startswith("✅ ")) + self.assertIn("Buildkite CI #123", github.comments[0]) + + def test_unapproved_authors_are_denied_without_buildkite(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(), + review_decision="REVIEW_REQUIRED", + ) + buildkite = FakeBuildkite() + run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) + + self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.reactions, ["eyes"]) + self.assertTrue(github.comments[0].startswith("❌ ")) + self.assertIn("approve the PR", github.comments[0]) + + run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) + self.assertEqual(len(github.comments), 1) + + def test_untrusted_approval_cannot_launch_ci(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(), + review_decision="APPROVED", + reviews=[ + { + "state": "APPROVED", + "user": {"login": "untrusted-reviewer"}, + } + ], + ) + buildkite = FakeBuildkite() + + run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) + + self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.reactions, ["eyes"]) + self.assertTrue(github.comments[0].startswith("❌ ")) + + def test_ready_label_notifies_author_once(self) -> None: + pr = make_pr(labels=[{"name": "ready"}]) + event = { + "action": "labeled", + "label": {"name": "ready"}, + "pull_request": pr, + } + github = FakeGitHub(permission="read", pr=pr) + + notify_authorized(event, github) + notify_authorized(event, github) + + self.assertEqual(len(github.comments), 1) + self.assertTrue(github.comments[0].startswith("✅ @author")) + self.assertIn("`/ci run`", github.comments[0]) + self.assertIn("`/ci retry`", github.comments[0]) + self.assertIn(CI_AUTHORIZED_COMMENT_MARKER, github.comments[0]) + + def test_ready_label_does_not_notify_after_trusted_approval(self) -> None: + pr = make_pr(labels=[{"name": "ready"}]) + event = { + "action": "labeled", + "label": {"name": "ready"}, + "pull_request": pr, + } + github = FakeGitHub( + permission="read", + permissions={"reviewer": "write"}, + pr=pr, + review_decision="APPROVED", + reviews=[ + { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + ], + ) + + notify_authorized(event, github) + + self.assertEqual(github.comments, []) + + def test_trusted_approval_notifies_author_once(self) -> None: + event = { + "action": "submitted", + "pull_request": make_pr(), + "review": {"state": "approved"}, + } + github = FakeGitHub( + permission="read", + permissions={"reviewer": "write"}, + review_decision="APPROVED", + reviews=[ + { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + ], + ) + + notify_authorized(event, github) + notify_authorized(event, github) + + self.assertEqual(len(github.comments), 1) + self.assertIn("@author", github.comments[0]) + + def test_approval_does_not_notify_when_ready_label_exists(self) -> None: + pr = make_pr(labels=[{"name": "ready"}]) + event = { + "action": "submitted", + "pull_request": pr, + "review": {"state": "approved"}, + } + github = FakeGitHub( + permission="read", + permissions={"reviewer": "write"}, + pr=pr, + review_decision="APPROVED", + reviews=[ + { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + ], + ) + + notify_authorized(event, github) + + self.assertEqual(github.comments, []) + + def test_untrusted_approval_does_not_notify_author(self) -> None: + event = { + "action": "submitted", + "pull_request": make_pr(), + "review": {"state": "approved"}, + } + github = FakeGitHub( + permission="read", + review_decision="APPROVED", + reviews=[ + { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + ], + ) + + notify_authorized(event, github) + + self.assertEqual(github.comments, []) + + def test_notification_skips_authors_who_already_have_write(self) -> None: + pr = make_pr(labels=[{"name": "ready"}]) + event = { + "action": "labeled", + "label": {"name": "ready"}, + "pull_request": pr, + } + github = FakeGitHub(permission="write", pr=pr) + + notify_authorized(event, github) + + self.assertEqual(github.comments, []) + + def test_notification_skips_draft_prs(self) -> None: + pr = make_pr(draft=True, labels=[{"name": "ready"}]) + event = { + "action": "labeled", + "label": {"name": "ready"}, + "pull_request": pr, + } + github = FakeGitHub(permission="read", pr=pr) + + notify_authorized(event, github) + + self.assertEqual(github.comments, []) + + def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite( + [ + [ + { + "created_at": "2026-07-28T01:00:00Z", + "number": 123, + "pull_request": {"id": 42}, + "state": "failing", + "web_url": "https://buildkite.example/builds/123", + } + ] + ] + ) + run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)]) + self.assertIn("Queued 3 failed job", github.comments[0]) + + def test_ci_retry_creates_filtered_build_for_new_head(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + source_build = { + "commit": "old-commit", + "created_at": "2026-07-28T01:00:00Z", + "finished_at": "2026-07-28T02:00:00Z", + "number": 122, + "pull_request": {"id": 42}, + "state": "failed", + "web_url": "https://buildkite.example/builds/122", + } + buildkite = FakeBuildkite( + [[], [source_build]], + [ + [ + { + "state": "failed", + "step_key": "basic-models-test-other-cpu", + "type": "script", + }, + { + "state": "failed", + "step_key": "basic-models-test-other-cpu", + "type": "script", + }, + { + "state": "timed_out", + "step_key": "distributed-tests-2xh100-2xmi300", + "type": "script", + }, + ] + ], + ) + + run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.job_list_calls, [122]) + self.assertEqual(len(buildkite.created_builds), 1) + payload = buildkite.created_builds[0] + self.assertEqual(payload["commit"], "0123456789abcdef") + self.assertEqual(payload["message"], "PR #42 /ci retry by @author") + self.assertEqual( + json.loads(payload["env"]["VLLM_CI_ONLY_STEP_KEYS"]), + [ + "basic-models-test-other-cpu", + "distributed-tests-2xh100-2xmi300", + ], + ) + self.assertEqual( + payload["meta_data"]["github-retry-source-build"], + "122", + ) + self.assertEqual( + payload["meta_data"]["github-retry-source-commit"], + "old-commit", + ) + self.assertIn("running 2 failed step", github.comments[0]) + self.assertIn("Buildkite CI #122", github.comments[0]) + + def test_ci_retry_new_head_requires_stable_step_keys(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite( + [ + [], + [ + { + "commit": "old-commit", + "created_at": "2026-07-28T01:00:00Z", + "finished_at": "2026-07-28T02:00:00Z", + "number": 122, + "pull_request": {"id": 42}, + "state": "failed", + "web_url": "https://buildkite.example/builds/122", + } + ], + ], + [[{"state": "failed", "step_key": None, "type": "script"}]], + ) + + run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.created_builds, []) + self.assertIn("without stable step keys", github.comments[0]) + self.assertIn("Use `/ci run`", github.comments[0]) + + def test_ci_retry_new_head_rejects_incomplete_setup_failure(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite( + [ + [], + [ + { + "commit": "old-commit", + "created_at": "2026-07-28T01:00:00Z", + "finished_at": "2026-07-28T02:00:00Z", + "number": 122, + "pull_request": {"id": 42}, + "state": "failed", + "web_url": "https://buildkite.example/builds/122", + } + ], + ], + [ + [ + { + "state": "failed", + "step_key": "image-build", + "type": "script", + } + ] + ], + ) + + run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.created_builds, []) + self.assertIn("failed during CI setup", github.comments[0]) + self.assertIn("Use `/ci run`", github.comments[0]) + + def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None: + transport = FakeTransport({"retried_jobs_count": 2}) + client = BuildkiteClient( + "secret", + "vllm", + "ci", + transport=transport, + ) + client.retry_failed_jobs(123, RETRY_STATES) + + call = transport.calls[0] + self.assertEqual(call["method"], "PUT") + self.assertTrue(call["url"].endswith("/123/retry_failed_jobs")) + self.assertEqual(call["body"], {"states": RETRY_STATES}) + + def test_buildkite_list_builds_allows_query_on_builds_endpoint(self) -> None: + transport = FakeTransport([]) + client = BuildkiteClient( + "secret", + "vllm", + "ci", + transport=transport, + ) + + builds = client.list_builds( + "current-commit", + metadata=("github-pr-number", "42"), + ) + + self.assertEqual(builds, []) + url = transport.calls[0]["url"] + self.assertIn("?exclude_jobs=true", url) + self.assertIn("commit=current-commit", url) + self.assertIn("meta_data%5Bgithub-pr-number%5D=42", url) + + def test_buildkite_failed_jobs_follow_cursor_pagination(self) -> None: + next_url = ( + "https://api.buildkite.com/v2/organizations/vllm/pipelines/ci/" + "builds/123/jobs?after=cursor" + ) + transport = FakeTransport( + ( + { + "items": [{"id": "first"}], + "links": {"next": next_url}, + }, + { + "items": [{"id": "second"}], + "links": {"next": None}, + }, + ) + ) + client = BuildkiteClient( + "secret", + "vllm", + "ci", + transport=transport, + ) + + jobs = client.list_failed_jobs(123) + + self.assertEqual(jobs, [{"id": "first"}, {"id": "second"}]) + self.assertIn("state%5B%5D=failed", transport.calls[0]["url"]) + self.assertIn("include_retried_jobs=false", transport.calls[0]["url"]) + self.assertEqual(transport.calls[1]["url"], next_url) + + +if __name__ == "__main__": + unittest.main() diff --git a/.gitignore b/.gitignore index 43787b8cd283..46418c03924e 100644 --- a/.gitignore +++ b/.gitignore @@ -173,9 +173,6 @@ venv.bak/ # mkdocs documentation /site -docs/argparse -docs/examples/* -!docs/examples/README.md # mypy .mypy_cache/ diff --git a/.markdownlint.yaml b/.markdownlint.yaml index 937487f47364..9140af925eea 100644 --- a/.markdownlint.yaml +++ b/.markdownlint.yaml @@ -3,6 +3,9 @@ MD007: MD013: false MD024: siblings_only: true +MD025: + # Allow front matter title to be different from the first heading in the document. + front_matter_title: "" MD031: list_items: false MD033: false diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3d26a51bbacc..35b53a416de4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_install_hook_types: default_stages: - pre-commit # Run locally - manual # Run in CI -exclude: 'vllm/third_party/.*' +exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*' repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.0 @@ -260,10 +260,6 @@ repos: files: ^docker/(Dockerfile|versions\.json)$ pass_filenames: false additional_dependencies: [dockerfile-parse] - - id: attention-backend-docs - name: Check attention backend documentation is up to date - entry: python tools/pre_commit/generate_attention_backend_docs.py --check - language: python - id: check-boolean-context-manager name: Check for boolean ops in with-statements entry: python tools/pre_commit/check_boolean_context_manager.py diff --git a/CMakeLists.txt b/CMakeLists.txt index f514ba41aa52..54374c2d01c0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. @@ -68,8 +68,8 @@ endif() # requirements.txt files and should be kept consistent. The ROCm torch # versions are derived from docker/Dockerfile.rocm # -set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0") -set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0") +set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0") +set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0") # TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported- # version check would always warn. Only treat it as a nightly build when the # value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which @@ -114,6 +114,11 @@ find_package(Torch REQUIRED) # Supported NVIDIA architectures. # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4) + # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports + # targeting it directly. + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0") +elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc) # to support the whole generation without specifying all sub-architectures @@ -214,10 +219,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # the set of architectures we want to compile for and remove the from the # CMAKE_CUDA_FLAGS so that they are not applied globally. # - # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch - # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only - # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's - # component-specific arch list below. + # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really + # needs PTX, add `+PTX` to that kernel's component-specific arch list below. # clear_cuda_arches(CUDA_ARCH_FLAGS) extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") @@ -227,6 +230,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") cuda_archs_loose_intersection(CUDA_ARCHS "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}") message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}") + if(NOT CUDA_ARCHS) + message(FATAL_ERROR + "No supported CUDA architectures; the build would produce a binary " + "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; " + "supported: ${CUDA_SUPPORTED_ARCHS}. " + "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).") + endif() else() # # For other GPU targets override the GPU architectures detected by cmake/torch @@ -411,8 +421,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -420,7 +433,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS - "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") + "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") @@ -695,7 +708,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() @@ -815,7 +828,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x) # require CUDA 12.8 or later if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -899,7 +912,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -924,7 +937,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # moe_data.cu is used by all CUTLASS MoE kernels. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() @@ -981,7 +994,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -1047,7 +1060,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Runtime dispatch is gated in # vllm/v1/attention/backends/mla/cutlass_mla.py. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -1069,6 +1082,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") set(MLA_ARCHS) endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS + "9.0a;10.0f;12.0f" "${CUDA_ARCHS}") + endif() + if(FUSED_KDA_DECODE_ARCHS) + set(FUSED_KDA_DECODE_SRC + "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_KDA_DECODE_SRC}" + CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") + set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY + COMPILE_OPTIONS "$<$:--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}") + message(STATUS + "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}") + endif() + + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS + "10.0f" "${CUDA_ARCHS}") + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + set(KIMI_K3_ATTN_RES_SRC + "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${KIMI_K3_ATTN_RES_SRC}" + CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") + set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY + COMPILE_OPTIONS + "$<$:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}") + message(STATUS + "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}") + endif() + # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) @@ -1110,6 +1158,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_COOPERATIVE_TOPK=1) endif() + if(FUSED_KDA_DECODE_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_FUSED_KDA_DECODE=1) + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_KIMI_K3_ATTN_RES=1) + endif() # Needed by CUTLASS kernels target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) @@ -1366,6 +1422,38 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms.cu" "csrc/rocm/skinny_gemms_int4.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS}) + + # skinny_gemms*.cu are built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy + # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude them from the + # gfx1250 build and disable their op registrations (VLLM_SKIP_SKINNY_GEMMS); + # vLLM falls back to default/Triton GEMM for those ops on gfx1250. + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms_int4.cu") + set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES}) + list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250") + if(VLLM_SKINNY_ARCHES) + message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}") + set(VLLM_SKINNY_SRC + "csrc/rocm/skinny_gemms.cu" + "csrc/rocm/skinny_gemms_int4.cu") + hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "${VLLM_SKINNY_SRC}") + unset(_VLLM_LAST_HIPIFY_TARGET) + add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS}) + add_dependencies(_rocm_C_skinny hipify_all) + set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG}) + set_target_properties(_rocm_C_skinny PROPERTIES + ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}" + POSITION_INDEPENDENT_CODE ON) + target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc) + target_compile_options(_rocm_C_skinny PRIVATE + $<$:${VLLM_ROCM_EXT_FLAGS}>) + target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C") + target_link_libraries(_rocm_C_skinny PRIVATE torch) + else() + message(STATUS "Only gfx1250, skipping skinny_gemms") + list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") + endif() set(VLLM_ROCM_HAS_GFX1100 OFF) if(VLLM_GPU_ARCHES MATCHES "gfx1100") @@ -1381,10 +1469,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP") DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_ROCM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} + COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + + if(TARGET _rocm_C_skinny) + target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny) + else() + target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) + endif() if(VLLM_ROCM_HAS_GFX1100) target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) @@ -1407,6 +1501,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) + include(cmake/external_projects/flashkda.cmake) include(cmake/external_projects/qutlass.cmake) include(cmake/external_projects/tml_fa4.cmake) diff --git a/README.md b/README.md index 42777436c63d..ca31a8cce94d 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ vLLM is flexible and easy to use with: - Tool calling and reasoning parsers - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support - Efficient multi-LoRA support for dense and MoE layers -- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. +- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. vLLM seamlessly supports 200+ model architectures on Hugging Face, including: diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 2fdd5ce88dc9..8f01be94d2d3 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -617,7 +617,7 @@ def main(): "--sparse-mla-mha-variants", nargs="+", default=None, - choices=["dense_mha", "mqa"], + choices=["dense_mha", "masked_mha", "mqa"], help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.", ) @@ -748,6 +748,9 @@ def main(): args.sparse_mla_mha_variants = yaml_config.get( "sparse_mla_mha_variants", args.sparse_mla_mha_variants ) + args.sparse_mla_masked_mha_max_seq_len = yaml_config.get( + "sparse_mla_masked_mha_max_seq_len", None + ) # Parameter sweep configuration if "parameter_sweep" in yaml_config: @@ -1119,6 +1122,7 @@ def main(): console.print(f"Prefill backend: {prefill_backend}") available_variants = [ ("dense_mha", False, "dense"), + ("masked_mha", False, "masked"), ("mqa", True, "auto"), ] requested_variants = getattr(args, "sparse_mla_mha_variants", None) @@ -1139,6 +1143,9 @@ def main(): ] else: variants = available_variants + masked_mha_max_seq_len = getattr( + args, "sparse_mla_masked_mha_max_seq_len", None + ) formatter = ResultsFormatter(console) total = 0 for spec in args.batch_specs: @@ -1148,6 +1155,10 @@ def main(): variant_label == "dense_mha" and dense_mha_max_seq_len is not None and q_len > dense_mha_max_seq_len + ) or ( + variant_label == "masked_mha" + and masked_mha_max_seq_len is not None + and q_len > masked_mha_max_seq_len ): continue total += len(backends) @@ -1161,6 +1172,10 @@ def main(): variant_label == "dense_mha" and dense_mha_max_seq_len is not None and q_len > dense_mha_max_seq_len + ) or ( + variant_label == "masked_mha" + and masked_mha_max_seq_len is not None + and q_len > masked_mha_max_seq_len ): continue config = BenchmarkConfig( @@ -1190,6 +1205,7 @@ def main(): sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len, sparse_mla_topk_pattern=sparse_mla_topk_pattern, prefill_backend=prefill_backend, + sparse_mla_masked_mha_max_seq_len=masked_mha_max_seq_len, ) # run_mla_benchmark needs the real backend name @@ -1358,6 +1374,10 @@ def main(): profile_memory=args.profile_memory, warmup_ms=args.warmup_ms, prefill_backend=pb, + kv_lora_rank=args.kv_lora_rank, + qk_nope_head_dim=args.qk_nope_head_dim, + qk_rope_head_dim=args.qk_rope_head_dim, + v_head_dim=args.v_head_dim, ) result = run_benchmark(config) diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index b8ca5739c5c5..45858f9148b3 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -299,8 +299,9 @@ class BenchmarkConfig: num_kv_splits: int | None = None # CUTLASS MLA reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA sparse_mla_force_mqa: bool = False # Force MQA path for sparse MLA - sparse_mla_mha_mode: str = "auto" # "auto" or "dense" + sparse_mla_mha_mode: str = "auto" # "auto", "dense", or "masked" sparse_mla_dense_mha_max_seq_len: int | None = None + sparse_mla_masked_mha_max_seq_len: int | None = None sparse_mla_topk_pattern: str = "random" # "random", "prefix", "sliding_window" num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled) diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa.yaml new file mode 100644 index 000000000000..a02b7f5ba0a0 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa.yaml @@ -0,0 +1,123 @@ +# Sparse MLA benchmark: forward_mha vs forward_mqa crossover plot +# +# Dense pure-prefill sweep to generate a smooth speedup curve. The 100 points +# run from q32 through q262144, approximately even on the log2 scale, and include +# routing boundaries plus extra samples around the 6k-8k crossover region. +# Dense MHA is only benchmarked through q32768 because it dominates runtime above +# that range. +# +# Usage: +# python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml + +mode: mha_vs_mqa + +model: + name: "deepseek-v3" + num_layers: 60 + num_q_heads: 16 + num_kv_heads: 1 + head_dim: 576 + kv_lora_rank: 512 + qk_nope_head_dim: 128 + qk_rope_head_dim: 64 + v_head_dim: 128 + block_size: 128 + max_model_len: 262144 + +batch_specs: + - "q32" + - "q36" + - "q40" + - "q45" + - "q50" + - "q56" + - "q62" + - "q70" + - "q78" + - "q87" + - "q97" + - "q109" + - "q122" + - "q136" + - "q152" + - "q170" + - "q190" + - "q212" + - "q237" + - "q255" + - "q256" + - "q257" + - "q265" + - "q296" + - "q331" + - "q370" + - "q413" + - "q462" + - "q516" + - "q577" + - "q645" + - "q721" + - "q806" + - "q901" + - "q1007" + - "q1125" + - "q1257" + - "q1405" + - "q1571" + - "q1756" + - "q1962" + - "q2048" + - "q2049" + - "q2193" + - "q2451" + - "q2740" + - "q3062" + - "q3422" + - "q3825" + - "q4275" + - "q4778" + - "q5340" + - "q5854" + - "q5969" + - "q6000" + - "q6144" + - "q6367" + - "q6671" + - "q6925" + - "q7168" + - "q7456" + - "q7532" + - "q8192" + - "q8193" + - "q8333" + - "q9314" + - "q10410" + - "q11635" + - "q13004" + - "q14534" + - "q16244" + - "q18156" + - "q20292" + - "q22680" + - "q25349" + - "q28332" + - "q31665" + - "q32768" + - "q32769" + - "q35391" + - "q39556" + - "q44210" + - "q49413" + - "q55227" + - "q61726" + - "q65536" + +backends: + - FLASHMLA_SPARSE + +device: "cuda:0" +profile_memory: true +sparse_mla_dense_mha_max_seq_len: 25349 +sparse_mla_masked_mha_max_seq_len: 8192 +sparse_mla_mha_variants: ["dense_mha", "masked_mha", "mqa"] +sparse_mla_topk_pattern: "random" diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index f78aca3bce40..6d56fb11ace2 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -931,10 +931,17 @@ def _run_single_benchmark( prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) fused_output = output_scale is not None and fuse_quant_op + mha_mode = getattr(config, "sparse_mla_mha_mode", "auto") # Build forward function (runs a single decode/prefill pass) def forward_fn(): results = [] + impl._sparse_mla_force_masked_mha = ( # type: ignore[attr-defined] + mha_mode == "masked" + ) + impl._sparse_mla_force_dense_mha = ( # type: ignore[attr-defined] + mha_mode == "dense" + ) if has_decode: results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)) if has_prefill: diff --git a/benchmarks/kernels/benchmark_inkling_qkvr_prep.py b/benchmarks/kernels/benchmark_inkling_qkvr_prep.py new file mode 100644 index 000000000000..8aa1233c50a8 --- /dev/null +++ b/benchmarks/kernels/benchmark_inkling_qkvr_prep.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import statistics + +import torch +from tabulate import tabulate + +from vllm.models.inkling.nvidia.ops import qkvr_prep +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def make_inputs(tokens: int, tp_size: int, is_local: bool): + torch.manual_seed(0) + num_q_heads = 64 // tp_size + num_kv_heads = (16 if is_local else 8) // tp_size + head_dim = 128 + d_rel = 16 + rel_extent = 512 if is_local else 1024 + page_size = 16 + num_blocks = (tokens + page_size - 1) // page_size + q_width = num_q_heads * head_dim + kv_width = num_kv_heads * head_dim + r_width = num_q_heads * d_rel + device = "cuda" + + qkvr = torch.randn( + tokens, + q_width + 2 * kv_width + r_width, + device=device, + dtype=torch.bfloat16, + ) + k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16) + v_weight = torch.randn_like(k_weight) + q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) + k_norm_weight = torch.randn_like(q_norm_weight) + rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16) + conv_cache = torch.zeros( + num_blocks, + num_kv_heads, + page_size, + 2 * head_dim, + device=device, + dtype=torch.bfloat16, + ) + key_cache = torch.empty( + num_blocks, + page_size, + num_kv_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ) + value_cache = torch.empty_like(key_cache) + positions = torch.arange(tokens, device=device, dtype=torch.int64) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None] + seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32) + slots = torch.arange(tokens, device=device, dtype=torch.int64) + query_start = torch.zeros(tokens, device=device, dtype=torch.int32) + log_scaling = None + if not is_local: + effective_n = (positions + 1).to(torch.float32) + log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0)) + return ( + qkvr, + k_weight, + v_weight, + q_norm_weight, + k_norm_weight, + rel_proj, + 1e-6, + num_q_heads, + num_kv_heads, + head_dim, + d_rel, + conv_cache, + key_cache, + value_cache, + positions, + block_table, + seq_idx, + slots, + query_start, + slots, + 0, + head_dim, + page_size, + log_scaling, + ) + + +def capture(implementation, inputs): + outputs = [] + + def run(): + outputs[:] = implementation.fused_qkvr_prep(*inputs) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + torch.accelerator.synchronize() + return graph, outputs + + +def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float: + for _ in range(warmup): + graph.replay() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000 / repeats + + +def benchmark(inputs, args) -> float: + graph, _ = capture(qkvr_prep, inputs) + return statistics.median( + time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials) + ) + + +@torch.inference_mode() +def main(args): + rows = [] + for tp_size in args.tp_sizes: + for tokens in args.tokens: + for is_local in (True, False): + triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args) + rows.append( + [ + tp_size, + tokens, + "local" if is_local else "global", + triton_us, + ] + ) + + print("Inkling QKVR prep (CUDA graph, median latency)") + print( + tabulate( + rows, + headers=[ + "TP", + "tokens", + "scope", + "Triton (us)", + ], + floatfmt=("d", "d", "", ".2f"), + ) + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser() + parser.add_argument( + "--tokens", + type=int, + nargs="+", + default=[1 << power for power in range(15)], + ) + parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=200) + parser.add_argument("--trials", type=int, default=5) + main(parser.parse_args()) diff --git a/benchmarks/kernels/benchmark_k3_cutedsl_residual.py b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py new file mode 100644 index 000000000000..7b9e349a792c --- /dev/null +++ b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py @@ -0,0 +1,367 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM. + +The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]`` +with FP32 accumulation and BF16 output. Both backends execute through CUDA +Graph replay. Weights and residuals rotate across buffers exceeding L2 so the +comparison models the full latent MoE projection-and-add path. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import importlib.util +import json +import math +import statistics +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings import driver as cuda +from cuda.bindings.driver import CUstream +from quack.compile_utils import make_fake_tensor + +N = 7168 +K = 3584 + + +@dataclasses.dataclass(frozen=True, slots=True) +class Config: + block_size: int + outputs_per_block: int + k_unroll: int + vector_width: int = 8 + + +def parse_config(value: str) -> Config: + try: + parts = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" + ) from error + if len(parts) == 3: + return Config(*parts) + if len(parts) == 4: + return Config(*parts) + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" + ) + + +def production_residual_config(m: int) -> Config | None: + """The measured Latent-MoE residual config for M, from the K3 table.""" + from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + + spec = KIMI_K3_PROJECTIONS.get((N, K)) + config = spec.residual_config(m) if spec is not None else None + if config is None: + return None + return Config( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]: + if mode == "selected": + if selected is not None: + return [selected] + # No explicit --config: fall back to the production table for this M. + config = production_residual_config(m) + return [config] if config is not None else [] + if mode == "baseline": + return [Config(224, 4, 2)] + return [ + Config(block_size, outputs_per_block, k_unroll, vector_width) + for vector_width in (4, 8) + for block_size in (32, 64, 128, 224, 448) + if block_size % 32 == 0 and K % (block_size * vector_width) == 0 + for outputs_per_block in (1, 2, 4, 7, 8) + if N % outputs_per_block == 0 + for k_unroll in (1, 2, 4) + ] + + +def load_kernel_class(path: Path): + spec = importlib.util.spec_from_file_location("cute_skinny_device", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load CuTe kernel from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.CuteSkinnyGemm + + +def stream() -> CUstream: + return CUstream(torch.cuda.current_stream().cuda_stream) + + +def compile_kernel(kernel_class, m: int, config: Config, max_registers: int): + element_type = cutlass.BFloat16 + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width) + b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width) + residual = make_fake_tensor(element_type, (m, n), divisibility=1) + c = make_fake_tensor(element_type, (m, n), divisibility=1) + kernel = kernel_class( + element_type=element_type, + num_rows=m, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=True, + use_pdl=True, + ) + return cute.compile( + kernel, + a, + b, + residual, + c, + stream(), + options=( + "--enable-tvm-ffi --keep-cubin " + f"--ptxas-options -maxrregcount={max_registers} " + "--ptxas-options -lineinfo" + ), + ) + + +def resource_usage(compiled) -> dict[str, Any]: + executor = getattr(compiled, "_default_executor", None) + context = getattr(executor, "exec_context", None) + functions = getattr(context, "kernel_functions", None) + if not functions: + return {"resource_metrics_available": False} + + def attribute(name, function) -> int: + error, value = cuda.cuFuncGetAttribute(name, function) + if error != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuFuncGetAttribute failed with {error}") + return int(value) + + registers = [ + attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function) + for function in functions + ] + local_bytes = [ + attribute( + cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + function, + ) + for function in functions + ] + return { + "resource_metrics_available": True, + "registers_per_thread": max(registers, default=0), + "spill_bytes": max(local_bytes, default=0), + } + + +def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int: + properties = torch.cuda.get_device_properties(0) + bytes_per_pair = (N * K + m * N) * 2 + target = math.ceil(multiplier * properties.L2_cache_size) + return max(2, min(limit, math.ceil(target / bytes_per_pair))) + + +def graph_samples( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None], + activation: torch.Tensor, + weights: Sequence[torch.Tensor], + residuals: Sequence[torch.Tensor], + repeats: int, + replays: int, +) -> tuple[list[float], list[torch.Tensor]]: + outputs = [torch.empty_like(residual) for residual in residuals] + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + for _ in range(20): + graph.replay() + torch.accelerator.synchronize() + + samples = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights))) + return samples, outputs + + +def summarize(samples: Sequence[float]) -> dict[str, Any]: + ordered = sorted(samples) + + def percentile(fraction: float) -> float: + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + mean = statistics.mean(samples) + return { + "median_us": statistics.median(samples), + "p10_us": percentile(0.1), + "p90_us": percentile(0.9), + "mean_us": mean, + "cv_pct": statistics.pstdev(samples) / mean * 100.0, + "samples_us": list(samples), + } + + +def correctness( + output: torch.Tensor, + activation: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> dict[str, Any]: + actual = output.float() + reference = activation.float() @ weight.float().t() + residual.float() + error = (actual - reference).abs() + scaled_error = error / (reference.abs() + 1.0) + cosine = torch.nn.functional.cosine_similarity( + actual.flatten(), reference.flatten(), dim=0 + ).item() + return { + "valid": cosine > 0.999, + "cosine": cosine, + "max_abs_error": error.max().item(), + "max_scaled_error": scaled_error.max().item(), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kernel", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--mode", choices=("baseline", "sweep", "selected"), default="baseline" + ) + parser.add_argument("--config", type=parse_config) + parser.add_argument("--m", type=int, action="append") + parser.add_argument("--config-shard", type=int, default=0) + parser.add_argument("--num-config-shards", type=int, default=1) + parser.add_argument("--repeats", type=int, default=21) + parser.add_argument("--replays", type=int, default=200) + parser.add_argument("--cache-multiplier", type=float, default=3.0) + parser.add_argument("--max-buffers", type=int, default=32) + parser.add_argument("--max-registers", type=int, default=64) + args = parser.parse_args() + + token_counts = args.m or list(range(1, 17)) + if any(not 1 <= m <= 16 for m in token_counts): + raise ValueError("expected 1 <= M <= 16") + if not 0 <= args.config_shard < args.num_config_shards: + raise ValueError("config shard must be in [0, num_config_shards)") + torch.accelerator.set_device_index(0) + if torch.cuda.get_device_capability() != (10, 3): + raise RuntimeError("this benchmark requires SM103") + + kernel_class = load_kernel_class(args.kernel) + properties = torch.cuda.get_device_properties(0) + metadata = { + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability()), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as output_file: + for m in token_counts: + configs = candidate_configs(args.mode, args.config, m) + torch.manual_seed(20260722 + m) + count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers) + activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn((N, K), device="cuda", dtype=torch.bfloat16) + for _ in range(count) + ] + residuals = [ + torch.randn((m, N), device="cuda", dtype=torch.bfloat16) + for _ in range(count) + ] + candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)] + candidates.extend( + ("cute_residual", config) + for index, config in enumerate(configs) + if index % args.num_config_shards == args.config_shard + ) + for backend, config in candidates: + row: dict[str, Any] = { + "m": m, + "n": N, + "k": K, + "backend": backend, + "mode": args.mode, + "config": dataclasses.asdict(config) if config else {}, + "num_buffers": count, + "cache_multiplier": args.cache_multiplier, + **metadata, + } + try: + if backend == "cublas_addmm": + launch = lambda a, b, residual, c: torch.addmm( + residual, a, b.t(), out=c + ) + else: + if config is None: + raise AssertionError("missing CuTe config") + compiled = compile_kernel( + kernel_class, m, config, args.max_registers + ) + launch = lambda a, b, residual, c, fn=compiled: fn( + a, b, residual, c, stream() + ) + row.update(resource_usage(compiled)) + samples, outputs = graph_samples( + launch, + activation, + weights, + residuals, + args.repeats, + args.replays, + ) + row.update( + correctness(outputs[0], activation, weights[0], residuals[0]) + ) + row.update(summarize(samples)) + except Exception as error: # noqa: BLE001 + row.update( + { + "valid": False, + "error": f"{type(error).__name__}: {error}", + } + ) + output_file.write(json.dumps(row, sort_keys=True) + "\n") + output_file.flush() + print(json.dumps(row, sort_keys=True), flush=True) + + del activation, weights, residuals + torch.accelerator.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py new file mode 100644 index 000000000000..3c1391919f88 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -0,0 +1,806 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels. + +The ``up-projection`` subcommand isolates the TP-local dynamic and static-M +skinny GEMMs. It rotates weights through a working set larger than L2 to model +successive model layers. + +The ``whole-tail`` subcommand measures the distributed operator. Its reference +path includes two AllReduces, RMSNorm, the replicated up-projection, and the +final add. CUDA-event samples report the slowest rank so cross-rank skew is +included. + +Examples: + +.. code-block:: console + + .venv/bin/python \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection + + torchrun --nproc-per-node=8 \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail + +For multi-node runs, launch one ``torchrun`` agent per node and use a shared +rendezvous endpoint. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +from collections.abc import Callable, Sequence +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.nn.functional as F +from cuda.bindings import driver as cuda + +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, + set_custom_all_reduce, +) +from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import ( + fused_add_multicast_gemm, + fused_add_multicast_skinny_gemm, +) + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +RMS_EPS = 0.1 +MAX_NUM_TOKENS = 16 +MMA_TILER_MN = (64, 32) +CLUSTER_SHAPE_MN = (1, 8) +B_PRIME_STAGES = 2 + + +def parse_up_projection_config( + value: str, +) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]" + ) from error + if len(values) in (3, 4): + return fused_add_multicast_skinny_gemm.SkinnyConfig(*values) + if len(values) == 5 and values[4] in (0, 1): + return fused_add_multicast_skinny_gemm.SkinnyConfig( + *values[:4], + prefetch_b_before_pdl=bool(values[4]), + ) + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL" + "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_tail_skinny_config( + value: str, +) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]" + ) from error + if len(values) == 4: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 5: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 6 and values[5] in (0, 1): + num_tokens, block, outputs, unroll, vector_width, prefetch = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig( + block, + outputs, + unroll, + vector_width, + bool(prefetch), + ) + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL" + "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="scope", required=True) + + up_projection = subparsers.add_parser( + "up-projection", + help="Benchmark the isolated TP-local up-projection kernels.", + ) + up_projection.add_argument( + "--backend", + choices=("dynamic", "skinny", "both"), + default="both", + ) + up_projection.add_argument("--tp-size", type=int, default=16) + up_projection.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[*range(1, 9), 16], + ) + up_projection.add_argument( + "--skinny-config", + type=parse_up_projection_config, + action="append", + help="Benchmark a static-M config for every selected token count.", + ) + up_projection.add_argument("--cache-multiplier", type=float, default=2.0) + up_projection.add_argument("--max-weights", type=int, default=64) + up_projection.add_argument("--warmup-replays", type=int, default=10) + up_projection.add_argument("--samples", type=int, default=31) + up_projection.add_argument("--output", type=Path) + + whole_tail = subparsers.add_parser( + "whole-tail", + help="Benchmark the distributed latent-MoE tail operator.", + ) + whole_tail.add_argument( + "--backend", + choices=("reference", "fused", "both"), + default="both", + ) + whole_tail.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[1, 5, 8, 16], + ) + whole_tail.add_argument("--warmup-replays", type=int, default=20) + whole_tail.add_argument("--samples", type=int, default=51) + whole_tail.add_argument( + "--skinny-max-num-tokens", + type=int, + nargs="+", + help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.", + ) + whole_tail.add_argument( + "--skinny-config", + type=parse_tail_skinny_config, + action="append", + help="Override one static-M config for tuning.", + ) + whole_tail.add_argument("--output", type=Path) + return parser.parse_args() + + +def percentile(samples: Sequence[float], fraction: float) -> float: + ordered = sorted(samples) + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + upper_weight = position - lower + return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight + + +def summarize(samples_us: Sequence[float]) -> dict[str, Any]: + mean_us = statistics.mean(samples_us) + return { + "median_us": statistics.median(samples_us), + "p10_us": percentile(samples_us, 0.1), + "p90_us": percentile(samples_us, 0.9), + "mean_us": mean_us, + "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0, + "samples_us": list(samples_us), + } + + +def rotating_weight_count( + shard_size: int, + cache_multiplier: float, + limit: int, +) -> int: + properties = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ) + weight_bytes = shard_size * LATENT_SIZE * 2 + target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier) + return max(2, min(limit, math.ceil(target_bytes / weight_bytes))) + + +def capture_up_projection_graph( + launches: Sequence[Callable[[], None]], +) -> torch.cuda.CUDAGraph: + for launch in launches: + launch() + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for launch in launches: + launch() + torch.accelerator.synchronize() + return graph + + +def benchmark_up_projection_graph( + graph: torch.cuda.CUDAGraph, + *, + operations_per_replay: int, + warmup_replays: int, + samples: int, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + samples_us = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + start.record() + graph.replay() + end.record() + end.synchronize() + samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay) + return summarize(samples_us) + + +class DynamicKernel: + def __init__( + self, + shard_size: int, + mailbox: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + self.shard_size = shard_size + self.mailbox = mailbox + self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox) + compile_latent = torch.empty( + (1, MAX_NUM_TOKENS, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + compile_weight = torch.empty( + (1, shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + cluster_size = math.prod(CLUSTER_SHAPE_MN) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + self.compiled = fused_add_multicast_gemm.compile_kernel( + (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1), + fused_add_multicast_gemm._as_cute( + compile_latent, + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(compile_weight), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + HIDDEN_SIZE, + shard_size, + MMA_TILER_MN, + CLUSTER_SHAPE_MN, + max_active_clusters, + B_PRIME_STAGES, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + self.compiled( + fused_add_multicast_gemm._as_cute( + latent.unsqueeze(0), + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + cutlass.Int64(latent.shape[0]), + cutlass.Int64(self.mailbox.data_ptr()), + stream, + ) + + +class SkinnyKernel: + def __init__( + self, + num_tokens: int, + shard_size: int, + config: fused_add_multicast_skinny_gemm.SkinnyConfig, + ) -> None: + self.compiled = fused_add_multicast_skinny_gemm.compile_kernel( + num_rows=num_tokens, + latent_dim=LATENT_SIZE, + hidden_dim=HIDDEN_SIZE, + shard_dim=shard_size, + config=config, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + ) -> None: + self.compiled( + fused_add_multicast_skinny_gemm._as_cute(latent), + fused_add_multicast_skinny_gemm._as_cute(weight), + fused_add_multicast_skinny_gemm._as_cute(shared_shard), + cutlass.Int64(mailbox.data_ptr()), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def check_up_projection_output( + actual: torch.Tensor, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, +) -> None: + gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16) + expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def make_up_projection_launches( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None], + latent: torch.Tensor, + weights: Sequence[torch.Tensor], + shared_shard: torch.Tensor, +) -> list[Callable[[], None]]: + return [ + lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights + ] + + +def benchmark_up_projection(args: argparse.Namespace) -> None: + if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size: + raise ValueError("TP size must be positive and divide the hidden size") + if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.cache_multiplier <= 0 or args.max_weights <= 0: + raise ValueError("cache multiplier and max weights must be positive") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + + torch.accelerator.set_device_index(0) + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + shard_size = HIDDEN_SIZE // args.tp_size + weight_count = rotating_weight_count( + shard_size, + args.cache_multiplier, + args.max_weights, + ) + torch.manual_seed(20260726) + weights = [ + torch.randn( + (shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + for _ in range(weight_count) + ] + mailbox = torch.empty( + (1, MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared = torch.randn( + (MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared_shard = shared[:, :shard_size] + use_dynamic = args.backend in ("dynamic", "both") + use_skinny = args.backend in ("skinny", "both") + dynamic_kernel = ( + DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None + ) + + results = [] + for num_tokens in args.num_tokens: + latent = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if dynamic_kernel is not None: + launches = make_up_projection_launches( + dynamic_kernel.launch, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + result["dynamic"] = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + if use_skinny: + configs = args.skinny_config or [ + fused_add_multicast_skinny_gemm.config_for_m( + num_tokens, + shard_size, + ) + ] + skinny_results = [] + for config in configs: + skinny_kernel = SkinnyKernel(num_tokens, shard_size, config) + + def launch_skinny( + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + *, + skinny_kernel: SkinnyKernel = skinny_kernel, + num_tokens: int = num_tokens, + ) -> None: + skinny_kernel.launch( + latent, + weight, + shared_shard[:num_tokens], + mailbox, + ) + + launches = make_up_projection_launches( + launch_skinny, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + timing = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + skinny_results.append( + { + "config": asdict(config), + **timing, + } + ) + result["skinny"] = skinny_results + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "up-projection", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "tp_size": args.tp_size, + "shard_size": shard_size, + "weight_count": weight_count, + "cache_multiplier": args.cache_multiplier, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "results": results, + } + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +def capture_tail_graph( + operation: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]: + for _ in range(3): + dist.barrier(group=cpu_group) + output = operation() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = operation() + torch.accelerator.synchronize() + return graph, output + + +def benchmark_tail_graph( + graph: torch.cuda.CUDAGraph, + *, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + for start, end in zip(starts, ends): + start.record() + graph.replay() + end.record() + torch.accelerator.synchronize() + + samples_us = torch.tensor( + [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)], + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group) + return summarize(samples_us[1:].tolist()) + + +def make_inputs( + num_tokens: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(20260726 + 100 * num_tokens + rank) + routed = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ).mul_(0.01) + shared = torch.randn( + (num_tokens, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + return routed, shared + + +def make_reference( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + device_group: dist.ProcessGroup, +) -> Callable[[], torch.Tensor]: + routed_workspace = torch.empty_like(routed) + shared_workspace = torch.empty_like(shared) + + def reference() -> torch.Tensor: + routed_workspace.copy_(routed) + dist.all_reduce(routed_workspace, group=device_group) + normalized = F.rms_norm( + routed_workspace, + (LATENT_SIZE,), + rms_weight, + RMS_EPS, + ) + projected = F.linear(normalized, up_weight) + shared_workspace.copy_(shared) + dist.all_reduce(shared_workspace, group=device_group) + return projected.add(shared_workspace) + + return reference + + +def check_fused_output( + fused_output: torch.Tensor, + reference: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> None: + dist.barrier(group=cpu_group) + expected = reference() + torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2) + + +def benchmark_whole_tail(args: argparse.Namespace) -> None: + if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + if args.skinny_max_num_tokens is not None and any( + not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens + ): + raise ValueError("--skinny-max-num-tokens must be in [0, 8]") + skinny_configs = dict(args.skinny_config or ()) + if len(skinny_configs) != len(args.skinny_config or ()): + raise ValueError("--skinny-config must not repeat an M value") + if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs): + raise ValueError("--skinny-config M values must be in [1, 8]") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys(): + raise RuntimeError("launch this benchmark with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + device = torch.device("cuda", local_rank) + torch.accelerator.set_device_index(device) + init_distributed_environment() + if world_size > 8: + set_custom_all_reduce(False) + initialize_model_parallel(tensor_model_parallel_size=world_size) + device_group = get_tp_group().device_group + cpu_group = dist.new_group(backend="gloo") + + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + torch.manual_seed(20260726) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + ) + up_weight = ( + torch.randn( + (HIDDEN_SIZE, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + ) + + use_reference = args.backend in ("reference", "both") + use_fused = args.backend in ("fused", "both") + fused_ops = [] + if use_fused: + production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m + + def config_for_m( + num_rows: int, + shard_dim: int = 896, + ) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + config = skinny_configs.get(num_rows) + if config is not None: + return config + return production_config_for_m(num_rows, shard_dim) + + fused_add_multicast_skinny_gemm.config_for_m = config_for_m + cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS] + for cutoff in cutoffs: + latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff + latent_moe_tail.KimiK3LatentMoETailOp._instances.clear() + fused_ops.append( + ( + cutoff, + latent_moe_tail.KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=RMS_EPS, + ), + ) + ) + cutedsl_warmup() + + results = [] + for num_tokens in args.num_tokens: + routed, shared = make_inputs(num_tokens, rank, device) + reference = make_reference( + routed, + shared, + rms_weight, + up_weight, + device_group, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if use_reference: + reference_graph, _ = capture_tail_graph(reference, cpu_group) + result["reference"] = benchmark_tail_graph( + reference_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + for cutoff, fused_op in fused_ops: + + def fused( + routed: torch.Tensor = routed, + shared: torch.Tensor = shared, + fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op, + ) -> torch.Tensor: + return fused_op(routed, shared, rms_weight, up_weight) + + fused_graph, fused_output = capture_tail_graph(fused, cpu_group) + fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}" + result[fused_key] = benchmark_tail_graph( + fused_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + check_fused_output(fused_output, reference, cpu_group) + if "reference" in result: + speedup = ( + result["reference"]["median_us"] / result[fused_key]["median_us"] + ) + if len(fused_ops) == 1: + result["speedup"] = speedup + else: + result[f"{fused_key}_speedup"] = speedup + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "whole-tail", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "world_size": world_size, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops], + "skinny_configs": { + str(num_tokens): asdict(config) + for num_tokens, config in skinny_configs.items() + }, + "timing_scope": { + "reference": ( + "two input copies, two AllReduces, RMSNorm, full replicated " + "up-projection GEMM, and final add" + ), + "fused": ( + "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded " + "up-projection/multicast, and Lamport copy" + ), + }, + "results": results, + } + if rank == 0: + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + dist.barrier(group=cpu_group) + + +def main() -> None: + args = parse_args() + if args.scope == "up-projection": + benchmark_up_projection(args) + return + + from vllm.config import VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + benchmark_whole_tail(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py new file mode 100644 index 000000000000..9a244e6c8265 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import os +import statistics +from collections.abc import Callable + +import torch +import torch.distributed as dist + +import vllm._custom_ops as ops +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024]) + parser.add_argument("--hidden-size", type=int, default=7168) + parser.add_argument("--graph-repeats", type=int, default=20) + parser.add_argument("--warmup-replays", type=int, default=5) + parser.add_argument("--samples", type=int, default=15) + return parser.parse_args() + + +def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + op() + stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(repeats): + op() + torch.cuda.current_stream().wait_stream(stream) + return graph + + +def max_rank_graph_time( + graph: torch.cuda.CUDAGraph, + repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> float: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + timings = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + dist.barrier(group=cpu_group) + start.record() + graph.replay() + end.record() + end.synchronize() + elapsed = torch.tensor( + start.elapsed_time(end) / repeats, + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) + timings.append(elapsed.item()) + return statistics.median(timings) + + +def check_outputs( + comm: CustomAllreduce, + local: torch.Tensor, + reduce_input: torch.Tensor, + device_group: dist.ProcessGroup, +) -> None: + expected_gather = torch.empty( + (local.shape[0] * dist.get_world_size(), local.shape[1]), + dtype=local.dtype, + device=local.device, + ) + dist.all_gather_into_tensor(expected_gather, local, group=device_group) + gathered = comm.custom_all_gather(local) + assert gathered is not None + torch.testing.assert_close(gathered, expected_gather) + + expected_scatter = torch.empty_like(local) + dist.reduce_scatter_tensor( + expected_scatter, + reduce_input.clone(), + group=device_group, + ) + scattered = comm.custom_reduce_scatter(reduce_input) + assert scattered is not None + torch.testing.assert_close(scattered, expected_scatter) + + +def benchmark_shape( + comm: CustomAllreduce, + global_tokens: int, + hidden_size: int, + graph_repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, float | int]: + world_size = dist.get_world_size() + rank = dist.get_rank() + padded_tokens = (global_tokens + world_size - 1) // world_size * world_size + local_tokens = padded_tokens // world_size + local = torch.full( + (local_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=torch.accelerator.current_device_index(), + ) + reduce_input = torch.full( + (padded_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=local.device, + ) + check_outputs(comm, local, reduce_input, device_group) + + custom_gather_out = torch.empty( + (padded_tokens, hidden_size), + dtype=local.dtype, + device=local.device, + ) + custom_scatter_out = torch.empty_like(local) + nccl_gather_out = torch.empty_like(custom_gather_out) + nccl_scatter_out = torch.empty_like(local) + + def custom_ag() -> None: + ops.mnnvl_lamport_all_gather( + comm._ptr, + local, + custom_gather_out, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_ag_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def custom_rs() -> None: + ops.mnnvl_lamport_reduce_scatter( + comm._ptr, + reduce_input, + custom_scatter_out, + comm.mnnvl_lamport_rs_local_ptr, + comm.mnnvl_lamport_rs_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def nccl_ag() -> None: + dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group) + + def nccl_rs() -> None: + dist.reduce_scatter_tensor( + nccl_scatter_out, + reduce_input, + group=device_group, + ) + + graphs = { + "custom_ag_us": capture_graph(custom_ag, graph_repeats), + "nccl_ag_us": capture_graph(nccl_ag, graph_repeats), + "custom_rs_us": capture_graph(custom_rs, graph_repeats), + "nccl_rs_us": capture_graph(nccl_rs, graph_repeats), + } + times = { + name: max_rank_graph_time( + graph, + graph_repeats, + warmup_replays, + samples, + device_group, + cpu_group, + ) + * 1000 + for name, graph in graphs.items() + } + torch.testing.assert_close(custom_gather_out, nccl_gather_out) + torch.testing.assert_close(custom_scatter_out, nccl_scatter_out) + return { + "global_tokens": global_tokens, + "padded_tokens": padded_tokens, + "local_bytes": local.nbytes, + "full_bytes": reduce_input.nbytes, + **times, + "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"], + "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"], + } + + +def main() -> None: + args = parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group("nccl") + device_group = dist.group.WORLD + cpu_group = dist.new_group(backend="gloo") + + comm = CustomAllreduce( + group=cpu_group, + device=torch.device("cuda", local_rank), + ) + assert not comm.disabled + assert comm.world_size == 16 + assert comm.mnnvl_only + assert comm.mnnvl_multicast_ptr + + results = [ + benchmark_shape( + comm, + tokens, + args.hidden_size, + args.graph_repeats, + args.warmup_replays, + args.samples, + device_group, + cpu_group, + ) + for tokens in args.tokens + ] + if dist.get_rank() == 0: + print(json.dumps(results, indent=2), flush=True) + + comm.close() + dist.destroy_process_group(cpu_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py new file mode 100644 index 000000000000..6bb504b7f090 --- /dev/null +++ b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Benchmark complete AITER BF16 and FP8 ViT attention calls on ROCm. + +The FP8 timing includes Q/K/V quantization and attention. Dynamic scales are +calibrated once per input before timing so the measured path matches serving +with a static scale file. + +Example: + python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \ + --seq-lens 2304 4096 8192 16384 --head-dim 72 +""" + +from types import SimpleNamespace + +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.multimodal import MultiModalConfig +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, +) +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +def make_attention(num_heads: int, head_dim: int, fp8: bool) -> MMEncoderAttention: + mm_config = MultiModalConfig( + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA, + mm_encoder_attn_dtype="fp8" if fp8 else None, + ) + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + with set_current_vllm_config(vllm_config): + return MMEncoderAttention(num_heads, head_dim).to("cuda") + + +def bench( + seq_lens: list[int], + num_heads: int, + head_dim: int, + warmup_ms: int, + repeat_ms: int, +) -> None: + if not current_platform.is_rocm(): + raise RuntimeError("This benchmark requires ROCm and AITER.") + + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + fp8_attention = make_attention(num_heads, head_dim, fp8=True) + bf16_attention = make_attention(num_heads, head_dim, fp8=False) + finally: + torch.set_default_dtype(old_dtype) + + print( + f"{'seq_len':>8} {'BF16 ms':>12} {'FP8 ms':>12} " + f"{'speedup':>10} {'FP8/BF16':>12}" + ) + print("-" * 60) + for seq_len in seq_lens: + torch.manual_seed(0) + qkv = torch.randn( + 1, + seq_len, + 3, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + q, k, v = qkv.unbind(dim=2) + cu_seqlens = torch.tensor([0, seq_len], device="cuda", dtype=torch.int32) + max_seqlen = torch.tensor(seq_len, device="cuda", dtype=torch.int32) + + # Calibrate per-tensor scales once, then benchmark the static-scale path. + fp8_attention._fp8_dynamic_scale = True + fp8_attention._forward_aiter_fp8(q, k, v, cu_seqlens, max_seqlen) + fp8_attention._fp8_dynamic_scale = False + + bf16_ms = triton.testing.do_bench( + lambda q=q, k=k, v=v, cu=cu_seqlens, ms=max_seqlen: ( + bf16_attention._forward_fa(q, k, v, cu, ms) + ), + warmup=warmup_ms, + rep=repeat_ms, + ) + fp8_ms = triton.testing.do_bench( + lambda q=q, k=k, v=v, cu=cu_seqlens, ms=max_seqlen: ( + fp8_attention._forward_aiter_fp8(q, k, v, cu, ms) + ), + warmup=warmup_ms, + rep=repeat_ms, + ) + speedup = bf16_ms / fp8_ms + ratio = fp8_ms / bf16_ms + print( + f"{seq_len:>8} {bf16_ms:>12.3f} {fp8_ms:>12.3f} " + f"{speedup:>9.2f}x {ratio:>12.3f}" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark AITER BF16 vs FP8 ViT attention." + ) + parser.add_argument( + "--seq-lens", type=int, nargs="+", default=[2304, 4096, 8192, 16384] + ) + parser.add_argument("--num-heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=72) + parser.add_argument("--warmup-ms", type=int, default=100) + parser.add_argument("--repeat-ms", type=int, default=500) + args = parser.parse_args() + bench( + args.seq_lens, + args.num_heads, + args.head_dim, + args.warmup_ms, + args.repeat_ms, + ) diff --git a/benchmarks/kernels/cpu/benchmark_cpu_attn.py b/benchmarks/kernels/cpu/benchmark_cpu_attn.py index 08afd693c333..cdfcb1d404ae 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_attn.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_attn.py @@ -154,7 +154,7 @@ def run_benchmark(iters: int) -> list[float]: scale=scale, causal=True, alibi_slopes=None, - sliding_window=window_size, + sliding_window=window_size if sliding_window is not None else -1, block_table=block_tables, softcap=0, scheduler_metadata=metadata, diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py new file mode 100644 index 000000000000..b7084a21315b --- /dev/null +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel. + +Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a +long greedy decode (CUDA graphs on) once with the standard kernel and once with +ReplaySSM, then reports the per-step / throughput speedup. The two modes run in +separate subprocesses so each gets a clean CUDA context. + +The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under +CUDA-graph capture on the pre-release Blackwell FP4 path); pass +--no-disable-flashinfer-autotune for non-FP4 models. + +Examples: + python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 +""" + +import argparse +import json +import os +import subprocess +import sys +import time + +DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and" + +MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"} + + +def parse_args(): + p = argparse.ArgumentParser( + description="E2E decode speedup: ReplaySSM vs the standard SSM kernel." + ) + p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16") + p.add_argument("--prompt", default=DEFAULT_PROMPT) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--num-steps", type=int, default=1000) + p.add_argument("--warmup-steps", type=int, default=128) + p.add_argument("--repeats", type=int, default=1) + p.add_argument( + "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length." + ) + p.add_argument( + "--dtype", + default="bfloat16", + choices=["bfloat16", "float16", "float32", "auto"], + ) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--max-model-len", type=int, default=None) + p.add_argument( + "--disable-flashinfer-autotune", + action=argparse.BooleanOptionalAction, + default=True, + help="Disable the FlashInfer FP4-MoE autotuner (default: on). " + "It is unstable under CUDA-graph capture on the " + "pre-release Blackwell FP4 path; pass " + "--no-disable-flashinfer-autotune for non-FP4 models.", + ) + p.add_argument( + "--mamba-ssm-cache-dtype", + default="auto", + choices=["auto", "float32", "float16", "bfloat16"], + help="SSM state dtype (both modes). 'auto' = config-driven; " + "'float32' = fp32 state, 'bfloat16' = s16 state.", + ) + p.add_argument( + "--baseline-ssm-config", + default="", + help="Pin the STANDARD baseline's SSM launch config as " + "'bsm,nw' via override_ssm_config (forces the in-process " + "engine so the override reaches the kernel). Empty = off.", + ) + p.add_argument( + "--worker", + choices=["standard", "replayssm"], + default=None, + help=argparse.SUPPRESS, + ) + return p.parse_args() + + +def resolve_max_model_len(args) -> int: + if args.max_model_len is not None: + return args.max_model_len + return args.num_steps + 256 + + +def run_worker(args): + # override_ssm_config is a module global; it only reaches the model if the + # engine runs in-process (default V1 spawns a separate EngineCore). Force it. + if args.worker == "standard" and args.baseline_ssm_config: + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + + import torch + + from vllm import LLM, SamplingParams + + mode = args.worker + max_model_len = resolve_max_model_len(args) + + llm_kwargs = dict( + model=args.model_id, + tensor_parallel_size=1, + dtype=args.dtype, + max_model_len=max_model_len, + trust_remote_code=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + max_num_seqs=args.batch_size, + max_num_batched_tokens=max(max_model_len, args.batch_size * 64), + enforce_eager=False, + disable_log_stats=True, + gpu_memory_utilization=args.gpu_memory_utilization, + # SSM state dtype (applies to both standard and ReplaySSM). + mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, + ) + if args.disable_flashinfer_autotune: + # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell; + # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models. + llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + if mode == "replayssm": + llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len) + + _ssm_cm = None + if mode == "standard" and args.baseline_ssm_config: + from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config + + _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(",")) + _ssm_cm = override_ssm_config((_bsm, _nw)) + _ssm_cm.__enter__() # active through LLM() graph capture + decode + print( + f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})", + flush=True, + ) + + llm = LLM(**llm_kwargs) + prompts = [args.prompt] * args.batch_size + + def timed_generate(n_tokens): + sp = SamplingParams( + n=1, + temperature=0.0, + ignore_eos=True, + min_tokens=n_tokens, + max_tokens=n_tokens, + ) + if torch.accelerator.is_available(): + torch.accelerator.synchronize() + t0 = time.perf_counter() + outs = llm.generate(prompts, sp, use_tqdm=False) + if torch.accelerator.is_available(): + torch.accelerator.synchronize() + elapsed = time.perf_counter() - t0 + produced = min(len(o.outputs[0].token_ids) for o in outs) + assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}" + return elapsed + + timed_generate(args.warmup_steps) + + best = None + for _ in range(args.repeats): + elapsed = timed_generate(args.num_steps) + tok_s = args.batch_size * args.num_steps / elapsed + per_step_ms = elapsed / args.num_steps * 1e3 + print( + f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step", + flush=True, + ) + if best is None or elapsed < best["elapsed_s"]: + best = { + "mode": mode, + "elapsed_s": elapsed, + "tok_s": tok_s, + "per_step_ms": per_step_ms, + } + + print("RESULT_JSON " + json.dumps(best), flush=True) + if _ssm_cm is not None: + _ssm_cm.__exit__(None, None, None) + + +def run_one_mode(args, mode) -> dict: + cmd = [ + sys.executable, + __file__, + "--worker", + mode, + "--model-id", + args.model_id, + "--prompt", + args.prompt, + "--batch-size", + str(args.batch_size), + "--num-steps", + str(args.num_steps), + "--warmup-steps", + str(args.warmup_steps), + "--repeats", + str(args.repeats), + "--buffer-len", + str(args.buffer_len), + "--dtype", + args.dtype, + "--gpu-memory-utilization", + str(args.gpu_memory_utilization), + "--mamba-ssm-cache-dtype", + args.mamba_ssm_cache_dtype, + "--baseline-ssm-config", + args.baseline_ssm_config, + ] + cmd.append( + "--disable-flashinfer-autotune" + if args.disable_flashinfer_autotune + else "--no-disable-flashinfer-autotune" + ) + if args.max_model_len is not None: + cmd += ["--max-model-len", str(args.max_model_len)] + + result = None + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1 + ) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + if line.startswith("RESULT_JSON "): + result = json.loads(line[len("RESULT_JSON ") :]) + proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}") + if result is None: + raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line") + return result + + +def main(): + args = parse_args() + if args.worker is not None: + run_worker(args) + return + + print( + f"model={args.model_id} batch_size={args.batch_size} " + f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}" + ) + + std = run_one_mode(args, "standard") + fla = run_one_mode(args, "replayssm") + speedup = std["per_step_ms"] / fla["per_step_ms"] + + print() + header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}" + print(header) + print("-" * len(header)) + for r in (std, fla): + print( + f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}" + f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}" + ) + print("-" * len(header)) + print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x") + + +if __name__ == "__main__": + main() diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 3aca9bcea91d..cdfc6cb6f84f 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_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM}) set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -96,12 +97,14 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") set(ENABLE_NUMA OFF) check_sysctl(hw.optional.neon ASIMD_FOUND) check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND) + check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND) else() find_isa(${CPUINFO} "Power11" POWER11_FOUND) find_isa(${CPUINFO} "POWER10" POWER10_FOUND) find_isa(${CPUINFO} "POWER9" POWER9_FOUND) find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support + find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support find_isa(${CPUINFO} "S390" S390_FOUND) find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support @@ -111,6 +114,11 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + if (ENABLE_ARM_I8MM) + set(ARM_I8MM_FOUND ON) + message(STATUS + "ARM I8MM support enabled via VLLM_CPU_ARM_I8MM 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. @@ -166,6 +174,11 @@ elseif (ASIMD_FOUND) message(WARNING "BF16 functionality is not available") set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16") endif() + if(ARM_I8MM_FOUND) + message(STATUS "I8MM extension detected") + string(APPEND MARCH_FLAGS "+i8mm") + add_compile_definitions(ARM_I8MM_SUPPORT) + endif() list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS}) elseif (S390_FOUND) message(STATUS "S390 detected") @@ -244,7 +257,7 @@ endif() # Build oneDNN for GEMM kernels -if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND) +if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND OR S390_FOUND) # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64 # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "") @@ -355,7 +368,23 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE}) set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size - FetchContent_MakeAvailable(oneDNN) + + if(S390_FOUND) + FetchContent_GetProperties(oneDNN) + if(NOT onednn_POPULATED) + FetchContent_Populate(oneDNN) + # Patch s390x helpers.h: ALWAYS_INLINE on operator+= breaks C++20/GCC14 + file(READ "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" _helpers_content) + string(REPLACE + "vec_type_t &ALWAYS_INLINE operator+=" + "ALWAYS_INLINE vec_type_t &operator+=" + _helpers_content "${_helpers_content}") + file(WRITE "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" "${_helpers_content}") + add_subdirectory("${onednn_SOURCE_DIR}" "${onednn_BINARY_DIR}") + endif() + else() + FetchContent_MakeAvailable(oneDNN) + endif() set(CMAKE_BUILD_TYPE ${VLLM_BUILD_TYPE}) add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp") target_include_directories( @@ -430,6 +459,7 @@ set(VLLM_EXT_SRC "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp") @@ -446,8 +476,13 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) "csrc/cpu/shm.cpp" "csrc/cpu/activation_lut_bf16.cpp" "csrc/cpu/cpu_tanhf_neon.hpp" - "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) + if (ARM_BF16_FOUND) + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) + if (ARM_I8MM_FOUND) + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC}) + endif() + endif() endif() if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) @@ -489,6 +524,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/dnnl_kernels.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" @@ -502,6 +538,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 38d218d00acb..be0eccd45aed 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # NOTE: This is currently targeting nv-dev branch due to sm120 support - set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # TODO: switch to nv_dev branch after it support situ + set(_DEEPGEMM_UPSTREAM_TAG "e21c821f39a2056d68067a466c64ddc942200106") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -68,6 +68,9 @@ endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f") + endif() else() list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") endif() diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake new file mode 100644 index 000000000000..28b471a9e746 --- /dev/null +++ b/cmake/external_projects/flashkda.cmake @@ -0,0 +1,74 @@ +include(FetchContent) + +if(DEFINED ENV{FLASH_KDA_SRC_DIR}) + set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR}) +endif() + +if(FLASH_KDA_SRC_DIR) + FetchContent_Declare( + flashkda + SOURCE_DIR ${FLASH_KDA_SRC_DIR} + ) +else() + FetchContent_Declare( + flashkda + GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git + GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f + GIT_PROGRESS TRUE + GIT_SUBMODULES cutlass + ) +endif() + +FetchContent_MakeAvailable(flashkda) +message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}") + +set(FLASH_KDA_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a") +endif() +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f") +elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a") +endif() + +cuda_archs_loose_intersection( + FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}") + +if(FLASH_KDA_ARCHS) + message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}") + + set(FLASH_KDA_SOURCES + csrc/flashkda_registration.cpp + ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp + ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu) + set(FLASH_KDA_INCLUDES + ${flashkda_SOURCE_DIR}/csrc + ${flashkda_SOURCE_DIR}/cutlass/include + ${flashkda_SOURCE_DIR}/cutlass/examples/common + ${flashkda_SOURCE_DIR}/cutlass/tools/util/include) + + set_gencode_flags_for_srcs( + SRCS "${FLASH_KDA_SOURCES}" + CUDA_ARCHS "${FLASH_KDA_ARCHS}") + + define_extension_target( + _flashkda_C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${FLASH_KDA_SOURCES} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + target_compile_options(_flashkda_C PRIVATE + $<$:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3> + $<$:-UPy_LIMITED_API>) +else() + message(STATUS + "FlashKDA will not compile: CUDA >=12.0 and a supported architecture " + "(SM90, SM10x, or SM12x) are required") + add_custom_target(_flashkda_C) +endif() diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 56f7a83f6786..4d8c2d74c20c 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -60,6 +60,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) # CUDA 12.9 has introduced "Family-Specific Architecture Features" # this supports all compute_10x family list(APPEND SUPPORT_ARCHS "10.0f") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + list(APPEND SUPPORT_ARCHS "10.7f") + endif() elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) list(APPEND SUPPORT_ARCHS "10.0a") endif() @@ -188,4 +191,3 @@ else() add_custom_target(_flashmla_C) add_custom_target(_flashmla_extension_C) endif() - diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 052966b27552..442097dec7d3 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73 + GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index 29c5c6528b9c..cd3665c889e9 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR) set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused") else() set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git") - set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65") + set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64") set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _qutlass_fc_root) @@ -55,7 +55,11 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}") - cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}") + endif() else() cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}") @@ -125,8 +129,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) CUDA_ARCHS "${QUTLASS_ARCHS}" ) - # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION. - # Keep it as its own extension (registers torch.ops._qutlass_C). define_extension_target( _qutlass_C DESTINATION vllm @@ -139,9 +141,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) WITH_SOABI) target_compile_definitions(_qutlass_C PRIVATE - QUTLASS_DISABLE_PYBIND=1 + QUTLASS_MINIMAL_BUILD=1 TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC} - CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1 + TORCH_TARGET_VERSION=0x020B000000000000ULL + USE_CUDA) set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS $<$:--expt-relaxed-constexpr --use_fast_math -O3> diff --git a/cmake/external_projects/triton_kernels.cmake b/cmake/external_projects/triton_kernels.cmake index 2966c78030bd..8312c1f7ca22 100644 --- a/cmake/external_projects/triton_kernels.cmake +++ b/cmake/external_projects/triton_kernels.cmake @@ -1,7 +1,5 @@ # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels -set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1") - # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to # be directly set to the triton_kernels python directory. if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) @@ -12,13 +10,20 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) ) else() - set(TRITON_GIT "https://github.com/triton-lang/triton.git") - message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}") + if (VLLM_TARGET_DEVICE STREQUAL "rocm") + set(TRITON_GIT "https://github.com/ROCm/triton.git") + # Pinned from release/internal/3.6.x + set(TRITON_KERNELS_TAG "0f380657dbf3ee86eb57558ff71df24f03b5d4e7") + else() + set(TRITON_GIT "https://github.com/triton-lang/triton.git") + set(TRITON_KERNELS_TAG "v3.5.1") + endif() + message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${TRITON_KERNELS_TAG}") FetchContent_Declare( triton_kernels # TODO (varun) : Fetch just the triton_kernels directory from Triton - GIT_REPOSITORY https://github.com/triton-lang/triton.git - GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG} + GIT_REPOSITORY ${TRITON_GIT} + GIT_TAG ${TRITON_KERNELS_TAG} GIT_PROGRESS TRUE SOURCE_SUBDIR python/triton_kernels/triton_kernels ) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 28a6b336e27b..c7ebbb7fd0a3 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8 + GIT_TAG 28e862d21806bc3580207aa0ad4e2759151e9827 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/cmake/utils.cmake b/cmake/utils.cmake index e3e766541df1..bbae89c1f57c 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -241,14 +241,15 @@ endmacro() # `.`, dedupes them and then sorts them in ascending order and # stores them in `OUT_ARCHES`. # -# Example: -# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" -# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS) -# OUT_ARCHES="7.5;...;9.0" +# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags. +# This handles mismatches such as `arch=compute_20,code=sm_121`. function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) set(_CUDA_ARCHES) foreach(_ARCH ${CUDA_ARCH_FLAGS}) - string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + if (NOT _COMPUTE) + string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + endif() if (_COMPUTE) set(_COMPUTE ${CMAKE_MATCH_1}) endif() @@ -396,14 +397,24 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x # family. The output uses TGT's value to preserve the user's compilation flags. set(_CUDA_ARCHS) + # Resolve exact base matches before family fallbacks so a generic entry such + # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry. foreach(_arch ${_SRC_CUDA_ARCHS}) if(_arch MATCHES "[af]$") - list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") string(REGEX REPLACE "[af]$" "" _base "${_arch}") - if ("${_base}" IN_LIST TGT_CUDA_ARCHS) + if("${_base}" IN_LIST _TGT_CUDA_ARCHS) + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") list(APPEND _CUDA_ARCHS "${_arch}") - elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS) + endif() + endif() + endforeach() + + foreach(_arch ${_SRC_CUDA_ARCHS}) + if(_arch MATCHES "[af]$") + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") + string(REGEX REPLACE "[af]$" "" _base "${_arch}") + if("${_base}a" IN_LIST _TGT_CUDA_ARCHS) list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a") list(APPEND _CUDA_ARCHS "${_base}a") elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS) @@ -487,7 +498,7 @@ endfunction() function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") + cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") else() cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}") endif() diff --git a/csrc/core/scalar_type.hpp b/csrc/core/scalar_type.hpp index b6f39ed795f3..e1e3490c2d90 100644 --- a/csrc/core/scalar_type.hpp +++ b/csrc/core/scalar_type.hpp @@ -327,6 +327,8 @@ static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); +static inline constexpr auto kFE2M3f = + ScalarType::float_(2, 3, true, ScalarType::NAN_NONE); static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); static inline constexpr auto kFE8M0fnu = @@ -346,6 +348,7 @@ static inline constexpr auto kUint8b128 = kU8B128; static inline constexpr auto kFloat4_e2m1f = kFE2M1f; static inline constexpr auto kFloat6_e3m2f = kFE3M2f; +static inline constexpr auto kFloat6_e2m3f = kFE2M3f; static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; static inline constexpr auto kFloat8_e5m2 = kFE5M2; static inline constexpr auto kFloat16_e8m7 = kFE8M7; diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index 53ae70497c0f..ccec4c2e87c7 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -172,4 +172,26 @@ #endif // __riscv_v +// Power VSX +#ifdef __powerpc__ + // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which + // implements a vectorised 5-term minimax polynomial using VSX intrinsics. + #define DEFINE_FAST_EXP \ + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { return vec.exp(); }; \ + auto fast_exp_f16 = fast_exp; + +#endif // __powerpc__ + +// IBM Z (s390x) VXE +#ifdef __s390x__ + // FP32Vec16::exp() in cpu_types_vxe.hpp delegates to FP32Vec8::exp(), which + // implements a vectorised 5-term minimax polynomial using VXE intrinsics. + #define DEFINE_FAST_EXP \ + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { return vec.exp(); }; \ + auto fast_exp_f16 = fast_exp; + +#endif // __s390x__ + #endif diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index fa22861157e2..52c806f4fa37 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -30,7 +30,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool causal, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal) { + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype) { cpu_attention::ISA isa; if (isa_hint == "amx") { isa = cpu_attention::ISA::AMX; @@ -65,9 +66,11 @@ torch::Tensor get_scheduler_metadata( input.dynamic_causal = dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + const int64_t kv_cache_idx = + static_cast(parse_fp8_kv_dtype(kv_cache_dtype)); VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { - CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() { - input.elem_size = sizeof(scalar_t); + CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() { + input.elem_size = sizeof(attn_impl::kv_cache_t); input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); input.output_buffer_elem_size = diff --git a/csrc/cpu/cpu_attn_vec.hpp b/csrc/cpu/cpu_attn_vec.hpp index c3983e0578a5..0f3e92a2f4a6 100644 --- a/csrc/cpu/cpu_attn_vec.hpp +++ b/csrc/cpu/cpu_attn_vec.hpp @@ -102,7 +102,9 @@ class TileGemm82 { kv_cache_t* __restrict__ curr_b = b_tile; for (int32_t k = 0; k < dynamic_k_size; ++k) { - auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b); + auto fp32_b_regs = load_b_pair_vec(curr_b); + auto fp32_b_0_reg = fp32_b_regs.first; + auto fp32_b_1_reg = fp32_b_regs.second; float* __restrict__ curr_m_a = curr_a; vec_op::unroll_loop([&](int32_t i) { diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 07b0aaf86888..b7bbef4c2d77 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -1,5 +1,6 @@ #include "cpu/cpu_types.hpp" #include "cpu/utils.hpp" +#include "cpu/cpu_fused_moe_activations.hpp" #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp" #include "cpu/cpu_arch_macros.h" @@ -43,193 +44,9 @@ }() namespace { -enum class FusedMOEAct { - SiluAndMul, - SwigluOAIAndMul, - GeluAndMul, - GeluTanhAndMul, -}; - -FusedMOEAct get_act_type(const std::string& act) { - if (act == "silu") { - return FusedMOEAct::SiluAndMul; - } else if (act == "swigluoai") { - return FusedMOEAct::SwigluOAIAndMul; - } else if (act == "gelu") { - return FusedMOEAct::GeluAndMul; - } else if (act == "gelu_tanh") { - return FusedMOEAct::GeluTanhAndMul; - } else { - TORCH_CHECK(false, "Invalid act type: " + act); - } -} -template -void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, - const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; -#if !defined(__aarch64__) - // For GPT-OSS interleaved gate-up weights - alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, - 16, 18, 20, 22, 24, 26, 28, 30}; - vec_op::INT32Vec16 index_vec(index); -#endif - vec_op::FP32Vec16 gate_up_max_vec(7.0); - vec_op::FP32Vec16 up_min_vec(-7.0); - vec_op::FP32Vec16 alpha_vec(1.702); - vec_op::FP32Vec16 one_vec(1.0); - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < n_size; n += 32) { - // Note: AdvSIMD does not support gather loads -#if defined(__aarch64__) - vec_op::FP32Vec16 gate_vec(vec_op::uninit); - vec_op::FP32Vec16 up_vec(vec_op::uninit); - vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); -#else - vec_op::FP32Vec16 gate_vec(input + n, index_vec); - vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); -#endif - gate_vec = gate_vec.min(gate_up_max_vec); - up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec); - auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec)); - auto glu = gate_vec * sigmoid_vec; - auto gated_output_fp32 = (one_vec + up_vec) * glu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n / 2); - } - input += input_stride; - output += output_stride; - } -} - -template -void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec)); - auto silu = gate_vec * sigmoid_vec; - auto gated_output_fp32 = up_vec * silu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template -void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - vec_op::FP32Vec16 w1_vec(M_SQRT1_2); - vec_op::FP32Vec16 w2_vec(0.5); - alignas(64) float temp[16]; - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto er_input_vec = gate_vec * w1_vec; - - er_input_vec.save(temp); - for (int32_t i = 0; i < 16; ++i) { - temp[i] = std::erf(temp[i]); - } - vec_op::FP32Vec16 er_vec(temp); - auto gelu = gate_vec * w2_vec * (one_vec + er_vec); - auto gated_output_fp32 = up_vec * gelu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template -void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, - const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - vec_op::FP32Vec16 w1_vec(0.7978845608028654); - vec_op::FP32Vec16 w2_vec(0.5); - vec_op::FP32Vec16 w3_vec(0.044715); - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; - auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); - // Note: can't use fast_exp form because diffusiongemma will generate - // wrong results - auto tanh_vec = inner_vec.tanh(); - auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); - auto gated_output_fp32 = up_vec * gelu_tanh; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template -FORCE_INLINE void apply_gated_act(const FusedMOEAct act, - float* __restrict__ input, - scalar_t* __restrict__ output, - const int32_t m, const int32_t n, - const int32_t input_stride, - const int32_t output_stride) { - switch (act) { - case FusedMOEAct::SwigluOAIAndMul: - swigluoai_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::SiluAndMul: - silu_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::GeluAndMul: - gelu_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::GeluTanhAndMul: - gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride); - return; - default: - TORCH_CHECK(false, "Unsupported act type."); - } -} +using cpu_fused_moe_utils::apply_gated_act; +using cpu_fused_moe_utils::FusedMOEAct; template void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr, @@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, } } } + } // namespace void prepack_moe_weight( @@ -864,7 +682,7 @@ void cpu_fused_moe( const int32_t input_size_2 = w2.size(2); const int32_t output_size_2 = w2.size(1); const int32_t topk_num = topk_id.size(1); - const FusedMOEAct act_type = get_act_type(act); + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); TORCH_CHECK(!skip_weighted || topk_num == 1, "skip_weighted is only supported for topk=1 on CPU"); diff --git a/csrc/cpu/cpu_fused_moe_activations.hpp b/csrc/cpu/cpu_fused_moe_activations.hpp new file mode 100644 index 000000000000..31e8cb009b7d --- /dev/null +++ b/csrc/cpu/cpu_fused_moe_activations.hpp @@ -0,0 +1,204 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP +#define CPU_FUSED_MOE_ACTIVATIONS_HPP + +#include +#include +#include + +#include "cpu/cpu_arch_macros.h" +#include "cpu/utils.hpp" + +namespace cpu_fused_moe_utils { +enum class FusedMOEAct { + SiluAndMul, + SwigluOAIAndMul, + GeluAndMul, + GeluTanhAndMul, +}; + +inline FusedMOEAct get_act_type(const std::string& act) { + if (act == "silu") { + return FusedMOEAct::SiluAndMul; + } else if (act == "swigluoai") { + return FusedMOEAct::SwigluOAIAndMul; + } else if (act == "gelu") { + return FusedMOEAct::GeluAndMul; + } else if (act == "gelu_tanh") { + return FusedMOEAct::GeluTanhAndMul; + } else { + TORCH_CHECK(false, "Invalid act type: " + act); + } +} + +template +void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, + const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; +#if !defined(__aarch64__) + // For GPT-OSS interleaved gate-up weights + alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, + 16, 18, 20, 22, 24, 26, 28, 30}; + vec_op::INT32Vec16 index_vec(index); +#endif + vec_op::FP32Vec16 gate_up_max_vec(7.0); + vec_op::FP32Vec16 up_min_vec(-7.0); + vec_op::FP32Vec16 alpha_vec(1.702); + vec_op::FP32Vec16 one_vec(1.0); + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < n_size; n += 32) { + // Note: AdvSIMD does not support gather loads +#if defined(__aarch64__) + vec_op::FP32Vec16 gate_vec(vec_op::uninit); + vec_op::FP32Vec16 up_vec(vec_op::uninit); + vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); +#else + vec_op::FP32Vec16 gate_vec(input + n, index_vec); + vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); +#endif + gate_vec = gate_vec.min(gate_up_max_vec); + up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec); + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec)); + auto glu = gate_vec * sigmoid_vec; + auto gated_output_fp32 = (one_vec + up_vec) * glu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n / 2); + } + input += input_stride; + output += output_stride; + } +} + +template +void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec)); + auto silu = gate_vec * sigmoid_vec; + auto gated_output_fp32 = up_vec * silu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template +void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(M_SQRT1_2); + vec_op::FP32Vec16 w2_vec(0.5); + alignas(64) float temp[16]; + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto er_input_vec = gate_vec * w1_vec; + + er_input_vec.save(temp); + for (int32_t i = 0; i < 16; ++i) { + temp[i] = std::erf(temp[i]); + } + vec_op::FP32Vec16 er_vec(temp); + auto gelu = gate_vec * w2_vec * (one_vec + er_vec); + auto gated_output_fp32 = up_vec * gelu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template +void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, + const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(0.7978845608028654); + vec_op::FP32Vec16 w2_vec(0.5); + vec_op::FP32Vec16 w3_vec(0.044715); + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; + auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); + // Note: can't use fast_exp form because diffusiongemma will generate + // wrong results + auto tanh_vec = inner_vec.tanh(); + auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); + auto gated_output_fp32 = up_vec * gelu_tanh; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template +FORCE_INLINE void apply_gated_act(const FusedMOEAct act, + float* __restrict__ input, + scalar_t* __restrict__ output, + const int32_t m, const int32_t n, + const int32_t input_stride, + const int32_t output_stride) { + switch (act) { + case FusedMOEAct::SwigluOAIAndMul: + swigluoai_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::SiluAndMul: + silu_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::GeluAndMul: + gelu_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::GeluTanhAndMul: + gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride); + return; + default: + TORCH_CHECK(false, "Unsupported act type."); + } +} +} // namespace cpu_fused_moe_utils + +#endif diff --git a/csrc/cpu/cpu_fused_moe_int8.cpp b/csrc/cpu/cpu_fused_moe_int8.cpp new file mode 100644 index 000000000000..649317b4786a --- /dev/null +++ b/csrc/cpu/cpu_fused_moe_int8.cpp @@ -0,0 +1,647 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "cpu/cpu_arch_macros.h" + +#include +#include +#include +#include +#include + +#include "cpu/cpu_fused_moe_activations.hpp" +#include "cpu/cpu_types.hpp" +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" +#include "cpu/utils.hpp" + +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) + #include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp" + #define NEON_DISPATCH(SCALAR_TYPE, ...) \ + case cpu_utils::ISA::NEON: { \ + using gemm_t = \ + cpu_micro_gemm::MicroGemmINT8; \ + return __VA_ARGS__(); \ + } +#else + #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON: +#endif + +#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \ + [&] { \ + switch (ISA_TYPE) { \ + NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ + default: { \ + TORCH_CHECK(false, "Invalid CPU ISA type."); \ + } \ + } \ + }() + +namespace { +using cpu_fused_moe_utils::apply_gated_act; +using cpu_fused_moe_utils::FusedMOEAct; + +template +void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr, + int8_t* __restrict__ packed_weight_ptr, + const int32_t expert_num, + const int32_t output_size, + const int32_t input_size, + const int64_t expert_stride) { +#pragma omp parallel for + for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) { + gemm_t::pack_weight(weight_ptr + expert_stride * e_idx, + packed_weight_ptr + expert_stride * e_idx, output_size, + input_size); + } +} + +// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp +template +void fused_moe_int8_impl( + scalar_t* __restrict__ output, const scalar_t* __restrict__ input, + const int8_t* __restrict__ w13, const int8_t* __restrict__ w2, + const float* __restrict__ w13_scales, const float* __restrict__ w2_scales, + scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias, + const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id, + const FusedMOEAct act_type, const int32_t token_num, + const int32_t expert_num, const int32_t topk_num, + const int32_t input_size_13, const int32_t output_size_13, + const int32_t input_size_2, const int32_t output_size_2, + const bool skip_weighted) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + constexpr int32_t gemm_n_tile_size = gemm_t::NSize; + constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; + constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size; + + TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0); + TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0); + TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0); + TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); + TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); + + const int32_t thread_num = cpu_utils::get_max_threads(); + const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( + gemm_m_tile_size * input_size_13 * sizeof(int8_t)); + const int32_t w2_input_buffer_size = + cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t)); + + const int32_t w13_n_tile_size = [&]() { + const int64_t cache_size = cpu_utils::get_available_l2_size(); + const int32_t n_size_cache_limit = + (cache_size - w13_input_buffer_size) / + (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t)); + const int32_t n_size_thread_limit = + output_size_13 / std::max(1, thread_num / topk_num); + const int32_t n_size = cpu_utils::round_down( + std::min(n_size_cache_limit, n_size_thread_limit)); + return std::max(n_size, min_w13_n_tile_size); + }(); + + const int32_t w2_n_tile_size = [&]() { + const int64_t cache_size = cpu_utils::get_available_l2_size(); + const int32_t n_size_cache_limit = + (cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t)); + const int32_t n_size_thread_limit = + output_size_2 / std::max(1, thread_num / topk_num); + const int32_t n_size = cpu_utils::round_down( + std::min(n_size_cache_limit, n_size_thread_limit)); + return std::max(n_size, gemm_n_tile_size); + }(); + + int32_t common_buffer_offset = 0; + const int32_t token_num_per_group_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t)); + const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t)); + const int32_t expanded_token_num = token_num * topk_num; + const int32_t expand_token_id_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); + const int32_t expand_token_id_index_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); + const int32_t input_quant_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t)); + const int32_t input_scale_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float)); + const int32_t w13_gemm_output_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>( + expanded_token_num * input_size_2 * sizeof(scalar_t)); + const int32_t w13_output_scale_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(float)); + const int32_t w2_gemm_output_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>( + expanded_token_num * output_size_2 * sizeof(float)); + + int32_t gemm_thread_buffer_offset = 0; + const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset; + gemm_thread_buffer_offset += + std::max(w13_input_buffer_size, w2_input_buffer_size); + const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset; + gemm_thread_buffer_offset += cpu_utils::round_up<64>( + gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) * + sizeof(int32_t)); + + const int32_t ws_output_buffer_offset = 0; + const int32_t ws_thread_buffer_size = + cpu_utils::round_up<64>(output_size_2 * sizeof(float)); + const int32_t thread_buffer_size = + std::max(gemm_thread_buffer_offset, ws_thread_buffer_size); + const int32_t buffer_size = + common_buffer_offset + thread_buffer_size * thread_num; + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size); + uint8_t* common_buffer_start = + cpu_utils::ScratchPadManager::get_scratchpad_manager() + ->get_data(); + uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset; + + int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast( + common_buffer_start + token_num_per_group_buffer_offset); + int32_t* __restrict__ cu_token_num_per_group_buffer = + reinterpret_cast(common_buffer_start + + cu_token_num_per_group_buffer_offset); + int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast( + common_buffer_start + expand_token_id_buffer_offset); + int32_t* __restrict__ expand_token_id_index_buffer = + reinterpret_cast(common_buffer_start + + expand_token_id_index_buffer_offset); + int8_t* __restrict__ input_quant_buffer = reinterpret_cast( + common_buffer_start + input_quant_buffer_offset); + float* __restrict__ input_scale_buffer = + reinterpret_cast(common_buffer_start + input_scale_buffer_offset); + + std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t)); + for (int32_t i = 0; i < expanded_token_num; ++i) { + ++token_num_per_group_buffer[topk_id[i]]; + } + + int32_t token_num_sum = 0; + cu_token_num_per_group_buffer[0] = 0; + int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1; + for (int32_t i = 0; i < expert_num; ++i) { + token_index_buffer[i] = token_num_sum; + token_num_sum += token_num_per_group_buffer[i]; + } + + for (int32_t i = 0; i < token_num; ++i) { + const int32_t* curr_topk_id = topk_id + i * topk_num; + int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num; + for (int32_t j = 0; j < topk_num; ++j) { + const int32_t curr_expert_id = curr_topk_id[j]; + const int32_t curr_index = token_index_buffer[curr_expert_id]++; + expand_token_id_buffer[curr_index] = i; + curr_index_buffer[j] = curr_index; + } + } + +// quantize inputs +#pragma omp parallel for + for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) { + gemm_t::quantize_row(input + token_idx * input_size_13, + input_quant_buffer + token_idx * input_size_13, + input_scale_buffer[token_idx], input_size_13); + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +// w13 GEMM + act +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + const int32_t task_num_per_expert = + (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size; + const int32_t task_num = task_num_per_expert * expert_num; + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + int8_t* __restrict__ gemm_input_buffer = + reinterpret_cast(thread_buffer + gemm_input_buffer_offset); + float* __restrict__ gemm_output_buffer = + reinterpret_cast(thread_buffer + gemm_output_buffer_offset); + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast( + common_buffer_start + w13_gemm_output_buffer_offset); + gemm_t gemm; + + const int32_t w13_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_13; + const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13; + + for (;;) { + const int32_t task_id = counter_ptr->acquire_counter(); + if (task_id >= task_num) { + break; + } + const int32_t curr_expert_id = task_id / task_num_per_expert; + const int32_t curr_output_group_id = task_id % task_num_per_expert; + const int32_t curr_token_num = + token_num_per_group_buffer[curr_expert_id]; + if (curr_token_num == 0) { + continue; + } + + const int32_t actual_n_tile_size = + std::min(w13_n_tile_size, + output_size_13 - curr_output_group_id * w13_n_tile_size); + const int32_t* __restrict__ curr_expand_token_id_buffer = + expand_token_id_buffer + + cu_token_num_per_group_buffer[curr_expert_id]; + scalar_t* __restrict__ curr_w13_gemm_output_buffer = + w13_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 + + curr_output_group_id * w13_n_tile_size / 2; + + const int8_t* w13_weight_ptr_0 = nullptr; + const int8_t* w13_weight_ptr_1 = nullptr; + const float* w13_scale_ptr_0 = nullptr; + const float* w13_scale_ptr_1 = nullptr; + scalar_t* w13_bias_ptr_0 = nullptr; + scalar_t* w13_bias_ptr_1 = nullptr; + if (act_type == FusedMOEAct::SwigluOAIAndMul) { + const int32_t output_offset = curr_output_group_id * w13_n_tile_size; + w13_weight_ptr_0 = w13 + + curr_expert_id * input_size_13 * output_size_13 + + output_offset * input_size_13; + w13_weight_ptr_1 = + w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13; + w13_scale_ptr_0 = + w13_scales + curr_expert_id * output_size_13 + output_offset; + w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2; + if (w13_bias != nullptr) { + w13_bias_ptr_0 = + w13_bias + curr_expert_id * output_size_13 + output_offset; + w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2; + } + } else { + const int32_t output_offset = + curr_output_group_id * (w13_n_tile_size / 2); + w13_weight_ptr_0 = w13 + + curr_expert_id * input_size_13 * output_size_13 + + output_offset * input_size_13; + w13_weight_ptr_1 = + w13_weight_ptr_0 + output_size_13 / 2 * input_size_13; + w13_scale_ptr_0 = + w13_scales + curr_expert_id * output_size_13 + output_offset; + w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2; + if (w13_bias != nullptr) { + w13_bias_ptr_0 = + w13_bias + curr_expert_id * output_size_13 + output_offset; + w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2; + } + } + + for (int32_t token_idx = 0; token_idx < curr_token_num; + token_idx += gemm_m_tile_size) { + const int32_t actual_token_num = + std::min(gemm_m_tile_size, curr_token_num - token_idx); + const int8_t* input_rows[gemm_m_tile_size]; + alignas(64) float input_scales[gemm_m_tile_size]; + // gather and pack + for (int32_t i = 0; i < actual_token_num; ++i) { + const int32_t curr_token_id = curr_expand_token_id_buffer[i]; + input_rows[i] = input_quant_buffer + curr_token_id * input_size_13; + input_scales[i] = input_scale_buffer[curr_token_id]; + } + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, + actual_token_num, input_size_13); + curr_expand_token_id_buffer += actual_token_num; + + const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0; + const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1; + const float* w13_scale_ptr_0_iter = w13_scale_ptr_0; + const float* w13_scale_ptr_1_iter = w13_scale_ptr_1; + scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0; + scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1; + float* w13_output_buffer_0_iter = gemm_output_buffer; + float* w13_output_buffer_1_iter = + gemm_output_buffer + actual_n_tile_size / 2; + + for (int32_t i = 0; i < actual_n_tile_size; + i += min_w13_n_tile_size) { + auto* output_0_int32 = + reinterpret_cast(w13_output_buffer_0_iter); + gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32, + actual_token_num, input_size_13, w13_n_group_stride, + actual_n_tile_size); + gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter, + input_scales, w13_scale_ptr_0_iter, + actual_token_num, gemm_n_tile_size, + actual_n_tile_size); + if (w13_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue( + w13_output_buffer_0_iter, w13_output_buffer_0_iter, + w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size, + actual_n_tile_size); + w13_bias_ptr_0_iter += gemm_n_tile_size; + } + + auto* output_1_int32 = + reinterpret_cast(w13_output_buffer_1_iter); + gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32, + actual_token_num, input_size_13, w13_n_group_stride, + actual_n_tile_size); + gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter, + input_scales, w13_scale_ptr_1_iter, + actual_token_num, gemm_n_tile_size, + actual_n_tile_size); + if (w13_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue( + w13_output_buffer_1_iter, w13_output_buffer_1_iter, + w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size, + actual_n_tile_size); + w13_bias_ptr_1_iter += gemm_n_tile_size; + } + + w13_weight_ptr_0_iter += w13_n_tile_stride; + w13_weight_ptr_1_iter += w13_n_tile_stride; + w13_scale_ptr_0_iter += gemm_n_tile_size; + w13_scale_ptr_1_iter += gemm_n_tile_size; + w13_output_buffer_0_iter += gemm_n_tile_size; + w13_output_buffer_1_iter += gemm_n_tile_size; + } + + apply_gated_act(act_type, gemm_output_buffer, + curr_w13_gemm_output_buffer, actual_token_num, + actual_n_tile_size, actual_n_tile_size, input_size_2); + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; + } + } + } + } + + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast( + common_buffer_start + w13_gemm_output_buffer_offset); + float* __restrict__ w13_output_scale_buffer = reinterpret_cast( + common_buffer_start + w13_output_scale_buffer_offset); + +// quantize w2 inputs - in place +#pragma omp parallel for + for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) { + scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2; + int8_t* output_row = reinterpret_cast(input_row); + gemm_t::quantize_row(input_row, output_row, + w13_output_scale_buffer[token_idx], input_size_2); + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +// w2 gemm +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + const int32_t task_num_per_expert = + (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size; + const int32_t task_num = task_num_per_expert * expert_num; + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + int8_t* __restrict__ gemm_input_buffer = + reinterpret_cast(thread_buffer + gemm_input_buffer_offset); + float* __restrict__ gemm_output_buffer = + reinterpret_cast(thread_buffer + gemm_output_buffer_offset); + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast( + common_buffer_start + w2_gemm_output_buffer_offset); + gemm_t gemm; + + const int32_t w2_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_2; + const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2; + + for (;;) { + const int32_t task_id = counter_ptr->acquire_counter(); + if (task_id >= task_num) { + break; + } + const int32_t curr_expert_id = task_id / task_num_per_expert; + const int32_t curr_output_group_id = task_id % task_num_per_expert; + const int32_t curr_token_num = + token_num_per_group_buffer[curr_expert_id]; + if (curr_token_num == 0) { + continue; + } + + const int32_t actual_n_tile_size = + std::min(w2_n_tile_size, + output_size_2 - curr_output_group_id * w2_n_tile_size); + scalar_t* __restrict__ curr_w13_gemm_output_buffer = + w13_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2; + float* __restrict__ curr_w13_output_scale_buffer = + w13_output_scale_buffer + + cu_token_num_per_group_buffer[curr_expert_id]; + float* __restrict__ curr_w2_gemm_output_buffer = + w2_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 + + curr_output_group_id * w2_n_tile_size; + const int8_t* __restrict__ w2_weight_ptr = + w2 + curr_expert_id * output_size_2 * input_size_2 + + curr_output_group_id * w2_n_tile_size * input_size_2; + const float* __restrict__ w2_scale_ptr = + w2_scales + curr_expert_id * output_size_2 + + curr_output_group_id * w2_n_tile_size; + scalar_t* w2_bias_ptr = nullptr; + if (w2_bias != nullptr) { + w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 + + curr_output_group_id * w2_n_tile_size; + } + + for (int32_t token_idx = 0; token_idx < curr_token_num; + token_idx += gemm_m_tile_size) { + const int32_t actual_token_num = + std::min(gemm_m_tile_size, curr_token_num - token_idx); + const int8_t* input_rows[gemm_m_tile_size]; + alignas(64) float input_scales[gemm_m_tile_size]; + for (int32_t i = 0; i < actual_token_num; ++i) { + input_rows[i] = reinterpret_cast( + curr_w13_gemm_output_buffer + i * input_size_2); + input_scales[i] = curr_w13_output_scale_buffer[i]; + } + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, + actual_token_num, input_size_2); + + const int8_t* w2_weight_ptr_iter = w2_weight_ptr; + const float* w2_scale_ptr_iter = w2_scale_ptr; + scalar_t* w2_bias_ptr_iter = w2_bias_ptr; + float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer; + for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) { + auto* output_int32 = reinterpret_cast(gemm_output_buffer); + gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32, + actual_token_num, input_size_2, w2_n_group_stride, + gemm_n_tile_size); + gemm_t::dequantize_tile(output_int32, gemm_output_buffer, + input_scales, w2_scale_ptr_iter, + actual_token_num, gemm_n_tile_size, + gemm_n_tile_size); + if (w2_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue( + gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter, + actual_token_num, gemm_n_tile_size, gemm_n_tile_size); + w2_bias_ptr_iter += gemm_n_tile_size; + } + for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) { + std::memcpy( + curr_w2_gemm_output_buffer_iter + m_idx * output_size_2, + gemm_output_buffer + m_idx * gemm_n_tile_size, + gemm_n_tile_size * sizeof(float)); + } + + w2_weight_ptr_iter += w2_n_tile_stride; + w2_scale_ptr_iter += gemm_n_tile_size; + curr_w2_gemm_output_buffer_iter += gemm_n_tile_size; + } + + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; + curr_w13_output_scale_buffer += gemm_m_tile_size; + curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2; + } + } + } + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + float* __restrict__ ws_output_buffer = + reinterpret_cast(thread_buffer + ws_output_buffer_offset); + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast( + common_buffer_start + w2_gemm_output_buffer_offset); + + for (;;) { + const int32_t token_id = counter_ptr->acquire_counter(); + if (token_id >= token_num) { + break; + } + int32_t* __restrict__ curr_expand_token_id_index_buffer = + expand_token_id_index_buffer + token_id * topk_num; + const float* __restrict__ curr_weight = + topk_weights + token_id * topk_num; + const float first_weight = skip_weighted ? 1.0f : curr_weight[0]; + scalar_t* __restrict__ curr_output_buffer = + output + token_id * output_size_2; + + if (topk_num > 1) { + int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; + float* w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + float* ws_output_buffer_iter = ws_output_buffer; + vec_op::FP32Vec16 weight_vec(first_weight); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + (vec * weight_vec).save(ws_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + } + + for (int32_t idx = 1; idx < topk_num - 1; ++idx) { + w2_output_idx = curr_expand_token_id_index_buffer[idx]; + w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + ws_output_buffer_iter = ws_output_buffer; + weight_vec = vec_op::FP32Vec16(curr_weight[idx]); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + vec_op::FP32Vec16 sum(ws_output_buffer_iter); + (sum + vec * weight_vec).save(ws_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + } + } + + const int32_t last_idx = topk_num - 1; + w2_output_idx = curr_expand_token_id_index_buffer[last_idx]; + w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + ws_output_buffer_iter = ws_output_buffer; + scalar_t* curr_output_buffer_iter = curr_output_buffer; + weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + vec_op::FP32Vec16 sum(ws_output_buffer_iter); + scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + curr_output_buffer_iter += 16; + } + } else { + const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; + float* w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + scalar_t* curr_output_buffer_iter = curr_output_buffer; + vec_op::FP32Vec16 weight_vec(first_weight); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter); + w2_output_iter += 16; + curr_output_buffer_iter += 16; + } + } + } + } + } +} +} // namespace + +void prepack_moe_weight_int8( + const torch::Tensor& weight, // [expert_num, output_size, input_size] + torch::Tensor& packed_weight, const std::string& isa) { + TORCH_CHECK(weight.is_contiguous()); + const int32_t expert_num = weight.size(0); + const int32_t output_size = weight.size(1); + const int32_t input_size = weight.size(2); + const int64_t expert_stride = weight.stride(0); + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); + TORCH_CHECK_EQ(output_size % 32, 0); + + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() { + TORCH_CHECK_EQ(input_size % gemm_t::K, 0); + prepack_moe_weight_int8_impl( + weight.data_ptr(), packed_weight.data_ptr(), expert_num, + output_size, input_size, expert_stride); + }); +} + +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& w13, const torch::Tensor& w2, + const torch::Tensor& w13_scale, + const torch::Tensor& w2_scale, + const std::optional& w13_bias, + const std::optional& w2_bias, + const torch::Tensor& topk_weights, + const torch::Tensor& topk_id, const bool skip_weighted, + const std::string& act, const std::string& isa) { + const int32_t token_num = input.size(0); + const int32_t input_size_13 = input.size(1); + const int64_t input_stride = input.stride(0); + TORCH_CHECK_EQ(input_stride, input_size_13); + const int32_t expert_num = w13.size(0); + const int32_t output_size_13 = w13.size(1); + const int32_t input_size_2 = w2.size(2); + const int32_t output_size_2 = w2.size(1); + const int32_t topk_num = topk_id.size(1); + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); + TORCH_CHECK(!skip_weighted || topk_num == 1, + "skip_weighted is only supported for topk=1 on CPU"); + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "cpu_fused_moe_int8", [&]() { + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() { + fused_moe_int8_impl( + output.data_ptr(), input.data_ptr(), + w13.data_ptr(), w2.data_ptr(), + w13_scale.data_ptr(), w2_scale.data_ptr(), + w13_bias.has_value() ? w13_bias->data_ptr() : nullptr, + w2_bias.has_value() ? w2_bias->data_ptr() : nullptr, + topk_weights.data_ptr(), topk_id.data_ptr(), + act_type, token_num, expert_num, topk_num, input_size_13, + output_size_13, input_size_2, output_size_2, skip_weighted); + }); + }); +} diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 250c870dbe4b..42083bc3eb69 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec { explicit FP32Vec4(__vector float data) : reg(data) {} - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} + FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} }; struct FP32Vec8 : public Vec { @@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec { explicit FP32Vec8(f32x4x2_t data) : reg(data) {} - explicit FP32Vec8(const FP32Vec8& data) { + FP32Vec8(const FP32Vec8& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; } @@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec { reg.val[1] = fp16_to_fp32_bits(raw_lo); } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&result, &ar](int i) { result += ar.values[i]; }); - - return result; + // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds. + // Step 1: pairwise sum of the two 4-wide halves + __vector float s = vec_add(reg.val[0], reg.val[1]); + // Step 2: rotate by 8 bytes (2 floats) and add + s = vec_add(s, vec_sld(s, s, 8)); + // Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total + s = vec_add(s, vec_sld(s, s, 4)); + return vec_extract(s, 0); } FP32Vec8 exp() const { f32x4x2_t out; @@ -592,7 +593,7 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {} explicit FP32Vec16(f32x4x4_t data) : reg(data) {} - explicit FP32Vec16(const FP32Vec16& data) { + FP32Vec16(const FP32Vec16& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; reg.val[2] = data.reg.val[2]; @@ -746,6 +747,15 @@ struct FP32Vec16 : public Vec { vec_abs(reg.val[2]), vec_abs(reg.val[3])})); } + FP32Vec16 exp() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_e = lo.exp(); + auto hi_e = hi.exp(); + return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1], + hi_e.reg.val[0], hi_e.reg.val[1]}); + } + float reduce_max() { __vector float max01 = vec_max(reg.val[0], reg.val[1]); __vector float max23 = vec_max(reg.val[2], reg.val[3]); diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index bf96554a8dff..9b28300cedfb 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -140,10 +140,21 @@ struct BF16Vec16 : public Vec { explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { - // Save 256 bits in two parts vec_xst(reg.val[0], 0, (signed short*)ptr); vec_xst(reg.val[1], 16, (signed short*)ptr); } + + void save(void* ptr, const int elem_num) const { + auto* dst = reinterpret_cast(ptr); + union { + ss16x8x2_t r; + c10::BFloat16 values[16]; + } ar; + ar.r = reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + dst[i] = ar.values[i]; + } + } }; const static __vector signed short zero = vec_splats((signed short)0); @@ -269,7 +280,7 @@ struct FP32Vec4 : public Vec { explicit FP32Vec4(__vector float data) : reg(data) {} - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} + FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} }; struct FP32Vec8 : public Vec { @@ -298,7 +309,7 @@ struct FP32Vec8 : public Vec { explicit FP32Vec8(f32x4x2_t data) : reg(data) {} - explicit FP32Vec8(const FP32Vec8& data) { + FP32Vec8(const FP32Vec8& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; } @@ -322,13 +333,12 @@ struct FP32Vec8 : public Vec { } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&result, &ar](int i) { result += ar.values[i]; }); - - return result; + __vector float sum = vec_add(reg.val[0], reg.val[1]); + __vector float hi = vec_sld(sum, sum, 8); + sum = vec_add(sum, hi); + __vector float lo = vec_sld(sum, sum, 4); + sum = vec_add(sum, lo); + return vec_extract(sum, 0); } FP32Vec8 exp() const { @@ -643,7 +653,7 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(f32x4x4_t data) : reg(data) {} - explicit FP32Vec16(const FP32Vec16& data) { + FP32Vec16(const FP32Vec16& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; reg.val[2] = data.reg.val[2]; @@ -672,6 +682,8 @@ struct FP32Vec16 : public Vec { reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero); } + explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {} + explicit FP32Vec16(const FP16Vec16& v) { __vector unsigned int raw_hi_0 = (__vector unsigned int)vec_unpackh(v.reg.val[0]); @@ -723,13 +735,13 @@ struct FP32Vec16 : public Vec { } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&result, &ar](int i) { result += ar.values[i]; }); - - return result; + __vector float sum = vec_add(vec_add(reg.val[0], reg.val[1]), + vec_add(reg.val[2], reg.val[3])); + __vector float hi = vec_sld(sum, sum, 8); + sum = vec_add(sum, hi); + __vector float lo = vec_sld(sum, sum, 4); + sum = vec_add(sum, lo); + return vec_extract(sum, 0); } template @@ -754,13 +766,68 @@ struct FP32Vec16 : public Vec { } float reduce_max() const { - AliasReg ar; - ar.reg = reg; - float result = ar.values[0]; - unroll_loop([&result, &ar](int i) { - if (ar.values[i] > result) result = ar.values[i]; - }); - return result; + __vector float m = vec_max(vec_max(reg.val[0], reg.val[1]), + vec_max(reg.val[2], reg.val[3])); + __vector float hi = vec_sld(m, m, 8); + m = vec_max(m, hi); + __vector float lo = vec_sld(m, m, 4); + m = vec_max(m, lo); + return vec_extract(m, 0); + } + + FP32Vec16 exp() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_e = lo.exp(); + auto hi_e = hi.exp(); + return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1], + hi_e.reg.val[0], hi_e.reg.val[1]}); + } + + FP32Vec16 abs() const { + return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]), + vec_abs(reg.val[2]), vec_abs(reg.val[3])})); + } + + FP32Vec16 min(const FP32Vec16& b) const { + return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]), + vec_min(reg.val[1], b.reg.val[1]), + vec_min(reg.val[2], b.reg.val[2]), + vec_min(reg.val[3], b.reg.val[3])})); + } + + FP32Vec16 clamp(const FP32Vec16& lo, const FP32Vec16& hi) const { + return this->max(lo).min(hi); + } + + float reduce_min() const { + __vector float m = vec_min(vec_min(reg.val[0], reg.val[1]), + vec_min(reg.val[2], reg.val[3])); + __vector float h = vec_sld(m, m, 8); + m = vec_min(m, h); + __vector float l = vec_sld(m, m, 4); + m = vec_min(m, l); + return vec_extract(m, 0); + } + + FP32Vec16 min(const FP32Vec16& b, const int elem_num) const { + AliasReg ar_this, ar_b; + ar_this.reg = reg; + ar_b.reg = b.reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + ar_this.values[i] = std::min(ar_this.values[i], ar_b.values[i]); + } + return FP32Vec16(ar_this.reg); + } + + FP32Vec16 max(const FP32Vec16& b, const int elem_num) const { + AliasReg ar_this, ar_b; + ar_this.reg = reg; + ar_b.reg = b.reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + ar_this.values[i] = std::max(ar_this.values[i], ar_b.values[i]); + } + return FP32Vec16(ar_this.reg); } void save(float* ptr) const { @@ -769,6 +836,67 @@ struct FP32Vec16 : public Vec { vec_xst(reg.val[2], 32, ptr); vec_xst(reg.val[3], 48, ptr); } + + void save(float* ptr, const int elem_num) const { + AliasReg ar; + ar.reg = reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + ptr[i] = ar.values[i]; + } + } + + void save(c10::Half* ptr) const { + FP16Vec16 fp16(*this); + fp16.save(ptr); + } + + void save(c10::Half* ptr, const int elem_num) const { + FP16Vec16 fp16(*this); + union { + ss16x8x2_t r; + c10::Half values[16]; + } ar; + ar.r = fp16.reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + ptr[i] = ar.values[i]; + } + } +}; + +struct INT8Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + + union AliasReg { + __vector signed char reg; + int8_t values[VEC_ELEM_NUM]; + }; + + __vector signed char reg; + + explicit INT8Vec16(const FP32Vec16& vec) { + __vector signed int ret[4]; + ret[0] = vec_signed(vec.reg.val[0]); + ret[1] = vec_signed(vec.reg.val[1]); + ret[2] = vec_signed(vec.reg.val[2]); + ret[3] = vec_signed(vec.reg.val[3]); + + __vector signed short packed1 = vec_packs(ret[0], ret[1]); + __vector signed short packed2 = vec_packs(ret[2], ret[3]); + + reg = vec_packs(packed1, packed2); + } + + void save(void* ptr) const { + *reinterpret_cast<__vector signed char*>(ptr) = reg; + } + + void save(int8_t* ptr, const int elem_num) const { + AliasReg ar; + ar.reg = reg; + for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) { + ptr[i] = ar.values[i]; + } + } }; template diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp new file mode 100644 index 000000000000..54e4f99c2d6e --- /dev/null +++ b/csrc/cpu/mamba_cpu.cpp @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// CPU at::Tensor wrappers for Mamba decode-step kernels defined in +// mamba_kernels.hpp. + +#include "cpu/mamba_kernels.hpp" + +#include +#include +#include + +#include "cpu_types.hpp" + +// --------------------------------------------------------------------------- +// causal_conv1d_update +// --------------------------------------------------------------------------- +at::Tensor causal_conv1d_update_cpu_impl( + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, + const c10::optional& bias, + const c10::optional& activation, + const c10::optional& conv_state_indices, + const c10::optional& query_start_loc, int64_t pad_slot_id) { + bool do_silu = false; + if (activation.has_value()) { + const std::string& act = activation.value(); + do_silu = (act == "silu" || act == "swish"); + } + + at::ScalarType dtype = x.scalar_type(); + + // Input x: contiguous in native dtype. + at::Tensor x_c = x.is_contiguous() ? x : x.contiguous(); + + // conv_state: NEVER copy the full paged tensor just for layout reasons. + // If the dtype matches we work directly on conv_state (contiguous or not) + // by extracting strides and passing them to the kernel. + // Only a dtype-conversion copy is made when types differ (rare for BF16). + bool state_type_ok = (conv_state.scalar_type() == dtype); + at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype); + // state_c and conv_state may be non-contiguous — that is intentional. + + // Weight: coerce to same dtype if needed (should match in practice) + at::Tensor w_c = + (weight.scalar_type() != dtype) + ? weight.to(dtype).contiguous() + : (weight.is_contiguous() ? weight : weight.contiguous()); + + // Bias stays float32 (small scalar, used only for fp32 accumulation) + at::Tensor bias_f32; + if (bias.has_value() && bias.value().defined()) + bias_f32 = bias.value().to(at::kFloat).contiguous(); + + int64_t batch = x_c.size(0); + int64_t dim = x_c.size(1); + int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1; + int64_t width = w_c.size(1); + int64_t state_len = state_c.size(2); + + // Extract strides — works for contiguous AND non-contiguous (transposed) + // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in + // contiguous) stride(1): between conv channels (dim stride) stride(2): + // between state elements (=1 when contiguous, =dim when transposed) + int64_t stride_s_slot = state_c.stride(0); + int64_t stride_s_dim = state_c.stride(1); + int64_t stride_s_state = state_c.stride(2); + + at::Tensor out = x_c.clone(); // native dtype, no float32 alloc + + const int32_t* cache_idx_ptr = nullptr; + at::Tensor cache_idx_int; + if (conv_state_indices.has_value()) { + cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous(); + cache_idx_ptr = cache_idx_int.data_ptr(); + } + + VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] { + mamba_cpu::causal_conv1d_update_kernel( + x_c.data_ptr(), state_c.data_ptr(), stride_s_slot, + stride_s_dim, stride_s_state, w_c.data_ptr(), + bias_f32.defined() ? bias_f32.data_ptr() : nullptr, + out.data_ptr(), cache_idx_ptr, + static_cast(pad_slot_id), batch, dim, seqlen, width, state_len, + do_silu); + }); + + // Write back only when a type-conversion copy was made. + // Layout-only non-contiguity is handled via strides above — no copy needed. + if (!state_type_ok) conv_state.copy_(state_c); + + return out; +} + +// --------------------------------------------------------------------------- +// selective_state_update +// --------------------------------------------------------------------------- +void selective_state_update_cpu_impl( + at::Tensor& state, // (nstates, nheads, dim, dstate) + const at::Tensor& x, // (N, nheads, dim) + const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B, + const at::Tensor& C, const c10::optional& D, + const c10::optional& z, + const c10::optional& dt_bias, bool dt_softplus, + const c10::optional& state_batch_indices, + const c10::optional& dst_state_batch_indices, + int64_t null_block_id, at::Tensor& out, + const c10::optional& num_accepted_tokens, + const c10::optional& cu_seqlens) { + at::ScalarType state_type = state.scalar_type(); + at::ScalarType input_type = x.scalar_type(); + + // x, B, C must be contiguous and match input_type + auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor { + at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; + return r.is_contiguous() ? r : r.contiguous(); + }; + at::Tensor x_in = ensure_input(x); + at::Tensor B_in = ensure_input(B); + at::Tensor C_in = ensure_input(C); + at::Tensor z_in; + if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value()); + + // A, D, dt_bias are float32 model parameters that arrive here as expanded + // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0). + // We need just the scalar value per head as a (nheads,) 1-D array so that + // A_ptr[h] in the kernel correctly reads head h's value. + // + // Strategy: peel trailing expanded (stride=0) dims via .select(), which is + // a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0) + // → .select(2,0) → (nheads, head_dim) strides (1,0) + // → .select(1,0) → (nheads,) stride (1,) ← contiguous, free. + // No allocation, no type conversion (A is already float32). + auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor { + at::Tensor r = t; + // Peel trailing dimensions that are broadcast (stride=0 or size=1) + while (r.dim() > 1) r = r.select(r.dim() - 1, 0); + if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); + return r.is_contiguous() ? r : r.contiguous(); + }; + + at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32 + at::Tensor D_f32, dt_bias_f32; + if (D.has_value() && D.value().defined()) + D_f32 = to_per_head_1d_f32(D.value()); + if (dt_bias.has_value() && dt_bias.value().defined()) + dt_bias_f32 = to_per_head_1d_f32(dt_bias.value()); + + // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE + // the type conversion so we convert head_dim x fewer elements. + at::Tensor dt_f32; + { + // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2, + // take a zero-copy view of index 0 along that dim first. + at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads) + at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2; + dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous(); + } + + int64_t nheads = state.size(1); + int64_t dim = state.size(2); + int64_t dstate = state.size(3); + int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined()) + ? cu_seqlens.value().size(0) - 1 + : x_in.size(0); + int64_t ngroups = B_in.size(1); + + // Strides + int64_t stride_state_n = state.stride(0); + int64_t stride_state_h = state.stride(1); + int64_t stride_state_d = state.stride(2); + int64_t stride_x_n = x_in.stride(0); + int64_t stride_x_h = x_in.stride(1); + int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads) + int64_t stride_BC_n = B_in.stride(0); + int64_t stride_BC_g = B_in.stride(1); + int64_t stride_out_n = out.stride(0); + int64_t stride_out_h = out.stride(1); + + // Optional index pointers + auto get_int32_ptr = + [](const c10::optional& opt) -> const int32_t* { + return (opt.has_value() && opt.value().defined()) + ? opt.value().data_ptr() + : nullptr; + }; + const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); + const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices); + const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); + const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); + + // Dispatch on (state_t, input_t, out_t): write directly into `out` + // without any intermediate float32 buffer. + VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] { + using state_t = scalar_t; + VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { + using input_t = scalar_t; + VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] { + using out_t = scalar_t; + mamba_cpu::selective_state_update_kernel( + state.data_ptr(), stride_state_n, stride_state_h, + stride_state_d, x_in.data_ptr(), stride_x_n, stride_x_h, + dt_f32.data_ptr(), stride_dt_n, A_f32.data_ptr(), + B_in.data_ptr(), C_in.data_ptr(), stride_BC_n, + stride_BC_g, D_f32.defined() ? D_f32.data_ptr() : nullptr, + z_in.defined() ? z_in.data_ptr() : nullptr, + dt_bias_f32.defined() ? dt_bias_f32.data_ptr() : nullptr, + out.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, + dsbi_ptr, static_cast(null_block_id), nat_ptr, csl_ptr, N, + nheads, ngroups, dim, dstate, dt_softplus); + }); + }); + }); +} + +// --------------------------------------------------------------------------- +// mamba_chunk_scan_fwd_cpu +// --------------------------------------------------------------------------- +void mamba_chunk_scan_fwd_cpu_impl( + at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller + at::Tensor& + final_states, // [batch, nheads, headdim, dstate] float32 contiguous + const at::Tensor& x, // [seqlen, nheads, headdim] + const at::Tensor& + dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp) + const at::Tensor& A, // [nheads] float32 + const at::Tensor& B, // [seqlen, ngroups, dstate] + const at::Tensor& C, // [seqlen, ngroups, dstate] + const c10::optional& D, // [nheads] float32 (optional) + const c10::optional& z, // [seqlen, nheads, headdim] (optional) + const at::Tensor& cu_seqlens // [batch+1] int32 +) { + const at::ScalarType input_type = x.scalar_type(); + + auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor { + at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; + return r.is_contiguous() ? r : r.contiguous(); + }; + at::Tensor x_in = ensure_contig(x); + at::Tensor B_in = ensure_contig(B); + at::Tensor C_in = ensure_contig(C); + at::Tensor z_in; + if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value()); + + // A and D are float32 model parameters, potentially broadcast-expanded. + // Strip trailing broadcast dims to get a contiguous (nheads,) array. + auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor { + at::Tensor r = t; + while (r.dim() > 1) r = r.select(r.dim() - 1, 0); + if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); + return r.is_contiguous() ? r : r.contiguous(); + }; + at::Tensor A_f32 = to_per_head_f32(A); + at::Tensor D_f32; + if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value()); + + // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in + // Python. + at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous(); + if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat); + + at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous(); + + const int64_t batch = final_states.size(0); + const int64_t nheads = final_states.size(1); + const int64_t headdim = final_states.size(2); + const int64_t dstate = final_states.size(3); + const int64_t ngroups = B_in.size(1); + + TORCH_CHECK(final_states.is_contiguous(), + "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); + TORCH_CHECK(out.is_contiguous(), + "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via " + "raw data_ptr)"); + + VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] { + mamba_cpu::mamba_chunk_scan_fwd_kernel( + final_states.data_ptr(), x_in.data_ptr(), + dt_c.data_ptr(), A_f32.data_ptr(), + B_in.data_ptr(), C_in.data_ptr(), + D_f32.defined() ? D_f32.data_ptr() : nullptr, + z_in.defined() ? z_in.data_ptr() : nullptr, + out.data_ptr(), cu_int.data_ptr(), batch, nheads, + ngroups, headdim, dstate); + }); +} diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp new file mode 100644 index 000000000000..722dad97e3ca --- /dev/null +++ b/csrc/cpu/mamba_kernels.hpp @@ -0,0 +1,382 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Fused CPU vector kernels for Mamba decode-step hotspots: +// - causal_conv1d_update (depthwise 1-D conv state roll + compute) +// - selective_state_update (SSM recurrence, single-step) + +#pragma once + +#include "cpu_types.hpp" +#include +#include +#include +#include + +namespace mamba_cpu { + +// --------------------------------------------------------------------------- +// causal_conv1d_update — templated for native BF16/FP32 +// +// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor. +// Explicit strides are passed so the kernel writes directly into the +// correct memory locations without making a contiguous copy of the full +// paged tensor (which was the source of the 34-41% direct_copy_kernel). +// +// stride_s_slot = state.stride(0) — between cache slots +// stride_s_dim = state.stride(1) — between conv_dim channels +// stride_s_state = state.stride(2) — between state elements +// +// When stride_s_state == 1 (contiguous), the memmove fast path is used. +// --------------------------------------------------------------------------- +template +inline void causal_conv1d_update_kernel( + const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr, + int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state, + const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, + scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, + int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, + int64_t width, int64_t state_len, bool do_silu) { +#pragma omp parallel for + for (int64_t b = 0; b < batch; ++b) { + int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; + if (cache_idx == pad_slot_id) continue; + + for (int64_t t = 0; t < seqlen; ++t) { + const scalar_t* x_b = x_ptr + (b * dim * seqlen + t); + scalar_t* out_b = out_ptr + (b * dim * seqlen + t); + // Base of this slot in the (possibly non-contiguous) paged state + scalar_t* s_base = state_ptr + cache_idx * stride_s_slot; + + for (int64_t d = 0; d < dim; ++d) { + float x_val = static_cast(x_b[d * seqlen]); + scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state + const scalar_t* w = weight_ptr + d * width; + + // Accumulate in float32 for precision + float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; + for (int64_t k = 0; k < state_len; ++k) { + acc += static_cast(w[k]) * + static_cast(sd[k * stride_s_state]); + } + acc += static_cast(w[state_len]) * x_val; + + // Shift state left and append new input. + // Use memmove when contiguous (stride==1); element loop otherwise. + if (stride_s_state == 1) { + if (state_len > 1) + std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); + if (state_len > 0) sd[state_len - 1] = static_cast(x_val); + } else { + for (int64_t k = 0; k < state_len - 1; ++k) + sd[k * stride_s_state] = sd[(k + 1) * stride_s_state]; + if (state_len > 0) + sd[(state_len - 1) * stride_s_state] = static_cast(x_val); + } + + if (do_silu) { + float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc)) + : std::exp(acc) / (1.0f + std::exp(acc)); + acc *= sigmoid; + } + out_b[d * seqlen] = static_cast(acc); + } + } + } +} + +// --------------------------------------------------------------------------- +// selective_state_update +// +// Template parameters: +// state_t - dtype of ssm_state cache (typically BFloat16) +// input_t - dtype of x, B, C (typically BFloat16) +// out_t - dtype of output tensor (typically BFloat16) +// Write directly — no float32 intermediate buffer needed. +// +// A, D, dt_bias are accepted as const float* (they are always float32 +// model parameters in Mamba2). This eliminates the per-call float32→BF16 +// conversion and the .contiguous() materialisation of the broadcast-expand. +// +// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the +// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed. +// --------------------------------------------------------------------------- +template +inline void selective_state_update_kernel( + state_t* __restrict__ state_ptr, int64_t stride_state_n, + int64_t stride_state_h, int64_t stride_state_d, + const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h, + // dt: (N, nheads) — scalar per head, NOT expanded to head_dim + const float* __restrict__ dt_ptr, int64_t stride_dt_n, + // A: (nheads,) float32 — scalar per head + const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr, + const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g, + // D: (nheads,) float32 — scalar per head (nullptr if not used) + const float* __restrict__ D_ptr, + // z: same shape as x (optional) + const input_t* __restrict__ z_ptr, + // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used) + const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr, + int64_t stride_out_n, int64_t stride_out_h, + const int32_t* __restrict__ state_batch_indices, + const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, + const int32_t* __restrict__ num_accepted_tokens, + const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads, + int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { + using state_vec_t = vec_op::vec_t; + using input_vec_t = vec_op::vec_t; + constexpr int VEC_ELEM_NUM = 8; + + int64_t nheads_per_group = nheads / ngroups; + + for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) { + int64_t bos, seq_len; + if (cu_seqlens != nullptr) { + bos = cu_seqlens[seq_idx]; + seq_len = cu_seqlens[seq_idx + 1] - bos; + } else { + bos = seq_idx; + seq_len = 1; + } + + int64_t state_read_idx = (state_batch_indices != nullptr) + ? state_batch_indices[seq_idx] + : seq_idx; + if (state_read_idx == null_block_id) continue; + + int64_t state_write_idx = (num_accepted_tokens == nullptr) + ? ((dst_state_batch_indices != nullptr) + ? dst_state_batch_indices[seq_idx] + : state_read_idx) + : -1; + + state_t* s = state_ptr + state_read_idx * stride_state_n; + + for (int64_t t = 0; t < seq_len; ++t) { + int64_t token_idx = bos + t; + const input_t* x_tok = x_ptr + token_idx * stride_x_n; + // dt: (N, nheads) — one float per head per token + const float* dt_tok = dt_ptr + token_idx * stride_dt_n; + const input_t* B_tok = B_ptr + token_idx * stride_BC_n; + const input_t* C_tok = C_ptr + token_idx * stride_BC_n; + out_t* out_tok = out_ptr + token_idx * stride_out_n; + +#pragma omp parallel for + for (int64_t h = 0; h < nheads; ++h) { + int64_t g = h / nheads_per_group; + const input_t* x_h = x_tok + h * stride_x_h; + const input_t* B_g = B_tok + g * stride_BC_g; + const input_t* C_g = C_tok + g * stride_BC_g; + out_t* out_h = out_tok + h * stride_out_h; + state_t* s_h = s + h * stride_state_h; + + // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing + float dt_val = dt_tok[h]; + if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h]; + if (dt_softplus) { + dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; + } + const float A_val = A_ptr[h]; // scalar: same for all dim, dstate + const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; + + const input_t* z_h = + (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h + : nullptr; + + vec_op::FP32Vec8 dt_vec(dt_val); + // dA = exp(A * dt): A and dt are SCALARS per head, so compute once + // and broadcast. This saves 7 redundant std::exp() calls that + // FP32Vec8::exp() would otherwise make on the broadcast vector. + const float dA_scalar = std::exp(A_val * dt_val); + vec_op::FP32Vec8 dA(dA_scalar); // broadcast + + for (int64_t d = 0; d < dim; ++d) { + float x_val = static_cast(x_h[d]); + + vec_op::FP32Vec8 out_vec(0.0f); + state_t* s_hd = s_h + d * stride_state_d; + const input_t* B_g_base = B_g; + const input_t* C_g_base = C_g; + + vec_op::FP32Vec8 x_vec(x_val); + // dBx = B * x * dt — same dA for all dstate (A is scalar) + // s_new = s * dA + B * x * dt + + int64_t n = 0; + for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { + vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n))); + vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n))); + vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n))); + + vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec; + vec_op::FP32Vec8 s_new = s_v * dA + dBx; + + state_vec_t(s_new).save(s_hd + n); + out_vec = out_vec + s_new * C_v; + } + + float out_val = out_vec.reduce_sum(); + for (; n < dstate; ++n) { + // Reuse dA_scalar computed once per head — no exp() re-call + float dBx = static_cast(B_g[n]) * x_val * dt_val; + float s_new = static_cast(s_hd[n]) * dA_scalar + dBx; + s_hd[n] = static_cast(s_new); + out_val += s_new * static_cast(C_g[n]); + } + + if (D_ptr != nullptr) out_val += x_val * D_val; + if (z_h != nullptr) { + float z_val = static_cast(z_h[d]); + float sigmoid = (z_val >= 0) + ? 1.0f / (1.0f + std::exp(-z_val)) + : std::exp(z_val) / (1.0f + std::exp(z_val)); + out_val *= z_val * sigmoid; + } + out_h[d] = static_cast(out_val); + } + } + + if (num_accepted_tokens != nullptr && + dst_state_batch_indices != nullptr) { + int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; + if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { + state_t* dst_s = state_ptr + token_dst_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); + } + } + } + + if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && + state_write_idx != state_read_idx) { + state_t* dst_s = state_ptr + state_write_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); + } + } +} + +// --------------------------------------------------------------------------- +// mamba_chunk_scan_fwd +// +// Prefill SSM recurrence for Mamba2 / SSD models. +// +// Key difference from selective_state_update_kernel (decode path): +// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop. +// Each thread owns a (batch, head) slice and runs the entire token +// sequence without any per-token OpenMP synchronisation overhead. +// For seqlen=256, this eliminates 256 thread-barrier launches per batch. +// +// `dt` arrives already processed (float32, after bias + softplus + clamp) +// to keep this kernel simple. Preprocessing is done in the Python wrapper. +// +// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output +// tensor, pre-initialised by the caller (zero or from initial_states). +// Each (b, h) slice is private to exactly one thread via collapse(2), so +// there are no write conflicts. +// +// D is treated as a scalar per head ([nheads] float32). +// --------------------------------------------------------------------------- +template +inline void mamba_chunk_scan_fwd_kernel( + float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32 + const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim] + const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed) + const float* __restrict__ A_ptr, // [nheads] f32 + const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate] + const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate] + const float* __restrict__ D_ptr, // [nheads] f32 (nullable) + const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable) + input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim] + const int32_t* __restrict__ cu_seqlens, // [batch+1] int32 + int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim, + int64_t dstate) { + using input_vec_t = vec_op::vec_t; + constexpr int VEC_ELEM_NUM = 8; + + const int64_t nheads_per_group = nheads / ngroups; + // states layout: [batch, nheads, headdim, dstate] contiguous (caller + // guarantee) + const int64_t stride_s_b = nheads * headdim * dstate; + const int64_t stride_s_h = headdim * dstate; + // stride_s_d = dstate, stride_s_n = 1 + +#pragma omp parallel for collapse(2) schedule(static) + for (int64_t b = 0; b < batch; ++b) { + for (int64_t h = 0; h < nheads; ++h) { + const int64_t seq_start = cu_seqlens[b]; + const int64_t seq_end = cu_seqlens[b + 1]; + const int64_t g = h / nheads_per_group; + + const float A_val = A_ptr[h]; + const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; + + // Working state slice: states[b, h, :, :] — float32, headdim * dstate. + // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB). + float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h; + + for (int64_t t = seq_start; t < seq_end; ++t) { + const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim; + const float* dt_h = dt_ptr + t * nheads + h; + const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate; + const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate; + const input_t* z_h = (z_ptr != nullptr) + ? z_ptr + t * nheads * headdim + h * headdim + : nullptr; + input_t* out_h = out_ptr + t * nheads * headdim + h * headdim; + + const float dt_val = *dt_h; + const float dA_val = std::exp(A_val * dt_val); + const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar + const vec_op::FP32Vec8 dt_vec(dt_val); + + for (int64_t d = 0; d < headdim; ++d) { + const float x_val = static_cast(x_h[d]); + float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32 + + // Vectorised SSM update + readout over dstate: + // s_new = s * dA + x * dt * B + // y += s_new * C + int64_t n = 0; + vec_op::FP32Vec8 y_vec(0.0f); + const vec_op::FP32Vec8 x_vec(x_val); + + for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { + const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n))); + const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n))); + const vec_op::FP32Vec8 s_v(s_bhd + n); + + const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v; + s_new.save(s_bhd + n); + y_vec = y_vec + s_new * C_v; + } + + float y_val = y_vec.reduce_sum(); + + // Scalar tail for remaining dstate elements + for (; n < dstate; ++n) { + const float B_n = static_cast(B_g[n]); + const float C_n = static_cast(C_g[n]); + const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n; + s_bhd[n] = s_new; + y_val += s_new * C_n; + } + + // D skip connection (scalar per head) + if (D_ptr != nullptr) y_val += x_val * D_val; + + // z gating: out = y * z * sigmoid(z) (SiLU) + if (z_h != nullptr) { + const float z_val = static_cast(z_h[d]); + const float sigmoid = + (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val)) + : std::exp(z_val) / (1.0f + std::exp(z_val)); + y_val *= z_val * sigmoid; + } + + out_h[d] = static_cast(y_val); + } + } + } + } +} + +} // namespace mamba_cpu diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp index f0471f714703..75505176eeee 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp @@ -31,6 +31,9 @@ class MicroGemm { } }; +template +class MicroGemmINT8; + template FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr, scalar_t* __restrict__ d_ptr, diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp new file mode 100644 index 000000000000..2e194fd39037 --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp @@ -0,0 +1,424 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP +#define CPU_MICRO_GEMM_INT8_NEON_HPP + +#include +#include + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#include +#include +#include +#include +#include + +namespace cpu_micro_gemm { + +namespace neon_smmla { + +constexpr int32_t K = 8; +constexpr int32_t Cols = 2; +constexpr int32_t TileSize = K * Cols; + +FORCE_INLINE float32x4x2_t load_as_f32(const float* input) { + float32x4x2_t result; + result.val[0] = vld1q_f32(input); + result.val[1] = vld1q_f32(input + 4); + return result; +} + +FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) { + const auto input_vec = vld1q_f16(reinterpret_cast(input)); + float32x4x2_t result; + result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec)); + result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec)); + return result; +} + +FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) { + const auto input_vec = vld1q_bf16(reinterpret_cast(input)); + float32x4x2_t result; + result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec)); + result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec)); + return result; +} + +FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01, + const int32x4_t acc23, + const int32x4_t acc45, + const int32x4_t acc67, + int32_t* __restrict__ c_ptr, + const int64_t ldc, const int32_t m_rows) { + if (m_rows == 0) { + return; + } + + vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23))); + vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67))); + + if (m_rows == 2) { + vst1q_s32(c_ptr + ldc, + vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23))); + vst1q_s32(c_ptr + ldc + 4, + vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67))); + } +} + +FORCE_INLINE void gemm_micro_smmla_8x8_packed_a( + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed, + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size, + const int64_t ldc) { + const int32x4_t zero = vdupq_n_s32(0); + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero; + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero; + int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero; + int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero; + + const int8_t* __restrict__ a_tile = a_packed; + const int8_t* __restrict__ b_tile = b_packed; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const int8x16_t a_tile01 = vld1q_s8(a_tile); + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize); + const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize); + const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize); + + const int8x16_t b_tile01 = vld1q_s8(b_tile); + const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize); + const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize); + const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize); + + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01); + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01); + acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01); + acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01); + + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23); + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23); + acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23); + acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23); + + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45); + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45); + acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45); + acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45); + + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67); + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67); + acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67); + acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67); + + a_tile += 4 * TileSize; + b_tile += 4 * TileSize; + } + + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, + std::min(2, m)); + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + std::min(2, std::max(0, m - 2))); + store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc, + std::min(2, std::max(0, m - 4))); + store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc, + std::min(2, std::max(0, m - 6))); +} + +FORCE_INLINE void gemm_micro_smmla_4x16_packed_a( + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed, + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size, + const int64_t b_n_group_stride, const int64_t ldc) { + const int32_t m_rows_01 = std::min(2, m); + const int32_t m_rows_23 = std::min(2, std::max(0, m - 2)); + const int32x4_t zero = vdupq_n_s32(0); + + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero; + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero; + int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero, + acc011415 = zero; + int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero, + acc231415 = zero; + + const int8_t* __restrict__ a_tile = a_packed; + // note: b packs 8 panels contiguously, so we need 2 b_tile ptrs + // for the 4x16 microkernel + const int8_t* __restrict__ b_tile0 = b_packed; + const int8_t* __restrict__ b_tile1 = b_packed + b_n_group_stride; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const int8x16_t a_tile01 = vld1q_s8(a_tile); + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize); + const int8x16_t b_tile01 = vld1q_s8(b_tile0); + const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize); + const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize); + const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize); + const int8x16_t b_tile89 = vld1q_s8(b_tile1); + const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize); + const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize); + const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize); + + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01); + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01); + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23); + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23); + + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45); + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45); + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67); + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67); + + acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89); + acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89); + acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011); + acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011); + + acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213); + acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213); + acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415); + acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415); + + a_tile += 2 * TileSize; + b_tile0 += 4 * TileSize; + b_tile1 += 4 * TileSize; + } + + // rows 0-1, columns 0-7 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01); + // rows 0-1, columns 8-15 + store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc, + m_rows_01); + // rows 2-3, columns 0-7 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + m_rows_23); + // rows 2-3, columns 8-15 + store_acc_rowpair(acc2389, acc231011, acc231213, acc231415, + c_ptr + 2 * ldc + 8, ldc, m_rows_23); +} + +} // namespace neon_smmla + +template +class MicroGemmINT8 { + public: + static constexpr int32_t K = neon_smmla::K; + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = Nr; + static_assert(MaxMSize % Mr == 0); + + static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output, + float& scale, const int32_t size) { + TORCH_CHECK_EQ(size % K, 0); + float32x4_t max_vec = vdupq_n_f32(0.0f); + + for (int32_t i = 0; i < size; i += K) { + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i); + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0])); + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1])); + } + + const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f); + scale = abs_max / 127.0f; + const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max); + + for (int32_t i = 0; i < size; i += K) { + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i); + const int32x4_t output_low = + vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec)); + const int32x4_t output_high = + vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec)); + const int16x8_t output_s16 = + vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high)); + vst1_s8(output + i, vqmovn_s16(output_s16)); + } + } + + // with current code, fusing this into the gemm micro kernel didn't move the + // needle + static FORCE_INLINE void dequantize_tile( + int32_t* input, float* output, const float* __restrict__ input_scales, + const float* __restrict__ weight_scales, const int32_t m, const int32_t n, + const int32_t stride) { + TORCH_CHECK_EQ(n % 4, 0); + for (int32_t m_idx = 0; m_idx < m; ++m_idx) { + const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]); + for (int32_t n_idx = 0; n_idx < n; n_idx += 4) { + const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx); + const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx); + const float32x4_t output_vec = + vmulq_f32(vcvtq_f32_s32(input_vec), + vmulq_f32(input_scale_vec, weight_scale_vec)); + vst1q_f32(output + m_idx * stride + n_idx, output_vec); + } + } + } + + // physical layout [ + // M / (8 or 4); Mr is 8 or 4 + // K / 8; K for smmla is 8 + // 4, ; 4 row-pairs for each 8 rows + // 2, ; row-pair is 2 rows + // 4 ; 4 elements per row + // ] + static void pack_input_from_rows(const int8_t* const* __restrict__ rows, + int8_t* __restrict__ a_packed, + const int32_t m, const int32_t k) { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK(k % K == 0); + const int8x8_t zero = vdup_n_s8(0); + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* const* panel_rows = rows + row_base; + int8_t* __restrict__ out = a_packed + row_base * k; + + // fast path for full 8-row panels (fast path for 4-row panels didn't move + // the needle) + if (panel_m == Mr) { + const int8_t* __restrict__ row0 = panel_rows[0]; + const int8_t* __restrict__ row1 = panel_rows[1]; + const int8_t* __restrict__ row2 = panel_rows[2]; + const int8_t* __restrict__ row3 = panel_rows[3]; + const int8_t* __restrict__ row4 = panel_rows[4]; + const int8_t* __restrict__ row5 = panel_rows[5]; + const int8_t* __restrict__ row6 = panel_rows[6]; + const int8_t* __restrict__ row7 = panel_rows[7]; + int32_t k_idx = 0; + for (; k_idx + 2 * K <= k; k_idx += 2 * K) { + int8_t* __restrict__ block0 = out; + int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize; + + int8x16_t a0 = vld1q_s8(row0 + k_idx); + int8x16_t a1 = vld1q_s8(row1 + k_idx); + vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row2 + k_idx); + a1 = vld1q_s8(row3 + k_idx); + vst1q_s8(block0 + neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row4 + k_idx); + a1 = vld1q_s8(row5 + k_idx); + vst1q_s8(block0 + 2 * neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + 2 * neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row6 + k_idx); + a1 = vld1q_s8(row7 + k_idx); + vst1q_s8(block0 + 3 * neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + 3 * neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + out += 8 * neon_smmla::TileSize; + } + + for (; k_idx < k; k_idx += K) { + int8x8_t a0 = vld1_s8(row0 + k_idx); + int8x8_t a1 = vld1_s8(row1 + k_idx); + vst1q_s8(out, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row2 + k_idx); + a1 = vld1_s8(row3 + k_idx); + vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row4 + k_idx); + a1 = vld1_s8(row5 + k_idx); + vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row6 + k_idx); + a1 = vld1_s8(row7 + k_idx); + vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1)); + + out += 4 * neon_smmla::TileSize; + } + continue; + } + + const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2; + for (int32_t k_idx = 0; k_idx < k; k_idx += K) { + for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) { + const int32_t row_idx = pair_idx * 2; + const int8x8_t row0 = + (row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero; + const int8x8_t row1 = (row_idx + 1 < panel_m) + ? vld1_s8(panel_rows[row_idx + 1] + k_idx) + : zero; + vst1q_s8(out, vcombine_s8(row0, row1)); + out += neon_smmla::TileSize; + } + } + } + } + + // physical layout [ + // N / 8; Nr is 8 + // K / 8; K for smmla is 8 + // 4, ; 4 col-pairs for each 8 cols + // 2, ; col-pair is 2 cols + // 4 ; 4 elements per col + // ] + static void pack_weight(const int8_t* __restrict__ weight, + int8_t* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK(output_size % NSize == 0); + TORCH_CHECK(input_size % K == 0); + + for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) { + int8_t* __restrict__ dst = packed_weight + o_idx * input_size; + for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) { + for (int32_t pair_idx = 0; pair_idx < Nr; + pair_idx += neon_smmla::Cols) { + const int8_t* __restrict__ row0 = + weight + (o_idx + pair_idx) * input_size + k_idx; + const int8_t* __restrict__ row1 = row0 + input_size; + vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1))); + dst += neon_smmla::TileSize; + } + } + } + } + + void gemm(const int8_t* __restrict__ a_packed, + const int8_t* __restrict__ b_packed, int32_t* __restrict__ c, + const int32_t m, const int32_t k, const int64_t b_n_group_stride, + const int64_t ldc) const { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK(k % K == 0); + + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) { + const int8_t* __restrict__ b_panel = b_packed + n_idx * k; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* __restrict__ a_panel = a_packed + row_base * k; + int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx; + + if (panel_m <= 4) { + neon_smmla::gemm_micro_smmla_4x16_packed_a( + a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc); + } else { + neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel, + panel_m, k, ldc); + neon_smmla::gemm_micro_smmla_8x8_packed_a( + a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k, + ldc); + } + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp index 7d4898852bb3..b38337956f6b 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + #ifndef CPU_MICRO_GEMM_NEON_HPP #define CPU_MICRO_GEMM_NEON_HPP @@ -16,9 +19,6 @@ namespace { constexpr int32_t K = 4; constexpr int32_t Cols = 2; constexpr int32_t TileSize = K * Cols; -constexpr int32_t Mr = 8; -constexpr int32_t Nr = 8; -constexpr int32_t Nr_gemv = 16; // a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1] FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) { @@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a( acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67); a_tile += 4 * TileSize; - b_tile += Nr * K; + b_tile += 4 * TileSize; } store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, @@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a( acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415); a_tile += 2 * TileSize; - b_tile0 += Nr * K; - b_tile1 += Nr * K; + b_tile0 += 4 * TileSize; + b_tile1 += 4 * TileSize; } store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01); @@ -223,6 +223,9 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a( template class MicroGemm { public: + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; static constexpr int32_t MaxMSize = 8; static constexpr int32_t NSize = 32; static constexpr int32_t WeightOCGroupSize = Nr; @@ -246,6 +249,9 @@ class MicroGemm { public: using scalar_t = c10::BFloat16; + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; static constexpr int32_t MaxMSize = 8; static constexpr int32_t NSize = 32; static constexpr int32_t WeightOCGroupSize = Nr; @@ -253,7 +259,7 @@ class MicroGemm { public: // physical layout [ - // M / 8; Mr is 8 + // M / (8 or 4); Mr is 8 or 4 // K / 4; K for bfmmla is 4 // 4, ; 4 row-pairs for each 8 rows // 2, ; row-pair is 2 rows @@ -439,7 +445,7 @@ class MicroGemm { (void)lda; // A is packed, so lda is not needed TORCH_CHECK_EQ(k % K, 0); - for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) { + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) { const bfloat16_t* __restrict__ b_panel = reinterpret_cast(b_ptr) + n_idx * k; diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index d501d2adb811..cbd7e691f779 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -1,11 +1,16 @@ // Adapted from // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu +// +// Synced from +// https://github.com/sgl-project/sglang/tree/7c248dde7fe1f3b5100966f8143f97a9932c22a4/sgl-kernel/csrc/cpu +// Sync date: 2026-07-29 // clang-format off #pragma once #include +#include #include #if defined(_OPENMP) @@ -49,6 +54,14 @@ namespace { } \ }() +// Half + BFloat16, plus one extra scalar type +#define AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, ...) \ + AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES(__VA_ARGS__) \ + AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__) + +#define AT_DISPATCH_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, __VA_ARGS__)) + // dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4) #define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \ [&] { \ diff --git a/csrc/cpu/sgl-kernels/conv.cpp b/csrc/cpu/sgl-kernels/conv.cpp index b918aed8bff8..38079ba95011 100644 --- a/csrc/cpu/sgl-kernels/conv.cpp +++ b/csrc/cpu/sgl-kernels/conv.cpp @@ -182,14 +182,13 @@ struct tinygemm_kernel { using fVec = at::vec::Vectorized; using bVec = at::vec::Vectorized; - const fVec one = fVec(1.f); auto storec = [&](auto i, int64_t m) { constexpr int col = i; fVec x0 = fVec(vc[col * 2 + 0]); fVec x1 = fVec(vc[col * 2 + 1]); if constexpr (has_silu) { - x0 = x0 / (one + x0.neg().exp_u20()); - x1 = x1 / (one + x1.neg().exp_u20()); + x0 = fast_silu(x0); + x1 = fast_silu(x1); } bVec out_vec = convert_from_float_ext(x0, x1); out_vec.store(C + m * lda + col * 32); @@ -451,6 +450,90 @@ void causal_conv1d_update_kernel_impl( }); } +template +void causal_conv1d_update_multi_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + scalar_t* __restrict__ conv_states, + const scalar_t* __restrict__ weight, + const scalar_t* __restrict__ bias, + const int32_t* __restrict__ num_accepted_tokens, + const int32_t* __restrict__ conv_indices, + bool silu_activation, + int64_t batch, + int64_t dim, + int64_t seqlen, + int64_t width, + int64_t state_len, + int64_t conv_state_slot_stride) { + constexpr int64_t BLOCK_N = block_size_n() * 2; + const int64_t NB = div_up(dim, BLOCK_N); + + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] { + at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, nb{0}; + data_index_init(begin, bs, batch, nb, NB); + + for (int64_t i = begin; i < end; ++i) { + const int64_t nb_start = nb * BLOCK_N; + const int64_t nb_size = std::min(dim - nb_start, BLOCK_N); + const int32_t conv_state_index = conv_indices[bs]; + const int32_t history_offset = num_accepted_tokens[bs] - 1; + + switch (width << 4 | nb_size >> 4) { + case 0x42: + tinygemm_kernel::apply( + input + bs * seqlen * dim + nb_start, + weight + nb_start * width, + out + bs * seqlen * dim + nb_start, + has_bias ? bias + nb_start : nullptr, + conv_states + conv_state_index * conv_state_slot_stride + + history_offset * dim + nb_start, + true, + seqlen, + dim, + true); + break; + case 0x44: + tinygemm_kernel::apply( + input + bs * seqlen * dim + nb_start, + weight + nb_start * width, + out + bs * seqlen * dim + nb_start, + has_bias ? bias + nb_start : nullptr, + conv_states + conv_state_index * conv_state_slot_stride + + history_offset * dim + nb_start, + true, + seqlen, + dim, + true); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size); + } + + data_index_step(bs, batch, nb, NB); + } + }); + }); + + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { + for (int64_t bs = begin; bs < end; ++bs) { + const int32_t conv_state_index = conv_indices[bs]; + const int32_t num_accepted = num_accepted_tokens[bs]; + scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride; + + std::memmove( + state, + state + num_accepted * dim, + (state_len - seqlen) * dim * sizeof(scalar_t)); + std::memcpy( + state + (state_len - seqlen) * dim, + input + bs * seqlen * dim, + seqlen * dim * sizeof(scalar_t)); + } + }); +} + } // anonymous namespace // from [dim, width] or [N, K] @@ -545,7 +628,7 @@ at::Tensor get_block_indices(const std::optional& offsets, int64_t n // query_start_loc: (batch + 1) int32 // cache_indices: (batch) int32 // has_initial_state: (batch) bool -// conv_states: (..., dim, width - 1) itype +// conv_states: (..., dim, state_len) itype, where state_len >= width - 1 // activation: either None or "silu" or "swish" // pad_slot_id: int // @@ -586,11 +669,14 @@ at::Tensor causal_conv1d_fwd_cpu( CHECK_EQ(conv_states_val.scalar_type(), scalar_type); CHECK_GE(padded_batch, batch); CHECK_EQ(conv_states_val.size(1), dim); - CHECK_EQ(conv_states_val.size(2), width - 1); + const int64_t state_len = conv_states_val.size(2); + CHECK_GE(state_len, width - 1); // adjust `conv_states` to be contiguous on `dim` // should happen only once if (conv_states_val.stride(-2) != 1) { + TORCH_CHECK(state_len == width - 1, + "causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim."); auto conv_states_copy = conv_states_val.clone(); conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim}); conv_states_val.copy_(conv_states_copy); @@ -651,14 +737,14 @@ at::Tensor causal_conv1d_fwd_cpu( // API aligned with GPUs // -// x: (batch, dim) or (batch, dim, seqlen) +// x: (batch, dim) or (batch, seqlen, dim) // conv_state: (..., dim, state_len), where state_len >= width - 1 // weight: (dim, width) // bias: (dim,) -// cache_seqlens: (batch,), dtype int32. +// num_accepted_tokens: (batch,), dtype int32. // conv_state_indices: (batch,), dtype int32 // pad_slot_id: int -// out: (batch, dim) or (batch, dim, seqlen) +// out: (batch, dim) or (batch, seqlen, dim) // at::Tensor causal_conv1d_update_cpu( const at::Tensor& x, @@ -666,7 +752,7 @@ at::Tensor causal_conv1d_update_cpu( const at::Tensor& weight, const std::optional& bias, bool silu_activation, - const std::optional& cache_seqlens, + const std::optional& num_accepted_tokens, const std::optional& conv_state_indices, int64_t pad_slot_id, bool is_vnni) { @@ -674,13 +760,13 @@ at::Tensor causal_conv1d_update_cpu( CHECK_CONTIGUOUS(weight); auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight); - // TODO: add multi-token prediction support - TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor."); - TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens."); + TORCH_CHECK( + x.dim() == 2 || x.dim() == 3, + "causal_conv1d_update_cpu: expect x to be 2D or 3D tensor."); int64_t batch = x.size(0); - int64_t dim = x.size(1); - int64_t seqlen = 1; + int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2); + int64_t seqlen = x.dim() == 2 ? 1 : x.size(1); int64_t width = weight.size(-1); const auto scalar_type = x.scalar_type(); @@ -690,10 +776,84 @@ at::Tensor causal_conv1d_update_cpu( CHECK_EQ(conv_states.scalar_type(), scalar_type); CHECK_EQ(conv_states.size(1), dim); - CHECK_EQ(conv_states.size(2), width - 1); + const int64_t state_len = conv_states.size(2); + CHECK_GE(state_len, width - 1); + + if (x.dim() == 3) { + TORCH_CHECK( + num_accepted_tokens.has_value(), + "causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x."); + TORCH_CHECK( + conv_state_indices.has_value(), + "causal_conv1d_update_cpu: conv_state_indices is required for 3D x."); + CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt); + TORCH_CHECK( + width == 4, + "causal_conv1d_update_cpu: support only width of 4 for 3D x."); + TORCH_CHECK( + seqlen > 0, + "causal_conv1d_update_cpu: expect non-empty sequence for 3D x."); + TORCH_CHECK( + state_len >= seqlen, + "causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x."); + TORCH_CHECK( + conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim, + "causal_conv1d_update_cpu: 3D x requires SD conv_states layout."); + + const int32_t* accepted_counts = + num_accepted_tokens.value().data_ptr(); + const int32_t* indices = conv_state_indices.value().data_ptr(); + const int64_t num_slots = conv_states.size(0); + for (int64_t bs = 0; bs < batch; ++bs) { + const int32_t num_accepted = accepted_counts[bs]; + const int32_t conv_state_index = indices[bs]; + TORCH_CHECK( + conv_state_index != pad_slot_id, + "causal_conv1d_update_cpu: 3D x does not support pad slots."); + TORCH_CHECK( + conv_state_index >= 0 && conv_state_index < num_slots, + "causal_conv1d_update_cpu: conv_state_indices out of range."); + TORCH_CHECK( + num_accepted >= 1 && num_accepted <= seqlen, + "causal_conv1d_update_cpu: num_accepted_tokens must be in [1, " + "seqlen]."); + TORCH_CHECK( + num_accepted - 1 + width - 1 <= state_len, + "causal_conv1d_update_cpu: history window exceeds conv_states."); + } + + int64_t conv_state_slot_stride = conv_states.stride(0); + at::Tensor out = at::empty_like(x); + AT_DISPATCH_REDUCED_FLOATING_TYPES( + scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] { + causal_conv1d_update_multi_kernel_impl( + out.data_ptr(), + x.data_ptr(), + conv_states.data_ptr(), + packed_w.data_ptr(), + conditional_data_ptr(bias), + accepted_counts, + indices, + silu_activation, + batch, + dim, + seqlen, + width, + state_len, + conv_state_slot_stride); + }); + return out; + } + + TORCH_CHECK( + !num_accepted_tokens.has_value(), + "causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D " + "x."); // adjust `conv_states` to be contiguous on `dim` if (conv_states.stride(-2) != 1) { + TORCH_CHECK(state_len == width - 1, + "causal_conv1d_update_cpu: wide conv_states must be contiguous on dim."); int64_t num_cache_lines = conv_states.size(0); auto conv_states_copy = conv_states.clone(); conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim}); diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index b9a793c599f5..fec6648cec66 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -9,929 +9,1550 @@ #include "vec_pack.h" namespace { -// For this cpu kernel, we have some innovations aside from the existing gpu kernels: -// 1) Use less parallel loops, i.e. 4 including l2_norm. -// 2) Fuse part of l2_norm with the rest of the computation. -#define THREAD_BUFFER_ALLOC(dst, base_ptr, offset, type, size) \ - type* dst = reinterpret_cast((base_ptr) + (offset)); \ - offset += (size); +// [NOTE] GDN Optimizations on AMX CPU +// * intra loop: fuse `kkt_solve` and `recompute_w_u` so as to avoid materialize `A`. +// * inter loop: fuse `recompute_w_u` and `update_v` so as to avoid materialize `h` and `v_new`. +// * intra loop parallel on H instead of Hv, remove duplicated key @ key.T +// * fuse format pack with elemwise OP as much as possible. +// * update state (FP32) with amx-bf16 where C(FP32) += A(BF16) * B(BF16) +// * compile time mask out upper triangular part in decay mask and tril solve, reduce fma needed. + +// * convert to vnni format， expect contiguous input and output +// from [K/2, 2, N] FP32 to [K/2, N, 2] BF16 +// * update src = src * exp(g_last) +template +void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float g_last, int ld_src, int ld_dst) { + static_assert(K % 32 == 0); + static_assert(N % 32 == 0); + + const float scale = std::exp(g_last); +#if defined(CPU_CAPABILITY_AVX512) + constexpr int KB = K / 2; + constexpr int NB = N / 32; + + __m512i s0, s1, d0, d1; + __m512 vd = _mm512_set1_ps(scale); + + const auto trans = [&](auto i) { + constexpr int kb = i / NB; + constexpr int nb = i % NB; + + // [K/2, 2, N/32, 32] -> [K/2, N/32, 32, 2] + constexpr int k0 = kb * 2 + 0; + constexpr int k1 = kb * 2 + 1; + __m512 v00 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32); + __m512 v01 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32 + 16); + __m512 v10 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32); + __m512 v11 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32 + 16); + s0 = (__m512i)_mm512_cvtne2ps_pbh(v01, v00); + s1 = (__m512i)_mm512_cvtne2ps_pbh(v11, v10); + + std::tie(d0, d1) = transpose_2x32_16bit(s0, s1); + _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2, d0); + _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2 + 32, d1); + + // update src = src * exp(g_last) + _mm512_storeu_ps(src + k0 * ld_src + nb * 32, _mm512_mul_ps(v00, vd)); + _mm512_storeu_ps(src + k0 * ld_src + nb * 32 + 16, _mm512_mul_ps(v01, vd)); + _mm512_storeu_ps(src + k1 * ld_src + nb * 32, _mm512_mul_ps(v10, vd)); + _mm512_storeu_ps(src + k1 * ld_src + nb * 32 + 16, _mm512_mul_ps(v11, vd)); + }; + Unroll{}(trans); +#else + // [K/2, 2, N] -> [K/2, N, 2] + for (int k = 0; k < K; k += 2) { + for (int n = 0; n < N; ++n) { + const float v0 = src[(k + 0) * ld_src + n]; + const float v1 = src[(k + 1) * ld_src + n]; + dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = static_cast(v0); + dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = static_cast(v1); + src[(k + 0) * ld_src + n] = v0 * scale; + src[(k + 1) * ld_src + n] = v1 * scale; + } + } +#endif +} -template -inline void fill_stub(scalar_t* __restrict__ out, float val, int size) { +template +inline void fill_stub(scalar_t* __restrict__ out, float val) { using Vec = at::vec::Vectorized; constexpr int kVecSize = Vec::size(); + static_assert(SIZE % kVecSize == 0); const Vec data_vec = Vec(static_cast(val)); - int d = 0; -#pragma GCC unroll 4 - for (; d <= size - kVecSize; d += kVecSize) { +#pragma GCC unroll 8 + for (int d = 0; d < SIZE; d += kVecSize) { data_vec.store(out + d); } - if (size - d > 0) { - data_vec.store(out + d, size - d); - } } -template -void chunk_gated_delta_rule_kernel_impl( - scalar_t* __restrict__ out, // [B, T, HV, EV] - float* __restrict__ final_state_data, // [N, HV, EK, EV] - const scalar_t* __restrict__ q_orig, // [B, T, HK, EK] - const scalar_t* __restrict__ k_orig, // [B, T, HK, EK] - const scalar_t* __restrict__ v_orig, // [B, T, HV, EV] - const float* __restrict__ g_orig, // [B, T, HV] FP32 - const scalar_t* __restrict__ b_orig, // [B, T, HV] - const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] INT32 - float* __restrict__ buff, - scalar_t* __restrict__ reduced_buff, - scalar_t* __restrict__ thread_buff, - const int32_t* __restrict__ chunk_offsets_ptr, - const int32_t* __restrict__ chunk_indices_ptr, - bool use_qk_l2norm_in_kernel, - const int64_t& batch_size, - const int64_t& global_seq_len, - const int64_t& qk_num_head, - const int64_t& v_num_head, - const int64_t& qk_head_size, - const int64_t& v_head_size, - const int64_t& qStrideH, - const int64_t& qStrideT, - const int64_t& kStrideH, - const int64_t& kStrideT, - const int64_t& vStrideH, - const int64_t& vStrideT, - const int64_t& oStrideH, - const int64_t& oStrideT, - const int64_t& global_total_seq_length, - const int64_t& global_num_chunk, - const int64_t& buff_size_16bit_per_thread, - double eps = 1e-5) { - int64_t gStrideH = 1; - int64_t gStrideT = v_num_head; - int64_t bStrideH = 1; - int64_t bStrideT = v_num_head; - int64_t final_state_StrideN = v_num_head * qk_head_size * v_head_size; - int64_t final_state_StrideH = qk_head_size * v_head_size; - int64_t final_state_StrideE = v_head_size; - int64_t head_group = v_num_head / qk_num_head; - float scale = 1.0 / std::sqrt(qk_head_size); - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int64_t VecSize = bVec::size(); - constexpr int64_t fVecSize = fVec::size(); - - // Data pointers - float* g_pad = buff; - float* core_attn_out = g_pad + v_num_head * global_total_seq_length; - float* decay_mask = core_attn_out + batch_size * v_num_head * global_total_seq_length * v_head_size; - float* v_beta_attn = decay_mask + v_num_head * global_total_seq_length * chunk_size; - - scalar_t* q_pad = reduced_buff; - scalar_t* k_pad = q_pad + qk_num_head * global_total_seq_length * qk_head_size; - scalar_t* v_pad = k_pad + qk_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_beta = v_pad + v_num_head * global_total_seq_length * v_head_size; - scalar_t* v_beta = k_beta + v_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_cumdecay_reduced = v_beta + v_num_head * global_total_seq_length * v_head_size; - scalar_t* q_norm_sum = k_cumdecay_reduced + v_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_norm_sum = q_norm_sum + qk_num_head * global_seq_len; +// Portable fallback for non-AVX512 builds (ARM/NEON, old x86 without +// AVX512BF16), vectorized via at::vec::Vectorized (portable across +// AVX2/NEON/generic) mirroring the idioms used elsewhere in this file (see +// l2norm_fwd_kernel_impl's predecessor and fused_gdn_gating_kernel_impl): +// bVec/fVec pairs with convert_to_float/convert_from_float for bf16<->float, +// plain scalar tails for remainders. Two kernels (cumsum_kernel, +// update_key_kernel) write a transposed layout relative to their vectorized +// read axis; those vectorize the load/compute and unpack lanes for the +// (unavoidably strided) store. +template +struct l2norm_kernel { + static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float eps) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int bVecSize = bVec::size(); + const float scale = 1.f / std::sqrt(static_cast(D)); + + fVec sum_fvec0(0.f), sum_fvec1(0.f); + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + bVec in_bvec = bVec::loadu(input + d); + fVec in0, in1; + std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); + sum_fvec0 = sum_fvec0 + in0 * in0; + sum_fvec1 = sum_fvec1 + in1 * in1; + } + float sqsum = vec_reduce_sum(sum_fvec0 + sum_fvec1); + for (; d < D; ++d) { + float v = static_cast(input[d]); + sqsum += v * v; + } - if (use_qk_l2norm_in_kernel) { - at::parallel_for(0, qk_num_head * global_seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t h_qk = 0, l = 0; - data_index_init(begin, h_qk, qk_num_head, l, global_seq_len); - for (int64_t i = begin; i < end; ++i) { - auto q_norm_sum_ptr = q_norm_sum + h_qk * global_seq_len + l; - auto k_norm_sum_ptr = k_norm_sum + h_qk * global_seq_len + l; - float sum_q = float(0); - float sum_k = float(0); - fVec sum_q_fvec = fVec(float(0)); - fVec sum_k_fvec = fVec(float(0)); - int64_t q_offset = l * qStrideT + h_qk * qStrideH; - int64_t k_offset = l * qStrideT + h_qk * qStrideH; - int64_t d; - for (d = 0; d <= qk_head_size - VecSize; d += VecSize) { - bVec q_bvec = bVec::loadu(q_orig + q_offset + d); - fVec q_fvec0, q_fvec1; - std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec); - sum_q_fvec += q_fvec0 * q_fvec0; - sum_q_fvec += q_fvec1 * q_fvec1; - bVec k_bvec = bVec::loadu(k_orig + k_offset + d); - fVec k_fvec0, k_fvec1; - std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec); - sum_k_fvec += k_fvec0 * k_fvec0; - sum_k_fvec += k_fvec1 * k_fvec1; + float rscale = 1.f / std::sqrt(sqsum + eps); + fVec rscale_fvec(rscale); + fVec scale_fvec(scale); + d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + bVec in_bvec = bVec::loadu(input + d); + fVec in0, in1; + std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); + in0 = in0 * rscale_fvec; + in1 = in1 * rscale_fvec; + if constexpr (has_scale) { + in0 = in0 * scale_fvec; + in1 = in1 * scale_fvec; + } + bVec out_bvec = at::vec::convert_from_float(in0, in1); + out_bvec.store(out + d); + } + for (; d < D; ++d) { + float v = static_cast(input[d]) * rscale; + if constexpr (has_scale) { + v *= scale; + } + out[d] = static_cast(v); + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct l2norm_kernel { + static inline void apply(at::BFloat16* __restrict__ out, const at::BFloat16* __restrict__ input, float eps) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 32; + + __m512bh va[COLS]; + __m512 vrscale; + + const float scale = 1.f / std::sqrt(D); + __m512 vscale = _mm512_set1_ps(scale); + + // step 1: load input and do reduce with avx512-bf16 + __m512 vsum = _mm512_set1_ps(0.f); + auto reduce = [&](auto col) { + va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32)); + vsum = _mm512_dpbf16_ps(vsum, va[col], va[col]); + }; + Unroll{}(reduce); + + float sqsum = _mm512_reduce_add_ps(vsum); + float rscale = 1.f / std::sqrt(sqsum + eps); + vrscale = _mm512_set1_ps(rscale); + + // step 2: apply scale to output + auto map = [&](auto col) { + __m512i a16 = (__m512i)va[col]; + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + va0 = _mm512_mul_ps(va0, vrscale); + va1 = _mm512_mul_ps(va1, vrscale); + // keep the mul order same as torch code: + // query = l2norm(query) * scale + if constexpr (has_scale) { + va0 = _mm512_mul_ps(va0, vscale); + va1 = _mm512_mul_ps(va1, vscale); + } + _mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); + }; + Unroll{}(map); + } +}; +#endif + +template +struct cumsum_kernel { + static inline void apply( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + int mb_size, + int hb_size, + int ld_src, + int ld_dst) { + // out: [hb_size valid rows, CHUNK_SIZE] within a [BLOCK_H, ld_dst] buffer + // input: [mb_size valid rows, hb_size] within a [CHUNK_SIZE(padded), ld_src] buffer + // input is contiguous along j (vectorize the load/accumulate); out is + // contiguous along i instead (transposed), so the store is per-lane. + using Vec = at::vec::Vectorized; + constexpr int VecSize = Vec::size(); + alignas(64) scalar_t lane_buf[VecSize]; + int j = 0; + for (; j <= hb_size - VecSize; j += VecSize) { + Vec running(static_cast(0)); + for (int i = 0; i < CHUNK_SIZE; ++i) { + if (i < mb_size) { + running = running + Vec::loadu(input + i * ld_src + j); + } + running.store(lane_buf); + for (int lane = 0; lane < VecSize; ++lane) { + out[(j + lane) * ld_dst + i] = lane_buf[lane]; } - sum_q += vec_reduce_sum(sum_q_fvec); - sum_k += vec_reduce_sum(sum_k_fvec); - q_norm_sum_ptr[0] = static_cast(float(1) / std::sqrt(sum_q + eps)); - k_norm_sum_ptr[0] = static_cast(float(1) / std::sqrt(sum_k + eps)); - data_index_step(h_qk, qk_num_head, l, global_seq_len); } - }); + } + for (; j < hb_size; ++j) { + float running = 0.f; + for (int i = 0; i < CHUNK_SIZE; ++i) { + if (i < mb_size) { + running += static_cast(input[i * ld_src + j]); + } + out[j * ld_dst + i] = static_cast(running); + } + } } - - // query = query * scale - // k_beta = key * beta.unsqueeze(-1) - // v_beta = value * beta.unsqueeze(-1) - // Padding for q/k/v/beta - at::parallel_for(0, qk_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) { - int ompIdx = at::get_thread_num(); - int64_t h_qk = 0, c = 0; - data_index_init(begin, h_qk, qk_num_head, c, global_num_chunk); - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch - int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk - int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size; - int64_t l = c * chunk_size; - bool is_tail = (c + 1 == chunk_offsets_ptr[ib + 1]); - int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib]; - int64_t real_chunk_size = is_tail ? seq_len - ic * chunk_size : chunk_size; - auto q_orig_ptr = q_orig + h_qk * qStrideH + l_orig * qStrideT; - auto k_orig_ptr = k_orig + h_qk * kStrideH + l_orig * kStrideT; - auto v_orig_ptr = v_orig + l_orig * vStrideT; - auto b_orig_ptr = b_orig + l_orig * bStrideT; - auto q_pad_ptr = q_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size; - auto k_pad_ptr = k_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size; - auto v_pad_ptr = v_pad + l * v_head_size; - auto k_beta_ptr = k_beta + l * qk_head_size; - auto v_beta_ptr = v_beta + l * v_head_size; - - for (int64_t j = 0; j < real_chunk_size; j++) { - auto curr_q_orig = q_orig_ptr + j * qStrideT; - auto curr_k_orig = k_orig_ptr + j * kStrideT; - auto curr_q_pad = q_pad_ptr + j * qk_head_size; - auto curr_k_pad = k_pad_ptr + j * qk_head_size; - auto q_scale = - use_qk_l2norm_in_kernel ? *(q_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast(1); - auto k_scale = - use_qk_l2norm_in_kernel ? *(k_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast(1); - auto q_scale_vec = bVec(q_scale); - auto k_scale_vec = bVec(k_scale); - int64_t i = 0; - scalar_t scale_reduced = static_cast(scale); - auto vec_scale_reduced = bVec(scale_reduced); - for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_q_orig + i, fVecSize); - auto tmp1 = tmp0 * q_scale_vec * vec_scale_reduced; - tmp1.store(curr_q_pad + i, fVecSize); - auto tmp3 = bVec::loadu(curr_k_orig + i, fVecSize); - auto tmp4 = tmp3 * k_scale_vec; - tmp4.store(curr_k_pad + i, fVecSize); +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct cumsum_kernel { + static inline void + apply(float* __restrict__ out, const float* __restrict__ input, int mb_size, int hb_size, int ld_src, int ld_dst) { + // vector length of fp32 for avx512 + static_assert(BLOCK_H == 16); + TORCH_CHECK(hb_size > 0 && hb_size <= BLOCK_H); + const __mmask16 vmask = static_cast<__mmask16>((1u << hb_size) - 1u); + + __m512i va[16]; + __m512 vsum = _mm512_set1_ps(0.f); + + for (int i = 0; i < CHUNK_SIZE; i += 16) { + // load input data + Unroll<16>{}([&](auto j) { + __m512 v; + if (i + j < mb_size) { + v = _mm512_maskz_loadu_ps(vmask, input + (i + j) * ld_src); + } else { + v = _mm512_setzero_ps(); } - - for (auto hi = 0; hi < head_group; hi++) { - int64_t h = h_qk * head_group + hi; - auto curr_v_orig = v_orig_ptr + h * vStrideH + j * vStrideT; - auto curr_b_orig = b_orig_ptr + h * bStrideH + j * bStrideT; - scalar_t b_orig_val_reduced = *(curr_b_orig); - auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size; - auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - - // query = query * scale - // k_beta = key * beta.unsqueeze(-1) - int64_t i = 0; - auto vec_b_reduced = bVec(b_orig_val_reduced); - for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_k_orig + i, fVecSize); - auto tmp2 = tmp0 * k_scale_vec * vec_b_reduced; - tmp2.store(curr_k_beta + i, fVecSize); - } - // v_beta = value * beta.unsqueeze(-1) - i = 0; - for (; i < VecSize * (v_head_size / VecSize); i += VecSize) { - auto tmp3 = bVec::loadu(curr_v_orig + i); - tmp3.store(curr_v_pad + i); - auto tmp5 = tmp3 * vec_b_reduced; - tmp5.store(curr_v_beta + i); - } + vsum = _mm512_add_ps(vsum, v); + va[j] = _mm512_castps_si512(vsum); + }); + // transpose + transpose_16x16_32bit(va); + // store output data + Unroll<16>{}([&](auto j) { + if (j < hb_size) { + _mm512_storeu_si512(out + j * ld_dst + i, va[j]); } + }); + } + } +}; +#endif + +template +struct decay_mask_kernel { + // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input) { + using Vec = at::vec::Vectorized; + constexpr int VecSize = Vec::size(); + const Vec zero(static_cast(0)); + for (int row = 0; row < CHUNK_SIZE; ++row) { + Vec g_row(input[row]); + Vec limit_vec(static_cast(row)); + int col = 0; + for (; col <= CHUNK_SIZE - VecSize; col += VecSize) { + Vec g_col = Vec::loadu(input + col); + Vec vc = (g_row - g_col).exp_u20(); + Vec idx = Vec::arange(static_cast(col), static_cast(1)); + Vec result = Vec::blendv(zero, vc, idx <= limit_vec); + result.store(out + row * CHUNK_SIZE + col); + } + for (; col < CHUNK_SIZE; ++col) { + out[row * CHUNK_SIZE + col] = + col <= row + ? static_cast(std::exp(static_cast(input[row]) - static_cast(input[col]))) + : static_cast(0); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct decay_mask_kernel { + static inline void apply(float* __restrict__ out, const float* __restrict__ input) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int ROWS = CHUNK_SIZE; + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 va; + __m512 vb[COLS]; + + // step 1: load g[j] + auto loadb = [&](auto i) { vb[i] = _mm512_loadu_ps(input + i * 16); }; + Unroll{}(loadb); + + // step2: exp(g[i] - g[j]) + auto compute = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(input[row]); } - for (int64_t j = real_chunk_size; j < chunk_size; j++) { - auto curr_q_pad = q_pad_ptr + j * qk_head_size; - auto curr_k_pad = k_pad_ptr + j * qk_head_size; - int64_t i = 0; - auto vec_zero = bVec(0.0); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_q_pad + i); - vec_zero.store(curr_k_pad + i); + // mask vb[col] (already loaded in step 1) for the lower-triangular region + constexpr int len = std::max(0, std::min(row + 1 - col * 16, 16)); + + __m512 vc; + if constexpr (len == 16) { + vc = _mm512_fexp_u20_ps(va - vb[col]); + } else if constexpr (len == 0) { + vc = _mm512_setzero_ps(); + } else { + vc = _mm512_fexp_u20_ps(va - vb[col]); + // do mask for vc + constexpr __mmask16 vmask = (1 << len) - 1; + vc = _mm512_mask_blend_ps(vmask, _mm512_setzero_ps(), vc); + } + _mm512_storeu_ps(out + row * CHUNK_SIZE + col * 16, vc); + }; + Unroll{}(compute); + } +}; +#endif + +template +struct apply_mask_kernel { + // has_beta: attn2 = -attn * beta * d (strict lower, col < row) + // !has_beta: attn2 = attn * d (lower incl. diagonal, col <= row) + static inline void apply( + scalar_t* __restrict__ attn2, + const float* __restrict__ attn, + const scalar_t* __restrict__ beta, + const float* __restrict__ d, + int size, + int b_stride = 0) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + const fVec zero(0.f); + for (int row = 0; row < size; ++row) { + int col_limit = has_beta ? row : row + 1; + float beta_val = 1.f; + if constexpr (has_beta) { + beta_val = -static_cast(beta[row * b_stride]); + } + fVec beta_fvec(beta_val); + fVec limit_fvec(static_cast(col_limit)); + int col = 0; + for (; col <= CHUNK_SIZE - bVecSize; col += bVecSize) { + fVec a0 = fVec::loadu(attn + row * CHUNK_SIZE + col); + fVec a1 = fVec::loadu(attn + row * CHUNK_SIZE + col + fVecSize); + fVec d0 = fVec::loadu(d + row * CHUNK_SIZE + col); + fVec d1 = fVec::loadu(d + row * CHUNK_SIZE + col + fVecSize); + fVec v0 = a0 * beta_fvec * d0; + fVec v1 = a1 * beta_fvec * d1; + fVec idx0 = fVec::arange(static_cast(col), 1.f); + fVec idx1 = fVec::arange(static_cast(col + fVecSize), 1.f); + v0 = fVec::blendv(zero, v0, idx0 < limit_fvec); + v1 = fVec::blendv(zero, v1, idx1 < limit_fvec); + bVec out_bvec = at::vec::convert_from_float(v0, v1); + out_bvec.store(attn2 + row * CHUNK_SIZE + col); + } + for (; col < CHUNK_SIZE; ++col) { + float v = 0.f; + if (col < col_limit) { + v = attn[row * CHUNK_SIZE + col] * beta_val * d[row * CHUNK_SIZE + col]; } - for (auto hi = 0; hi < head_group; hi++) { - int64_t h = h_qk * head_group + hi; - auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size; - auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - int64_t i = 0; - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_k_beta + i); + attn2[row * CHUNK_SIZE + col] = static_cast(v); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct apply_mask_kernel { + static inline void apply( + at::BFloat16* __restrict__ attn2, + const float* __restrict__ attn, + const at::BFloat16* __restrict__ beta, + const float* __restrict__ d, + int size, + int b_stride = 0) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int ROWS = CHUNK_SIZE; + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 vbeta; + + // has_beta: attn2 = -attn * beta * d (strict lower) + // !has_beta: attn2 = attn * d (lower incl. diagonal) + auto compute = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + constexpr int len = + has_beta ? std::max(0, std::min(row - col * 16, 16)) : std::max(0, std::min(row + 1 - col * 16, 16)); + if (row < size) { + if constexpr (has_beta) { + if constexpr (col == 0) { + vbeta = _mm512_set1_ps(-static_cast(beta[row * b_stride])); } - i = 0; - for (; i < VecSize * (v_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_v_pad + i); - vec_zero.store(curr_v_beta + i); + } + + __m512 vc; + if constexpr (len == 0) { + vc = _mm512_setzero_ps(); + } else { + constexpr __mmask16 vmask = (1 << len) - 1; + __m512 va = _mm512_maskz_loadu_ps(vmask, attn + row * CHUNK_SIZE + col * 16); + __m512 vd = _mm512_maskz_loadu_ps(vmask, d + row * CHUNK_SIZE + col * 16); + if constexpr (has_beta) { + vc = _mm512_mul_ps(_mm512_mul_ps(va, vbeta), vd); + } else { + vc = _mm512_mul_ps(va, vd); } } + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(attn2 + row * CHUNK_SIZE + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc))); + } + }; + Unroll{}(compute); + } +}; +#endif + +template +struct solve_tril_kernel { + // (I + L)^{-1} via forward substitution, L = strict-lower part of attn2. + static inline void apply(scalar_t* __restrict__ attn2, int size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + + // for len == 0 and row < size, we don't have to write back zero again + // as in `apply_mask_kernel`, we already set zero for the upper-triangular region + for (int i = 1; i < size; ++i) { + scalar_t* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE; + float vsum[CHUNK_SIZE]; + int j = 0; + for (; j <= i - bVecSize; j += bVecSize) { + bVec row_bvec = bVec::loadu(row_ptr + j); + fVec f0, f1; + std::tie(f0, f1) = at::vec::convert_to_float(row_bvec); + f0.store(vsum + j); + f1.store(vsum + j + fVecSize); + } + for (; j < i; ++j) { + vsum[j] = static_cast(row_ptr[j]); } - // Move to the next query - data_index_step(h_qk, qk_num_head, c, global_num_chunk); - } - }); - at::parallel_for(0, v_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) { - int64_t h = 0, c = 0; - data_index_init(begin, h, v_num_head, c, global_num_chunk); - int ompIdx = at::get_thread_num(); - int64_t offset = 0; - scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread; - THREAD_BUFFER_ALLOC(k_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(v_pack, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(k_beta_g, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(k_beta_g_pack, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(curr_attn, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2); - THREAD_BUFFER_ALLOC(curr_attn_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size); - THREAD_BUFFER_ALLOC(k_cumdecay, thread_buff_ptr, offset, float, chunk_size* qk_head_size * 2); - THREAD_BUFFER_ALLOC(row, thread_buff_ptr, offset, float, chunk_size * 2); - THREAD_BUFFER_ALLOC(updated, thread_buff_ptr, offset, float, chunk_size * 2); - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch - int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk - int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size; - int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib]; - int64_t h_qk = h / head_group; - auto curr_g_orig = g_orig + h * gStrideH + l_orig * gStrideT; - auto curr_g_pad = g_pad + h * global_total_seq_length + c * chunk_size; - auto curr_decay_mask = decay_mask + h * global_total_seq_length * chunk_size + c * chunk_size * chunk_size; - auto curr_k_pad = k_pad + h_qk * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_k_beta = k_beta + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_k_cumdecay_reduced = - k_cumdecay_reduced + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_v_beta = v_beta + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size; - auto curr_value = v_beta_attn + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size; - - float acc_val = 0; - for (int64_t i = 0; i < chunk_size; i++) { - // Padding for g - // g = g.cumsum(dim=-1) - // g: [B, HV, num_chunk, chunk_size] - if (ic * chunk_size + i < seq_len) { - acc_val += curr_g_orig[i * gStrideT]; + // row = attn[..., i, :i].clone() + // sub = attn[..., :i, :i].clone() + // vsum = row + (row.unsqueeze(-1) * sub).sum(-2) + for (int k = 0; k < i; ++k) { + // read BEFORE row_ptr is written back below (row k was finalized in + // an earlier outer iteration; row i itself is untouched until the + // final write-back after this loop) + float va = static_cast(row_ptr[k]); + fVec va_vec(va); + const scalar_t* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE; + int jj = 0; + for (; jj <= k - bVecSize; jj += bVecSize) { + bVec rk_bvec = bVec::loadu(row_k_ptr + jj); + fVec rk0, rk1; + std::tie(rk0, rk1) = at::vec::convert_to_float(rk_bvec); + fVec vsum0 = fVec::loadu(vsum + jj); + fVec vsum1 = fVec::loadu(vsum + jj + fVecSize); + vsum0 = vsum0 + va_vec * rk0; + vsum1 = vsum1 + va_vec * rk1; + vsum0.store(vsum + jj); + vsum1.store(vsum + jj + fVecSize); } - curr_g_pad[i] = acc_val; - // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - // decay_mask: [B, HV, num_chunk, chunk_size, chunk_size] - float curr_g_pad_i = static_cast(curr_g_pad[i]); - auto vec_curr_g_pad_i = fVec(curr_g_pad_i); - int64_t j = 0; - int64_t len = i + 1; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_g_pad + j); - auto tmp1 = vec_curr_g_pad_i - tmp0; - auto tmp2 = tmp1.exp_u20(); - tmp2.store(curr_decay_mask + i * chunk_size + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(curr_g_pad + j, len - j); - auto tmp1 = vec_curr_g_pad_i - tmp0; - auto tmp2 = tmp1.exp_u20(); - tmp2.store(curr_decay_mask + i * chunk_size + j, len - j); + for (; jj < k; ++jj) { + vsum[jj] += va * static_cast(row_k_ptr[jj]); } } - // attn = k_beta @ key.transpose(-1, -2) - // attn: [B, HV, num_chunk, chunk_size, chunk_size] - // transpose and pack for key - if constexpr (brgemm_supported()) { - pack_vnni( - /* dst */ k_transpose, - /* src */ curr_k_pad, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // k_beta @ key.transpose(-1, -2) - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, - /* add_C */ false, - /* A */ curr_k_beta, - /* B */ k_transpose, - /* C */ curr_attn); - } else { - blas_gemm( - at::native::TransposeType::Transpose, - at::native::TransposeType::NoTranspose, - chunk_size, - chunk_size, - qk_head_size, - 1.0f, - curr_k_pad, - qk_head_size, - curr_k_beta, - qk_head_size, - 0.0f, - curr_attn, - chunk_size); + j = 0; + for (; j <= i - bVecSize; j += bVecSize) { + fVec f0 = fVec::loadu(vsum + j); + fVec f1 = fVec::loadu(vsum + j + fVecSize); + bVec out_bvec = at::vec::convert_from_float(f0, f1); + out_bvec.store(row_ptr + j); } - // attn = attn * decay_mask - for (int64_t m = 0; m < chunk_size; m++) { - at::vec::map2( - [](fVec x, fVec y) { return fVec(0) - x * y; }, - curr_attn + m * chunk_size, - curr_attn + m * chunk_size, - curr_decay_mask + m * chunk_size, - chunk_size); + for (; j < i; ++j) { + row_ptr[j] = static_cast(vsum[j]); } + } - // chunk decay - // attn: [B, HV, num_chunk, chunk_size, chunk_size] - // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) - // attn = -attn.masked_fill(mask, 0) - // attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) [B, HV, num_chunk, i] - // attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - // attn = -attn.masked_fill(mask, 0) - for (int i = 0; i < chunk_size; i++) { - const auto vec_zero = fVec(0); - int64_t len = chunk_size - i; - int64_t front = len % fVecSize; - int64_t j = i; - // first masked vec for alignment - if (front > 0) { - vec_zero.store(curr_attn + i * chunk_size + j, front); - j += front; - } - for (; j < fVecSize * (chunk_size / fVecSize); j += fVecSize) { - vec_zero.store(curr_attn + i * chunk_size + j); + // attn = attn + torch.eye(chunk_size) + for (int i = 0; i < size; ++i) { + attn2[i * CHUNK_SIZE + i] = static_cast(static_cast(attn2[i * CHUNK_SIZE + i]) + 1.f); + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct solve_tril_kernel { + static inline void apply(at::BFloat16* __restrict__ attn2, int size) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vsum[COLS]; + + // for len == 0 and row < size, we don't have to write back zero again + // as in `apply_mask_kernel`, we already set zero for the upper-triangular region + for (int i = 1; i < size; ++i) { + // load row attn[..., i, :i] + at::BFloat16* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE; + Unroll{}([&](auto col) { + int len = std::min(i - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + vsum[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_ptr + col * 16)); } + }); + + // row = attn[..., i, :i].clone() + // sub = attn[..., :i, :i].clone() + // vsum = row + (row.unsqueeze(-1) * sub).sum(-2) + for (int k = 0; k < i; ++k) { + va = _mm512_set1_ps(static_cast(row_ptr[k])); + + const at::BFloat16* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE; + Unroll{}([&](auto col) { + int len = std::min(k - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + vb[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_k_ptr + col * 16)); + vsum[col] = _mm512_fmadd_ps(va, vb[col], vsum[col]); + } + }); } - for (int i = 1; i < chunk_size; i++) { - // row = attn[..., i, :i] [B, HK, num_chunk, i] - int64_t j = 0; - int64_t len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j); - tmp0.store(row + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j, len - j); - tmp0.store(row + j, len - j); + + // attn[..., i, :i] = vsum + Unroll{}([&](auto col) { + int len = std::min(i - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + _mm256_mask_storeu_epi16(row_ptr + col * 16, vmask, (__m256i)(_mm512_cvtneps_pbh(vsum[col]))); } - // (row.unsqueeze(-1) * sub).sum(-2) - fill_stub(updated, 0, i); - for (int k = 0; k < i; k++) { - float row_k = row[k]; - auto vec_row_k = fVec(row_k); - int64_t j = 0; - int64_t len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j); - auto tmp1 = vec_row_k * tmp0; - auto tmp2 = fVec::loadu(updated + j); - auto tmp3 = tmp1 + tmp2; - tmp3.store(updated + j); + }); + } + + // attn = attn + torch.eye(chunk_size) + for (int i = 0; i < size; ++i) { + attn2[i * CHUNK_SIZE + i] += 1.f; + } + } +}; +#endif + +template +struct apply_beta_kernel { + static inline void apply( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const scalar_t* __restrict__ beta, + const float* __restrict__ g, + int size, + int ld_src, + int ld_dst, + int b_stride) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + + for (int i = 0; i < size; ++i) { + float scale = 1.f; + if constexpr (has_beta) { + scale *= static_cast(beta[i * b_stride]); + } + if constexpr (has_g) { + scale *= std::exp(g[i]); + } + fVec scale_fvec(scale); + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + bVec in_bvec = bVec::loadu(input + i * ld_src + d); + fVec in0, in1; + std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); + in0 = in0 * scale_fvec; + in1 = in1 * scale_fvec; + bVec out_bvec = at::vec::convert_from_float(in0, in1); + out_bvec.store(out + i * ld_dst + d); + } + for (; d < D; ++d) { + out[i * ld_dst + d] = static_cast(static_cast(input[i * ld_src + d]) * scale); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct apply_beta_kernel { + static inline void apply( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ input, + const at::BFloat16* __restrict__ beta, + const float* __restrict__ g, + int size, + int ld_src, + int ld_dst, + int b_stride) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + // get g.exp() and g is padded to CHUNK_SIZE + alignas(64) float g_arr[CHUNK_SIZE]; + if constexpr (has_g) { + Unroll{}([&](auto col) { + __m512 vg = _mm512_loadu_ps(g + col * 16); + __m512 vg_exp = _mm512_fexp_u20_ps(vg); + _mm512_storeu_ps(g_arr + col * 16, vg_exp); + }); + } + + for (int i = 0; i < size; ++i) { + __m512 vbeta; + if constexpr (has_beta) { + vbeta = _mm512_set1_ps(static_cast(beta[i * b_stride])); + } + __m512 vg; + if constexpr (has_g) { + vg = _mm512_set1_ps(g_arr[i]); + } + + Unroll{}([&](auto col) { + // load for 0, 2, 4, 6 + if constexpr (col % 2 == 0) { + __m512i a16 = _mm512_loadu_si512(input + i * ld_src + col * 16); + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + if constexpr (has_beta) { + va0 = _mm512_mul_ps(va0, vbeta); + va1 = _mm512_mul_ps(va1, vbeta); } - if (j < len) { - auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j, len - j); - auto tmp1 = vec_row_k * tmp0; - auto tmp2 = fVec::loadu(updated + j); - auto tmp3 = tmp1 + tmp2; - tmp3.store(updated + j, len - j); + if constexpr (has_g) { + va0 = _mm512_mul_ps(va0, vg); + va1 = _mm512_mul_ps(va1, vg); } + _mm512_storeu_si512(out + i * ld_dst + col * 16, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); } - // attn[..., i, :i] = row + sum(...) - j = 0; - len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(row + j); - auto tmp1 = fVec::loadu(updated + j); - auto tmp2 = tmp0 + tmp1; - tmp2.store(curr_attn + i * chunk_size + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(row + j, len - j); - auto tmp1 = fVec::loadu(updated + j, len - j); - auto tmp2 = tmp0 + tmp1; - tmp2.store(curr_attn + i * chunk_size + j, len - j); + }); + } + } +}; +#endif + +template +struct update_kernel { + static inline void + apply(scalar_t* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + + for (int i = 0; i < size; ++i) { + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + fVec f0 = fVec::loadu(input + i * ld_src + d); + fVec f1 = fVec::loadu(input + i * ld_src + d + fVecSize); + bVec out_bvec = at::vec::convert_from_float(f0, f1); + out_bvec.store(out + i * ld_dst + d); + } + for (; d < D; ++d) { + out[i * ld_dst + d] = static_cast(input[i * ld_src + d]); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_kernel { + static inline void + apply(at::BFloat16* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + for (int i = 0; i < size; ++i) { + Unroll{}([&](auto col) { + if constexpr (col % 2 == 0) { + __m512 va0 = _mm512_loadu_ps(input + i * ld_src + (col + 0) * 16); + __m512 va1 = _mm512_loadu_ps(input + i * ld_src + (col + 1) * 16); + __m512i a16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(out + i * ld_dst + col * 16, a16); } + }); + } + } +}; +#endif + +template +struct update_value_kernel { + static inline void apply( + scalar_t* __restrict__ v_prime2, + const scalar_t* __restrict__ v, + const float* __restrict__ v_prime, + int size, + int padded_size, + int v_strideT) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + + // v2' = v - v' + for (int i = 0; i < size; ++i) { + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + bVec v_bvec = bVec::loadu(v + i * v_strideT + d); + fVec v0, v1; + std::tie(v0, v1) = at::vec::convert_to_float(v_bvec); + fVec vp0 = fVec::loadu(v_prime + i * D + d); + fVec vp1 = fVec::loadu(v_prime + i * D + d + fVecSize); + v0 = v0 - vp0; + v1 = v1 - vp1; + bVec out_bvec = at::vec::convert_from_float(v0, v1); + out_bvec.store(v_prime2 + i * D + d); } - for (int i = 0; i < chunk_size; i++) { - curr_attn[i * chunk_size + i] += 1.0f; - at::vec::map( - [](fVec x) { return x; }, curr_attn_reduced + i * chunk_size, curr_attn + i * chunk_size, chunk_size); + for (; d < D; ++d) { + float val = static_cast(v[i * v_strideT + d]) - v_prime[i * D + d]; + v_prime2[i * D + d] = static_cast(val); } + } + // pad the last chunk + const bVec zero_bvec(static_cast(0)); + for (int i = size; i < padded_size; ++i) { + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + zero_bvec.store(v_prime2 + i * D + d); + } + for (; d < D; ++d) { + v_prime2[i * D + d] = static_cast(0); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_value_kernel { + static inline void apply( + at::BFloat16* __restrict__ v_prime2, + const at::BFloat16* __restrict__ v, + const float* __restrict__ v_prime, + int size, + int padded_size, + int v_strideT) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + // v2' = v - v' + for (int i = 0; i < size; ++i) { + Unroll{}([&](auto col) { + // load for 0, 2, 4, 6 + if constexpr (col % 2 == 0) { + __m512i v16 = _mm512_loadu_si512(v + i * v_strideT + col * 16); + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 1)); + + __m512 v_prime0 = _mm512_loadu_ps(v_prime + i * D + col * 16); + __m512 v_prime1 = _mm512_loadu_ps(v_prime + i * D + col * 16 + 16); + va0 = _mm512_sub_ps(va0, v_prime0); + va1 = _mm512_sub_ps(va1, v_prime1); + __m512i o16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(v_prime2 + i * D + col * 16, o16); + } + }); + } - // v_beta_attn = attn @ v_beta - // k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) - // v_beta_attn: [B, HV, num_chunk, chunk_size, EV] - // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK] - // k_cumdecay: [B, HV, num_chunk, chunk_size, EK] - // pack for value - if constexpr (brgemm_supported()) { - pack_vnni2( - /* dst */ v_pack, - /* src */ curr_v_beta, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // value = attn @ v_beta - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ v_pack, - /* C */ curr_value); - } else { - blas_gemm( - at::native::TransposeType::NoTranspose, - at::native::TransposeType::NoTranspose, - v_head_size, - chunk_size, - chunk_size, - 1.0f, - curr_v_beta, - v_head_size, - curr_attn_reduced, - chunk_size, - 0.0f, - curr_value, - v_head_size); + // pad the last chunk + for (int i = size; i < padded_size; ++i) { + Unroll{}([&](auto col) { + if constexpr (col % 2 == 0) { + __m512i v16 = _mm512_setzero_si512(); + _mm512_storeu_si512(v_prime2 + i * D + col * 16, v16); + } + }); + } + } +}; +#endif + +template +struct update_key_kernel { + static inline void apply( + scalar_t* __restrict__ k_updated, + const scalar_t* __restrict__ k, + const float* __restrict__ g, + int size, + int k_strideT) { + // k_updated is transposed: [D, CHUNK_SIZE], k_updated[d, t] = k[t, d] * exp(g_last - g[t]). + // k's D dim is contiguous (vectorize load/compute); k_updated's D dim is + // strided (CHUNK_SIZE apart), so the store is unpacked per lane. + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int fVecSize = fVec::size(); + constexpr int bVecSize = bVec::size(); + alignas(64) scalar_t lane_buf[bVecSize]; + + const float g_last = g[size - 1]; + for (int t = 0; t < size; ++t) { + float scale = std::exp(g_last - g[t]); + fVec scale_fvec(scale); + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + bVec k_bvec = bVec::loadu(k + t * k_strideT + d); + fVec k0, k1; + std::tie(k0, k1) = at::vec::convert_to_float(k_bvec); + k0 = k0 * scale_fvec; + k1 = k1 * scale_fvec; + bVec out_bvec = at::vec::convert_from_float(k0, k1); + out_bvec.store(lane_buf); + for (int lane = 0; lane < bVecSize; ++lane) { + k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane]; + } + } + for (; d < D; ++d) { + k_updated[d * CHUNK_SIZE + t] = static_cast(static_cast(k[t * k_strideT + d]) * scale); } - // k_beta_g = k_beta * g.exp().unsqueeze(-1) - for (int64_t j = 0; j < chunk_size; j++) { - int64_t i = 0; - float g_exp = std::exp(curr_g_pad[j]); - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(curr_k_beta + j * qk_head_size + i); - auto tmp1 = tmp0 * vec_g_exp_reduced; - tmp1.store(k_beta_g + j * qk_head_size + i); + } + const bVec zero_bvec(static_cast(0)); + for (int t = size; t < CHUNK_SIZE; ++t) { + int d = 0; + for (; d <= D - bVecSize; d += bVecSize) { + zero_bvec.store(lane_buf); + for (int lane = 0; lane < bVecSize; ++lane) { + k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane]; } } - // pack for k_beta_g + for (; d < D; ++d) { + k_updated[d * CHUNK_SIZE + t] = static_cast(0); + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_key_kernel { + static inline void apply( + at::BFloat16* __restrict__ k_updated, + const at::BFloat16* __restrict__ k, + const float* __restrict__ g, + int size, + int k_strideT) { + static_assert(D % 32 == 0); + const int MB = div_up(size, 16); + const int KB = D / 16; + + const float g_last = g[size - 1]; + const __m512 vg_last = _mm512_set1_ps(g_last); + + float scale_arr[16]; + __m256i va[16]; + + // from [C, D](MB, KB) to [D, C](KB, MB) + // pad size to 16 in this kernel so that transpose can be done in one loop + for (int mb = 0; mb < MB; ++mb) { + const int mb_size = std::min(size - mb * 16, 16); + // prepare exp(g_last - g) + __m512 vg = _mm512_loadu_ps(g + mb * 16); + _mm512_storeu_ps(scale_arr, _mm512_fexp_u20_ps(_mm512_sub_ps(vg_last, vg))); + for (int kb = 0; kb < KB; ++kb) { + const at::BFloat16* __restrict__ k_ptr = k + mb * 16 * k_strideT + kb * 16; + at::BFloat16* __restrict__ k_updated_ptr = k_updated + kb * 16 * CHUNK_SIZE + mb * 16; + // load 16 regs + Unroll<16>{}([&](auto m) { + if (m < mb_size) { + __m256i v16 = _mm256_loadu_si256(reinterpret_cast(k_ptr + m * k_strideT)); + __m512 v32 = _mm512_mul_ps(CVT_BF16_TO_FP32(v16), _mm512_set1_ps(scale_arr[m])); + va[m] = (__m256i)_mm512_cvtneps_pbh(v32); + } else { + va[m] = _mm256_setzero_si256(); + } + }); + // transpose 16x16 + transpose_16x16_16bit(va); + // store 16 regs + Unroll<16>{}( + [&](auto k) { _mm256_storeu_si256(reinterpret_cast<__m256i*>(k_updated_ptr + k * CHUNK_SIZE), va[k]); }); + } + } + } +}; +#endif + + +// template head_dim here to reduce extra read +// * normal approach: read inputs 2 times: +// - reduce: 1R +// - scale: 1R + 1W +// * keep input data in register: +// - reduce: 1R +// - scale: 1W +template +void l2norm_fwd_kernel_impl( + scalar_t* __restrict__ query_norm, + scalar_t* __restrict__ key_norm, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ key, + float eps, + int64_t T, + int64_t H, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH) { + // expected to be contuguous + int64_t qn_strideH = D; + int64_t kn_strideH = D; + + // parallel on [B, T, H] + at::parallel_for(0, T * H, 0, [&](int64_t begin, int64_t end) { + int64_t t{0}, h{0}; + data_index_init(begin, t, T, h, H); + + for (int64_t i = begin; i < end; ++i) { + const scalar_t* __restrict__ q_ptr = query + t * q_strideT + h * q_strideH; + const scalar_t* __restrict__ k_ptr = key + t * k_strideT + h * k_strideH; + scalar_t* __restrict__ qn_ptr = query_norm + i * qn_strideH; + scalar_t* __restrict__ kn_ptr = key_norm + i * kn_strideH; + + l2norm_kernel::apply(qn_ptr, q_ptr, eps); + l2norm_kernel::apply(kn_ptr, k_ptr, eps); + + // move to the next index + data_index_step(t, T, h, H); + } + }); +} + +// g : [B, T, Hv] +// g_ : [B, Hv, NT, C] -> [B, NT, HB, BLOCK_H, C] +// cu_seqlens : [num_seqs + 1] +// chunk_indices : [NT * 2] +template +void chunk_local_cumsum_kernel_impl( + scalar_t* __restrict__ g_, + const scalar_t* __restrict__ g, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_indices, + int64_t Hv, + int64_t NT) { + constexpr int BLOCK_H = 16; + int64_t HB = div_up(Hv, int64_t(BLOCK_H)); + + // parallel on [NT * HB] to increase parallelism + at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) { + int64_t nt{0}, hb{0}; + data_index_init(begin, nt, NT, hb, HB); + + for (int64_t i = begin; i < end; ++i) { + int32_t bs = chunk_indices[nt * 2 + 0]; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + int64_t hb_size = std::min(Hv - hb * BLOCK_H, int64_t(BLOCK_H)); + + const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H; + scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE); + cumsum_kernel::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE); + + // move to the next index + data_index_step(nt, NT, hb, HB); + } + }); +} + +#define DECL_BUF(type, name, size_expr) alignas(64) type name[(size_expr)] +#define DECL_ZERO_BUF(type, name, size_expr) \ + DECL_BUF(type, name, size_expr); \ + fill_stub(name, 0.f) + +// w : [B, T, Hv, D] +// u : [B, T, Hv, Dv] +// d : [B, NT, Hv, C, C] +// k : [B, T, H, D] +// v : [B, T, Hv, Dv] +// g : [B, NT, Hv, C] +// beta : [B, T, Hv] +// cu_seqlens : [num_seqs + 1] +// chunk_indices : [NT * 2] +template +void chunk_gated_delta_rule_fwd_intra_kernel_impl( + scalar_t* __restrict__ w, + scalar_t* __restrict__ u, + float* __restrict__ d, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const float* __restrict__ g, + const scalar_t* __restrict__ beta, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_indices, + int64_t H, + int64_t Hv, + int64_t NT, + int64_t k_strideT, + int64_t k_strideH, + int64_t v_strideT, + int64_t v_strideH) { + // head group, expect to be 1，2，4 for qwen3.5 + const int64_t HG = Hv / H; + + // strides + const int64_t w_strideT = Hv * D; + const int64_t w_strideH = D; + const int64_t u_strideT = Hv * D; + const int64_t u_strideH = D; + + // [NB]: parallel on [NT, H] + // * parallel on num_heads and go sequential on num_heads_v, + // * avoid instantialize k_beta (beta * k) + // * compute key @ key^T * beta instead of k_beta @ key^T, same as triton impl + // * compute key @ key^T once for each k head index and reuse for v head index + at::parallel_for(0, NT * H, 0, [&](int64_t begin, int64_t end) { + int64_t nt{0}, h{0}; + data_index_init(begin, nt, NT, h, H); + + // thread local temp buffer + DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); + DECL_ZERO_BUF(scalar_t, tmp2, CHUNK_SIZE * D); + DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); + DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); + DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); + + // alias + scalar_t* __restrict__ k_packed = tmp; + scalar_t* __restrict__ k_beta = tmp; + scalar_t* __restrict__ v_beta = tmp; + scalar_t* __restrict__ k_beta_packed = tmp2; + scalar_t* __restrict__ v_beta_packed = tmp2; + float* __restrict__ k_updated = tmp3; + float* __restrict__ v_updated = tmp3; + + for (int64_t i = begin; i < end; ++i) { + int32_t bs = chunk_indices[nt * 2 + 0]; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + + // mb_size` is K in 5.c, 5.g, pad to TILE_K; + const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; + + // step 1: decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; + float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + decay_mask_kernel::apply(d_ptr, g_ptr); + } + + // step 2: attn = key @ key^T + const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; if constexpr (brgemm_supported()) { - pack_vnni2( - /* dst */ k_beta_g_pack, - /* src */ k_beta_g, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ qk_head_size); - // k_cumdecay = attn @ k_beta_g + pack_vnni( + /* dst */ k_packed, + /* src */ k_ptr, + /* N */ mb_size, + /* K */ D, + /* ld_src */ k_strideT, + /* ld_dst */ CHUNK_SIZE); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ qk_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ qk_head_size, - /* ldc */ qk_head_size, + /* M */ mb_size, + /* N */ mb_size, + /* K */ D, + /* lda */ k_strideT, + /* ldb */ CHUNK_SIZE, + /* ldc */ CHUNK_SIZE, /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ k_beta_g_pack, - /* C */ k_cumdecay); + /* A */ k_ptr, + /* B */ k_packed, + /* C */ attn); } else { blas_gemm( + at::native::TransposeType::Transpose, at::native::TransposeType::NoTranspose, - at::native::TransposeType::NoTranspose, - qk_head_size, - chunk_size, - chunk_size, + mb_size, + mb_size, + D, 1.0f, - k_beta_g, - qk_head_size, - curr_attn_reduced, - chunk_size, + k_ptr, + k_strideT, + k_ptr, + k_strideT, 0.0f, - k_cumdecay, - qk_head_size); + attn, + CHUNK_SIZE); } - for (int i = 0; i < chunk_size; i++) { - at::vec::map( - [](fVec x) { return x; }, - curr_k_cumdecay_reduced + i * qk_head_size, - k_cumdecay + i * qk_head_size, - qk_head_size); + + for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { + // step 3: attn2 = -attn * beta * d + const scalar_t* __restrict__ beta_ptr = beta + (batch_offset + mb_start) * Hv + hv; + const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + apply_mask_kernel::apply(attn2, attn, beta_ptr, d_ptr, mb_size, Hv); + + // step 4: solve_tril(attn2) -> (I + L)^{-1}, L = strict-lower from step 3 + // for i in 1..C-1: attn2[i, :i] += (attn2[i, :i] * attn2[:i, :i]).sum(-1) + // attn2 += eye(C) + solve_tril_kernel::apply(attn2, mb_size); + + // step 5: recompute_w_u + // w = attn2 @ (k_beta * g.exp().unsqueeze(-1)) + // u = attn2 @ value * beta.unsqueeze(-1) + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; + const scalar_t* __restrict__ v_ptr = v + (batch_offset + mb_start) * v_strideT + hv * v_strideH; + + // 5.a key = key * beta * g.exp + apply_beta_kernel::apply( + k_beta, k_ptr, beta_ptr, g_ptr, mb_size, k_strideT, D, Hv); + + // 5.b pack key + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ k_beta_packed, + /* src */ k_beta, + /* K */ mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); + + // 5.c w = attn2 @ k_beta + at::native::cpublas::brgemm( + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, + /* add_C */ false, + /* A */ attn2, + /* B */ k_beta_packed, + /* C */ k_updated); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + D, + mb_size, + padded_mb_size, + 1.0f, + k_beta, + D, + attn2, + CHUNK_SIZE, + 0.0f, + k_updated, + D); + } + + // 5.d k_updated -> w + scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; + update_kernel::apply(w_ptr, k_updated, mb_size, D, w_strideT); + + // 5.e value = value * beta + apply_beta_kernel::apply( + v_beta, v_ptr, beta_ptr, nullptr, mb_size, v_strideT, D, Hv); + + // 5.f pack value + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ v_beta_packed, + /* src */ v_beta, + /* K */ mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); + + // 5.g u = attn2 @ v_beta + at::native::cpublas::brgemm( + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, + /* add_C */ false, + /* A */ attn2, + /* B */ v_beta_packed, + /* C */ v_updated); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + D, + mb_size, + padded_mb_size, + 1.0f, + v_beta, + D, + attn2, + CHUNK_SIZE, + 0.0f, + v_updated, + D); + } + + // 5.h v_updated -> u + scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; + update_kernel::apply(u_ptr, v_updated, mb_size, D, u_strideT); } - // Move to the next query - data_index_step(h, v_num_head, c, global_num_chunk); + // move to the next index + data_index_step(nt, NT, h, H); } + at::native::cpublas::brgemm_release(); }); +} - // for each chunk - at::parallel_for(0, batch_size * v_num_head, 1, [&](int64_t begin, int64_t end) { - int64_t b = 0, h = 0; - data_index_init(begin, b, batch_size, h, v_num_head); - int ompIdx = at::get_thread_num(); - int64_t offset = - /* k_transpose */ qk_head_size * chunk_size + - /* v_pack */ chunk_size * v_head_size + - /* k_beta_g */ chunk_size * qk_head_size + - /* k_beta_g_pack */ chunk_size * qk_head_size + - /* attn */ chunk_size * chunk_size * 2 + - /* attn_reduced */ chunk_size * chunk_size + - /* k_cumdecay */ chunk_size * qk_head_size * 2 + - /* row */ chunk_size * 2 + - /* updated */ chunk_size * 2; - scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread; - THREAD_BUFFER_ALLOC( - curr_last_recurrent_state_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size); - THREAD_BUFFER_ALLOC( - curr_last_recurrent_state_pack_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size); - THREAD_BUFFER_ALLOC(k_transpose_i, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(attn_i, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2); - THREAD_BUFFER_ALLOC(attn_i_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size); - THREAD_BUFFER_ALLOC(v_prime, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2); - THREAD_BUFFER_ALLOC(v_prime_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(v_prime_pack_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(qg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(attn_inter, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2); - THREAD_BUFFER_ALLOC(kg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(kg_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(kgv, thread_buff_ptr, offset, float, qk_head_size* v_head_size * 2); - - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t start_q = cu_seqlens_ptr[b]; - int64_t seq_len = cu_seqlens_ptr[b + 1] - start_q; - int64_t num_chunk = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b]; - int64_t chunk_offset = chunk_offsets_ptr[b]; - int64_t len_offset = chunk_offset * chunk_size; - - int64_t h_qk = h / head_group; - auto out_ptr = out + start_q * oStrideT; - auto curr_q = q_pad + len_offset * qk_head_size + - h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_k = k_pad + len_offset * qk_head_size + - h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_v = v_beta_attn + h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV] - auto curr_decay_mask = - decay_mask + h * global_total_seq_length * chunk_size; // [num_chunk, chunk_size, chunk_size] - auto curr_k_cumdecay_reduced = - k_cumdecay_reduced + h * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_last_recurrent_state = - final_state_data + b * final_state_StrideN + h * final_state_StrideH; // [EK, EV] - auto curr_g_pad = g_pad + len_offset + h * global_total_seq_length; // [num_chunk, chunk_size] - auto curr_core_attn_out = core_attn_out + len_offset * v_head_size + - h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV] - for (int64_t c = 0; c < num_chunk; c++) { - for (int i = 0; i < qk_head_size; i++) { - at::vec::map( - [](fVec x) { return x; }, - curr_last_recurrent_state_reduced + i * v_head_size, - curr_last_recurrent_state + i * v_head_size, - v_head_size); - } - auto q_i = curr_q + c * chunk_size * qk_head_size; // [chunk_size, EK] - auto k_i = curr_k + c * chunk_size * qk_head_size; // [chunk_size, EK] - auto v_i = curr_v + (chunk_offset + c) * chunk_size * v_head_size; // [chunk_size, EV] - auto decay_mask_i = curr_decay_mask + (chunk_offset + c) * chunk_size * chunk_size; // [chunk_size, chunk_size] - auto k_cumdecay_i_reduced = - curr_k_cumdecay_reduced + (chunk_offset + c) * chunk_size * qk_head_size; // [chunk_size, EK] - auto g_pad_i = curr_g_pad + c * chunk_size; // [chunk_size] - auto core_attn_out_i = curr_core_attn_out + c * chunk_size * v_head_size; // [chunk_size, EV] +// +// out : [B, T, Hv, Dv] +// state : [num_seqs, Hv, Dv, D] +// q : [B, T, H, D] +// k : [B, T, H, D] +// w : [B, T, Hv, D] +// u : [B, T, Hv, Dv] +// g : [B, NT, Hv, C] +// d : [B, NT, Hv, C, C] +// cu_seqlens : [num_seqs + 1] +// chunk_offsets : [num_seqs + 1] +template +void chunk_gated_delta_rule_fwd_inter_kernel_impl( + scalar_t* __restrict__ out, + float* __restrict__ state, + const int32_t* __restrict__ indices, + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ w, + const scalar_t* __restrict__ u, + const float* __restrict__ g, + const float* __restrict__ d, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_offsets, + int64_t H, + int64_t Hv, + int64_t num_seqs, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH, + int64_t state_strideS) { + // head group, expect to be 1，2，4 for qwen3.5 + const int64_t HG = Hv / H; + + // strides + const int64_t w_strideT = Hv * D; + const int64_t w_strideH = D; + const int64_t u_strideT = Hv * D; + const int64_t u_strideH = D; + const int64_t o_strideT = Hv * D; + const int64_t o_strideH = D; + + // [NB]: parallel on [num_seqs, Hv] + // * choose to parallel on Hv instead of H, though this means q @ kT has duplicated compute + // * H might be 16 which is not enough to use 32C when num_seqs is small + at::parallel_for(0, num_seqs * Hv, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, hv{0}; + data_index_init(begin, bs, num_seqs, hv, Hv); + + // thread local temp buffer + DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); + DECL_ZERO_BUF(scalar_t, tmp2, D * D); + DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); + DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D); + DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); + DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); + + // alias + scalar_t* __restrict__ k_packed = tmp; + scalar_t* __restrict__ s_packed = tmp2; + float* __restrict__ v_prime = tmp3; + scalar_t* __restrict__ v_prime2 = tmp; + float* __restrict__ attn_inter = tmp3; + scalar_t* __restrict__ qg_exp = tmp4; + scalar_t* __restrict__ v_packed = tmp4; + scalar_t* __restrict__ k_updated = tmp; + + for (int64_t i = begin; i < end; ++i) { + int64_t h = hv / HG; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t nt = chunk_offsets[bs]; + + for (int64_t mb_start = 0; mb_start < seqlen; mb_start += CHUNK_SIZE, ++nt) { + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + // mb_size` is K in 4.a, pad to TILE_K; + const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; + + // step 1.a: attn = query @ key^T // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - // k_transpose_i = k_i.transpose(-1, -2) + const scalar_t* __restrict__ q_ptr = q + (batch_offset + mb_start) * q_strideT + h * q_strideH; + const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; if constexpr (brgemm_supported()) { pack_vnni( - /* dst */ k_transpose_i, - /* src */ k_i, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // attn_i = q_i @ k_transpose_i + /* dst */ k_packed, + /* src */ k_ptr, + /* N */ mb_size, + /* K */ D, + /* ld_src */ k_strideT, + /* ld_dst */ CHUNK_SIZE); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, + /* M */ mb_size, + /* N */ mb_size, + /* K */ D, + /* lda */ q_strideT, + /* ldb */ CHUNK_SIZE, + /* ldc */ CHUNK_SIZE, /* add_C */ false, - /* A */ q_i, - /* B */ k_transpose_i, - /* C */ attn_i); + /* A */ q_ptr, + /* B */ k_packed, + /* C */ attn); } else { blas_gemm( at::native::TransposeType::Transpose, at::native::TransposeType::NoTranspose, - chunk_size, - chunk_size, - qk_head_size, + mb_size, + mb_size, + D, 1.0f, - k_i, - qk_head_size, - q_i, - qk_head_size, + k_ptr, + k_strideT, + q_ptr, + q_strideT, 0.0f, - attn_i, - chunk_size); - } - // attn_i = attn_i * decay_mask_i - for (int64_t m = 0; m < chunk_size; m++) { - auto attn_i_m = attn_i + m * chunk_size; - auto attn_i_reduced_m = attn_i_reduced + m * chunk_size; - auto decay_mask_i_m = decay_mask_i + m * chunk_size; - int64_t n = 0; - for (; n < fVecSize * (chunk_size / fVecSize); n += fVecSize) { - auto tmp0 = fVec::loadu(attn_i_m + n); - auto tmp1 = fVec::loadu(decay_mask_i_m + n); - auto tmp2 = tmp0 * tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(attn_i_reduced_m + n, fVecSize); - } - if (n < chunk_size) { - auto tmp0 = fVec::loadu(attn_i_m + n, chunk_size - n); - auto tmp1 = fVec::loadu(decay_mask_i_m + n, chunk_size - n); - auto tmp2 = tmp0 * tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(attn_i_reduced_m + n, chunk_size - n); - } - } - // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) - // attn_i = attn_i.masked_fill_(mask, 0) - for (int i = 0; i < chunk_size - 1; i++) { - const auto vec_zero = bVec(0); - int64_t len = chunk_size - i - 1; - int64_t front = len % VecSize; - int64_t j = i + 1; - // first masked vec for alignment - if (front > 0) { - vec_zero.store(attn_i_reduced + i * chunk_size + j, front); - j += front; - } - for (; j < VecSize * (chunk_size / VecSize); j += VecSize) { - vec_zero.store(attn_i_reduced + i * chunk_size + j); - } + attn, + CHUNK_SIZE); } - // pack for curr_last_recurrent_state + // step 1.b: attn = attn * decay_mask.masked_fill_(mask, 0) + const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + apply_mask_kernel::apply(attn2, attn, nullptr, d_ptr, mb_size); + + // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing) + float* __restrict__ s_ptr = state + indices[bs] * state_strideS + hv * (D * D); + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE); + float g_last = g_ptr[mb_size - 1]; + const scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; if constexpr (brgemm_supported()) { - pack_vnni2( - /* dst */ curr_last_recurrent_state_pack_reduced, - /* src */ curr_last_recurrent_state_reduced, - /* N */ qk_head_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - - // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV] - // k_cumdecay_i: [chunk_size, EK] - // curr_last_recurrent_state: [EK, EV] + pack_vnni2( + /* dst */ s_packed, + /* src */ s_ptr, + /* g_last */ g_last, + /* ld_src */ D, + /* ld_dst */ D); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ D, + /* lda */ w_strideT, + /* ldb */ D, + /* ldc */ D, /* add_C */ false, - /* A */ k_cumdecay_i_reduced, - /* B */ curr_last_recurrent_state_pack_reduced, + /* A */ w_ptr, + /* B */ s_packed, /* C */ v_prime); } else { + // brgemm_supported()==false path: pack_vnni2 above packs the + // *unscaled* state into its dst (for this GEMM's B operand) while + // separately scaling src in place by exp(g_last) (consumed later, + // at step 5.3's state accumulation). Replicate both halves in the + // same order: snapshot s_ptr into s_packed BEFORE scaling s_ptr, + // not after, since the GEMM below needs the pre-scale state. + float g_last_scale = std::exp(g_last); + for (int64_t d0 = 0; d0 < D * D; ++d0) { + s_packed[d0] = static_cast(s_ptr[d0]); + s_ptr[d0] *= g_last_scale; + } blas_gemm( at::native::TransposeType::NoTranspose, at::native::TransposeType::NoTranspose, - v_head_size, - chunk_size, - qk_head_size, + D, + mb_size, + D, 1.0f, - curr_last_recurrent_state_reduced, - v_head_size, - k_cumdecay_i_reduced, - qk_head_size, + s_packed, + D, + w_ptr, + w_strideT, 0.0f, v_prime, - v_head_size); + D); } - // v_new = v_prime = v_i - v_prime - // v_i: [chunk_size, EV] - for (int64_t m = 0; m < chunk_size; m++) { - int64_t i = 0; - for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) { - auto tmp0 = fVec::loadu(v_i + m * v_head_size + i); - auto tmp1 = fVec::loadu(v_prime + m * v_head_size + i); - auto tmp2 = tmp0 - tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(v_prime_reduced + m * v_head_size + i, fVecSize); - } - } + // step 2.b: v2' = u - v' + const scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; + update_value_kernel::apply(v_prime2, u_ptr, v_prime, mb_size, padded_mb_size, u_strideT); - // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - // qg = q_i * g[:, :, i, :, None].exp(): [chunk_size, EK] - // q_i: [chunk_size, EK] - // g[:, :, i, :, None]: [chunk_size, 1] - for (int64_t m = 0; m < chunk_size; m++) { - auto g_pad_i_m = g_pad_i + m; - auto g_exp = std::exp(*g_pad_i_m); - int64_t i = 0; - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(q_i + m * qk_head_size + i); - auto tmp2 = tmp0 * vec_g_exp_reduced; - tmp2.store(qg + m * qk_head_size + i); - } - } - // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV] - // curr_last_recurrent_state: [EK, EV] + // step 3.a: qg_exp = q * exp(g) + apply_beta_kernel::apply( + qg_exp, q_ptr, nullptr, g_ptr, mb_size, q_strideT, D, /*b_stride*/ 0); + + // step 3.b: attn_inter = qg_exp @ state if constexpr (brgemm_supported()) { at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ D, + /* lda */ D, + /* ldb */ D, + /* ldc */ D, /* add_C */ false, - /* A */ qg, - /* B */ curr_last_recurrent_state_pack_reduced, - /* C */ attn_inter); + /* A */ qg_exp, + /* B */ s_packed, + /* C */ attn_inter); } else { blas_gemm( at::native::TransposeType::NoTranspose, at::native::TransposeType::NoTranspose, - v_head_size, - chunk_size, - qk_head_size, + D, + mb_size, + D, 1.0f, - curr_last_recurrent_state_reduced, - v_head_size, - qg, - qk_head_size, + s_packed, + D, + qg_exp, + D, 0.0f, attn_inter, - v_head_size); + D); } - // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new - // pack for v_prime + // step 4.a: attn_inter += attn2 @ v2' if constexpr (brgemm_supported()) { pack_vnni2( - /* dst */ v_prime_pack_reduced, - /* src */ v_prime_reduced, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV] - // attn_i: [chunk_size, chunk_size] - // v_new: [chunk_size, EV] + /* dst */ v_packed, + /* src */ v_prime2, + /* K */ padded_mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, /* add_C */ true, - /* A */ attn_i_reduced, - /* B */ v_prime_pack_reduced, - /* C */ attn_inter); + /* A */ attn2, + /* B */ v_packed, + /* C */ attn_inter); } else { blas_gemm( at::native::TransposeType::NoTranspose, at::native::TransposeType::NoTranspose, - v_head_size, - chunk_size, - chunk_size, + D, + mb_size, + padded_mb_size, 1.0f, - v_prime_reduced, - v_head_size, - attn_i_reduced, - chunk_size, + v_prime2, + D, + attn2, + CHUNK_SIZE, 1.0f, attn_inter, - v_head_size); + D); } - // core_attn_out[:, :, i] = attn_inter - for (int64_t m = 0; m < chunk_size; m++) { - at::vec::map( - [](fVec x) { return x; }, core_attn_out_i + m * v_head_size, attn_inter + m * v_head_size, v_head_size); + // step 4.b: write attn_inter -> out + scalar_t* __restrict__ o_ptr = out + (batch_offset + mb_start) * o_strideT + hv * o_strideH; + update_kernel::apply(o_ptr, attn_inter, mb_size, D, o_strideT); + + // brgemm_supported()==false path: step 5.2 below overwrites k_updated, + // which aliases the same buffer as v_prime2 (both are `tmp`). Snapshot + // v_prime2 into v_packed (otherwise brgemm-only, and free by now since + // qg_exp was consumed at step 3.b) before that happens, so step 5.3 + // doesn't read k_updated's data under the v_prime2 name. + if constexpr (!brgemm_supported()) { + std::copy(v_prime2, v_prime2 + padded_mb_size * D, v_packed); } - // last_recurrent_state = ( - // last_recurrent_state * g[:, :, i, -1, None, None].exp() - // + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new - // ) - // 1) last_recurrent_state * g[:, :, i, -1, None, None].exp() - // curr_last_recurrent_state: [EK, EV] - // g[:, :, i, -1, None, None]: [1, 1] - // last_recurrent_state * g[:, :, i, -1, None, None].exp(): [EK, EV] - auto g_pad_i_last = g_pad_i + chunk_size - 1; - auto g_exp_last = std::exp(g_pad_i_last[0]); - for (int64_t m = 0; m < qk_head_size; m++) { - int64_t i = 0; - auto vec_g_exp_last = fVec(g_exp_last); - for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i); - auto tmp1 = at::vec::convert(tmp0); - auto tmp2 = tmp1 * vec_g_exp_last; - tmp2.store(curr_last_recurrent_state + m * v_head_size + i); - } - if (i < v_head_size) { - auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i, v_head_size - i); - auto tmp1 = at::vec::convert(tmp0); - auto tmp2 = tmp1 * vec_g_exp_last; - tmp2.store(curr_last_recurrent_state + m * v_head_size + i, v_head_size - i); - } - } - // 2) (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new - // k_i: [chunk_size, EK] - // g[:, :, i, -1, None]: [1] - // g[:, :, i]: [chunk_size] - // (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, 1] - // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, EK] - // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2): [EK, chunk_size] - // v_new: [chunk_size, EV] - // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new: [EK, EV] - // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None] - for (int64_t m = 0; m < chunk_size; m++) { - auto g_exp = std::exp((g_pad_i_last[0] - g_pad_i[m])); - int64_t i = 0; - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(k_i + m * qk_head_size + i); - auto tmp2 = tmp0 * vec_g_exp_reduced; - tmp2.store(kg + m * qk_head_size + i); - } - } - // kg.transpose(-1, -2): [EK, chunk_size] - at::native::utils::transpose( - /* M */ chunk_size, - /* N */ qk_head_size, - /* src */ kg, - /* ld_src */ qk_head_size, - /* dst */ kg_transpose, - /* ld_dst */ chunk_size); - // kgv = kg.transpose(-1, -2) @ v_new - // v_new: [chunk_size, EV] + // step 5: update state + // state_new = state * exp(g_last) + (k * exp(g_last - g)).T @ v2' + + // step 5.1 state *= exp(g_last) fused with step 2.a + + // step 5.2 k' = k * exp(g_last - g).T; TODO: fuse this with 1.a + update_key_kernel::apply(k_updated, k_ptr, g_ptr, mb_size, k_strideT); + + // step 5.3 state += k' @ v2' if constexpr (brgemm_supported()) { at::native::cpublas::brgemm( - /* M */ qk_head_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ kg_transpose, - /* B */ v_prime_pack_reduced, - /* C */ kgv); + /* M */ D, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, + /* add_C */ true, + /* A */ k_updated, + /* B */ v_packed, + /* C */ s_ptr); } else { blas_gemm( at::native::TransposeType::NoTranspose, at::native::TransposeType::NoTranspose, - v_head_size, - qk_head_size, - chunk_size, + D, + D, + padded_mb_size, 1.0f, - v_prime_reduced, - v_head_size, - kg_transpose, - chunk_size, - 0.0f, - kgv, - v_head_size); - } - // last_recurrent_state = 1) + 2) - for (int64_t m = 0; m < qk_head_size; m++) { - at::vec::map2( - [](fVec x, fVec y) { return x + y; }, - curr_last_recurrent_state + m * v_head_size, - curr_last_recurrent_state + m * v_head_size, - kgv + m * v_head_size, - v_head_size); + v_packed, + D, + k_updated, + CHUNK_SIZE, + 1.0f, + s_ptr, + D); } } - // core_attn_out -> output - // output: [B, T, HV, EV] - // core_attn_out: [B, HV, padded_T, EV] - auto curr_out = out_ptr + h * oStrideH; - for (int64_t m = 0; m < seq_len; m++) { - at::vec::map( - [](fVec x) { return x; }, curr_out + m * oStrideT, curr_core_attn_out + m * v_head_size, v_head_size); - } - - // Move to the next query - data_index_step(b, batch_size, h, v_num_head); + // move to the next index + data_index_step(bs, num_seqs, hv, Hv); } + at::native::cpublas::brgemm_release(); }); } @@ -1011,14 +1632,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( int64_t d; #pragma GCC unroll 4 for (d = 0; d <= head_dim - VecSize; d += VecSize) { - bVec q_bvec = bVec::loadu(q_ptr + q_offset + d); - fVec q_fvec0, q_fvec1; - std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec); + auto [q_fvec0, q_fvec1] = load_float_vec2(q_ptr + q_offset + d); sum_q_fvec += q_fvec0 * q_fvec0; sum_q_fvec += q_fvec1 * q_fvec1; - bVec k_bvec = bVec::loadu(k_ptr + k_offset + d); - fVec k_fvec0, k_fvec1; - std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec); + auto [k_fvec0, k_fvec1] = load_float_vec2(k_ptr + k_offset + d); sum_k_fvec += k_fvec0 * k_fvec0; sum_k_fvec += k_fvec1 * k_fvec1; } @@ -1066,14 +1683,11 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( fVec kv_mem_vec1 = fVec(float(0)); for (int di = 0; di < head_dim; ++di) { fVec k_val_vec = fVec(k_ptr[k_offset + di] * k_scale); - fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi); - fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize); + auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi); kv_mem_vec0 = kv_mem_vec0 + state_vec0 * g_val_exp_vec * k_val_vec; kv_mem_vec1 = kv_mem_vec1 + state_vec1 * g_val_exp_vec * k_val_vec; } - bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi); - fVec v_vec0, v_vec1; - std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec); + auto [v_vec0, v_vec1] = load_float_vec2(v_ptr + v_offset + dvi); fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec; fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec; fVec o_vec0 = fVec(float(0)); @@ -1081,8 +1695,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( for (int di = 0; di < head_dim; ++di) { fVec q_vec = fVec(q_ptr[q_offset + di] * q_scale); fVec k_vec = fVec(k_ptr[k_offset + di] * k_scale); - fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi); - fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize); + auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi); state_vec0 = state_vec0 * g_val_exp_vec + k_vec * dt_vec0; state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1; o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec; @@ -1289,25 +1902,19 @@ void fused_gdn_gating_kernel_impl( constexpr int vec_size = bVec::size(); constexpr int fvec_size = fVec::size(); const fVec neg_one(-1.0f); - const fVec one(1.0f); at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { int64_t j = 0; for (; j < num_heads - (num_heads % vec_size); j += vec_size) { - fVec A_log_vec0 = fVec::loadu(A_log + j); - fVec A_log_vec1 = fVec::loadu(A_log + j + fvec_size); - bVec dt_bias_vec = bVec::loadu(dt_bias + j); - bVec a_bvec = bVec::loadu(a + i * num_heads + j); - bVec b_bvec = bVec::loadu(b + i * num_heads + j); - fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1; - std::tie(a0, a1) = at::vec::convert_to_float(a_bvec); - std::tie(b0, b1) = at::vec::convert_to_float(b_bvec); - std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec); + auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j); + auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j); + auto [a0, a1] = load_float_vec2(a + i * num_heads + j); + auto [b0, b1] = load_float_vec2(b + i * num_heads + j); fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0); fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1); - fVec beta0 = one / (one + (neg_one * b0).exp_u20()); - fVec beta1 = one / (one + (neg_one * b1).exp_u20()); + fVec beta0 = fast_sigmoid(b0); + fVec beta1 = fast_sigmoid(b1); g0.store(out + i * num_heads + j); g1.store(out + i * num_heads + j + fvec_size); @@ -1337,26 +1944,19 @@ void fused_gdn_gating_kernel_impl( constexpr int vec_size = bVec::size(); constexpr int fvec_size = fVec::size(); const fVec neg_one(-1.0f); - const fVec one(1.0f); at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { int64_t j = 0; for (; j < num_heads - (num_heads % vec_size); j += vec_size) { - bVec A_log_bvec = bVec::loadu(A_log + j); - fVec A_log_vec0, A_log_vec1; - std::tie(A_log_vec0, A_log_vec1) = at::vec::convert_to_float(A_log_bvec); - bVec dt_bias_vec = bVec::loadu(dt_bias + j); - bVec a_bvec = bVec::loadu(a + i * num_heads + j); - bVec b_bvec = bVec::loadu(b + i * num_heads + j); - fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1; - std::tie(a0, a1) = at::vec::convert_to_float(a_bvec); - std::tie(b0, b1) = at::vec::convert_to_float(b_bvec); - std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec); + auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j); + auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j); + auto [a0, a1] = load_float_vec2(a + i * num_heads + j); + auto [b0, b1] = load_float_vec2(b + i * num_heads + j); fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0); fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1); - fVec beta0 = one / (one + (neg_one * b0).exp_u20()); - fVec beta1 = one / (one + (neg_one * b1).exp_u20()); + fVec beta0 = fast_sigmoid(b0); + fVec beta1 = fast_sigmoid(b1); g0.store(out + i * num_heads + j); g1.store(out + i * num_heads + j + fvec_size); @@ -1373,29 +1973,208 @@ void fused_gdn_gating_kernel_impl( } // anonymous namespace -template -inline void -CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const int64_t& dim, const at::IntArrayRef& sizes, at::ScalarType st) { - TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes()); - TORCH_CHECK(tensor.dtype() == st, "Input tensor dtype mismatch"); - CHECK_DIM(dim, tensor); - if (is_last_dim_contiguous) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor); - } else { - CHECK_CONTIGUOUS(tensor); +template +std::tuple prepare_chunk_indices(const at::Tensor& cu_seqlens) { + int64_t num_seqs = cu_seqlens.size(0) - 1; + at::Tensor chunk_offsets = at::empty({num_seqs + 1}, cu_seqlens.options()); + // get number of chunks and chunk offsets + const int32_t* offsets_data = cu_seqlens.data_ptr(); + int32_t num_chunks = 0; + chunk_offsets[0] = 0; + for (int64_t row = 0; row < num_seqs; ++row) { + num_chunks += div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); + chunk_offsets[row + 1] = num_chunks; } + // get chunk indices + at::Tensor chunk_indices = at::empty({num_chunks, 2}, cu_seqlens.options()); + int32_t* indices_data = chunk_indices.data_ptr(); + + int64_t idx = 0; + for (int32_t row = 0; row < num_seqs; ++row) { + int32_t num_chunks = div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); + + for (int32_t col = 0; col < num_chunks; ++col) { + indices_data[idx * 2 + 0] = row; + indices_data[idx * 2 + 1] = col; + idx++; + } + } + return std::make_tuple(chunk_indices, chunk_offsets); } -// query: [B, T, HK, EK] -// key: [B, T, HK, EK] -// value: [B, T, HV, EV] -// g: [B, T, HV] FP32 -// beta: [B, T, HV] -// initial_state: [N, HV, EK, EV] FP32 -// output_final_state: bool -// cu_seqlens: [N + 1] INT32 -// head_first: bool -// use_qk_l2norm_in_kernel: bool +#define DISPATCH_HEAD_DIM_CASE(launch_macro, hd) \ + case hd: { \ + launch_macro(hd); \ + break; \ + } + +// [NB]: add new head_dim support here +#define DISPATCH_HEAD_DIM(dim, launch_macro) \ + switch (dim) { \ + DISPATCH_HEAD_DIM_CASE(launch_macro, 64) \ + DISPATCH_HEAD_DIM_CASE(launch_macro, 128) \ + default: \ + TORCH_CHECK(false, "Unexpected head dim size, ", dim); \ + } + +#define LAUNCH_L2NORM_KERNEL(HD) \ + l2norm_fwd_kernel_impl( \ + query_norm.data_ptr(), \ + key_norm.data_ptr(), \ + query.data_ptr(), \ + key.data_ptr(), \ + eps, \ + T, \ + H, \ + query.stride(1), \ + query.stride(2), \ + key.stride(1), \ + key.stride(2)); + +std::tuple l2norm_fwd(const at::Tensor& query, const at::Tensor& key, double eps) { + int64_t B = query.size(0); + int64_t T = query.size(1); + int64_t H = query.size(2); + int64_t D = query.size(3); + + at::Tensor query_norm = at::empty_like(query); + at::Tensor key_norm = at::empty_like(key); + + AT_DISPATCH_REDUCED_FLOATING_TYPES( + query.scalar_type(), "l2norm_fwd", [&] { DISPATCH_HEAD_DIM(D, LAUNCH_L2NORM_KERNEL); }); + + return std::make_tuple(query_norm, key_norm); +} + +// [NB]: instantiate decay_mask to avoid heavy recomputation in the kernel with exp +template +at::Tensor chunk_local_cumsum(const at::Tensor& g, const at::Tensor& cu_seqlens, const at::Tensor& chunk_indices) { + int64_t B = g.size(0); + // int64_t T = g.size(1); + int64_t Hv = g.size(2); + int64_t NT = chunk_indices.size(0); + + at::Tensor g_ = at::empty({B, NT, Hv, CHUNK_SIZE}, g.options()); + AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "chunk_local_cumsum", [&] { + chunk_local_cumsum_kernel_impl( + g_.data_ptr(), + g.data_ptr(), + cu_seqlens.data_ptr(), + chunk_indices.data_ptr(), + Hv, + NT); + }); + return g_; +} + +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL(HD) \ + chunk_gated_delta_rule_fwd_intra_kernel_impl( \ + w.data_ptr(), \ + u.data_ptr(), \ + decay_mask.data_ptr(), \ + k.data_ptr(), \ + v.data_ptr(), \ + g.data_ptr(), \ + beta.data_ptr(), \ + cu_seqlens.data_ptr(), \ + chunk_indices.data_ptr(), \ + H, \ + Hv, \ + NT, \ + k.stride(1), \ + k.stride(2), \ + v.stride(1), \ + v.stride(2)); + +template +std::tuple chunk_gated_delta_rule_fwd_intra( + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& g, + const at::Tensor& beta, + const at::Tensor& cu_seqlens, + const at::Tensor& chunk_indices) { + int64_t B = k.size(0); + int64_t T = k.size(1); + int64_t H = k.size(2); + int64_t D = k.size(3); + int64_t Hv = v.size(2); + int64_t Dv = v.size(3); + int64_t NT = chunk_indices.size(0); + + at::Tensor w = at::empty({B, T, Hv, D}, k.options()); // BFloat16 + at::Tensor u = at::empty({B, T, Hv, Dv}, k.options()); // BFloat16 + at::Tensor decay_mask = at::empty({B, NT, Hv, CHUNK_SIZE, CHUNK_SIZE}, g.options()); // Float + AT_DISPATCH_REDUCED_FLOATING_TYPES(k.scalar_type(), "chunk_gated_delta_rule_fwd_intra", [&] { + DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL); + }); + + return std::make_tuple(w, u, decay_mask); +} + +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \ + chunk_gated_delta_rule_fwd_inter_kernel_impl( \ + o.data_ptr(), \ + initial_state.data_ptr(), \ + initial_state_indices.data_ptr(), \ + q.data_ptr(), \ + k.data_ptr(), \ + w.data_ptr(), \ + u.data_ptr(), \ + g.data_ptr(), \ + decay_mask.data_ptr(), \ + cu_seqlens.data_ptr(), \ + chunk_offsets.data_ptr(), \ + H, \ + Hv, \ + num_seqs, \ + q.stride(1), \ + q.stride(2), \ + k.stride(1), \ + k.stride(2), \ + initial_state.stride(0)); + +template +std::tuple chunk_gated_delta_rule_fwd_inter( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& g, + const at::Tensor& decay_mask, + const at::Tensor& initial_state, + bool output_final_state, + const at::Tensor& cu_seqlens, + const at::Tensor& chunk_offsets, + const at::Tensor& initial_state_indices) { + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t H = q.size(2); + const int64_t D = q.size(3); + const int64_t Hv = w.size(2); + const int64_t Dv = u.size(3); + const int64_t num_seqs = initial_state_indices.size(0); + + at::Tensor o = at::empty({B, T, Hv, Dv}, q.options()); + AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "chunk_gated_delta_rule_fwd_inter", [&] { + DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL); + }); + + return std::make_tuple(o, initial_state); +} + +// [NB]: Support only varlen inputs +// B: packed batch dim of q/k/v (== 1) +// num_seqs: number of variable-length sequences +// +// query: [B, T, H, D] +// key: [B, T, H, D] +// value: [B, T, Hv, Dv] +// g: [B, T, Hv] FP32 +// beta: [B, T, Hv] +// initial_state: [num_seqs, Hv, Dv, D] FP32 +// cu_seqlens: [num_seqs + 1] INT32 +// std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& query, const at::Tensor& key, @@ -1407,154 +2186,71 @@ std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, - double eps = 1e-5) { - TORCH_CHECK(head_first == false, "chunk_gated_delta_rule_cpu does not support head first"); - int64_t B = query.size(0); - int64_t global_seq_len = query.size(1); - int64_t qk_num_head = query.size(2); - int64_t qk_head_size = query.size(3); - int64_t v_num_head = value.size(2); - int64_t v_head_size = value.size(3); - int64_t batch_size = initial_state.size(0); - CHECK_EQ(B, 1); - TORCH_CHECK(v_num_head % qk_num_head == 0, "expect v_num_head multiple of qk_num_head."); - TORCH_CHECK(qk_head_size % 32 == 0, "expect qk_head_size to be multiples of 32."); - TORCH_CHECK(v_head_size % 32 == 0, "expect v_head_size to be multiples of 32."); - CHECK_INPUT_SHAPE_DTYPE(query, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(key, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(value, 4, {B, global_seq_len, v_num_head, v_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(g, 3, {B, global_seq_len, v_num_head}, at::kFloat); - CHECK_INPUT_SHAPE_DTYPE(beta, 3, {B, global_seq_len, v_num_head}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, 1, {batch_size + 1}, at::kInt); - CHECK_INPUT_SHAPE_DTYPE(initial_state, 4, {batch_size, v_num_head, qk_head_size, v_head_size}, at::kFloat); - - at::Tensor output = at::empty_like(value, value.options()); // [B, T, HV, EV] - at::Tensor final_state = initial_state.to(at::kFloat); // [N, HV, EK, EV] - - // Strides - int64_t qStrideH = query.stride(2); - int64_t qStrideT = query.stride(1); - int64_t kStrideH = key.stride(2); - int64_t kStrideT = key.stride(1); - int64_t vStrideH = value.stride(2); - int64_t vStrideT = value.stride(1); - int64_t oStrideH = output.stride(2); - int64_t oStrideT = output.stride(1); - - constexpr int64_t chunk_size = 64; - // Deduce the global chunks - // e.g. cu_seqlens: [0, 5, 13, 16], chunk_size = 4 - // chunk_offsets: [0, 2, 4, 5] - // chunk_indices (batch_id, local_chunk_id): [[0, 0], [0, 1], [1, 0], [1, 1], [2, 0]] - at::Tensor chunk_offsets = at::empty(batch_size + 1, cu_seqlens.options()); - auto chunk_offsets_ptr = chunk_offsets.data_ptr(); - chunk_offsets_ptr[0] = 0; - int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr(); - int64_t s = 0; - int64_t e = 0; - int64_t s_pad = 0; - int64_t e_pad = 0; - for (int64_t b = 0; b < batch_size; b++) { - e = cu_seqlens_ptr[b + 1]; - int64_t seq_len = e - s; - int64_t pad_size = (chunk_size - seq_len % chunk_size) % chunk_size; - int64_t total_seq_length = seq_len + pad_size; - e_pad = s_pad + total_seq_length; - chunk_offsets[b + 1] = e_pad / chunk_size; - s = e; - s_pad = e_pad; - } - int64_t global_total_seq_length = e_pad; - int64_t global_num_chunk = chunk_offsets_ptr[batch_size]; - at::Tensor chunk_indices = at::empty(global_num_chunk * 2, cu_seqlens.options()); - auto chunk_indices_ptr = chunk_indices.data_ptr(); - int64_t curr_c = 0; - for (int64_t b = 0; b < batch_size; b++) { - int64_t batch_chunk_num = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b]; - for (int64_t c = 0; c < batch_chunk_num; c++) { - chunk_indices_ptr[curr_c * 2] = b; - chunk_indices_ptr[curr_c * 2 + 1] = c; - curr_c += 1; - } - } + const at::Tensor& initial_state_indices, + double eps = 1e-6) { + TORCH_CHECK(!head_first, "chunk_gated_delta_rule_cpu: does not support head first"); - // Allocate buffer - int64_t buff_size = v_num_head * global_total_seq_length // g_pad_data - + batch_size * v_num_head * global_total_seq_length * v_head_size // core_attn - + v_num_head * global_total_seq_length * chunk_size // decay_mask - + v_num_head * global_total_seq_length * v_head_size; // v_beta_attn - at::Tensor buff_data = at::empty({buff_size}, query.options().dtype(at::kFloat)); - int64_t reduced_buff_size = qk_num_head * global_total_seq_length * qk_head_size // q_pad_data - + qk_num_head * global_total_seq_length * qk_head_size // k_pad_data - + v_num_head * global_total_seq_length * v_head_size // v_pad_data - + v_num_head * global_total_seq_length * qk_head_size // k_beta_data - + v_num_head * global_total_seq_length * v_head_size // v_beta_data - + v_num_head * global_total_seq_length * qk_head_size // k_cumdecay_reduced - + qk_num_head * global_seq_len // q_norm_sum - + qk_num_head * global_seq_len; // k_norm_sum - at::Tensor reduced_buff_data = at::empty({reduced_buff_size}, query.options()); - int64_t num_thread = at::get_num_threads(); - int64_t buff_size_16bit_per_thread = - /* k_transpose */ qk_head_size * chunk_size + - /* v_pack */ chunk_size * v_head_size + - /* k_beta_g */ chunk_size * qk_head_size + - /* k_beta_g_pack */ chunk_size * qk_head_size + - /* attn */ chunk_size * chunk_size * 2 + - /* attn_reduced */ chunk_size * chunk_size + - /* k_cumdecay */ chunk_size * qk_head_size * 2 + - /* row */ chunk_size * 2 + - /* updated */ chunk_size * 2 + - /* curr_last_recurrent_state_reduced */ qk_head_size * v_head_size + - /* curr_last_recurrent_state_pack_reduced */ qk_head_size * v_head_size + - /* k_transpose_i */ qk_head_size * chunk_size + - /* attn_i */ chunk_size * chunk_size * 2 + - /* attn_i_reduced */ chunk_size * chunk_size + - /* v_prime */ chunk_size * v_head_size * 2 + - /* v_prime_reduced */ chunk_size * v_head_size + - /* v_prime_pack_reduced */ chunk_size * v_head_size + - /* qg */ chunk_size * qk_head_size + - /* attn_inter */ chunk_size * v_head_size * 2 + - /* kg */ chunk_size * qk_head_size + - /* kg_transpose */ qk_head_size * chunk_size + - /* kgv */ qk_head_size * v_head_size * 2; - at::Tensor thread_buff_data = at::empty({num_thread, buff_size_16bit_per_thread}, query.options()); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "chunk_gated_delta_rule_kernel", [&] { - chunk_gated_delta_rule_kernel_impl( - output.data_ptr(), - final_state.data_ptr(), - query.data_ptr(), - key.data_ptr(), - value.data_ptr(), - g.data_ptr(), - beta.data_ptr(), - cu_seqlens_ptr, - buff_data.data_ptr(), - reduced_buff_data.data_ptr(), - thread_buff_data.data_ptr(), - chunk_offsets_ptr, - chunk_indices_ptr, - use_qk_l2norm_in_kernel, - batch_size, - global_seq_len, - qk_num_head, - v_num_head, - qk_head_size, - v_head_size, - qStrideH, - qStrideT, - kStrideH, - kStrideT, - vStrideH, - vStrideT, - oStrideH, - oStrideT, - global_total_seq_length, - global_num_chunk, - buff_size_16bit_per_thread, - eps); - }); - return std::make_tuple(std::move(output), std::move(final_state)); + int64_t B = query.size(0); + int64_t T = query.size(1); + int64_t H = query.size(2); + int64_t D = query.size(3); + int64_t Hv = value.size(2); + int64_t Dv = value.size(3); + int64_t num_seqs = initial_state_indices.size(0); + + TORCH_CHECK(B == 1, __func__, ": expect batch size to be 1"); + TORCH_CHECK(Hv % H == 0, __func__, ": expect num_heads_kv multiple of num_heads."); + TORCH_CHECK(D % 32 == 0, __func__, ": expect head_dim to be multiples of 32."); + TORCH_CHECK(Dv % 32 == 0, __func__, ": expect head_dim_v to be multiples of 32."); + TORCH_CHECK(D == Dv, __func__, ": expect head_dim to be equal to head_dim_v."); + CHECK_INPUT_SHAPE_DTYPE(query, {B, T, H, D}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(key, {B, T, H, D}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(value, {B, T, Hv, Dv}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(g, {B, T, Hv}, at::kFloat); + CHECK_INPUT_SHAPE_DTYPE(beta, {B, T, Hv}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {num_seqs + 1}, at::kInt); + TORCH_CHECK(initial_state.sizes() == at::IntArrayRef({initial_state.size(0), Hv, Dv, D}), + "chunk_gated_delta_rule_cpu: initial_state shape mismatch, got ", initial_state.sizes()); + TORCH_CHECK(initial_state.scalar_type() == at::kFloat, "chunk_gated_delta_rule_cpu: initial_state dtype mismatch"); + CHECK_CPU(initial_state); + // initial_state may be a pooled/paged buffer with padding between slots + // (e.g. mamba cache-align mode), so only the per-slot (Hv, Dv, D) layout + // needs to be densely packed; dim 0's stride is read at runtime instead of + // assumed, mirroring conv.cpp's conv_state_slot_stride handling. + TORCH_CHECK(initial_state.stride(-1) == 1 && initial_state.stride(-2) == D && + initial_state.stride(-3) == Dv * D, + "chunk_gated_delta_rule_cpu: expect initial_state to be contiguous per pool slot."); + CHECK_INPUT_SHAPE_DTYPE(initial_state_indices, {num_seqs}, at::kInt); + + constexpr int CHUNK_SIZE = 64; + + // prepare chunk indices + auto [chunk_indices, chunk_offsets] = prepare_chunk_indices(cu_seqlens); + + float scale = 1.0 / std::sqrt(D); + auto [query_, key_] = use_qk_l2norm_in_kernel ? l2norm_fwd(query, key, eps) : std::make_tuple(query.mul(scale), key); + + auto g_ = chunk_local_cumsum(g, cu_seqlens, chunk_indices); + + // fused kkt + solve_tril + recompute_w_u + auto [w, u, decay_mask] = + chunk_gated_delta_rule_fwd_intra(key_, value, g_, beta, cu_seqlens, chunk_indices); + + // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o` + auto [output, final_state] = chunk_gated_delta_rule_fwd_inter( + query_, + key_, + w, + u, + g_, + decay_mask, + initial_state, + output_final_state, + cu_seqlens, + chunk_offsets, + initial_state_indices); + + return std::make_tuple(output, final_state); } // A_log: [v_num_heads] diff --git a/csrc/cpu/sgl-kernels/gemm.cpp b/csrc/cpu/sgl-kernels/gemm.cpp index 38e6d9f4ce9d..212e1db5297e 100644 --- a/csrc/cpu/sgl-kernels/gemm.cpp +++ b/csrc/cpu/sgl-kernels/gemm.cpp @@ -116,8 +116,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d); - fVec data1 = fVec::loadu(input + d + fVec::size()); + auto [data0, data1] = load_float_vec2(input + d); bVec out_vec = convert_from_float_ext(data0, data1); out_vec.store(out + d); } @@ -135,9 +134,7 @@ inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ inpu int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0, data1; - bVec b_vec = bVec::loadu(input + d); - std::tie(data0, data1) = at::vec::convert_to_float(b_vec); + auto [data0, data1] = load_float_vec2(input + d); data0.store(out + d); data1.store(out + d + fVec::size()); } @@ -156,9 +153,9 @@ inline void copy_add_stub( int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d); - fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size()); - bVec out_vec = convert_from_float_ext(data0, data1); + auto [data0, data1] = load_float_vec2(input + d); + auto [bias0, bias1] = load_float_vec2(bias + d); + bVec out_vec = convert_from_float_ext(data0 + bias0, data1 + bias1); out_vec.store(out + d); } for (; d < size; ++d) { @@ -176,7 +173,6 @@ inline void scalar_sigmoid_and_mul( using bVec = at::vec::Vectorized; using fVec = at::vec::Vectorized; // scalar sigmoid - const fVec one = fVec(1.f); fVec X; if constexpr (has_bias) { assert(bias != nullptr); @@ -184,18 +180,13 @@ inline void scalar_sigmoid_and_mul( } else { X = fVec(input[0]); } - X = one / (one + X.neg().exp_u20()); + X = fast_sigmoid(X); // vec mul constexpr int kVecSize = bVec::size(); for (int d = 0; d < SIZE; d += kVecSize) { - bVec m_bvec = bVec::loadu(mul + d); - fVec m_fvec0, m_fvec1; - std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec); - m_fvec0 = m_fvec0 * X; - m_fvec1 = m_fvec1 * X; - - bVec out_vec = convert_from_float_ext(m_fvec0, m_fvec1); + auto [m_fvec0, m_fvec1] = load_float_vec2(mul + d); + bVec out_vec = convert_from_float_ext(m_fvec0 * X, m_fvec1 * X); out_vec.store(out + d); } } @@ -727,10 +718,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) { return packed_scale; } -// mat1 : [M, K] +// mat1 : [*, K] // mat2 : [N, K] ([K, N] if use_fma_gemm) // bias : [N] -// out : [M, N] +// out : [*, N] // at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, bool is_vnni) { @@ -740,23 +731,25 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional(a) == b; -} - -constexpr bool operator==(int a, CPUAcTMethod b) { - return a == static_cast(b); -} +enum class CPUActMethod : int { + silu_and_mul = 0, + swiglu = 1, + gelu_and_mul = 2, +}; enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 }; @@ -113,6 +109,17 @@ constexpr bool operator==(int64_t a, CPUQuantAlgo b) { return a == static_cast(b); } +inline int64_t get_row_size(CPUQuantMethod quant, int64_t K) { + switch (quant) { + case CPUQuantMethod::INT8_W8A8: + return K + sizeof(int32_t); + case CPUQuantMethod::MXFP4: + return K >> 1; + default: + return K; + } +} + inline int64_t get_4bit_block_k_size(int64_t group_size) { return group_size > 128 ? 128 : group_size; } @@ -184,7 +191,7 @@ void fused_experts_fp_kernel_impl( int64_t num_tokens_post_pad, float alpha, float limit, - CPUAcTMethod act_func, + CPUActMethod act_func, bool with_bias); // shared expert implementation for int8 w8a8 diff --git a/csrc/cpu/sgl-kernels/gemm_fp8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8.cpp index b47eb9256e03..8509af08a03e 100644 --- a/csrc/cpu/sgl-kernels/gemm_fp8.cpp +++ b/csrc/cpu/sgl-kernels/gemm_fp8.cpp @@ -18,8 +18,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d); - fVec data1 = fVec::loadu(input + d + fVec::size()); + auto [data0, data1] = load_float_vec2(input + d); bVec out_vec = convert_from_float_ext(data0, data1); out_vec.store(out + d); } @@ -38,9 +37,9 @@ inline void copy_add_stub( int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d); - fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size()); - bVec out_vec = convert_from_float_ext(data0, data1); + auto [data0, data1] = load_float_vec2(input + d); + auto [bias0, bias1] = load_float_vec2(bias + d); + bVec out_vec = convert_from_float_ext(data0 + bias0, data1 + bias1); out_vec.store(out + d); } for (; d < size; ++d) { @@ -57,9 +56,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ int d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d) * vscale; - fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale; - bVec out_vec = convert_from_float_ext(data0, data1); + auto [data0, data1] = load_float_vec2(input + d); + bVec out_vec = convert_from_float_ext(data0 * vscale, data1 * vscale); out_vec.store(out + d); } for (; d < size; ++d) { diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp index 06f9f7c37362..5a810d8f38ed 100644 --- a/csrc/cpu/sgl-kernels/moe.cpp +++ b/csrc/cpu/sgl-kernels/moe.cpp @@ -123,46 +123,6 @@ int moe_align_block_size( return num_tokens_post_pad; } -// silu : shape leading dimension -// input0 [m_size, BLOCK_N] BLOCK_N -// input1 [m_size, BLOCK_N] BLOCK_N -// output [M * topk, N] N -template -inline void silu_and_mul( - scalar_t* __restrict__ output, - const float* __restrict__ input0, // x: x0, x1 - const float* __restrict__ input1, // y: y0, y1 - int64_t m_size, - int64_t N) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - - const fVec one = fVec(1.f); - - // no remainder - for (int64_t m = 0; m < m_size; ++m) { - scalar_t* __restrict__ out = output + m * N; - const float* __restrict__ x = input0 + m * BLOCK_N; - const float* __restrict__ y = input1 + m * BLOCK_N; - - for (int64_t d = 0; d < BLOCK_N; d += bVec::size()) { - fVec x0 = fVec::loadu(x + d); - fVec x1 = fVec::loadu(x + d + fVec::size()); - fVec y0 = fVec::loadu(y + d); - fVec y1 = fVec::loadu(y + d + fVec::size()); - // silu - x0 = x0 / (one + x0.neg().exp_u20()); - x1 = x1 / (one + x1.neg().exp_u20()); - // mul - x0 = x0 * y0; - x1 = x1 * y1; - // convert - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - } -} - template struct tinygemm_kernel_nn2 { static inline void apply( @@ -239,23 +199,17 @@ struct tinygemm_kernel_nn2 { Unroll{}(compute, k); } - using Vec = at::vec::Vectorized; - const Vec one = Vec(1.f); auto storec = [&](auto i) { constexpr int row = i / COLS; constexpr int col = i % COLS; // for COLS = 2, 4 use 512bit store if constexpr (col % 2 == 0) { - Vec x0 = vc0[row * COLS + col + 0]; - Vec x1 = vc0[row * COLS + col + 1]; - Vec y0 = vc1[row * COLS + col + 0]; - Vec y1 = vc1[row * COLS + col + 1]; - // silu - x0 = x0 / (one + x0.neg().exp_u20()); - x1 = x1 / (one + x1.neg().exp_u20()); - // mul - x0 = x0 * y0; - x1 = x1 * y1; + __m512 x0 = vc0[row * COLS + col + 0]; + __m512 x1 = vc0[row * COLS + col + 1]; + __m512 y0 = vc1[row * COLS + col + 0]; + __m512 y1 = vc1[row * COLS + col + 1]; + x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0); + x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1); _mm512_storeu_si512( reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), @@ -455,6 +409,8 @@ void fused_experts_kernel_impl( const scalar_t* __restrict__ input, const scalar_t* __restrict__ packed_w1, const scalar_t* __restrict__ packed_w2, + const float* __restrict__ w1_bias, + const float* __restrict__ w2_bias, const float* __restrict__ topk_weights, const int32_t* __restrict__ sorted_ids, const int32_t* __restrict__ expert_ids, @@ -464,7 +420,11 @@ void fused_experts_kernel_impl( int64_t K, int64_t E, int64_t topk, - int64_t num_tokens_post_pad) { + int64_t num_tokens_post_pad, + float alpha, + float limit, + CPUActMethod act_func, + bool with_bias) { // handle 2 tiles per block constexpr int64_t BLOCK_M = block_size_m(); constexpr int64_t BLOCK_N = block_size_n(); @@ -499,6 +459,8 @@ void fused_experts_kernel_impl( int32_t expert_id = expert_ids[mb]; const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n; const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n; + const float* __restrict__ B0_bias = w1_bias + expert_id * 2 * N + nb_upper * BLOCK_N; + const float* __restrict__ B1_bias = w1_bias + expert_id * 2 * N + nb_lower * BLOCK_N; int64_t m_size = offsets[mb + 1] - offsets[mb]; @@ -538,23 +500,62 @@ void fused_experts_kernel_impl( /* B */ B1, /* C */ C1); - // 1.d silu and mul - const int64_t offset = offsets[mb]; - silu_and_mul(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N); } else { - // fused 1.bcd: silu_and_mul(A @ B0, A @ B1) const int64_t offset = offsets[mb]; - tinygemm_kernel( - /* A */ A, - /* B0 */ B0, - /* B1 */ B1, - /* C */ ic1 + offset * N + nb * BLOCK_N, - /* M */ m_size, - /* N */ n_size, - /* K */ K, - /* lda */ K, - /* ldb */ n_size, - /* ldc */ N); + if (act_func == CPUActMethod::swiglu) { + tinygemm_kernel( + /* A */ A, + /* B */ B0, + /* C */ C0, + /* M */ m_size, + /* N */ n_size, + /* K */ K, + /* lda */ K, + /* ldb */ n_size, + /* ldc */ BLOCK_N); + tinygemm_kernel( + /* A */ A, + /* B */ B1, + /* C */ C1, + /* M */ m_size, + /* N */ n_size, + /* K */ K, + /* lda */ K, + /* ldb */ n_size, + /* ldc */ BLOCK_N); + } else { + // fused 1.bcd: silu_and_mul(A @ B0, A @ B1) + tinygemm_kernel( + /* A */ A, + /* B0 */ B0, + /* B1 */ B1, + /* C */ ic1 + offset * N + nb * BLOCK_N, + /* M */ m_size, + /* N */ n_size, + /* K */ K, + /* lda */ K, + /* ldb */ n_size, + /* ldc */ N); + } + } + if (with_bias) { + for (int64_t m = 0; m < m_size; ++m) { + add_bias_stub(C0 + m * BLOCK_N, B0_bias, n_size); + add_bias_stub(C1 + m * BLOCK_N, B1_bias, n_size); + } + } + // 1.d silu and mul + const int64_t offset = offsets[mb]; + if (act_func == CPUActMethod::silu_and_mul && use_brgemm) { + for (int64_t m = 0; m < m_size; ++m) { + silu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N); + } + } else if (act_func == CPUActMethod::swiglu) { + for (int64_t m = 0; m < m_size; ++m) { + scalar_t* __restrict__ ic1_row = ic1 + (offset + m) * N; + clamp_sigmoid_and_mul_stub(ic1_row + nb * BLOCK_N / 2, C0 + m * BLOCK_N, BLOCK_N / 2, alpha, limit); + clamp_sigmoid_and_mul_stub(ic1_row + N / 2 + nb * BLOCK_N / 2, C1 + m * BLOCK_N, BLOCK_N / 2, alpha, limit); + } } }); @@ -591,6 +592,7 @@ void fused_experts_kernel_impl( // B shape [IC, n_size] in vnni format int32_t expert_id = expert_ids[mb]; const scalar_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc; + const float* __restrict__ B_bias = w2_bias + expert_id * OC + nb * BLOCK_N; // 2.a gemm: C = A @ B if (use_brgemm) { @@ -618,6 +620,11 @@ void fused_experts_kernel_impl( /* ldc */ BLOCK_N); } + if (with_bias) { + for (int64_t m = 0; m < m_size; ++m) { + add_bias_stub(C + m * BLOCK_N, B_bias, n_size); + } + } // 2.b copy from C to ic2 in original order // and also mul topk_weights in float32 for (int64_t m = 0; m < m_size; ++m) { @@ -717,7 +724,9 @@ void shared_expert_kernel_impl( /* C */ C1); // 1.d silu and mul - silu_and_mul(ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, m_size, N); + for (int64_t m = 0; m < m_size; ++m) { + silu_and_mul_stub(ic1 + (mb * BLOCK_M + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N); + } } else { // fused 1.bcd: silu_and_mul(A @ B0, A @ B1) tinygemm_kernel( @@ -809,24 +818,20 @@ void shared_expert_kernel_impl( } // anonymous namespace // common checks +template static inline void check_moe_scales( - bool use_int8_w8a8, - bool use_fp8_w8a16, - bool use_mxfp4, const std::optional& w1_scale, const std::optional& w2_scale, const std::optional> block_size) { - if (use_int8_w8a8) { + if constexpr (quant == CPUQuantMethod::INT8_W8A8) { TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8."); TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8."); - } - if (use_fp8_w8a16) { + } else if constexpr (quant == CPUQuantMethod::FP8_W8A16) { TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16."); TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a16."); TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16."); TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2."); - } - if (use_mxfp4) { + } else if constexpr (quant == CPUQuantMethod::MXFP4) { TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4."); TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4."); TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8."); @@ -834,6 +839,20 @@ static inline void check_moe_scales( } } +static inline void check_moe_scales( + int64_t moe_comp_method, + const std::optional& w1_scale, + const std::optional& w2_scale, + const std::optional> block_size) { + if (moe_comp_method == CPUQuantMethod::INT8_W8A8) { + check_moe_scales(w1_scale, w2_scale, block_size); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) { + check_moe_scales(w1_scale, w2_scale, block_size); + } else if (moe_comp_method == CPUQuantMethod::MXFP4) { + check_moe_scales(w1_scale, w2_scale, block_size); + } +} + #define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \ auto w1s = w1_scale.value(); \ auto w2s = w2_scale.value(); \ @@ -877,10 +896,14 @@ at::Tensor fused_experts_cpu( constexpr int64_t BLOCK_N = block_size_n(); const auto st = hidden_states.scalar_type(); + // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax) + // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided + auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt); + CHECK_INPUT(hidden_states); CHECK_INPUT(w1); CHECK_INPUT(w2); - CHECK_EQ(topk_weights.sizes(), topk_ids.sizes()); + CHECK_EQ(topk_weights.sizes(), topk_ids_.sizes()); CHECK_DIM(2, hidden_states); if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) { CHECK_DIM(4, w1); @@ -890,9 +913,8 @@ at::Tensor fused_experts_cpu( CHECK_DIM(3, w2); } CHECK_DIM(2, topk_weights); - CHECK_DIM(2, topk_ids); - - CHECK_EQ(topk_ids.scalar_type(), at::kInt); + CHECK_DIM(2, topk_ids_); + CHECK_EQ(topk_ids_.scalar_type(), at::kInt); // TODO: support topk_weights to be bf16 or fp16 in the kernel. // The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bf16/fp16 @@ -908,12 +930,8 @@ at::Tensor fused_experts_cpu( int64_t topk = topk_weights_.size(1); // we use int32_t compensation for int8 w8a8 - int64_t packed_K = moe_comp_method == CPUQuantMethod::MXFP4 - ? get_row_size(K) - : get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8); - int64_t packed_N = moe_comp_method == CPUQuantMethod::MXFP4 - ? get_row_size(N) - : get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8); + int64_t packed_K = get_row_size(static_cast(moe_comp_method), K); + int64_t packed_N = get_row_size(static_cast(moe_comp_method), N); // check weight shapes CHECK_EQ(w2.size(0), E); @@ -923,13 +941,7 @@ at::Tensor fused_experts_cpu( CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); } // check scales - check_moe_scales( - moe_comp_method == CPUQuantMethod::INT8_W8A8, - moe_comp_method == CPUQuantMethod::FP8_W8A16, - moe_comp_method == CPUQuantMethod::MXFP4, - w1_scale, - w2_scale, - block_size); + check_moe_scales(moe_comp_method, w1_scale, w2_scale, block_size); at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states); @@ -945,7 +957,7 @@ at::Tensor fused_experts_cpu( int64_t max_num_blocks = div_up(max_num_tokens_padded, BLOCK_M); auto buffer = at::empty( {max_num_tokens_padded + max_num_blocks + (num_threads + 1) * E + (E + 1) + (max_num_blocks + 1)}, - topk_ids.options()); + topk_ids_.options()); int32_t* __restrict__ sorted_ids = buffer.data_ptr(); int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded; @@ -969,7 +981,7 @@ at::Tensor fused_experts_cpu( // align experts index int64_t num_tokens_post_pad = moe_align_block_size( - sorted_ids, expert_ids, topk_ids.data_ptr(), total_cnts, cumsums, offsets, E, numel, num_threads); + sorted_ids, expert_ids, topk_ids_.data_ptr(), total_cnts, cumsums, offsets, E, numel, num_threads); // unlike triton kernel, we fuse silu with gemm1 so only need 2 intermediate_caches: // 1. intermediate_cache1 : [M * topk, N] @@ -1048,7 +1060,7 @@ at::Tensor fused_experts_cpu( scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); bool with_bias = w1_bias.has_value(); - auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul; + auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; CHECK_MOE_SCALES_FP8(1, 2); fused_experts_fp_kernel_impl( @@ -1088,7 +1100,7 @@ at::Tensor fused_experts_cpu( scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); bool with_bias = w1_bias.has_value(); - auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul; + auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; // mxfp4 supports only group size of 32 (2^5) constexpr int64_t group_size = 32; @@ -1172,6 +1184,8 @@ at::Tensor fused_experts_cpu( } else { scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K; float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); + bool with_bias = w1_bias.has_value(); + auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; fused_experts_kernel_impl( out_hidden_states.data_ptr(), @@ -1182,6 +1196,8 @@ at::Tensor fused_experts_cpu( hidden_states.data_ptr(), packed_w1.data_ptr(), packed_w2.data_ptr(), + with_bias ? w1_bias.value().data_ptr() : nullptr, + with_bias ? w2_bias.value().data_ptr() : nullptr, topk_weights_.data_ptr(), sorted_ids, expert_ids, @@ -1191,7 +1207,11 @@ at::Tensor fused_experts_cpu( K, E, topk, - num_tokens_post_pad); + num_tokens_post_pad, + alpha.has_value() ? float(alpha.value()) : 0, + limit.has_value() ? float(limit.value()) : 0, + act_func, + with_bias); } }); return out_hidden_states; @@ -1254,7 +1274,11 @@ at::Tensor shared_expert_cpu( CHECK_EQ(packed_w2.size(1), packed_N); // check scales - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, false, w1_scale, w2_scale, block_size); + if (use_int8_w8a8) { + check_moe_scales(w1_scale, w2_scale, block_size); + } else if (use_fp8_w8a16) { + check_moe_scales(w1_scale, w2_scale, block_size); + } at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states); diff --git a/csrc/cpu/sgl-kernels/moe.h b/csrc/cpu/sgl-kernels/moe.h index b6d2e9e7f6db..876c2d65db6a 100644 --- a/csrc/cpu/sgl-kernels/moe.h +++ b/csrc/cpu/sgl-kernels/moe.h @@ -64,9 +64,7 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict_ #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { auto [x0, x1] = load_float_vec2(input + d); - x0 = x0 * weight_vec; - x1 = x1 * weight_vec; - bVec out_vec = convert_from_float_ext(x0, x1); + bVec out_vec = convert_from_float_ext(x0 * weight_vec, x1 * weight_vec); out_vec.store(out + d); } for (; d < size; ++d) { @@ -91,10 +89,7 @@ inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ in fVec sum_fvec0 = fVec(0.f); fVec sum_fvec1 = fVec(0.f); for (int t = 0; t < topk; ++t) { - bVec x_bvec = bVec::loadu(input + t * K + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - + auto [x_fvec0, x_fvec1] = load_float_vec2(input + t * K + d); sum_fvec0 += x_fvec0; sum_fvec1 += x_fvec1; } @@ -137,11 +132,7 @@ inline void add_mul_stub( #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { auto [x0, x1] = load_float_vec2(input + d); - - bVec y_bvec = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); - + auto [y0, y1] = load_float_vec2(input2 + d); x0 = x0 + y0 * s_vec; x1 = x1 + y1 * s_vec; bVec out_vec = convert_from_float_ext(x0, x1); @@ -152,31 +143,52 @@ inline void add_mul_stub( } } -template +template inline void silu_and_mul_stub( - scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) { + scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) { + static_assert( + std::is_same_v || std::is_same_v, + "silu_and_mul_stub only supports input_t == float or input_t == scalar_t"); using bVec = at::vec::Vectorized; using fVec = at::vec::Vectorized; - const fVec one = fVec(1.f); // no remainder #pragma GCC unroll 4 for (int64_t d = 0; d < size; d += bVec::size()) { - bVec x = bVec::loadu(input + d); - fVec x0, x1; - std::tie(x0, x1) = at::vec::convert_to_float(x); - bVec y = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y); - x0 = x0 / (one + x0.neg().exp_u20()); - x1 = x1 / (one + x1.neg().exp_u20()); - x0 = x0 * y0; - x1 = x1 * y1; + auto [x0, x1] = load_float_vec2(input + d); + auto [y0, y1] = load_float_vec2(input2 + d); + x0 = fast_silu(x0) * y0; + x1 = fast_silu(x1) * y1; bVec out_vec = convert_from_float_ext(x0, x1); out_vec.store(out + d); } } +template +inline void clamp_sigmoid_and_mul_stub( + scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) { + static_assert( + std::is_same_v || std::is_same_v, + "clamp_sigmoid_and_mul_stub only supports input_t == float or input_t == scalar_t"); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + const fVec one = fVec(1.f); + const fVec limit_v = fVec(limit); + const fVec nlimit_v = fVec(-limit); + const fVec alpha_v = fVec(alpha); + +#pragma GCC unroll 4 + for (int64_t d = 0; d < 2 * size; d += bVec::size()) { + auto [x0_, y0_] = load_float_vec2(input + d); + auto [x0, y0] = at::vec::deinterleave2(x0_, y0_); + + x0 = at::vec::minimum(x0, limit_v); + y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0)); + x0 = fast_sigmoid_glu(x0, alpha_v) * (y0 + one); + store_from_float_ext(out + d / 2, x0); + } +} + template inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { using bVec = at::vec::Vectorized; @@ -186,9 +198,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec data0 = fVec::loadu(input + d) * weight_vec; - fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec; - bVec out_vec = convert_from_float_ext(data0, data1); + auto [x0, x1] = load_float_vec2(input + d); + bVec out_vec = convert_from_float_ext(x0 * weight_vec, x1 * weight_vec); out_vec.store(out + d); } for (; d < size; ++d) { @@ -222,63 +233,11 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict int64_t d; #pragma GCC unroll 4 for (d = 0; d <= size - kVecSize; d += kVecSize) { - bVec x = bVec::loadu(input + d); - fVec x0, x1; - std::tie(x0, x1) = at::vec::convert_to_float(x); - x0 = x0 * weight_vec; - x1 = x1 * weight_vec; - bVec out_vec = convert_from_float_ext(x0, x1); + auto [x0, x1] = load_float_vec2(input + d); + bVec out_vec = convert_from_float_ext(x0 * weight_vec, x1 * weight_vec); out_vec.store(out + d); } for (; d < size; ++d) { out[d] = static_cast(input[d] * weight); } } - -template -inline void clamp_sigmoid_and_mul_stub( - scalar_t* __restrict__ out, - const scalar_t* __restrict__ input, - int64_t size, - const float alpha, - const float limit) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - const fVec one = fVec(1.f); - const fVec zero = fVec(0.f); - const fVec limit_v = fVec(limit); - const fVec nlimit_v = fVec(-limit); - const fVec alpha_v = fVec(alpha); - - // no remainder -#pragma GCC unroll 4 - for (int64_t d = 0; d < size; d += bVec::size()) { - bVec x = bVec::loadu(input + d); - fVec x0_, y0_; - std::tie(x0_, y0_) = at::vec::convert_to_float(x); - float tmp_buffer[fVec::size() * 2]; // 32 - float tmp_glu[fVec::size()]; // 16 - float tmp_linear[fVec::size()]; // 16 - x0_.store(tmp_buffer); - y0_.store(tmp_buffer + fVec::size()); - // interleaved: x[2i] = glu, x[2i+1] = linear - for (int j = 0; j < fVec::size(); ++j) { - // x0 [0,2,..30] - tmp_glu[j] = tmp_buffer[j * 2]; - // y0 [1,3,...31] - tmp_linear[j] = tmp_buffer[j * 2 + 1]; - } - fVec x0 = fVec::loadu(tmp_glu); - fVec y0 = fVec::loadu(tmp_linear); - - // clamp - x0 = at::vec::minimum(x0, limit_v); - y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0)); - // x * sigmoid(x * alpha) - x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20()); - // (y + 1) * x - y0 = y0 + one; - x0 = x0 * y0; - convert_from_float_and_store(out + d / 2, x0); - } -} diff --git a/csrc/cpu/sgl-kernels/moe_fp8.cpp b/csrc/cpu/sgl-kernels/moe_fp8.cpp index 7b33a6585594..5c707c67aaed 100644 --- a/csrc/cpu/sgl-kernels/moe_fp8.cpp +++ b/csrc/cpu/sgl-kernels/moe_fp8.cpp @@ -37,7 +37,7 @@ void fused_experts_fp_kernel_impl( int64_t num_tokens_post_pad, float alpha, float limit, - CPUAcTMethod act_func, + CPUActMethod act_func, bool with_bias) { constexpr int64_t BLOCK_M = block_size_m(); constexpr int64_t BLOCK_N = block_size_n(); @@ -122,17 +122,17 @@ void fused_experts_fp_kernel_impl( }); // stage 1.5: intermediate_cache1 = silu(intermediate_cache0) - if (act_func == CPUAcTMethod::silu_and_mul) { + if (act_func == CPUActMethod::silu_and_mul) { at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { for (int64_t m = begin; m < end; ++m) { silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); } }); - } else if (act_func == CPUAcTMethod::swiglu) { + } else if (act_func == CPUActMethod::swiglu) { at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { for (int64_t m = begin; m < end; ++m) { - clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit); - clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit); + clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N / 2, alpha, limit); + clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N / 2, alpha, limit); } }); } @@ -243,7 +243,7 @@ void fused_experts_fp_kernel_impl( int64_t num_tokens_post_pad, \ float alpha, \ float limit, \ - CPUAcTMethod act_func, \ + CPUActMethod act_func, \ bool with_bias) INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false); diff --git a/csrc/cpu/sgl-kernels/moe_int8.cpp b/csrc/cpu/sgl-kernels/moe_int8.cpp index 3bdd5892d0b4..4bdddb9c53e0 100644 --- a/csrc/cpu/sgl-kernels/moe_int8.cpp +++ b/csrc/cpu/sgl-kernels/moe_int8.cpp @@ -53,16 +53,14 @@ inline void silu_and_mul( vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]); }; - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - const fVec one = fVec(1.f); auto silu_and_mul = [&](auto col) { - fVec x = fVec(vc0[col]); - fVec y = fVec(vc1[col]); - x = x / (one + x.neg().exp_u20()); - vc0[col] = x * y; + __m512 x = vc0[col]; + __m512 y = vc1[col]; + vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y); }; + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; auto storec = [&](auto col, int64_t m) { if constexpr (col % 2 == 0) { fVec x0 = fVec(vc0[col + 0]); @@ -229,23 +227,17 @@ struct tinygemm_kernel_vnni { }; Unroll{}(scalec); - using Vec = at::vec::Vectorized; - const Vec one = Vec(1.f); auto storec = [&](auto i) { constexpr int row = i / COLS; constexpr int col = i % COLS; // for COLS = 2, 4 use 512bit store if constexpr (col % 2 == 0) { - Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]); - Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]); - Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]); - Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]); - // silu - x0 = x0 / (one + x0.neg().exp_u20()); - x1 = x1 / (one + x1.neg().exp_u20()); - // mul - x0 = x0 * y0; - x1 = x1 * y1; + __m512 x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]); + __m512 x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]); + __m512 y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]); + __m512 y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]); + x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0); + x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1); _mm512_storeu_si512( reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 407cfe604343..04b696abf50f 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -32,11 +32,11 @@ inline Vectorized convert_from_float_ext(const Vectorized& a, c } template -inline void convert_from_float_and_store(scalar_t* out, const Vectorized& a) { - float out_buffer[at::vec::Vectorized::size()]; +inline void store_from_float_ext(scalar_t* out, const Vectorized& a) { + float out_buffer[Vectorized::size()]; a.store(out_buffer); - for (int i = 0; i < 16; i++) { - out[i] = (scalar_t)out_buffer[i]; + for (int i = 0; i < Vectorized::size(); ++i) { + out[i] = static_cast(out_buffer[i]); } } @@ -59,6 +59,17 @@ inline std::tuple, Vectorized> load_float_vec2(const fl return std::make_tuple(x0, x1); } +template , int> = 1> +inline at::vec::Vectorized load_float_vec(const scalar_t* __restrict__ data) { + at::vec::Vectorized out; + if constexpr (std::is_same_v) { + at::vec::load_fp32_from_bf16(data, out); + } else { + at::vec::load_fp32_from_fp16(data, out); + } + return out; +} + #if defined(CPU_CAPABILITY_AVX512) // `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics @@ -70,8 +81,14 @@ convert_from_float_ext(const Vectorized& a, const Vectorize } template <> -inline void convert_from_float_and_store(at::BFloat16* out, const Vectorized& a) { - _mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a)))); +inline void store_from_float_ext(at::BFloat16* out, const Vectorized& a) { + _mm256_storeu_si256(reinterpret_cast<__m256i*>(out), (__m256i)(_mm512_cvtneps_pbh(__m512(a)))); +} + +template <> +inline void store_from_float_ext(at::Half* out, const Vectorized& a) { + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(out), _mm512_cvtps_ph(__m512(a), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); } #define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16)) @@ -301,6 +318,77 @@ inline void quantize_row_int8( // transpose utils // taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998 #if defined(CPU_CAPABILITY_AVX512) +inline void transpose_16x16_16bit(__m256i* v) { + __m256i v1[16]; + v1[0] = _mm256_unpacklo_epi16(v[0], v[1]); + v1[1] = _mm256_unpackhi_epi16(v[0], v[1]); + v1[2] = _mm256_unpacklo_epi16(v[2], v[3]); + v1[3] = _mm256_unpackhi_epi16(v[2], v[3]); + v1[4] = _mm256_unpacklo_epi16(v[4], v[5]); + v1[5] = _mm256_unpackhi_epi16(v[4], v[5]); + v1[6] = _mm256_unpacklo_epi16(v[6], v[7]); + v1[7] = _mm256_unpackhi_epi16(v[6], v[7]); + v1[8] = _mm256_unpacklo_epi16(v[8], v[9]); + v1[9] = _mm256_unpackhi_epi16(v[8], v[9]); + v1[10] = _mm256_unpacklo_epi16(v[10], v[11]); + v1[11] = _mm256_unpackhi_epi16(v[10], v[11]); + v1[12] = _mm256_unpacklo_epi16(v[12], v[13]); + v1[13] = _mm256_unpackhi_epi16(v[12], v[13]); + v1[14] = _mm256_unpacklo_epi16(v[14], v[15]); + v1[15] = _mm256_unpackhi_epi16(v[14], v[15]); + + v[0] = _mm256_unpacklo_epi32(v1[0], v1[2]); + v[1] = _mm256_unpackhi_epi32(v1[0], v1[2]); + v[2] = _mm256_unpacklo_epi32(v1[1], v1[3]); + v[3] = _mm256_unpackhi_epi32(v1[1], v1[3]); + v[4] = _mm256_unpacklo_epi32(v1[4], v1[6]); + v[5] = _mm256_unpackhi_epi32(v1[4], v1[6]); + v[6] = _mm256_unpacklo_epi32(v1[5], v1[7]); + v[7] = _mm256_unpackhi_epi32(v1[5], v1[7]); + v[8] = _mm256_unpacklo_epi32(v1[8], v1[10]); + v[9] = _mm256_unpackhi_epi32(v1[8], v1[10]); + v[10] = _mm256_unpacklo_epi32(v1[9], v1[11]); + v[11] = _mm256_unpackhi_epi32(v1[9], v1[11]); + v[12] = _mm256_unpacklo_epi32(v1[12], v1[14]); + v[13] = _mm256_unpackhi_epi32(v1[12], v1[14]); + v[14] = _mm256_unpacklo_epi32(v1[13], v1[15]); + v[15] = _mm256_unpackhi_epi32(v1[13], v1[15]); + + v1[0] = _mm256_unpacklo_epi64(v[0], v[4]); + v1[1] = _mm256_unpackhi_epi64(v[0], v[4]); + v1[2] = _mm256_unpacklo_epi64(v[1], v[5]); + v1[3] = _mm256_unpackhi_epi64(v[1], v[5]); + v1[4] = _mm256_unpacklo_epi64(v[2], v[6]); + v1[5] = _mm256_unpackhi_epi64(v[2], v[6]); + v1[6] = _mm256_unpacklo_epi64(v[3], v[7]); + v1[7] = _mm256_unpackhi_epi64(v[3], v[7]); + v1[8] = _mm256_unpacklo_epi64(v[8], v[12]); + v1[9] = _mm256_unpackhi_epi64(v[8], v[12]); + v1[10] = _mm256_unpacklo_epi64(v[9], v[13]); + v1[11] = _mm256_unpackhi_epi64(v[9], v[13]); + v1[12] = _mm256_unpacklo_epi64(v[10], v[14]); + v1[13] = _mm256_unpackhi_epi64(v[10], v[14]); + v1[14] = _mm256_unpacklo_epi64(v[11], v[15]); + v1[15] = _mm256_unpackhi_epi64(v[11], v[15]); + + v[0] = _mm256_permute2x128_si256(v1[0], v1[8], 0x20); + v[1] = _mm256_permute2x128_si256(v1[1], v1[9], 0x20); + v[2] = _mm256_permute2x128_si256(v1[2], v1[10], 0x20); + v[3] = _mm256_permute2x128_si256(v1[3], v1[11], 0x20); + v[4] = _mm256_permute2x128_si256(v1[4], v1[12], 0x20); + v[5] = _mm256_permute2x128_si256(v1[5], v1[13], 0x20); + v[6] = _mm256_permute2x128_si256(v1[6], v1[14], 0x20); + v[7] = _mm256_permute2x128_si256(v1[7], v1[15], 0x20); + v[8] = _mm256_permute2x128_si256(v1[0], v1[8], 0x31); + v[9] = _mm256_permute2x128_si256(v1[1], v1[9], 0x31); + v[10] = _mm256_permute2x128_si256(v1[2], v1[10], 0x31); + v[11] = _mm256_permute2x128_si256(v1[3], v1[11], 0x31); + v[12] = _mm256_permute2x128_si256(v1[4], v1[12], 0x31); + v[13] = _mm256_permute2x128_si256(v1[5], v1[13], 0x31); + v[14] = _mm256_permute2x128_si256(v1[6], v1[14], 0x31); + v[15] = _mm256_permute2x128_si256(v1[7], v1[15], 0x31); +} + inline void transpose_16x16_32bit(__m512i* v) { __m512i v1[16]; v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); @@ -394,6 +482,62 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) } #pragma GCC diagnostic pop +// Note: mapped from aten exp_u20 +inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) { + const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); + const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); + const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); + const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); + const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); + const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e) + const __m512 vec_half = _mm512_set1_ps(0.5f); + const __m512 vec_one = _mm512_set1_ps(1.f); + const __m512 vec_zero = _mm512_set1_ps(0.f); + const __m512 vec_two = _mm512_set1_ps(2.f); + const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); + const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); + const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); + const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); + const int n_mantissa_bits = 23; + + // exp(x) = + // = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem + // = 2^n * exp(r) // simplify the exp(n*ln(2)) expression + + auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); + auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); + vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); + + // fx = floorf(x * log2ef + 0.5) + auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); + auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); + vec_fx = _mm512_cvtepi32_ps(vec_fx_i); + + // x = x - fx * ln2 + auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); + + // compute polynomial + auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); + + // compute 2^(n-1) + auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); + auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); + auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); + vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); + auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); + vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero); + + // y = y * 2^n + vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); + vec_res = _mm512_mul_ps(vec_res, vec_two); + return vec_res; +} + +// Note: mapped from aten fexp_u20 inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) { const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f); const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f); @@ -444,6 +588,51 @@ inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 val // final interpretation to float return _mm512_castsi512_ps(casted_integer); } + +// sigmoid(x) = 1 / (1 + exp(-x)); avoid vdivps via rcp14 +inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_ps(__m512 x) { + __m512 minus_x = _mm512_xor_ps(_mm512_set1_ps(-0.f), x); + __m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_x), _mm512_set1_ps(1.f)); + return _mm512_rcp14_ps(denom); +} + +// SiLU(x) = x * sigmoid(x) +inline __attribute__((always_inline)) __m512 _mm512_rcp14_silu_ps(__m512 x) { + return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(x)); +} + +// x * sigmoid(x * alpha) for clamped SwiGLU +inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512 x, __m512 alpha) { + __m512 xa = _mm512_mul_ps(x, alpha); + return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa)); +} + +#endif + +inline at::vec::Vectorized fast_sigmoid(const at::vec::Vectorized& x) { +#if defined(CPU_CAPABILITY_AVX512) + return at::vec::Vectorized(_mm512_rcp14_sigmoid_ps(x)); +#else + const auto one = at::vec::Vectorized(1.f); + return one / (one + x.neg().exp_u20()); +#endif +} + +inline at::vec::Vectorized fast_silu(const at::vec::Vectorized& x) { +#if defined(CPU_CAPABILITY_AVX512) + return at::vec::Vectorized(_mm512_rcp14_silu_ps(x)); +#else + return x * fast_sigmoid(x); #endif +} + +inline at::vec::Vectorized +fast_sigmoid_glu(const at::vec::Vectorized& x, const at::vec::Vectorized& alpha) { +#if defined(CPU_CAPABILITY_AVX512) + return at::vec::Vectorized(_mm512_rcp14_sigmoid_glu_ps(x, alpha)); +#else + return x * fast_sigmoid(x * alpha); +#endif +} } // anonymous namespace diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 88a593725eeb..e6130473e30c 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -110,7 +110,7 @@ std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& g, const at::Tensor& beta, const at::Tensor& initial_state, bool output_final_state, const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, - double eps = 1e-5); + const at::Tensor& initial_state_indices, double eps = 1e-5); at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q, @@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu( at::Tensor causal_conv1d_update_cpu( const at::Tensor& x, const at::Tensor& conv_states, const at::Tensor& weight, const std::optional& bias, - bool silu_activation, const std::optional& cache_seqlens, + bool silu_activation, const std::optional& num_accepted_tokens, const std::optional& conv_state_indices, int64_t pad_slot_id, bool is_vnni); @@ -163,7 +163,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool casual, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal); + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype); void cpu_attn_reshape_and_cache(const torch::Tensor& key, const torch::Tensor& value, @@ -207,12 +208,52 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input, const torch::Tensor& topk_id, const bool skip_weighted, const std::string& act, const std::string& isa); +void prepack_moe_weight_int8(const torch::Tensor& weight, + torch::Tensor& packed_weight, + const std::string& isa); + +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& w13, const torch::Tensor& w2, + const torch::Tensor& w13_scale, + const torch::Tensor& w2_scale, + const std::optional& w13_bias, + const std::optional& w2_bias, + const torch::Tensor& topk_weights, + const torch::Tensor& topk_id, const bool skip_weighted, + const std::string& act, const std::string& isa); + void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, const torch::Tensor positions, const torch::Tensor block_table, torch::Tensor slot_mapping, const int64_t block_size); +at::Tensor causal_conv1d_update_cpu_impl( + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, + const c10::optional& bias, + const c10::optional& activation, + const c10::optional& conv_state_indices, + const c10::optional& query_start_loc, int64_t pad_slot_id); + +void selective_state_update_cpu_impl( + at::Tensor& state, const at::Tensor& x, const at::Tensor& dt, + const at::Tensor& A, const at::Tensor& B, const at::Tensor& C, + const c10::optional& D, const c10::optional& z, + const c10::optional& dt_bias, bool dt_softplus, + const c10::optional& state_batch_indices, + const c10::optional& dst_state_batch_indices, + int64_t null_block_id, at::Tensor& out, + const c10::optional& num_accepted_tokens, + const c10::optional& cu_seqlens); + +void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states, + const at::Tensor& x, const at::Tensor& dt, + const at::Tensor& A, const at::Tensor& B, + const at::Tensor& C, + const c10::optional& D, + const c10::optional& z, + const at::Tensor& cu_seqlens); + void init_cpu_memory_env(std::vector node_ids); namespace cpu_utils { @@ -357,7 +398,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Quantization #if defined(__AVX512F__) || defined(__AVX2__) || \ (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \ - defined(__riscv_v) + defined(__riscv_v) || defined(__s390x__) // Helper function to release oneDNN handlers ops.def("release_dnnl_matmul_handler(int handler) -> ()", &release_dnnl_matmul_handler); @@ -476,7 +517,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " "weight, Tensor? bias, bool silu_activation," - "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " + "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int " + "pad_slot_id, " "bool is_vnni) -> Tensor"); ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif @@ -504,7 +546,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor g, Tensor beta, " "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool " "head_first, " - "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)"); + "bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float " + "eps=1e-5) -> (Tensor, Tensor)"); ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU, &chunk_gated_delta_rule_cpu); ops.def( @@ -536,7 +579,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " "query_start_loc, bool casual, int window_size, str isa_hint, bool " - "enable_kv_split, Tensor? dynamic_causal) -> Tensor", + "enable_kv_split, Tensor? dynamic_causal, " + "str kv_cache_dtype=\"auto\") -> Tensor", &get_scheduler_metadata); ops.def( "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) " @@ -570,8 +614,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif // fused moe -#if defined(__AVX512F__) || \ - (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT)) +#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__)) ops.def( "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) " "-> ()"); @@ -582,7 +625,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#endif +#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && + // !defined(__APPLE__)) +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ + !defined(__APPLE__) + ops.def( + "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, " + "str isa) -> ()"); + ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8); + ops.def( + "cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, " + "Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, " + "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool " + "skip_weighted, str act, str isa) -> ()"); + ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8); +#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && + // !defined(__APPLE__) ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," @@ -595,6 +653,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "block_size) -> ()", &compute_slot_mapping_kernel_impl); + // Mamba CPU kernels + ops.def( + "causal_conv1d_update_cpu_vec(" + "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, " + "Tensor? bias, str? activation, Tensor? conv_state_indices, " + "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor", + &causal_conv1d_update_cpu_impl); + + ops.def( + "selective_state_update_cpu(" + "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " + "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, " + "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, " + "SymInt null_block_id, Tensor(a13!) out, " + "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()", + &selective_state_update_cpu_impl); + + ops.def( + "mamba_chunk_scan_fwd_cpu(" + "Tensor(a0!) out, Tensor(a1!) final_states, " + "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " + "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()", + &mamba_chunk_scan_fwd_cpu_impl); + ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); // Speculative decoding kernels diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 2329d51a149e..720c990450b6 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -212,6 +212,29 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, if (error_code != 0) { return; } + +#ifdef USE_ROCM + // ROCm-only: zero the freshly (re)mapped region. + // + // On some amdgpu driver versions the physical VRAM backing a buffer is not + // scrubbed on release, so the pages handed back here (e.g. on wake_up after a + // sleep that discarded a tag) can still contain stale data from a previously + // freed allocation -- possibly from another process. Callers of the device + // allocator expect newly mapped memory to be well defined; in particular + // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent + // state cache, and a new sequence that reads an unzeroed slot will propagate + // garbage/NaN, collapsing generation to token id 0 (decoded as "!"). + // + // This is scoped to the ROCm path on purpose: the CUDA path is unchanged. + // Use CUDA_CHECK so a failure sets the global error_code; the Python caller + // (python_create_and_map) only inspects error_code, so returning without + // setting it would make a failed memset look successful and still hand back + // unzeroed memory. + CUDA_CHECK((CUresult)hipMemset(reinterpret_cast(d_mem), 0, size)); + if (error_code != 0) { + return; + } +#endif // std::cout << "create_and_map: device=" << device << ", size=" << size << ", // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; } diff --git a/csrc/custom_all_gather_reduce_scatter.cuh b/csrc/custom_all_gather_reduce_scatter.cuh new file mode 100644 index 000000000000..ac990d8106ab --- /dev/null +++ b/csrc/custom_all_gather_reduce_scatter.cuh @@ -0,0 +1,327 @@ +#pragma once + +#include "custom_collective_common.cuh" + +namespace vllm { + +constexpr int kMnnvlLamportAgThreads = 128; +constexpr int kMnnvlLamportRsThreads = 256; +constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192; + +using CopyPack = array_t; + +template +__global__ void __launch_bounds__(512, 1) + cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg, + CopyPack* __restrict__ result, int rank, + int size_per_rank) { + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + barrier_at_start(sg, self_sg, rank); +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + auto src = reinterpret_cast(dp.ptrs[src_rank]); + auto dst = result + src_rank * size_per_rank; + for (int idx = tid; idx < size_per_rank; idx += stride) { + dst[idx] = src[idx]; + } + } + barrier_at_end(sg, self_sg, rank); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, + int size_per_rank) { + using P = typename packed_t::P; + using A = typename packed_t::A; + auto dp = *_dp; + auto offset = rank * size_per_rank; + barrier_at_start(sg, self_sg, rank); + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank; + idx += gridDim.x * blockDim.x) { + reinterpret_cast(result)[idx] = + packed_reduce((const P**)&dp.ptrs[0], offset + idx); + } + barrier_at_end(sg, self_sg, rank); +} + +template +union LamportPack { + P packed; + uint32_t words[sizeof(P) / sizeof(uint32_t)]; +}; + +template +DINLINE LamportPack

load_lamport_pack(const P* ptr) { + static_assert(sizeof(P) == 16); + LamportPack

value; +#if !defined(USE_ROCM) + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(value.words[0]), "=r"(value.words[1]), + "=r"(value.words[2]), "=r"(value.words[3]) + : "l"(ptr) + : "memory"); +#else + const volatile uint32_t* src = + reinterpret_cast(ptr); + #pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = src[i]; + } +#endif + return value; +} + +template +DINLINE bool is_lamport_dirty(const LamportPack

& value) { +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) return true; + } + return false; +} + +template +DINLINE P lamport_sentinel() { + LamportPack

value; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = 0x80000000U; + } + return value.packed; +} + +template +DINLINE P sanitize_lamport_payload(P packed) { + LamportPack

value{.packed = packed}; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) value.words[i] = 0; + } + return value.packed; +} + +template +DINLINE P wait_lamport_payload(const P* ptr) { + auto value = load_lamport_pack(ptr); + while (is_lamport_dirty(value)) value = load_lamport_pack(ptr); + return value.packed; +} + +template +DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride, + P local_value, P (&values)[ngpus]) { + bool ready[ngpus]; +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + ready[src_rank] = src_rank == rank; + if (src_rank == rank) values[src_rank] = local_value; + } + + int remaining = ngpus - 1; + while (remaining != 0) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + if (!ready[src_rank]) { + auto value = load_lamport_pack(base + src_rank * rank_stride); + if (!is_lamport_dirty(value)) { + values[src_rank] = value.packed; + ready[src_rank] = true; + --remaining; + } + } + } + } +} + +template +DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input, + int rank, int size_per_rank, int idx) { + P source_zero = + rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx); + A tmp = upcast(source_zero); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + P value = src_rank == rank + ? packed_input[rank * size_per_rank + idx] + : wait_lamport_payload(current_local + + src_rank * size_per_rank + idx); + packed_assign_add(tmp, upcast(value)); + } + return sanitize_lamport_payload(downcast

(tmp)); +} + +DINLINE void lamport_cta_arrive(uint32_t* counter) { +#if !defined(USE_ROCM) + if (threadIdx.x < 32) { + asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory"); + if (threadIdx.x == 0) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("red.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #else + atomicAdd(counter, 1); + #endif + } + } else { + asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory"); + } +#else + __syncthreads(); + if (threadIdx.x == 0) atomicAdd(counter, 1); +#endif +} + +template +__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1) + mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input, + T* __restrict__ result, + T* __restrict__ multicast_buffer, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto current_multicast = + reinterpret_cast(multicast_buffer) + current_stage * stage_size; + auto packed_input = reinterpret_cast(input); + auto packed_result = reinterpret_cast(result); + + int total_size = size_per_rank * ngpus; + P local_value; + if (tid < size_per_rank) { + local_value = packed_input[tid]; + current_multicast[rank * size_per_rank + tid] = + sanitize_lamport_payload(local_value); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (tid < size_per_rank) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + int output_idx = src_rank * size_per_rank + tid; + P value = src_rank == rank + ? local_value + : wait_lamport_payload(current_local + output_idx); + packed_result[output_idx] = value; + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = total_size; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +template +__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1) + mnnvl_lamport_reduce_scatter_kernel(RankData* _dp, + const T* __restrict__ input, + T* __restrict__ result, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; + using A = typename packed_t::A; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int dst_rank = blockIdx.x % ngpus; + int tile = blockIdx.x / ngpus; + int idx = tile * blockDim.x + threadIdx.x; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto packed_input = reinterpret_cast(input); + + if (idx < size_per_rank && dst_rank != rank) { + auto dst = reinterpret_cast(const_cast(dp.ptrs[dst_rank])) + + current_stage * stage_size + rank * size_per_rank; + auto src = packed_input + dst_rank * size_per_rank; + dst[idx] = sanitize_lamport_payload(src[idx]); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (idx < size_per_rank && dst_rank == rank) { + if constexpr (ngpus == 4) { + if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) { + reinterpret_cast(result)[idx] = + reduce_lamport_payloads(current_local, packed_input, + rank, size_per_rank, idx); + } else { + P values[ngpus]; + wait_lamport_payloads( + current_local + idx, rank, size_per_rank, + packed_input[rank * size_per_rank + idx], values); + A tmp = upcast(values[0]); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + packed_assign_add(tmp, upcast(values[src_rank])); + } + reinterpret_cast(result)[idx] = + sanitize_lamport_payload(downcast

(tmp)); + } + } else { + reinterpret_cast(result)[idx] = reduce_lamport_payloads( + current_local, packed_input, rank, size_per_rank, idx); + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = size_per_rank * ngpus; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +} // namespace vllm diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 58926f6429dd..385a0e0ae79c 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -1,299 +1,8 @@ #pragma once -#include -#include -#include -#include - -#if defined(USE_ROCM) -typedef __hip_bfloat16 nv_bfloat16; -#endif - -#include -#include -#include -#include -#include -#include -#include -#include +#include "custom_collective_common.cuh" namespace vllm { -#define CUDACHECK(cmd) \ - do { \ - cudaError_t e = cmd; \ - if (e != cudaSuccess) { \ - printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ - cudaGetErrorString(e)); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -// Maximal number of blocks in allreduce kernel. -constexpr int kMaxBlocks = 36; - -// Default number of blocks in allreduce kernel. -#ifndef USE_ROCM -const int defaultBlockLimit = 36; -CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#else -const int defaultBlockLimit = 16; -hipPointer_attribute rangeStartAddrAttr = - HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#endif - -// Counter may overflow, but it's fine since unsigned int overflow is -// well-defined behavior. -using FlagType = uint32_t; - -// Two sets of peer counters are needed for two syncs: starting and ending an -// operation. The reason is that it's possible for peer GPU block to arrive at -// the second sync point while the current GPU block haven't passed the first -// sync point. Thus, peer GPU may write counter+1 while current GPU is busy -// waiting for counter. We use alternating counter array to avoid this -// possibility. -struct Signal { - alignas(128) FlagType start[kMaxBlocks][8]; - alignas(128) FlagType end[kMaxBlocks][8]; - alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank -}; - -struct __align__(16) RankData { - const void* ptrs[8]; -}; - -struct __align__(16) RankSignals { - Signal* signals[8]; -}; - -// like std::array, but aligned -template -struct __align__(alignof(T) * sz) array_t { - T data[sz]; - using type = T; - static constexpr int size = sz; -}; - -// use packed type to maximize memory efficiency -// goal: generate ld.128 and st.128 instructions -template -struct packed_t { - // the (P)acked type for load/store - using P = array_t; - // the (A)ccumulator type for reduction - using A = array_t; -}; - -#define DINLINE __device__ __forceinline__ - -// scalar cast functions -DINLINE float upcast_s(half val) { return __half2float(val); } - -template -DINLINE T downcast_s(float val); -template <> -DINLINE half downcast_s(float val) { - return __float2half(val); -} - -// scalar add functions -// for some reason when compiling with Pytorch, the + operator for half and -// bfloat is disabled so we call the intrinsics directly -DINLINE half& assign_add(half& a, half b) { - a = __hadd(a, b); - return a; -} -DINLINE float& assign_add(float& a, float b) { return a += b; } - -#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } -template <> -DINLINE nv_bfloat16 downcast_s(float val) { - return __float2bfloat16(val); -} -DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { - a = __hadd(a, b); - return a; -} -#endif - -template -DINLINE array_t& packed_assign_add(array_t& a, array_t b) { -#pragma unroll - for (int i = 0; i < N; i++) { - assign_add(a.data[i], b.data[i]); - } - return a; -} - -template -DINLINE array_t upcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - array_t out; -#pragma unroll - for (int i = 0; i < N; i++) { - out.data[i] = upcast_s(val.data[i]); - } - return out; - } -} - -template -DINLINE O downcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - O out; -#pragma unroll - for (int i = 0; i < O::size; i++) { - out.data[i] = downcast_s(val.data[i]); - } - return out; - } -} - -#if !defined(USE_ROCM) - -static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #else - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #endif -} - -static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { - FlagType flag; - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("ld.acquire.sys.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - #else - asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" - : "=r"(flag) - : "l"(flag_addr)); - #endif - return flag; -} - -static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { - asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); -} - -static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { - FlagType flag; - asm volatile("ld.volatile.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - return flag; -} - -// This function is meant to be used as the first synchronization in the all -// reduce kernel. Thus, it doesn't need to make any visibility guarantees for -// prior memory accesses. Note: volatile writes will not be reordered against -// other volatile writes. -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value - // from peer. - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -// This function is meant to be used as the second or the final -// synchronization barrier in the all reduce kernel. If it's the final -// synchronization barrier, we don't need to make any visibility guarantees -// for prior memory accesses. -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value from - // peer. - if constexpr (!final_sync) { - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } else { - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - } - if constexpr (!final_sync) __syncthreads(); - - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#else - -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_RELAXED, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], - flag, - final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, - __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while ( - __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], - final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - if constexpr (!final_sync) __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#endif - -template -DINLINE P packed_reduce(const P* ptrs[], int idx) { - A tmp = upcast(ptrs[0][idx]); -#pragma unroll - for (int i = 1; i < ngpus; i++) { - packed_assign_add(tmp, upcast(ptrs[i][idx])); - } - return downcast

(tmp); -} template __global__ void __launch_bounds__(512, 1) @@ -616,6 +325,21 @@ class CustomAllreduce { #undef KL } + void allgather(cudaStream_t stream, void* input, void* output, int size_bytes, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output, + void* local_buffer, void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes); + template + void reduce_scatter(cudaStream_t stream, T* input, T* output, int size, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output, + void* local_buffer, uint32_t* epochs, + int size, int stage_size_bytes); + ~CustomAllreduce() { for (auto [_, ptr] : ipc_handles_) { CUDACHECK(cudaIpcCloseMemHandle(ptr)); @@ -625,8 +349,8 @@ class CustomAllreduce { /** * To inspect PTX/SASS, copy paste this header file to compiler explorer and - add a template instantiation: + * add a template instantiation: * template void vllm::CustomAllreduce::allreduce(cudaStream_t, half *, - half *, int, int, int); -*/ -} // namespace vllm \ No newline at end of file + * half *, int, int, int); + */ +} // namespace vllm diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh new file mode 100644 index 000000000000..353dcd07e8b4 --- /dev/null +++ b/csrc/custom_collective_common.cuh @@ -0,0 +1,332 @@ +#pragma once + +#include +#include +#include +#include + +#if defined(USE_ROCM) +typedef __hip_bfloat16 nv_bfloat16; +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm { +constexpr int kMaxCustomCollectiveRanks = 16; + +#define CUDACHECK(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ + cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +// Maximal number of blocks in allreduce kernel. +constexpr int kMaxBlocks = 36; + +// Default number of blocks in allreduce kernel. +#ifndef USE_ROCM +inline constexpr int defaultBlockLimit = 36; +inline CUpointer_attribute rangeStartAddrAttr = + CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#else +inline constexpr int defaultBlockLimit = 16; +inline hipPointer_attribute rangeStartAddrAttr = + HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#endif + +// Counter may overflow, but it's fine since unsigned int overflow is +// well-defined behavior. +using FlagType = uint32_t; + +// Two sets of peer counters are needed for two syncs: starting and ending an +// operation. The reason is that it's possible for peer GPU block to arrive at +// the second sync point while the current GPU block haven't passed the first +// sync point. Thus, peer GPU may write counter+1 while current GPU is busy +// waiting for counter. We use alternating counter array to avoid this +// possibility. +struct Signal { + alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank +}; + +struct __align__(16) RankData { + const void* ptrs[kMaxCustomCollectiveRanks]; +}; + +struct __align__(16) RankSignals { + Signal* signals[kMaxCustomCollectiveRanks]; +}; + +// like std::array, but aligned +template +struct __align__(alignof(T) * sz) array_t { + T data[sz]; + using type = T; + static constexpr int size = sz; +}; + +// use packed type to maximize memory efficiency +// goal: generate ld.128 and st.128 instructions +template +struct packed_t { + // the (P)acked type for load/store + using P = array_t; + // the (A)ccumulator type for reduction + using A = array_t; +}; + +#define DINLINE __device__ __forceinline__ + +// scalar cast functions +DINLINE float upcast_s(half val) { return __half2float(val); } + +template +DINLINE T downcast_s(float val); +template <> +DINLINE half downcast_s(float val) { + return __float2half(val); +} + +// scalar add functions +// for some reason when compiling with Pytorch, the + operator for half and +// bfloat is disabled so we call the intrinsics directly +DINLINE half& assign_add(half& a, half b) { + a = __hadd(a, b); + return a; +} +DINLINE float& assign_add(float& a, float b) { return a += b; } + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } +template <> +DINLINE nv_bfloat16 downcast_s(float val) { + return __float2bfloat16(val); +} +DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { + a = __hadd(a, b); + return a; +} +#endif + +template +DINLINE array_t& packed_assign_add(array_t& a, array_t b) { +#pragma unroll + for (int i = 0; i < N; i++) { + assign_add(a.data[i], b.data[i]); + } + return a; +} + +template +DINLINE array_t upcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + array_t out; +#pragma unroll + for (int i = 0; i < N; i++) { + out.data[i] = upcast_s(val.data[i]); + } + return out; + } +} + +template +DINLINE O downcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + O out; +#pragma unroll + for (int i = 0; i < O::size; i++) { + out.data[i] = downcast_s(val.data[i]); + } + return out; + } +} + +#if !defined(USE_ROCM) + +static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #else + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #endif +} + +static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { + FlagType flag; + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + #else + asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" + : "=r"(flag) + : "l"(flag_addr)); + #endif + return flag; +} + +static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { + asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +} + +static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { + FlagType flag; + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + return flag; +} + +// This function is meant to be used as the first synchronization in the all +// reduce kernel. Thus, it doesn't need to make any visibility guarantees for +// prior memory accesses. Note: volatile writes will not be reordered against +// other volatile writes. +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value + // from peer. + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +// This function is meant to be used as the second or the final +// synchronization barrier in the all reduce kernel. If it's the final +// synchronization barrier, we don't need to make any visibility guarantees +// for prior memory accesses. +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value from + // peer. + if constexpr (!final_sync) { + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } else { + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + } + if constexpr (!final_sync) __syncthreads(); + + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#else + +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_RELAXED, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], + flag, + final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, + __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while ( + __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], + final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + if constexpr (!final_sync) __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#endif + +template +DINLINE P packed_reduce(const P* ptrs[], int idx) { + A tmp = upcast(ptrs[0][idx]); +#pragma unroll + for (int i = 1; i < ngpus; i++) { + packed_assign_add(tmp, upcast(ptrs[i][idx])); + } + return downcast

(tmp); +} + +} // namespace vllm diff --git a/csrc/flashkda_registration.cpp b/csrc/flashkda_registration.cpp new file mode 100644 index 000000000000..9058119e1193 --- /dev/null +++ b/csrc/flashkda_registration.cpp @@ -0,0 +1,17 @@ +#include "core/registration.h" +#include "flash_kda.h" + +TORCH_LIBRARY(_flashkda_C, m) { + m.def("get_workspace_size(int T_total, int H, int N=1) -> int", + &get_workspace_size); + m.def( + "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, " + "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, " + "float lower_bound, " + "Tensor? initial_state=None, Tensor(b!)? final_state=None, " + "Tensor? cu_seqlens=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); } + +REGISTER_EXTENSION(_flashkda_C) diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp index fdf3e614e644..6b7c8d7fe1e5 100644 --- a/csrc/fs_io.cpp +++ b/csrc/fs_io.cpp @@ -3,19 +3,155 @@ #include +#include +#include #include +#include +#include #include +#if defined(O_DIRECT) +constexpr int kODirectFlag = O_DIRECT; +#else +constexpr int kODirectFlag = 0; +#endif + extern "C" { -static void _batch_lookup(const std::vector& paths, +namespace { + +// Returns 0 on success, or the std::error_code's POSIX-compatible value on +// failure, mirroring the errno convention used by the syscalls below. +inline int ensure_parent_dirs(const std::string& path) { + const auto parent = std::filesystem::path(path).parent_path(); + if (parent.empty()) { + return 0; + } + std::error_code ec; + std::filesystem::create_directories(parent, ec); + return ec ? ec.value() : 0; +} + +// Core single-block store: src/size are raw pointer + byte count. Returns 0 +// on success, or the errno of the failing step on failure -- captured +// before any subsequent cleanup call can overwrite it. On failure, the temp +// file is removed. +inline int _store_block(const char* tmp_path, const char* dest_path, + const char* src, size_t size, bool use_o_direct) { + if (access(dest_path, F_OK) == 0) { + return 0; // Already present. + } + + if (const int err = ensure_parent_dirs(dest_path); err != 0) { + return err; + } + + const int o_direct_flag = use_o_direct ? kODirectFlag : 0; + const int fd = open( + tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644); + if (fd < 0) { + return errno; + } + + const ssize_t written = write(fd, src, size); + if (written < 0 || static_cast(written) != size) { + const int err = written < 0 ? errno : EIO; + close(fd); // Best-effort cleanup; the real error is already captured. + unlink(tmp_path); + return err; + } + + if (close(fd) != 0) { + const int err = errno; + unlink(tmp_path); + return err; + } + + if (rename(tmp_path, dest_path) != 0) { + const int err = errno; + unlink(tmp_path); + return err; + } + + return 0; +} + +// Core single-block load: dst/size are raw pointer + byte count. Returns 0 +// on success, or the errno of the failing step on failure. On failure, +// the source file is removed since a partially-read block should not be reused. +inline int _load_block(const char* source_path, char* dst, size_t size, + bool use_o_direct) { + const int o_direct_flag = use_o_direct ? kODirectFlag : 0; + const int fd = open(source_path, O_RDONLY | o_direct_flag, 0); + if (fd < 0) { + const int err = errno; + unlink(source_path); + return err; + } + + const ssize_t bytes_read = read(fd, dst, size); + if (bytes_read < 0 || static_cast(bytes_read) != size) { + const int err = bytes_read < 0 ? errno : EIO; + close(fd); + unlink(source_path); + return err; + } + + if (close(fd) != 0) { + const int err = errno; + unlink(source_path); + return err; + } + + return 0; +} + +inline void _batch_lookup(const std::vector& paths, std::vector& exists_flags) { for (size_t i = 0; i < paths.size(); i++) { exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; } } +// Helper: extract a list[str] of length n into a vector. +// Returns false and sets a Python exception on error. +inline bool extract_str_list(PyObject* list, Py_ssize_t n, + std::vector& out) { + for (Py_ssize_t i = 0; i < n; i++) { + out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr); + if (out[i] == nullptr) { + return false; + } + } + return true; +} + +// Helper: extract a Py_buffer per element of a list[bytes-like] of length n. +// On success, `out` holds n acquired buffers (caller must PyBuffer_Release +// each). On failure, any buffers already acquired are released before +// returning false, and a Python exception is set. +inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags, + std::vector& out) { + for (Py_ssize_t i = 0; i < n; i++) { + if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) { + for (Py_ssize_t j = 0; j < i; j++) { + PyBuffer_Release(&out[j]); + } + return false; + } + } + return true; +} + +inline void release_buffer_list(std::vector& buffers) { + for (auto& buf : buffers) { + PyBuffer_Release(&buf); + } +} + +} // namespace + /// @brief Check file existence for a batch of paths. /// @param paths list[str] – absolute paths to check. /// @return list[bool] – True if the corresponding path exists, False otherwise. @@ -51,11 +187,157 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { return result; } +/// @brief Store a batch of blocks, each from its own buffer, to disk. +/// @param tmp_paths list[str] – one temp path per block. +/// @param dest_paths list[str] – one destination path per block. +/// @param buffers list[bytes-like] – one source buffer per block. +/// @param use_o_direct bool – whether to open files with O_DIRECT +/// (default True). Ignored where O_DIRECT is unsupported +/// by the platform. +/// @note Releases the GIL for the entire batch. Raises on first error. +static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) { + PyObject* tmp_paths_obj = nullptr; + PyObject* dest_paths_obj = nullptr; + PyObject* buffers_obj = nullptr; + int use_o_direct = 1; + + if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj, + &PyList_Type, &dest_paths_obj, &PyList_Type, + &buffers_obj, &use_o_direct)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(tmp_paths_obj); + if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) { + PyErr_SetString( + PyExc_ValueError, + "tmp_paths, dest_paths and buffers must have the same length"); + return nullptr; + } + + std::vector tmp_paths(n); + std::vector dest_paths(n); + + if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr; + if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr; + + std::vector buffers(n); + if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) { + return nullptr; + } + + Py_ssize_t failed_index = -1; + int failure_errno = 0; + + { + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { + const char* buf = static_cast(buffers[i].buf); + const int err = + _store_block(tmp_paths[i], dest_paths[i], buf, + static_cast(buffers[i].len), use_o_direct); + if (err != 0) { + failed_index = i; + failure_errno = err; + break; + } + } + Py_END_ALLOW_THREADS + } + + release_buffer_list(buffers); + + if (failed_index >= 0) { + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. + errno = failure_errno; + return PyErr_SetFromErrnoWithFilename(PyExc_OSError, + dest_paths[failed_index]); + } + + Py_RETURN_NONE; +} + +/// @brief Load a batch of blocks from disk, each into its own buffer. +/// @param source_paths list[str] – one source path per block. +/// @param buffers list[writable bytes-like] – one destination buffer +/// per block. +/// @param use_o_direct bool – whether to open files with O_DIRECT +/// (default True). Ignored where O_DIRECT is unsupported +/// by the platform. +/// @note Releases the GIL for the entire batch. Raises on first error. +static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) { + PyObject* source_paths_obj = nullptr; + PyObject* buffers_obj = nullptr; + int use_o_direct = 1; + + if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj, + &PyList_Type, &buffers_obj, &use_o_direct)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(source_paths_obj); + if (PyList_Size(buffers_obj) != n) { + PyErr_SetString(PyExc_ValueError, + "source_paths and buffers must have the same length"); + return nullptr; + } + + std::vector source_paths(n); + if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr; + + std::vector buffers(n); + if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) { + return nullptr; + } + + Py_ssize_t failed_index = -1; + int failure_errno = 0; + + { + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { + char* buf = static_cast(buffers[i].buf); + const int err = + _load_block(source_paths[i], buf, static_cast(buffers[i].len), + use_o_direct); + if (err != 0) { + failed_index = i; + failure_errno = err; + break; + } + } + Py_END_ALLOW_THREADS + } + + release_buffer_list(buffers); + + if (failed_index >= 0) { + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. + errno = failure_errno; + return PyErr_SetFromErrnoWithFilename(PyExc_OSError, + source_paths[failed_index]); + } + + Py_RETURN_NONE; +} + static PyMethodDef fs_io_C_methods[] = { {"batch_lookup", batch_lookup, METH_VARARGS, "batch_lookup(paths: list[str]) -> list[bool]\n" "\n" "Check file existence for a batch of paths."}, + {"batch_store_block", batch_store_block, METH_VARARGS, + "batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n" + " buffers: list[bytes-like],\n" + " use_o_direct: bool = True) -> None\n" + "\n" + "Store a batch of blocks, each from its own buffer, to disk. Raises on " + "first error."}, + {"batch_load_block", batch_load_block, METH_VARARGS, + "batch_load_block(source_paths: list[str],\n" + " buffers: list[writable bytes-like],\n" + " use_o_direct: bool = True) -> None\n" + "\n" + "Load a batch of blocks from disk into corresponding buffers. " + "Raises on first error."}, {nullptr, nullptr, 0, nullptr}, }; diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index 60b8ca5f382d..aa2c1b6b92f4 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -103,9 +103,10 @@ __global__ void act_and_mul_kernel( scalar_t* __restrict__ out, // [..., d] const scalar_t* __restrict__ input, // [..., 2, d] const int d, const float limit, const float alpha, const float beta) { - const scalar_t* x_ptr = input + blockIdx.x * 2 * d; + const int64_t token_idx = blockIdx.x; + const scalar_t* x_ptr = input + token_idx * 2 * d; const scalar_t* y_ptr = x_ptr + d; - scalar_t* out_ptr = out + blockIdx.x * d; + scalar_t* out_ptr = out + token_idx * d; if constexpr (use_vec) { using cuda_t = typename CUDATypeConverter::Type; @@ -353,9 +354,10 @@ template ::Type; @@ -464,6 +466,66 @@ __global__ void swigluoai_and_mul_kernel( } } +// SITU (Kimi SituGLU) gated activation. Non-interleaved layout: +// input = [gate(d), up(d)] per token. +// gate_out = beta * tanh(gate / beta) * sigmoid(gate) +// up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up +// out = gate_out * up_out +// Compute is done in fp32 and written straight to `out` -- no intermediate +// tensors and no full-tensor fp32 upcast (the pure-torch forward_native +// allocated ~8 fp32 temporaries per call, which blows up MoE profiling). +template +__global__ void situ_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2, d] + const int d, const float beta, const float linear_beta) { + const int64_t row = blockIdx.x; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const float g = (float)VLLM_LDG(&gate_ptr[idx]); + const float u = (float)VLLM_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + +template +__global__ void masked_situ_and_mul_kernel( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, + const int* __restrict__ expert_num_tokens, const int max_num_tokens, + const int d, const float beta, const float linear_beta) { + const int expert = blockIdx.y; + const int num_tokens = expert_num_tokens[expert]; + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= d || num_tokens == 0) { + return; + } + + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + const int64_t expert_row = static_cast(expert) * max_num_tokens; + for (int token = 0; token < num_tokens; ++token) { + const int64_t row = expert_row + token; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const float g = (float)VLLM_LDG(&gate_ptr[idx]); + const float u = (float)VLLM_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + } // namespace vllm #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \ @@ -553,6 +615,54 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d] double alpha, double limit) { LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit); } + +// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed +// through), matching SituAndMul(linear_beta=None) on the Python side. +void situ_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + double beta, double linear_beta) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + if (num_tokens == 0) { + return; + } + dim3 grid(num_tokens); + dim3 block(std::min(d, 1024)); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "situ_and_mul_kernel", [&] { + vllm::situ_and_mul_kernel<<>>( + out.mutable_data_ptr(), input.const_data_ptr(), + d, (float)beta, (float)linear_beta); + }); +} + +void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d] + torch::stable::Tensor& input, // [E, T, 2 * d] + const torch::stable::Tensor& expert_num_tokens, + double beta, double linear_beta) { + int num_experts = input.size(0); + int max_num_tokens = input.size(1); + int d = input.size(2) / 2; + if (num_experts == 0 || max_num_tokens == 0) { + return; + } + constexpr int block_size = 256; + dim3 grid((d + block_size - 1) / block_size, num_experts); + dim3 block(block_size); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "masked_situ_and_mul_kernel", [&] { + vllm::masked_situ_and_mul_kernel<<>>( + out.mutable_data_ptr(), input.const_data_ptr(), + expert_num_tokens.const_data_ptr(), max_num_tokens, d, + (float)beta, (float)linear_beta); + }); +} namespace vllm { // Element-wise activation kernel template. @@ -562,8 +672,9 @@ __global__ void activation_kernel( scalar_t* __restrict__ out, // [..., d] const scalar_t* __restrict__ input, // [..., d] const int d) { - const scalar_t* in_ptr = input + blockIdx.x * d; - scalar_t* out_ptr = out + blockIdx.x * d; + const int64_t token_idx = blockIdx.x; + const scalar_t* in_ptr = input + token_idx * d; + scalar_t* out_ptr = out + token_idx * d; if constexpr (use_vec) { // Fast path: 128-bit/256-bit vectorized loop diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index b132e82e253f..cc89397f68a1 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -21,7 +21,10 @@ __global__ void merge_attn_states_kernel( const float* prefix_lse, const scalar_t* suffix_output, const float* suffix_lse, const uint num_tokens, const uint num_heads, const uint head_size, const uint prefix_head_stride, - const uint output_head_stride, const uint prefix_num_tokens, + const uint output_head_stride, const uint prefix_lse_head_stride, + const uint prefix_lse_token_stride, const uint suffix_lse_head_stride, + const uint suffix_lse_token_stride, const uint output_lse_head_stride, + const uint output_lse_token_stride, const uint prefix_num_tokens, const float* output_scale) { // Inputs always load 128-bit packs (pack_size elements of scalar_t). // Outputs store pack_size elements of output_t, which is smaller for FP8. @@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel( } } if (output_lse != nullptr && pack_idx == 0) { - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; - output_lse[head_idx * num_tokens + token_idx] = s_lse; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = s_lse; } return; } // For tokens within prefix range, merge prefix and suffix - float p_lse = prefix_lse[head_idx * num_tokens + token_idx]; - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; + float p_lse = prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; @@ -132,7 +139,8 @@ __global__ void merge_attn_states_kernel( } // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { - output_lse[head_idx * num_tokens + token_idx] = max_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = max_lse; } return; } @@ -187,7 +195,8 @@ __global__ void merge_attn_states_kernel( // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { float out_lse = logf(out_se) + max_lse; - output_lse[head_idx * num_tokens + token_idx] = out_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = out_lse; } } @@ -221,6 +230,9 @@ __global__ void merge_attn_states_kernel( reinterpret_cast(suffix_output.data_ptr()), \ reinterpret_cast(suffix_lse.data_ptr()), num_tokens, \ num_heads, head_size, prefix_head_stride, output_head_stride, \ + prefix_lse_head_stride, prefix_lse_token_stride, \ + suffix_lse_head_stride, suffix_lse_token_stride, \ + output_lse_head_stride, output_lse_token_stride, \ prefix_num_tokens, output_scale_ptr); \ } @@ -259,6 +271,19 @@ void merge_attn_states_launcher( const uint head_size = output.size(2); const uint prefix_head_stride = prefix_output.stride(1); const uint output_head_stride = output.stride(1); + // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views + // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index + // them by their actual strides rather than assuming a contiguous layout. + const uint prefix_lse_head_stride = prefix_lse.stride(0); + const uint prefix_lse_token_stride = prefix_lse.stride(1); + const uint suffix_lse_head_stride = suffix_lse.stride(0); + const uint suffix_lse_token_stride = suffix_lse.stride(1); + uint output_lse_head_stride = 0; + uint output_lse_token_stride = 0; + if (output_lse.has_value()) { + output_lse_head_stride = output_lse.value().stride(0); + output_lse_token_stride = output_lse.value().stride(1); + } // Thread mapping is based on input BF16 pack_size const uint pack_size = 16 / sizeof(scalar_t); STD_TORCH_CHECK(head_size % pack_size == 0, diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 2d4b47b4b43c..c35a516b5c95 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel( copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank); } +// Grouped variant of concat_and_cache_mla: inserts the context K/V for every +// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer +// reads its own cache base pointer from kv_cache_ptrs (same pointer-array +// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no +// scaling or quantization; scalar_t is uint16_t for portability. +template +__global__ void concat_and_cache_mla_grouped_kernel( + const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens, + // kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim] + const int64_t* __restrict__ kv_cache_ptrs, // [num_layers] + const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens] + const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride, + const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride, + const int64_t slot_layer_stride, const int64_t block_stride, + const int64_t entry_stride, const int kv_lora_rank, const int pe_dim, + const int block_size) { + const int64_t token_idx = blockIdx.x; + const int64_t layer_idx = blockIdx.y; + const int64_t slot_idx = + slot_mapping[layer_idx * slot_layer_stride + token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + scalar_t* __restrict__ kv_cache = + reinterpret_cast(kv_cache_ptrs[layer_idx]); + const scalar_t* __restrict__ kv_c_layer = + kv_c + layer_idx * kv_c_layer_stride; + const scalar_t* __restrict__ k_pe_layer = + k_pe + layer_idx * k_pe_layer_stride; + + auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride, + int size, int offset) { + for (int i = threadIdx.x; i < size; i += blockDim.x) { + const int64_t src_idx = token_idx * src_token_stride + i; + const int64_t dst_idx = + block_idx * block_stride + block_offset * entry_stride + i + offset; + kv_cache[dst_idx] = src[src_idx]; + } + }; + + copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0); + copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank); +} + template __global__ void concat_and_cache_ds_mla_kernel( const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] @@ -902,6 +951,53 @@ void concat_and_cache_mla( } } +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim] + torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device + torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64 + int64_t block_size, int64_t block_stride, int64_t entry_stride) { + int num_layers = kv_c.size(0); + int num_tokens = kv_c.size(1); + int kv_lora_rank = kv_c.size(2); + int pe_dim = k_pe.size(2); + + STD_TORCH_CHECK( + kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 && + k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=", + kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type()); + STD_TORCH_CHECK( + kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "kv_cache_ptrs must be int64"); + + if (num_tokens == 0 || num_layers == 0) { + return; + } + + const int64_t kv_c_layer_stride = kv_c.stride(0); + const int64_t kv_c_token_stride = kv_c.stride(1); + const int64_t k_pe_layer_stride = k_pe.stride(0); + const int64_t k_pe_token_stride = k_pe.stride(1); + const int64_t slot_layer_stride = slot_mapping.stride(0); + + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + dim3 grid(num_tokens, num_layers); + dim3 block(std::min(kv_lora_rank, 512)); + vllm::concat_and_cache_mla_grouped_kernel + <<>>( + reinterpret_cast(kv_c.data_ptr()), + reinterpret_cast(k_pe.data_ptr()), + kv_cache_ptrs.const_data_ptr(), + slot_mapping.const_data_ptr(), kv_c_layer_stride, + kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, + slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, + block_size); +} + namespace vllm { template @@ -1025,6 +1121,9 @@ __global__ void gather_and_maybe_dequant_cache( batch_offset += offset; int32_t block_table_id = batch_offset / block_size; int32_t slot_id = batch_offset % block_size; + // seq_starts may push the block index past the end of the batch's block + // table row. + if (block_table_id >= block_table_stride) continue; int32_t block_table_offset = batch_id * block_table_stride + block_table_id; int32_t block_id = block_table[block_table_offset]; int64_t cache_offset = diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu new file mode 100644 index 000000000000..88aed49d181a --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu @@ -0,0 +1,362 @@ +#include "torch_utils.h" + +#include +#include +#include +#include + +#include "custom_all_reduce.cuh" +#include "custom_all_gather_reduce_scatter.cuh" + +namespace vllm { + +void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output, + int size_bytes, int threads, int block_limit) { + if (size_bytes % sizeof(CopyPack) != 0) + throw std::runtime_error( + "custom allgather requires input byte size to be a multiple of " + + std::to_string(sizeof(CopyPack))); + + auto ptrs = buffers_.at(input); + int size_per_rank = size_bytes / sizeof(CopyPack); + int total_size = size_per_rank * world_size_; + int blocks = std::min(block_limit, (total_size + threads - 1) / threads); + +#define AG_CASE(ngpus) \ + case ngpus: \ + cross_device_all_gather<<>>( \ + ptrs, sg_, self_sg_, reinterpret_cast(output), rank_, \ + size_per_rank); \ + break; + + switch (world_size_) { + AG_CASE(2) + AG_CASE(4) + AG_CASE(6) + AG_CASE(8) + default: + throw std::runtime_error( + "custom allgather only supports num gpus in (2,4,6,8)"); + } +#undef AG_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input, + T* output, void* local_buffer, + void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes) { + if (size_bytes % sizeof(typename packed_t::P) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport allgather requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size_bytes / sizeof(typename packed_t::P); + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks = + (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportAgThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather, \ + ptrs, input, output, \ + reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + mnnvl_lamport_all_gather \ + <<>>( \ + ptrs, input, output, reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_AG_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_AG_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_AG_CASE(2) + MNNVL_LAMPORT_AG_CASE(4) + MNNVL_LAMPORT_AG_CASE(6) + MNNVL_LAMPORT_AG_CASE(8) + MNNVL_LAMPORT_AG_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_AG_CASE +#undef MNNVL_LAMPORT_AG_LAUNCH +} + +template +void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output, + int size, int threads, int block_limit) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0) + throw std::runtime_error( + "custom reduce-scatter requires each output shard byte size to be " + "a multiple of 16"); + + auto ptrs = buffers_.at(input); + int size_per_rank = size / packed_size / world_size_; + int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads); + +#define RS_CASE(ngpus) \ + case ngpus: \ + cross_device_reduce_scatter<<>>( \ + ptrs, sg_, self_sg_, output, rank_, size_per_rank); \ + break; + + switch (world_size_) { + RS_CASE(2) + RS_CASE(4) + RS_CASE(6) + RS_CASE(8) + default: + throw std::runtime_error( + "custom reduce-scatter only supports num gpus in (2,4,6,8)"); + } +#undef RS_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream, + T* input, T* output, + void* local_buffer, + uint32_t* epochs, int size, + int stage_size_bytes) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size / packed_size / world_size_; + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks_per_rank = + (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads; + int blocks = blocks_per_rank * world_size_; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportRsThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx( \ + &config, &mnnvl_lamport_reduce_scatter_kernel, ptrs, input, \ + output, epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + mnnvl_lamport_reduce_scatter_kernel \ + <<>>( \ + ptrs, input, output, epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_RS_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_RS_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_RS_CASE(2) + MNNVL_LAMPORT_RS_CASE(4) + MNNVL_LAMPORT_RS_CASE(6) + MNNVL_LAMPORT_RS_CASE(8) + MNNVL_LAMPORT_RS_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports num gpus in " + "(2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_RS_CASE +#undef MNNVL_LAMPORT_RS_LAUNCH +} + +} // namespace vllm + +using fptr_t = int64_t; +static_assert(sizeof(void*) == sizeof(fptr_t)); + +bool _is_weak_contiguous(torch::stable::Tensor& t); + +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size); +} + +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _local_buffer, + fptr_t _multicast_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto multicast_buffer = reinterpret_cast(_multicast_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports float32, float16 and " + "bfloat16"); + } +} + +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->reduce_scatter(stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), + inp.numel()); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } +#endif + default: + throw std::runtime_error( + "custom reduce-scatter only supports float32, float16 and bfloat16"); + } +} + +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t _local_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK(input_size <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, epochs, + inp.numel(), stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports float32, float16 and " + "bfloat16"); + } +} diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp new file mode 100644 index 000000000000..198b78d56dad --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp @@ -0,0 +1,29 @@ +#include "ops.h" +#include "core/registration.h" + +#include + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) { + custom_ag_rs.def( + "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int multicast_buffer, int epoch_buffer, int " + "stage_sz_bytes) -> ()"); + custom_ag_rs.def( + "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) { + custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather)); + custom_ag_rs.impl("mnnvl_lamport_all_gather", + TORCH_BOX(&mnnvl_lamport_all_gather)); + custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter)); + custom_ag_rs.impl("mnnvl_lamport_reduce_scatter", + TORCH_BOX(&mnnvl_lamport_reduce_scatter)); +} diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/libtorch_stable/custom_all_reduce.cu index 0f7f759949a8..07fca99ddd9e 100644 --- a/csrc/libtorch_stable/custom_all_reduce.cu +++ b/csrc/libtorch_stable/custom_all_reduce.cu @@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, torch::stable::Tensor& rank_data, int64_t rank, bool fully_connected) { int world_size = fake_ipc_ptrs.size(); - if (world_size > 8) - throw std::invalid_argument("world size > 8 is not supported"); + if (world_size > vllm::kMaxCustomCollectiveRanks) + throw std::invalid_argument("world size > 16 is not supported"); if (world_size % 2 != 0) throw std::invalid_argument("Odd num gpus is not supported for now"); if (rank < 0 || rank >= world_size) throw std::invalid_argument("invalid rank passed in"); - vllm::Signal* ipc_ptrs[8]; + vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks]; for (int i = 0; i < world_size; i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } @@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); } void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); - void* ipc_ptrs[8]; + void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks]; for (int i = 0; i < fake_ipc_ptrs.size(); i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu index 585004c047bf..d87b034be8cc 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu @@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( #endif } -template +template void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, - cudaStream_t const stream) { - constexpr int gemm_m = kHdOut; // 2112 - int const gemm_n = num_tokens; // 1-16 - constexpr int gemm_k = kHdIn; // 7168 + cudaStream_t const stream, bool enable_pdl) { + constexpr int gemm_m = kHdOut; + int const gemm_n = num_tokens; + constexpr int gemm_k = kHdIn; constexpr int batch_size = 1; std::swap(mat_a, mat_b); constexpr int tile_m = 16; - constexpr int tile_n = kTileN; // 8 or 16 - constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 + constexpr int tile_n = kTileN; + constexpr int tile_k = kTileK; constexpr int max_stage_cnt = 1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t)); constexpr int k_iter_cnt = gemm_k / tile_k; @@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, config.stream = stream; cudaLaunchAttribute attrs[1]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = + enable_pdl || getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; if (smem_bytes >= (48 * 1024)) { @@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, output, mat_a, mat_b, gemm_n); } -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); +template +void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b, + int num_tokens, cudaStream_t const stream, + bool enable_pdl) { + if (num_tokens <= 8) { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } else { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } +} void dsv3_fused_a_gemm(torch::stable::Tensor& output, torch::stable::Tensor const& mat_a, - torch::stable::Tensor const& mat_b) { + torch::stable::Tensor const& mat_b, bool enable_pdl) { STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2); int const num_tokens = mat_a.size(0); int const hd_in = mat_a.size(1); int const hd_out = mat_b.size(1); - constexpr int kHdIn = 7168; - constexpr int kHdOut = 2112; STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, "required 1 <= mat_a.shape[0] <= 16"); - STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168"); - STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112"); STD_TORCH_CHECK(output.size(0) == num_tokens, "required output.shape[0] == mat_a.shape[0]"); STD_TORCH_CHECK(output.size(1) == hd_out, "required output.shape[1] == mat_b.shape[1]"); + STD_TORCH_CHECK(mat_b.size(0) == hd_in, + "required mat_b.shape[0] == mat_a.shape[1]"); - STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); - STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor"); - STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); + STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() && + mat_a.get_device_index() == output.get_device_index(), + "mat_a, mat_b, and output must be on the same device"); + + // The kernels index global memory with raw pointers and packed strides, so + // reject any padded or transposed view rather than reading out of bounds. + STD_TORCH_CHECK( + mat_a.stride(0) == hd_in && mat_a.stride(1) == 1, + "mat_a must be a packed row-major [num_tokens, hd_in] tensor"); + STD_TORCH_CHECK( + output.stride(0) == hd_out && output.stride(1) == 1, + "output must be a packed row-major [num_tokens, hd_out] tensor"); + STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in, + "mat_b must be a packed column-major [hd_in, hd_out] tensor"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 && @@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); + auto* output_ptr = + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + auto const* mat_a_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); + auto const* mat_b_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); + +#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \ + if (hd_in == HD_IN && hd_out == HD_OUT) { \ + invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \ + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \ + return; \ } + + // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners + // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection. + DISPATCH_DSV3_SHAPE(7168, 1536) + DISPATCH_DSV3_SHAPE(7168, 2112) + DISPATCH_DSV3_SHAPE(1536, 2304) + DISPATCH_DSV3_SHAPE(1536, 4608) + DISPATCH_DSV3_SHAPE(7168, 3584) + DISPATCH_DSV3_SHAPE(768, 7168) + // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent + // because hd_in=2112 is not a multiple of any supported tile_k. + DISPATCH_DSV3_SHAPE(1536, 1152) + DISPATCH_DSV3_SHAPE(7168, 768) + DISPATCH_DSV3_SHAPE(7168, 3216) + DISPATCH_DSV3_SHAPE(7168, 4224) + +#ifdef VLLM_K3_BENCH_SHAPES + // The selector routes these shapes to CuTe or the default GEMM, so they are + // never reached in production. They are compiled only for offline + // DSV3-vs-CuTe benchmarking. + DISPATCH_DSV3_SHAPE(7168, 6288) + DISPATCH_DSV3_SHAPE(1536, 7168) + DISPATCH_DSV3_SHAPE(3584, 7168) + DISPATCH_DSV3_SHAPE(7168, 8448) + DISPATCH_DSV3_SHAPE(7168, 20480) + DISPATCH_DSV3_SHAPE(7168, 3072) + DISPATCH_DSV3_SHAPE(7168, 12448) + DISPATCH_DSV3_SHAPE(3072, 7168) + DISPATCH_DSV3_SHAPE(8448, 7168) + DISPATCH_DSV3_SHAPE(7168, 16896) + DISPATCH_DSV3_SHAPE(7168, 40960) +#endif + +#undef DISPATCH_DSV3_SHAPE + + if (hd_in == 128 && hd_out == 1536) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 128 && hd_out == 3072) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple + // of 256, so both need the 128 tile_k. + if (hd_in == 128 && hd_out == 768) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 384 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#ifdef VLLM_K3_BENCH_SHAPES + if (hd_in == 4224 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#endif + + STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 7bc435b8e0da..ae93266a4593 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -942,9 +942,10 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( +void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) + torch::stable::Tensor& q_out, // [N, q_head_padded, 512] torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8 torch::stable::Tensor const& slot_mapping, // [N] int64 torch::stable::Tensor const& position_ids, // [N] int64 @@ -970,8 +971,16 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(), "q_in and kv dtype must match"); + STD_TORCH_CHECK(q_out.device() == q_in.device() && q_out.is_contiguous(), + "q_out must be contiguous and on the same device as q_in"); + STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(), + "q_out dtype must match q_in"); STD_TORCH_CHECK(q_head_padded >= q_in.size(1), "q_head_padded must be >= q_in.size(1) (num_heads_q)"); + STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) && + q_out.size(1) == q_head_padded && + q_out.size(2) == q_in.size(2), + "q_out shape [N, q_head_padded, 512]"); STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte, "k_cache must be uint8"); STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, @@ -999,11 +1008,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q_in.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index()); - // Allocate the padded q output. The kernel writes every element (live - // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe. - auto q_out = torch::stable::new_empty( - q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); - VLLM_STABLE_DISPATCH_HALF_TYPES( q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { using qkv_scalar_t = scalar_t; @@ -1020,6 +1024,20 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( num_heads_q_padded, cache_block_size_i, kv_block_stride, stream); }); +} + +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size) { + auto q_out = torch::stable::new_empty( + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); + fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( + q_in, kv, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, + q_head_padded, eps, cache_block_size); return q_out; } diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu new file mode 100644 index 000000000000..b0d52e8b816b --- /dev/null +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -0,0 +1,1277 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Fused Kimi-K3 MLA prefill + decode epilogues with optional RoPE. + * + * Prefill: runs after the q_b_proj / kv_b_proj GEMMs, one launch per token + * slice. Decode: runs after BMM1 (q_nope x W_UK) right before forward_mqa, + * concatenating mqa_q = [ql_nope | q_pe] and inserting the latent cache + * (fused_kimi_k3_mla_decode_q_concat_kv_cache_{,fp8_,ds_mla_}insert). + * + * Prefill variants: + * + * bf16 (fused_kimi_k3_mla_key_concat_kv_cache_insert): + * - optional in-place q RoPE: rotate q[t, h, 128:192] + * - full key concat: k_out[t, h] = [k_nope[t, h] | k_pe[t]] (per head) + * - latent cache insert: cache[slot(t)] = [kv_c_normed[t] | k_pe[t]] + * (v is used as-is in bf16, so it is not touched here.) + * + * fp8 (fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert): + * - q_fp8[t, h] = quant(q[t, h], q_scale) + * - k_fp8[t, h] = quant([k_nope[t, h] | k_pe[t]], k_scale) + * - v_fp8[t, h] = quant(v[t, h], v_scale) + * - cache[slot(t)] = quant([kv_c_normed[t] | k_pe[t]], k_scale) + * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a + * separate cache scale). Per-tensor E4M3. + * + * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): + * - full key concat (bf16), and cache insert in DeepSeek's 656-byte + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic + * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel. + * + * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the + * producing GEMMs on sm_90+, and are structured after + * `fusedDeepseekV4FullCacheKernel`: one grid, one warp per (token, slot) with + * `slotsPerToken = num_heads + 1`. Slots [0, H) do the per-head work; the extra + * slot H does the per-token cache insert. All dims are multiples of 8, so bf16 + * copies move one uint4 (8 elems) per step and fp8 stores pack 8 elems into a + * uint2. + */ + +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + +#include "cuda_compat.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" + +#ifndef USE_ROCM + #include + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#else + #include + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#endif +#include +#include +#include + +#ifdef USE_ROCM +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { + #if defined(__gfx942__) + __hip_fp8_e4m3_fnuz fp8_val(val); + #else + __hip_fp8_e4m3 fp8_val(val); + #endif + return reinterpret_cast(fp8_val); +} +#endif + +namespace vllm { +namespace kimi_k3_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} +} // namespace + +// ──────────────────────────────────────────────────────────────────────────── +// Constants (Kimi-K3 MLA) +// ──────────────────────────────────────────────────────────────────────────── +constexpr int kKvLoraRank = 512; // L +constexpr int kQkNopeHeadDim = 128; // P +constexpr int kQkRopeHeadDim = 64; // R +constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192 +constexpr int kVHeadDim = 128; // V +constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576 +constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store + +#if defined(USE_ROCM) && defined(__gfx942__) +constexpr float kFp8Max = 224.0f; +#else +constexpr float kFp8Max = 448.0f; +#endif +// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu). +// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32 +// tile scales, [528,656) RoPE 64 bf16. +constexpr float kFp8ScaleDivisor = kFp8Max; + +// Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a +// uint4 (bf16); FP8=true decodes to fp32, scales by `scale_inv`, saturates to +// ±kFp8Max and packs into a uint2 of E4M3. +template +__device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, + float scale_inv, + const float* cos_sin = nullptr, + int rope_elem_base = 0) { + uint4 const v = *reinterpret_cast(src); + if constexpr (FP8 || APPLY_ROPE) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + // _typeConvert is unavailable on pre-Ampere. Kimi K3 uses + // bf16 inputs, so discard unsupported conversion paths in those builds. + if constexpr (std::is_same_v) { + return; + } else { +#endif + using Converter = vllm::_typeConvert; + auto const* p = + reinterpret_cast(&v); + float f[kVecElems]; +#pragma unroll + for (int i = 0; i < 4; i++) { + float2 x = Converter::convert(p[i]); + f[2 * i] = x.x; + f[2 * i + 1] = x.y; + } + if constexpr (APPLY_ROPE) { +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + int const pair_idx = rope_elem_base / 2 + i; + float const cos = static_cast(cos_sin[pair_idx]); + float const sin = + static_cast(cos_sin[pair_idx + kQkRopeHeadDim / 2]); + float const x = f[2 * i]; + float const y = f[2 * i + 1]; + f[2 * i] = x * cos - y * sin; + f[2 * i + 1] = x * sin + y * cos; + } + } + if constexpr (!FP8) { + uint4 out; + auto* o = reinterpret_cast(&out); +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1])); + } + *reinterpret_cast(dst) = out; + return; + } +#ifndef USE_ROCM + uint2 out; + auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out); + #pragma unroll + for (int i = 0; i < 4; i++) { + float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv); + s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max); + s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max); + o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3); + } + *reinterpret_cast(dst) = out; +#else + uint8_t out[kVecElems]; + #pragma unroll + for (int i = 0; i < kVecElems; i++) { + float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max); + out[i] = rocm_cvt_float_to_fp8_e4m3(s); + } + *reinterpret_cast(dst) = *reinterpret_cast(out); +#endif +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif + } else { + *reinterpret_cast(dst) = v; + } +} + +// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). +// FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size). +template +__device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope, + const scalar_t* k_pe, int laneId, + int dst_elem_size, float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, k_nope + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8( + d + e * dst_elem_size, k_pe + rope_e, scale_inv, cos_sin, rope_e); + } + } +} + +// Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias +// q (in-place); fp8 writes the quantized query directly to its output. +template +__device__ __forceinline__ void writePrefillQuery( + void* dst, const scalar_t* q, int laneId, int dst_elem_size, + float scale_inv, const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + if constexpr (FP8) { + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, q + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8(d + e * dst_elem_size, q + e, + scale_inv, cos_sin, rope_e); + } + } + } else if constexpr (APPLY_ROPE) { + for (int e = laneId * kVecElems; e < kQkRopeHeadDim; e += 32 * kVecElems) { + copyChunk8( + d + (kQkNopeHeadDim + e) * dst_elem_size, q + kQkNopeHeadDim + e, + scale_inv, cos_sin, e); + } + } +} + +// Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used +// for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache +// entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4. +template +__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512, + const scalar_t* b64, int laneId, + int dst_elem_size, + float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) { + if (e < kKvLoraRank) { + copyChunk8(d + e * dst_elem_size, a512 + e, scale_inv); + } else { + int const rope_e = e - kKvLoraRank; + copyChunk8(d + e * dst_elem_size, b64 + rope_e, + scale_inv, cos_sin, rope_e); + } + } +} + +// Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512 +// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at +// [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with +// concat_and_cache_ds_mla_kernel. +template +__device__ __forceinline__ void writeDsMlaCache( + uint8_t* row, const scalar_t* kvc, const scalar_t* pe, int laneId, + const float* cos_sin = nullptr) { + constexpr int kElemsPerLane = kKvLoraRank / 32; // 16 + int const tile = laneId >> 3; // 8 lanes per tile + scalar_t vals[kElemsPerLane]; + *reinterpret_cast(vals) = + *reinterpret_cast(kvc + laneId * kElemsPerLane); + *reinterpret_cast(vals + 8) = + *reinterpret_cast(kvc + laneId * kElemsPerLane + 8); + + float max_abs = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + max_abs = fmaxf(max_abs, fabsf(static_cast(vals[i]))); + } +#pragma unroll + for (int offset = 4; offset > 0; offset /= 2) { + max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8)); + } + float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + if ((laneId & 7) == 0) { + reinterpret_cast(row)[kKvLoraRank / 4 + tile] = tile_scale; + } + uint8_t res[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + res[i] = + fp8::scaled_convert( + vals[i], tile_scale); + } + *reinterpret_cast(row + laneId * kElemsPerLane) = + *reinterpret_cast(res); + scalar_t* row16 = reinterpret_cast(row); + scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2; + if constexpr (APPLY_ROPE) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + if constexpr (std::is_same_v) { + return; + } else { +#endif + using Converter = vllm::_typeConvert; + using packed_t = typename Converter::packed_hip_type; + packed_t const src = *reinterpret_cast(pe + laneId * 2); + float2 const xy = Converter::convert(src); + float const cos = static_cast(cos_sin[laneId]); + float const sin = + static_cast(cos_sin[laneId + kQkRopeHeadDim / 2]); + *reinterpret_cast(rope_dst) = Converter::convert( + make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos)); +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif + } else { + *reinterpret_cast(rope_dst) = + *reinterpret_cast(pe + laneId * 2); + } +} + +// ──────────────────────────────────────────────────────────────────────────── +// bf16 variant: optional q RoPE + full-key concat + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, scalar_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache, int const num_tokens, + int const num_heads, int const cache_block_size) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + scalar_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// fp8 variant: quant q / k / v + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( + const scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ v, int64_t const v_tok_stride, + int64_t const v_head_stride, uint8_t* __restrict__ q_fp8, + int64_t const qo_tok_stride, int64_t const qo_head_stride, + uint8_t* __restrict__ k_fp8, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8, + int64_t const vo_tok_stride, int64_t const vo_head_stride, + uint8_t* __restrict__ k_cache, int64_t const cache_block_stride, + int64_t const cache_token_stride, const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ k_scale_inv, + const float* __restrict__ v_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + int const h = slotIdx; + // q_fp8[t, h] = quant(q[t, h], q_scale) + float const qsi = __ldg(q_scale_inv); + const scalar_t* qh = q + tokenIdx * q_tok_stride + h * q_head_stride; + uint8_t* qo = q_fp8 + tokenIdx * qo_tok_stride + h * qo_head_stride; + writePrefillQuery(qo, qh, laneId, 1, qsi, + rope_cache); + // k_fp8[t, h] = quant([k_nope | k_pe], k_scale) + writeFullKey( + k_fp8 + tokenIdx * ko_tok_stride + h * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + h * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, __ldg(k_scale_inv), + rope_cache); + // v_fp8[t, h] = quant(v[t, h], v_scale) + float const vsi = __ldg(v_scale_inv); + const scalar_t* vh = v + tokenIdx * v_tok_stride + h * v_head_stride; + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + h * vo_head_stride; + for (int e = laneId * kVecElems; e < kVHeadDim; e += 32 * kVecElems) { + copyChunk8(vo + e, vh + e, vsi); + } + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + // The cache latent uses _k_scale (read back by decode / context); the + // attention key (k_fp8 above) uses its own k_scale. + float const ksi = __ldg(cache_scale_inv); + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert +// +// Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel: +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled +// [512, 528) 4 fp32 per-tile scales +// [528, 656) RoPE 64 vals as bf16 (unquantized) +// The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16) +// (tile = L>>3, absmax-reduced within its 8-lane tile group), then all 32 lanes +// write 2 RoPE bf16 each. +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + // Full key (bf16): k_out[t, h] = [k_nope[t, h] | k_pe[t]]. + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert, +// run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the +// plain per-tensor cache. (ds_mla cache uses the separate kernel below.) +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + void* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, void* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + constexpr int kMqElem = Q_FP8 ? 1 : sizeof(scalar_t); + constexpr int kCacheElem = KV_FP8 ? 1 : sizeof(scalar_t); + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(mqa_q) + + (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + kMqElem, qsi, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + float const ksi = KV_FP8 ? __ldg(cache_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(k_cache) + + (slot_id / cache_block_size * cache_block_stride + + slot_id % cache_block_size * cache_token_stride) * + kCacheElem, + kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride, + laneId, kCacheElem, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// Decode epilogue for fp8_ds_mla: concat mqa_q (bf16) + ds_mla cache insert. +template +__global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + scalar_t* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + writeLatent576( + mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + sizeof(scalar_t), 1.0f, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// PDL-aware launch of a (token, num_heads + 1)-warp grid. +template +static void launchPdl(KernelT kernel, int num_tokens, int num_heads, + cudaStream_t stream, Args... args) { + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int64_t const total_warps = + static_cast(num_tokens) * (num_heads + 1); + int const grid = + static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); +#ifndef USE_ROCM + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3(grid); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + cudaLaunchKernelEx(&config, kernel, args...); +#else + // clang-format off + // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the + // launch closer as ">>>". clang-format would otherwise re-split it (it does + // not parse the CUDA launch syntax in this libtorch_stable file). + kernel<<>>(args...); + // clang-format on +#endif +} + +void checkBfloat16Support(torch::headeronly::ScalarType dtype) { +#ifndef USE_ROCM + if (dtype == torch::headeronly::ScalarType::BFloat16) { + static int const sm_version = getSMVersion(); + STD_TORCH_CHECK( + sm_version >= 80, + "Kimi K3 fused MLA operations require sm_80+ (Ampere or newer); got " + "sm_", + sm_version); + } +#else + (void)dtype; +#endif +} + +} // namespace kimi_k3_fused_ops +} // namespace vllm + +// ──────────────────────────────────────────────────────────────────────────── +// Torch op wrappers +// ──────────────────────────────────────────────────────────────────────────── +namespace { +bool check_rope_inputs( + std::optional const& position_ids, + std::optional const& cos_sin_cache, + torch::stable::Tensor const& /*input*/, int64_t num_tokens) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(position_ids.has_value() == cos_sin_cache.has_value(), + "position_ids and cos_sin_cache must be provided together"); + if (!position_ids.has_value()) return false; + + auto const& positions = position_ids.value(); + auto const& rope_cache = cos_sin_cache.value(); + STD_TORCH_CHECK(positions.device().is_cuda() && positions.dim() == 1 && + positions.scalar_type() == ScalarType::Long && + positions.size(0) == num_tokens, + "position_ids must be int64 CUDA with shape [num_tokens]"); + STD_TORCH_CHECK(rope_cache.device().is_cuda() && rope_cache.dim() == 2 && + rope_cache.size(1) == 64 && rope_cache.stride(1) == 1 && + rope_cache.scalar_type() == ScalarType::Float, + "cos_sin_cache must have shape [max_position, 64], unit " + "last-dim stride, and be fp32 (RoPE math runs in fp32)"); + return true; +} +} // namespace + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] + torch::stable::Tensor const& k_pe, // [Tp, 64] + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192] CUDA"); + // k_pe is a strided view of the fused QKV-LoRA GEMM output; the kernel takes + // its row stride and reads the 64 cols contiguously, so unit last-dim stride + // (not full contiguity) is all that is required. + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride, CUDA"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && kv_c_normed.dim() == 2 && + kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous CUDA"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.dim() == 3 && k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous CUDA"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous CUDA"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK(q.scalar_type() == dt && k_pe.scalar_type() == dt && + kv_c_normed.scalar_type() == dt && + k_out.scalar_type() == dt && k_cache.scalar_type() == dt, + "all tensors must share k_nope's (bf16/fp16) dtype"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr, + num_tokens, num_heads, static_cast(cache_block_size)); + }; + if (apply_rope) { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.scalar_type() == dt && + q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.scalar_type() == dt && k_out.dim() == 3 && + k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous"); + // fp8_ds_mla entry is 656 bytes stored as uint8. + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, // [Tp, H, 192] bf16 + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor const& v, // [Tp, H, 128] bf16 + torch::stable::Tensor& q_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& k_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& v_fp8, // [Tp, H, 128] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& k_scale_inv, // scalar fp32 (1 / k scale) + torch::stable::Tensor const& v_scale_inv, // scalar fp32 (1 / v scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + auto check_in = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.scalar_type() == dt && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_in(q, 192, "q shape [Tp, H, 192]"); + check_in(k_nope, 128, "k_nope shape [Tp, H, 128]"); + check_in(v, 128, "v shape [Tp, H, 128]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + auto check_out = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.is_contiguous() && + t.scalar_type() == ScalarType::Float8_e4m3fn && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_out(q_fp8, 192, "q_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(k_fp8, 192, "k_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(v_fp8, 128, "v_fp8 shape [Tp, H, 128] fp8 contiguous"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(k_scale_inv, "k_scale_inv must be scalar float32 CUDA"); + check_scale(v_scale_inv, "v_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.const_data_ptr()), + q.stride(0), q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(v.const_data_ptr()), + v.stride(0), v.stride(1), + reinterpret_cast(q_fp8.mutable_data_ptr()), + q_fp8.stride(0), q_fp8.stride(1), + reinterpret_cast(k_fp8.mutable_data_ptr()), + k_fp8.stride(0), k_fp8.stride(1), + reinterpret_cast(v_fp8.mutable_data_ptr()), + v_fp8.stride(0), v_fp8.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + k_scale_inv.const_data_ptr(), + v_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } else { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } + }); +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue torch ops +// ──────────────────────────────────────────────────────────────────────────── +namespace { +// Shared shape checks for the decode ops. Verifies the query/latent inputs and +// slot_mapping; caller checks mqa_q / k_cache dtypes for its variant. +void check_decode_inputs(torch::stable::Tensor const& ql_nope, + torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& mqa_q, + torch::stable::Tensor const& slot_mapping) { + using torch::headeronly::ScalarType; + auto const dt = ql_nope.scalar_type(); + STD_TORCH_CHECK(ql_nope.device().is_cuda() && ql_nope.dim() == 3 && + ql_nope.size(2) == 512, + "ql_nope shape [B, H, 512] CUDA"); + STD_TORCH_CHECK(q_pe.device().is_cuda() && q_pe.scalar_type() == dt && + q_pe.dim() == 3 && q_pe.size(2) == 64, + "q_pe shape [B, H, 64]"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [B, 512] contiguous"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [B, 64], unit last-dim stride"); + STD_TORCH_CHECK(mqa_q.device().is_cuda() && mqa_q.is_contiguous() && + mqa_q.dim() == 3 && mqa_q.size(2) == 576, + "mqa_q shape [B, H, 576] contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); +} +} // namespace + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt && k_cache.scalar_type() == dt, + "mqa_q / k_cache must match ql_nope dtype (bf16)"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + nullptr, nullptr, num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == ScalarType::Float8_e4m3fn, + "mqa_q must be float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt, "mqa_q must be bf16 for ds_mla"); + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + kk3::checkBfloat16Support(dt); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(mqa_q.mutable_data_ptr()), + mqa_q.stride(0), mqa_q.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } + }); +} diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index d8460f032eb2..8769ba723bf9 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -269,6 +269,22 @@ __device__ __forceinline__ void storeElemsFp8( #endif } +// Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the +// inverse dequantization scale and converting it to E4M3. +template +__device__ __forceinline__ void storeScaledQElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane], + float const inv_scale) { + using Converter = vllm::_typeConvert; + float scaled[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + auto const rounded = Converter::convert(elems[i]); + scaled[i] = static_cast(rounded) * inv_scale; + } + storeElemsFp8(dst, scaled); +} + // ──────────────────────────────────────────────────────────────────────────── // Kernel // ──────────────────────────────────────────────────────────────────────────── @@ -297,6 +313,7 @@ template (store_ptr + dim_base, elems); } + if (isQ && q_fp8_out != nullptr) { + storeScaledQElemsFp8( + q_fp8_out + static_cast(tokenIdx) * nq * kHeadDim + + slot * kHeadDim + dim_base, + elems, q_fp8_inv_scale); + } } // ── Cache inserts (sparse serving only). ─────────────────────────────── @@ -493,13 +517,14 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( // ──────────────────────────────────────────────────────────────────────────── template void launchFusedMiniMaxM3( - scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w, - scalar_t const* k_norm_w, scalar_t const* iq_norm_w, - scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache, - int64_t const* positions, int64_t const* slot_mapping, - int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache, - float const eps, int const rotary_dim, int const num_tokens, int const nq, - int const nkv, int const niq, int const block_size, + scalar_t* qkv, scalar_t* q_out, uint8_t* q_fp8_out, void* index_q_out, + scalar_t const* q_norm_w, scalar_t const* k_norm_w, + scalar_t const* iq_norm_w, scalar_t const* ik_norm_w, + scalar_t const* cos_sin_cache, int64_t const* positions, + int64_t const* slot_mapping, int64_t const* index_slot_mapping, + cache_t* kv_cache, void* index_cache, float const eps, + float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens, + int const nq, int const nkv, int const niq, int const block_size, int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, int64_t const kv_s_dim, bool const has_index, bool const insert_kv, bool const process_index, bool const fp8_idx, cudaStream_t stream) { @@ -540,10 +565,11 @@ void launchFusedMiniMaxM3( fusedMiniMaxM3QNormRopeKVInsertKernel, \ - qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, k_norm_w, \ - iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \ - index_slot_mapping, kv_cache, reinterpret_cast(index_cache), \ - eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ + qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ + q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + slot_mapping, index_slot_mapping, kv_cache, \ + reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ + rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ kv_s_head, kv_s_token, kv_s_dim) #else // ROCm: standard kernel launch syntax (no PDL/stream serialization). @@ -552,12 +578,12 @@ void launchFusedMiniMaxM3( fusedMiniMaxM3QNormRopeKVInsertKernel< \ scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \ FP8><<>>( \ - qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, \ - k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ + q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ - reinterpret_cast(index_cache), eps, rotary_dim, num_tokens, \ - nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token, \ - kv_s_dim) + reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ + rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ + kv_s_head, kv_s_token, kv_s_dim) // clang-format on #endif @@ -599,6 +625,9 @@ void launchFusedMiniMaxM3( vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3( \ reinterpret_cast(qkv.data_ptr()), \ q_out.has_value() ? reinterpret_cast(q_out->data_ptr()) : nullptr, \ + q_fp8_out.has_value() \ + ? reinterpret_cast(q_fp8_out->data_ptr()) \ + : nullptr, \ index_q_out.has_value() \ ? reinterpret_cast(index_q_out->data_ptr()) \ : nullptr, \ @@ -622,10 +651,10 @@ void launchFusedMiniMaxM3( (insert_kv && process_index) \ ? reinterpret_cast(index_cache->data_ptr()) \ : nullptr, \ - static_cast(eps), static_cast(rotary_dim), num_tokens, nq, \ - nkv, niq, static_cast(block_size), kv_s_block, kv_s_head, \ - kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx, \ - stream) + static_cast(eps), 1.0f / static_cast(q_fp8_scale), \ + static_cast(rotary_dim), num_tokens, nq, nkv, niq, \ + static_cast(block_size), kv_s_block, kv_s_head, kv_s_token, \ + kv_s_dim, has_index, insert_kv, process_index, fp8_idx, stream) // clang-format on // ──────────────────────────────────────────────────────────────────────────── @@ -649,7 +678,10 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional q_out, // [N, nq*128] contiguous std::optional index_q_out, // [N, niq*128] contiguous - const std::string& kv_cache_dtype, bool skip_index_branch) { + const std::string& kv_cache_dtype, bool skip_index_branch, + std::optional + q_fp8_out, // [N, nq*128] contiguous E4M3 + double q_fp8_scale) { STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(), "qkv must be contiguous CUDA"); STD_TORCH_CHECK( @@ -780,6 +812,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert( q_out->numel() == static_cast(num_tokens) * nq * kHeadDim, "q_out must have num_tokens * num_heads * 128 elements"); } + if (q_fp8_out.has_value()) { + STD_TORCH_CHECK(q_fp8_out->is_cuda() && q_fp8_out->is_contiguous() && + q_fp8_out->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn, + "q_fp8_out must be a contiguous CUDA fp8 e4m3 tensor"); + STD_TORCH_CHECK( + q_fp8_out->numel() == static_cast(num_tokens) * nq * kHeadDim, + "q_fp8_out must have num_tokens * num_heads * 128 elements"); + STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0, + "q_fp8_scale must be finite and positive"); + } if (index_q_out.has_value()) { STD_TORCH_CHECK(process_index, "index_q_out requires index branch processing"); diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu new file mode 100644 index 000000000000..eb4dcf6bf5a2 --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu @@ -0,0 +1,954 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +// Production AttnRes forward for Blackwell (SM100). +// +// Warp-specialized online softmax + residual + RMSNorm: +// - 1 producer warp issues cp.async.bulk row loads into shared memory. +// - 8 consumer warps compute reductions and output. +// - Q=res_weight*rms_weight remains in registers across persistent tokens. +// - V rows are converted once and cached as FP32 in TMEM between passes. +// +// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major +// block residual storage. + +#include "../torch_utils.h" + +#include +#include +#include +#include +#include + +using bf16_t = __nv_bfloat16; + +namespace sm100 { +namespace fwd_prod_v2 { + +constexpr int K_TILE = 1024; +constexpr int N_CHUNK_DEFAULT = 4; +constexpr int CHUNK_DEPTH = 2; +constexpr int BLK = 288; // 1 producer warp + 8 consumer warps +constexpr int CONSUMER_THREADS = BLK - 32; // 256 +constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32; +constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups +constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS; +constexpr int FIRST_USER_NAMED_BARRIER = 8; + +__device__ __forceinline__ const bf16_t* residual_addr( + const bf16_t* block_res, const bf16_t* layer_res, int source, int N, + int token, int block_stride_m, int block_stride_r, int H) { + if (source < N - 1) { + return block_res + static_cast(token) * block_stride_m + + source * block_stride_r; + } + return layer_res + static_cast(token) * H; +} + +__device__ __forceinline__ uint32_t elect_one_sync() { + uint32_t pred = 0; + uint32_t laneid = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %2;\n" + "@%%px mov.s32 %1, 1;\n" + " mov.s32 %0, %%rx;\n" + "}\n" + : "+r"(laneid), "+r"(pred) + : "r"(0xffffffff)); + return pred; +} + +__device__ __forceinline__ void mbarrier_init(uint64_t& barrier, + int thread_count) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr), + "r"(thread_count)); +} + +__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier, + uint32_t bytes) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"( + barrier_addr), + "r"(bytes)); +} + +__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .pred p;\n" + "WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n" + "@p bra DONE;\n" + "bra WAIT;\n" + "DONE:\n" + "}\n" ::"r"(barrier_addr), + "r"(phase)); +} + +__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .b64 state;\n" + "mbarrier.arrive.shared::cta.b64 state, [%0];\n" + "}\n" ::"r"(barrier_addr)); +} + +__device__ __forceinline__ void fence_mbarrier_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads, + uint32_t user_barrier_id) { + asm volatile( + "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER), + "r"(num_threads) + : "memory"); +} + +__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) { + uint32_t const dst_addr = + static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"( + dst_addr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) { + asm volatile( + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_release_allocation_lock() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +__device__ __forceinline__ void tmem_store_wait() { + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); +} + +template +__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) { + uint32_t* values = reinterpret_cast(dst); + if constexpr (N == 8) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]), + "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7]) + : "r"(src_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]) + : "r"(src_addr)); + } +} + +template +__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) { + uint32_t* values = reinterpret_cast(src); + if constexpr (N == 8) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]), + "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr)); + } +} + +__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) { + float2 result; + asm volatile("add.rn.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b))); + return result; +} + +__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) { + float2 result; + asm volatile("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b))); + return result; +} + +__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b, + const float2& c) { + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +} + +template +struct FwdSmemPlan { + alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; + alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; + alignas(16) uint64_t bar_output_norm_ready; + alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; + uint32_t tmem_base; +}; + +__device__ __forceinline__ void cp_async_bulk(void* smem_dst, + const void* gmem_src, int bytes, + uint64_t& mbar) { + uint32_t const s = static_cast(__cvta_generic_to_shared(smem_dst)); + uint32_t const m = static_cast(__cvta_generic_to_shared(&mbar)); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], " + "[%1], %2, [%3];\n" ::"r"(s), + "l"(gmem_src), "r"(bytes), "r"(m) + : "memory"); +} + +template +__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( + const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res, + const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w, + const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T, + int B, int block_stride_m, int block_stride_r, float rms_eps, + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + constexpr float LOG2_E = 1.4426950408889634f; + constexpr int N_CHUNK = NC; + // The two-source specialization only consumes half of the TMEM columns. + constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256; + constexpr int NUM_BUFS = CHUNK_DEPTH * NC; + constexpr int NHT = H / K_TILE; + constexpr int SLICES_PER_GROUP = + (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS; + constexpr int VEC = 8; + constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC; + constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC; + constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; + static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC); + static_assert(H >= 4096 && H <= 8192); + static_assert(H % K_TILE == 0); + + const int tid = threadIdx.x; + const int wid = tid >> 5; + const int lane = tid & 31; + const int TB = T * B; + const int num_ctas = gridDim.x; + const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; + + const int comp_wid = wid - 1; + const int comp_tid = tid - 32; + const int group = (comp_wid >= 4) ? 1 : 0; + const int ct_in_group = + (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; + const int k_local = ct_in_group * VEC; + + constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); + constexpr size_t DELTA_BYTES = + HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0; + constexpr size_t OUTPUT_NORM_BYTES = + OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0; + extern __shared__ __align__(16) char smem_raw[]; + bf16_t* v_bufs = reinterpret_cast(smem_raw); // [NUM_BUFS][H] + bf16_t* delta_bufs = reinterpret_cast(smem_raw + V_BYTES); + bf16_t* output_norm_buf = + reinterpret_cast(smem_raw + V_BYTES + DELTA_BYTES); + FwdSmemPlan& plan = *reinterpret_cast*>( + smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES); + + auto slot_of = [](long long gci, int n) { + return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n; + }; + auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); }; + auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; }; + auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* { + return delta_bufs + chunk_slot * H; + }; + + if (wid == 0 && elect_one_sync()) { + #pragma unroll + for (int i = 0; i < CHUNK_DEPTH; i++) { + mbarrier_init(plan.bar_ready[i], 1); + mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS); + } + if constexpr (OUTPUT_NORM_IN_SMEM) { + mbarrier_init(plan.bar_output_norm_ready, 1); + } + fence_mbarrier_init(); + } + + // gdc wait BEFORE tmem alloc + cudaGridDependencySynchronize(); + + if (wid == 1) { + tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base); + if constexpr (RELEASE_TMEM) { + tmem_release_allocation_lock(); + } + } + __syncthreads(); + + if constexpr (OUTPUT_NORM_IN_SMEM) { + if (wid == 0 && elect_one_sync()) { + mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t)); + cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t), + plan.bar_output_norm_ready); + } + } + + const uint32_t my_v_tmem = + comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0; + float q_cache[ACC_PER_THREAD]; + if (comp_tid >= 0) { + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + #pragma unroll + for (int j = 0; j < 4; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + #pragma unroll + for (int j = 0; j < VEC; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + } + } + + if (wid == 0) { + if (elect_one_sync()) { + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + const int t = tb / B; + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pc = phase_of(gci); + mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1); + int transaction_bytes = an * H * (int)sizeof(bf16_t); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + transaction_bytes += H * (int)sizeof(bf16_t); + } + } + mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes); + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n >= an) continue; + int slot = slot_of(gci, n); + const bf16_t* src = + residual_addr(block_res, layer_res, ns + n, N, t, + block_stride_m, block_stride_r, H); + cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + cp_async_bulk(delta_buf_ptr(chunk_slot), + delta + (long long)tb * H, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + } + } + } + } + } else { + float acc32[ACC_PER_THREAD] = {}; + float eps_cache; + asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps)); + + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + float m_running = -FLT_MAX; + float s_running = 0.f; + #pragma unroll + for (int i = 0; i < ACC_PER_THREAD; i++) { + acc32[i] = 0.f; + } + + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pr = phase_of(gci); + mbarrier_wait(plan.bar_ready[chunk_slot], pr); + + float2 sq_local[N_CHUNK] = {}; + float2 dot_local[N_CHUNK] = {}; + + auto pass_A_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = + 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + const float* qv = &q_cache[si * VEC]; + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int2 vp = + *reinterpret_cast(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = + delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < 2; j++) { + auto delta2 = *reinterpret_cast( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[2] = {__bfloat1622float2(v2[0]), + __bfloat1622float2(v2[1])}; + tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); + sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); + dot_local[n] = + float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); + dot_local[n] = + float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + const float* qv = &q_cache[si * VEC]; + + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int4 vp = *reinterpret_cast(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + auto delta2 = *reinterpret_cast( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[4] = { + __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]), + __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])}; + tmem_store(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + sq_local[n] = float2_fma(f[j], f[j], sq_local[n]); + dot_local[n] = float2_fma( + f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]); + } + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_A_body(std::integral_constant{}); + break; + case 3: + pass_A_body(std::integral_constant{}); + break; + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_A_body(std::integral_constant{}); + break; + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } + if (lane == 0) { + mbarrier_arrive(plan.bar_consumed[chunk_slot]); + } + tmem_store_wait(); + + float2 reduce_pair[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y, + dot_local[n].x + dot_local[n].y); + } + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + uint64_t packed = reinterpret_cast(reduce_pair[n]); + packed = __shfl_xor_sync(0xffffffff, packed, offset); + float2 other = reinterpret_cast(packed); + reduce_pair[n] = float2_add(reduce_pair[n], other); + } + } + if (lane == 0) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + plan.ws_stats[comp_wid][n] = reduce_pair[n]; + } + } + named_barrier_sync(CONSUMER_THREADS, 0); + + float local_rsig = 0.f; + float local_logit = 0.f; + int stat_n = lane / CONSUMER_WARPS; + int stat_w = lane % CONSUMER_WARPS; + float2 totals = {}; + if (stat_n < N_CHUNK) { + totals = plan.ws_stats[stat_w][stat_n]; + } + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + totals.x += + __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS); + totals.y += + __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS); + } + if (stat_n < N_CHUNK && stat_w == 0) { + local_rsig = rsqrtf(totals.x / H + eps_cache); + local_logit = totals.y * local_rsig; + } + float logit_n[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS); + } + + float m_chunk = -FLT_MAX; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]); + } + float m_new = fmaxf(m_running, m_chunk); + float corr = exp2f((m_running - m_new) * LOG2_E); + float w_n[N_CHUNK] = {}; + float w_sum = 0.f; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) { + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); + w_sum += w_n[n]; + } + } + + auto pass_B_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + float2 corr2 = make_float2(corr, corr); + float2 a[2]; + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, + f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + float2 corr2 = make_float2(corr, corr); + float2 a[VEC / 2]; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][VEC / 2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_B_body(std::integral_constant{}); + break; + case 3: + pass_B_body(std::integral_constant{}); + break; + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_B_body(std::integral_constant{}); + break; + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } + + s_running = s_running * corr + w_sum; + m_running = m_new; + } + + float inv_s = 1.f / s_running; + bf16_t* out_ptr = output + (long long)tb * H; + float2 output_sq_pair = {}; + // When output RMSNorm is fused, the softmax denominator cancels: + // (acc / s) * rsqrt(mean((acc / s)^2) + eps) + // = acc * rsqrt(mean(acc^2) + eps * s^2). + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast(out_ptr + h_base) = packed; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = + make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast(out_ptr + h_base) = packed; + } + } + + if constexpr (HAS_OUTPUT_NORM) { + if constexpr (OUTPUT_NORM_IN_SMEM) { + // The immutable weight copy is acquired once, at its first use. + if (tb == blockIdx.x) { + mbarrier_wait(plan.bar_output_norm_ready, 0); + } + } + float output_sq = output_sq_pair.x + output_sq_pair.y; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset); + } + if (lane == 0) { + plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f); + } + named_barrier_sync(CONSUMER_THREADS, 0); + float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f; + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + total_sq += + __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS); + } + if (lane == 0) { + total_sq = + rsqrtf(total_sq / H + output_norm_eps * s_running * s_running); + } + float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0); + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* values = reinterpret_cast(&packed); + #pragma unroll + for (int j = 0; j < 4; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = __float2bfloat16(acc32[si * VEC + j] * + output_rsigma * weight); + } + *reinterpret_cast(out_ptr + h_base) = packed; + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* values = reinterpret_cast(&packed); + #pragma unroll + for (int j = 0; j < VEC; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = + __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); + } + *reinterpret_cast(out_ptr + h_base) = packed; + } + } + } + } + + cudaTriggerProgrammaticLaunchCompletion(); + __syncthreads(); + if (wid == 1) { + tmem_free(plan.tmem_base, TMEM_COLS_ALLOC); + } +#else + if (threadIdx.x == 0) { + printf("attn_res_fwd_online_v2_kernel requires sm_10x\n"); + } +#endif +} + +template +static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual, + const bf16_t* delta, const bf16_t* res_weight, + const bf16_t* rms_weight, bf16_t* output, int N, int T, + int B, float rms_eps, int num_sm, cudaStream_t stream, + const bf16_t* output_norm_weight = nullptr, + float output_norm_eps = 0.f, int block_stride_m = 0, + int block_stride_r = 0) { + constexpr size_t smem_size = + ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) + + (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) + + sizeof(FwdSmemPlan) + 15) & + ~size_t(15); + auto kernel = + &attn_res_fwd_online_v2_kernel; + static bool attrs_set = false; + if (!attrs_set) { + if (smem_size > 48 * 1024) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + } + attrs_set = true; + } + int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = BLK; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta, + res_weight, rms_weight, output, N, T, B, block_stride_m, + block_stride_r, rms_eps, output_norm_weight, + output_norm_eps); +} + +} // namespace fwd_prod_v2 +} // namespace sm100 + +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps) { + int const num_tokens = static_cast(prefix.size(0)); + int const device = prefix.get_device_index(); + torch::stable::accelerator::DeviceGuard const device_guard(device); + cudaDeviceProp const* properties = get_device_prop(); + STD_TORCH_CHECK(properties->major == 10, + "Kimi K3 AttnRes requires the SM100 family"); + + using namespace sm100::fwd_prod_v2; + // Two-source chunks and two resident CTAs are beneficial once setup is + // amortized by the long, full eight-block prefill workload. + if (num_blocks == 8 && num_tokens >= 4096) { + launch_fwd<7168, 2, true, true, true, true>( + static_cast(blocks.data_ptr()), + static_cast(prefix.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(qk_weight.data_ptr()), + static_cast(norm_weight.data_ptr()), + static_cast(output.data_ptr()), + static_cast(num_blocks) + 1, num_tokens, 1, + static_cast(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast(output_norm_weight.data_ptr()), + static_cast(output_norm_eps), static_cast(blocks.stride(0)), + static_cast(blocks.stride(1))); + } else { + launch_fwd<7168, 4, false, true, true, true>( + static_cast(blocks.data_ptr()), + static_cast(prefix.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(qk_weight.data_ptr()), + static_cast(norm_weight.data_ptr()), + static_cast(output.data_ptr()), + static_cast(num_blocks) + 1, num_tokens, 1, + static_cast(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast(output_norm_weight.data_ptr()), + static_cast(output_norm_eps), static_cast(blocks.stride(0)), + static_cast(blocks.stride(1))); + } + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu new file mode 100644 index 000000000000..0badeb9b772c --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu @@ -0,0 +1,1130 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include +#include +#include +#include +#include + +#include "../torch_utils.h" +#include "../../cuda_compat.h" + +namespace { + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; +constexpr int kConvStateWidth = kKernelWidth - 1; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; + +struct KdaDecodeStrides { + int64_t x_row; + int64_t beta_row; + int64_t onorm_row; + int64_t conv_slot; + int64_t state_slot; +}; + +__device__ __forceinline__ float bf16_load(const __nv_bfloat16* ptr, + int64_t idx) { + return __bfloat162float(ptr[idx]); +} + +__device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) { + return ptr[idx]; +} + +template +__device__ __forceinline__ float conv_weight_load(const float* ptr, int channel, + int width) { + return ptr[width * kChannels + channel]; +} + +__device__ __forceinline__ __nv_bfloat16 bf16_store(float value) { + return __float2bfloat16(value); +} + +template +__device__ __forceinline__ void store_state_float4(float* ptr, float4 value) { + if constexpr (kUseCacheGlobalStore) { + __stcg(reinterpret_cast(ptr), value); + } else { + *reinterpret_cast(ptr) = value; + } +} + +__device__ __forceinline__ float sigmoid_fast(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float silu_fast(float x) { + return x * sigmoid_fast(x); +} + +__device__ __forceinline__ float softplus_fast(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} + +__device__ __forceinline__ float warp_reduce_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +__device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr, + const float* gmem_ptr) { + uint32_t smem_addr = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem_addr), "l"(gmem_ptr)); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::); +} + +__device__ __forceinline__ void cp_async_wait_all() { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +} + +__device__ __forceinline__ void cp_async_wait_group_1() { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_state_chunk_for(float* s_state, + const float* state, + int slot, int i_hv, + int HV, int chunk) { + constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; + const int tid = threadIdx.x; + const int stage = chunk & 1; + const int v_base = chunk * kChunkV; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) { + const int elem = linear4 * 4; + const int row = elem / kDimK; + const int k = elem - row * kDimK; + float* dst = s_state + (stage * kChunkV + row) * kDimK + k; + const float* src = + state + ((slot * HV + i_hv) * kDimV + v_base + row) * kDimK + k; + cp_async_cg_16b(dst, src); + } + cp_async_commit(); +} + +__device__ __forceinline__ void cp_async_state_chunk(float* s_state, + const float* state, + int slot, int i_hv, int HV, + int chunk) { + cp_async_state_chunk_for(s_state, state, slot, i_hv, HV, chunk); +} + +__device__ __forceinline__ float block_reduce_sum(float value, float* scratch) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + float warp_total = warp_reduce_sum(value); + if (lane == 0) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +struct Sum2 { + float x; + float y; +}; + +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffffu, x, offset); + y += __shfl_xor_sync(0xffffffffu, y, offset); + } + return {x, y}; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + const float warp_x = warp_reduce_sum(x); + const float warp_y = warp_reduce_sum(y); + if (lane == 0) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +__device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y, + float* scratch) { + return block_reduce_sum2_for(x, y, scratch); +} + +template +__device__ __forceinline__ float block_reduce_sum_active_for(float value, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_total = 0.0f; + if (warp < kReduceWarps) { + warp_total = warp_reduce_sum(value); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_x = 0.0f; + float warp_y = 0.0f; + if (warp < kReduceWarps) { + warp_x = warp_reduce_sum(x); + warp_y = warp_reduce_sum(y); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +template +__global__ +__launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( + const __nv_bfloat16* __restrict__ x_q, + const __nv_bfloat16* __restrict__ x_k, + const __nv_bfloat16* __restrict__ x_v, const float* __restrict__ w_q_t, + const float* __restrict__ w_k_t, const float* __restrict__ w_v_t, + const float* __restrict__ bias_q, const float* __restrict__ bias_k, + const float* __restrict__ bias_v, __nv_bfloat16* __restrict__ cs_q, + __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, + const float* __restrict__ a_log, const __nv_bfloat16* __restrict__ g, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ onorm_g, + const float* __restrict__ onorm_weight, + const int* __restrict__ ssm_state_indices, + const int* __restrict__ cu_seqlens, float* __restrict__ state, + __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif + if constexpr (kUseStaticDecodeLayout) { + if constexpr (kUseHeadGrid) { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } else { + const int nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } else { + const int nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + const int hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } + + constexpr int kLocalDim = kFixedHeads * kDimK; + constexpr int kPackedDim = 3 * kLocalDim; + const int hk_off = i_h * kDimK; + const int hv_off = i_hv * kDimV; + constexpr int hv_count = kFixedValueHeads; + float* const state_for_slot = state + slot * strides.state_slot; + const int64_t conv_slot_offset = slot * strides.conv_slot; + __nv_bfloat16* const cs_q_for_slot = cs_q + conv_slot_offset; + __nv_bfloat16* const cs_k_for_slot = cs_k + conv_slot_offset; + __nv_bfloat16* const cs_v_for_slot = cs_v + conv_slot_offset; + + __shared__ float s_state[2][kChunkV][kDimK]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, 0); + + if constexpr (kUpdateConvState) { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim]; + const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim]; + q_acc += __bfloat162float(q_state) * + conv_weight_load(w_q_t, hk, w); + k_acc += __bfloat162float(k_state) * + conv_weight_load(w_k_t, hk, w); + if (w == 1) { + q_shift0 = q_state; + k_shift0 = k_state; + } else if (w == 2) { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += __bfloat162float(k_new) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + cs_q_for_slot[hk] = q_shift0; + cs_q_for_slot[hk + kPackedDim] = q_shift1; + cs_q_for_slot[hk + 2 * kPackedDim] = q_new; + cs_k_for_slot[hk] = k_shift0; + cs_k_for_slot[hk + kPackedDim] = k_shift1; + cs_k_for_slot[hk + 2 * kPackedDim] = k_new; + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } else { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hk + w * kPackedDim; + q_acc += bf16_load(cs_q_for_slot, cs_idx) * + conv_weight_load(w_q_t, hk, w); + k_acc += bf16_load(cs_k_for_slot, cs_idx) * + conv_weight_load(w_k_t, hk, w); + } + q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim]; + v_acc += __bfloat162float(v_state) * + conv_weight_load(w_v_t, hvv, w); + if (w == 1) { + v_shift0 = v_state; + } else if (w == 2) { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + cs_v_for_slot[hvv] = v_shift0; + cs_v_for_slot[hvv + kPackedDim] = v_shift1; + cs_v_for_slot[hvv + 2 * kPackedDim] = v_new; + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } else { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hvv + w * kPackedDim; + v_acc += bf16_load(cs_v_for_slot, cs_idx) * + conv_weight_load(w_v_t, hvv, w); + } + v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) { + const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv); + if constexpr (kApplyBetaSigmoid) { + s_beta = sigmoid_fast(beta_raw); + } else { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, + 1); + } + + const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } else { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + const int k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) { + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + if (chunk + 1 < kNumChunks) { + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + } else { + cp_async_wait_all(); + } + if constexpr (!kSkipWarpSync) { + __syncwarp(); + } + + if constexpr (!kPrefetchNextStateChunk) { + if (chunk + 1 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 1); + } + } + +#pragma unroll + for (int row = 0; row < kRowsPerWarp; row += 2) { + const int v_row_a = warp + row * kWarps; + const int v_row_b = warp + (row + 1) * kWarps; + const int v0 = chunk * kChunkV + v_row_a; + const int v1 = chunk * kChunkV + v_row_b; + float h_a_vals[4]; + float h_b_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + + const float4 raw_h_a = *reinterpret_cast( + &s_state[chunk & 1][v_row_a][k_base]); + const float4 raw_h_b = *reinterpret_cast( + &s_state[chunk & 1][v_row_b][k_base]); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + + const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); + const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta; + const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + const int state_idx_a = (i_hv * kDimV + v0) * kDimK + k_base; + const int state_idx_b = (i_hv * kDimV + v1) * kDimK + k_base; + const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + if constexpr (kComputeOutputBeforeStore) { + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + } else { + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + } + + const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); + if (lane == 0) { + s_o[v0] = dot_hq.x; + s_o[v1] = dot_hq.y; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) { + o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; + } + } + } + + if constexpr (kPrefetchNextStateChunk) { + if (chunk + 2 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 2); + } + } + } + __syncthreads(); + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + if constexpr (kApplyOnorm) { + if constexpr (kAccumulateOnormSumsq) { + if (lane == 0) { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = warp_reduce_sum(total_sumsq); + if (lane == 0) { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float raw_o = s_o[tid]; + const float rstd = + rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } else { + const float raw_o = tid < kDimV ? s_o[tid] : 0.0f; + const float o_sq = raw_o * raw_o; + float sumsq; + if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) { + sumsq = block_reduce_sum_active_for(o_sq, s_reduce); + } else { + sumsq = block_reduce_sum(o_sq, s_reduce); + } + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float rstd = + rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } else { + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +template +void launch_kda_decode_many_heads_raw( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, float lower_bound, float scale, + float onorm_eps, KdaDecodeStrides strides, cudaStream_t stream) { + auto kernel = &kda_decode_fusion_many_heads_kernel< + kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false, + false, true, true, true, kUpdateConvState, kUseLowerBound, + kApplyBetaSigmoid>; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(B, kHeads); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, + reinterpret_cast(x_q), + reinterpret_cast(x_k), + reinterpret_cast(x_v), + reinterpret_cast(w_q_t), + reinterpret_cast(w_k_t), + reinterpret_cast(w_v_t), + reinterpret_cast(bias_q), + reinterpret_cast(bias_k), + reinterpret_cast(bias_v), + reinterpret_cast<__nv_bfloat16*>(cs_q), + reinterpret_cast<__nv_bfloat16*>(cs_k), + reinterpret_cast<__nv_bfloat16*>(cs_v), a_log, + reinterpret_cast(g), dt_bias, + reinterpret_cast(beta), + reinterpret_cast(onorm_g), + onorm_weight, ssm_state_indices, cu_seqlens, state, + reinterpret_cast<__nv_bfloat16*>(out), B, H, HV, + lower_bound, scale, onorm_eps, strides); +} + +template +void launch_kda_decode_many_heads_selected( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides, + cudaStream_t stream) { +#define LAUNCH_KDA_DECODE(NUM_HEADS) \ + do { \ + if (update_conv_cache) { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } else { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } \ + } while (false) + + switch (H) { + case 12: + LAUNCH_KDA_DECODE(12); + break; + case 24: + LAUNCH_KDA_DECODE(24); + break; + case 48: + LAUNCH_KDA_DECODE(48); + break; + case 96: + LAUNCH_KDA_DECODE(96); + break; + default: + STD_TORCH_CHECK(false, "Unsupported number of heads: ", H); + } + +#undef LAUNCH_KDA_DECODE +} + +struct KdaDecodeLaunchParams { + const void* x_q; + const void* x_k; + const void* x_v; + const void* w_q_t; + const void* w_k_t; + const void* w_v_t; + const void* bias_q; + const void* bias_k; + const void* bias_v; + void* cs_q; + void* cs_k; + void* cs_v; + const float* a_log; + const void* g; + const float* dt_bias; + const void* beta; + const void* onorm_g; + const float* onorm_weight; + const int* ssm_state_indices; + const int* cu_seqlens; + float* state; + void* out; + int B; + int H; + int HV; + bool update_conv_cache; + float lower_bound; + float scale; + float onorm_eps; + KdaDecodeStrides strides; + cudaStream_t stream; +}; + +template +void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) { + launch_kda_decode_many_heads_selected( + p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, + p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale, + p.onorm_eps, p.strides, p.stream); +} + +template +void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p, + bool apply_beta_sigmoid) { + if (apply_beta_sigmoid) { + launch_kda_decode_selected_backend(p); + } else { + launch_kda_decode_selected_backend(p); + } +} + +template +void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p, + bool use_lower_bound, bool apply_beta_sigmoid) { + if (use_lower_bound) { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } else { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } +} + +void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p, + bool apply_onorm, bool use_lower_bound, + bool apply_beta_sigmoid) { + if (apply_onorm) { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } else { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } +} + +} // namespace + +extern "C" void launch_kda_decode_many_heads_cuda( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool apply_onorm, bool update_conv_cache, + bool use_lower_bound, bool apply_beta_sigmoid, float lower_bound, + float scale, float onorm_eps, const int64_t* raw_strides, + cudaStream_t stream) { + const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2], + raw_strides[3], raw_strides[4]}; + const KdaDecodeLaunchParams params{x_q, + x_k, + x_v, + w_q_t, + w_k_t, + w_v_t, + bias_q, + bias_k, + bias_v, + cs_q, + cs_k, + cs_v, + a_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + ssm_state_indices, + cu_seqlens, + state, + out, + B, + H, + HV, + update_conv_cache, + lower_bound, + scale, + onorm_eps, + strides, + stream}; + dispatch_kda_decode_features(params, apply_onorm, use_lower_bound, + apply_beta_sigmoid); +} + +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps) { + using torch::headeronly::ScalarType; + constexpr int kHeadDim = 128; + constexpr int kConvWidth = 4; + + STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16, + "x must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float, + "weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16, + "conv_state must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16, + "raw_g must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16, + "raw_beta must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, + "A_log must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float, + "dt_bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float, + "state must be a CUDA float32 tensor"); + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, + "out must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, + "state_indices must be a CUDA int32 tensor"); + + STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]"); + int const batch_size = static_cast(x.size(0)); + int64_t const qkv_width = x.size(1); + STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0, + "x must have shape [B, 3 * H * 128]"); + int64_t const num_heads = qkv_width / (3 * kHeadDim); + STD_TORCH_CHECK( + num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96, + "H must be 12, 24, 48, or 96, got ", num_heads); + STD_TORCH_CHECK(batch_size > 0, + "KDA decode fusion requires at least one row"); + int const dim = num_heads * kHeadDim; + + STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() && + weight.size(0) == 3 && weight.size(1) == kConvWidth && + weight.size(2) == dim, + "weight must have shape [3, 4, H * 128]"); + STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim && + conv_state.size(2) == kConvWidth - 1, + "conv_state must have shape [slots, 3 * H * 128, 3]"); + STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 && + raw_g.size(1) == batch_size && + raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim, + "raw_g must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 && + raw_beta.size(1) == batch_size && + raw_beta.size(2) == num_heads, + "raw_beta must have shape [1, B, H]"); + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads, + "A_log must be contiguous with H elements"); + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim, + "dt_bias must be contiguous with H * 128 elements"); + STD_TORCH_CHECK( + state_indices.is_contiguous() && state_indices.numel() == batch_size, + "state_indices must be contiguous with B elements"); + STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads && + state.size(2) == kHeadDim && state.size(3) == kHeadDim, + "state must have shape [slots, H, 128, 128]"); + STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 && + out.size(1) == batch_size && out.size(2) == num_heads && + out.size(3) == kHeadDim, + "out must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(x.stride(1) == 1, + "x must be contiguous in its channel dimension"); + STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && + conv_state.stride(1) == 1 && + conv_state.stride(2) == 3 * dim, + "conv_state must use the SD cache layout"); + STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && + state.stride(1) == kHeadDim * kHeadDim && + state.stride(2) == kHeadDim && state.stride(3) == 1, + "state must have contiguous [H, 128, 128] slot contents"); + STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous"); + STD_TORCH_CHECK(raw_beta.stride(2) == 1, + "raw_beta must be contiguous in its head dimension"); + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + + bool const apply_onorm = output_gate.has_value(); + STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(), + "output_gate and norm_weight must be provided together"); + void const* output_gate_ptr = nullptr; + float const* norm_weight_ptr = nullptr; + int64_t output_gate_row_stride = 0; + if (apply_onorm) { + STD_TORCH_CHECK(output_gate->is_cuda() && + output_gate->scalar_type() == ScalarType::BFloat16, + "output_gate must be a CUDA bfloat16 tensor"); + bool const gate_is_3d = + output_gate->dim() == 3 && output_gate->size(0) == batch_size && + output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim; + bool const gate_is_4d = + output_gate->dim() == 4 && output_gate->size(0) == 1 && + output_gate->size(1) == batch_size && + output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim; + STD_TORCH_CHECK(gate_is_3d || gate_is_4d, + "output_gate must have shape [B, H, 128] or " + "[1, B, H, 128]"); + int const row_dim = gate_is_3d ? 0 : 1; + STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1, + "output_gate must be contiguous in its last dimension"); + STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim, + "output_gate must have contiguous head rows"); + STD_TORCH_CHECK(norm_weight->is_cuda() && + norm_weight->scalar_type() == ScalarType::Float, + "norm_weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim, + "norm_weight must be contiguous with 128 elements"); + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); + output_gate_ptr = output_gate->data_ptr(); + norm_weight_ptr = static_cast(norm_weight->data_ptr()); + output_gate_row_stride = output_gate->stride(row_dim); + } + + void const* bias_ptr = nullptr; + if (bias.has_value()) { + STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float, + "bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim, + "bias must be contiguous with 3 * H * 128 elements"); + bias_ptr = bias->data_ptr(); + } + + auto const* x_ptr = static_cast(x.data_ptr()); + auto const* weight_ptr = static_cast(weight.data_ptr()); + auto* conv_ptr = static_cast(conv_state.data_ptr()); + auto const* bias_bytes = static_cast(bias_ptr); + int64_t const segment_bytes = dim * sizeof(__nv_bfloat16); + int64_t const weight_segment_bytes = + dim * kConvWidth * static_cast(sizeof(float)); + int64_t const conv_segment_bytes = + dim * conv_state.stride(1) * sizeof(__nv_bfloat16); + int64_t const bias_segment_bytes = dim * sizeof(float); + std::array const strides{ + x.stride(0), raw_beta.stride(1), output_gate_row_stride, + conv_state.stride(0), state.stride(0), + }; + bool const use_lower_bound = lower_bound.has_value(); + float const lower_bound_value = + use_lower_bound ? static_cast(*lower_bound) : 0.0f; + + torch::stable::accelerator::DeviceGuard const device_guard( + x.get_device_index()); + cudaStream_t const stream = get_current_cuda_stream(x.get_device_index()); + launch_kda_decode_many_heads_cuda( + x_ptr, x_ptr + segment_bytes, x_ptr + 2 * segment_bytes, weight_ptr, + weight_ptr + weight_segment_bytes, weight_ptr + 2 * weight_segment_bytes, + bias_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes, + conv_ptr, conv_ptr + conv_segment_bytes, + conv_ptr + 2 * conv_segment_bytes, + static_cast(a_log.data_ptr()), raw_g.data_ptr(), + static_cast(dt_bias.data_ptr()), raw_beta.data_ptr(), + output_gate_ptr, norm_weight_ptr, + static_cast(state_indices.data_ptr()), nullptr, + static_cast(state.data_ptr()), out.data_ptr(), batch_size, + num_heads, num_heads, apply_onorm, true, use_lower_bound, true, + lower_bound_value, 0.08838834764831845f, static_cast(norm_eps), + strides.data(), stream); + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 KDA decode kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 878b44df936c..b342c59f1801 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -110,7 +110,8 @@ fused_add_rms_norm_kernel( const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight - const float epsilon, const int num_tokens, const int hidden_size) { + const float epsilon, const int num_tokens, const int hidden_size, + const int64_t residual_stride) { // Sanity checks on our vector struct and type-punned pointer arithmetic static_assert(std::is_pod_v<_f16Vec>); static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); @@ -130,7 +131,7 @@ fused_add_rms_norm_kernel( reinterpret_cast*>(weight); for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { - int id = blockIdx.x * vec_hidden_size + idx; + int64_t id = blockIdx.x * residual_stride / width + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec temp = input_v[strided_id]; temp += residual_v[id]; @@ -148,7 +149,7 @@ fused_add_rms_norm_kernel( __syncthreads(); for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { - int id = blockIdx.x * vec_hidden_size + idx; + int64_t id = blockIdx.x * residual_stride / width + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec res = residual_v[id]; _f16Vec out; @@ -182,16 +183,17 @@ fused_add_rms_norm_kernel( const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight - const float epsilon, const int num_tokens, const int hidden_size) { + const float epsilon, const int num_tokens, const int hidden_size, + const int64_t residual_stride) { __shared__ float s_variance; float variance = 0.0f; for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { scalar_t z = input[blockIdx.x * input_stride + idx]; - z += residual[blockIdx.x * hidden_size + idx]; + z += residual[blockIdx.x * residual_stride + idx]; float x = (float)z; variance += x * x; - residual[blockIdx.x * hidden_size + idx] = z; + residual[blockIdx.x * residual_stride + idx] = z; } using BlockReduce = cub::BlockReduce; @@ -204,7 +206,7 @@ fused_add_rms_norm_kernel( __syncthreads(); for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { - float x = (float)residual[blockIdx.x * hidden_size + idx]; + float x = (float)residual[blockIdx.x * residual_stride + idx]; if constexpr (HasWeight) { float w = (float)weight[idx]; input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); @@ -249,7 +251,9 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -297,13 +301,13 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] input.mutable_data_ptr(), input_stride, \ residual.mutable_data_ptr(), \ weight->const_data_ptr(), epsilon, num_tokens, \ - hidden_size); \ + hidden_size, residual_stride); \ } else { \ vllm::fused_add_rms_norm_kernel \ <<>>( \ input.mutable_data_ptr(), input_stride, \ residual.mutable_data_ptr(), nullptr, epsilon, \ - num_tokens, hidden_size); \ + num_tokens, hidden_size, residual_stride); \ } \ }); @@ -312,21 +316,27 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] std::optional weight, double epsilon) { STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type()); - STD_TORCH_CHECK(residual.is_contiguous()); + STD_TORCH_CHECK(residual.stride(-1) == 1); if (weight.has_value()) { STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type()); STD_TORCH_CHECK(weight->is_contiguous()); } int hidden_size = input.size(-1); int64_t input_stride = input.stride(-2); + int64_t residual_stride = residual.stride(-2); int num_tokens = input.numel() / hidden_size; dim3 grid(num_tokens); /* This kernel is memory-latency bound in many scenarios. When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency - hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + hiding on global mem ops. In batch-invariant mode the block size must + not depend on num_tokens, otherwise the same token would use a different + reduction width (and thus a different floating-point summation order) + across batches; lock it to 1024 to keep results bit-exact. */ + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -336,8 +346,8 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] auto inp_ptr = reinterpret_cast(input.data_ptr()); auto res_ptr = reinterpret_cast(residual.data_ptr()); bool offsets_are_multiple_of_vector_width = - hidden_size % vector_width == 0 && input_stride % vector_width == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + hidden_size % vector_width == 0 && input_stride % vector_width == 0 && + residual_stride % vector_width == 0; const bool has_weight = weight.has_value(); if (has_weight) { auto wt_ptr = reinterpret_cast(weight->data_ptr()); diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index f3bf8882e775..f43be531de03 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant( int num_tokens = input.numel() / hidden_size; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant( When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant( auto wt_ptr = reinterpret_cast(weight.data_ptr()); bool ptrs_are_aligned = inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 && !batch_invariant_launch) { LAUNCH_FUSED_ADD_RMS_NORM(8); diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu index da9ef44d03f9..e329e067b266 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu @@ -25,6 +25,7 @@ #include "libtorch_stable/torch_utils.h" #include +#include #include #include #include @@ -448,7 +449,8 @@ enum ScoringFunc { SCORING_SIGMOID = 1 // apply sigmoid }; -// Efficient sigmoid approximation from TensorRT-LLM +// Adapted from +// https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu __device__ inline float sigmoid_accurate(float x) { return 0.5f * tanhf(0.5f * x) + 0.5f; } @@ -890,6 +892,434 @@ __global__ void grouped_topk_fused_small_expert_count_kernel( #endif } +// Adapted from +// https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh +namespace single_group_topk { +namespace detail { + +static constexpr int BlockDim = 256; +static constexpr uint32_t FullWarpMask = 0xffffffffU; +static constexpr float InvalidScore = -INFINITY; + +// TopK-only tuning: use wider workers and keep these tiers on the block path. +template +static constexpr bool UseTunedBlockPath = + MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024); + +template +__device__ __forceinline__ void preprocess_score(T input, BiasT correction_bias, + float& unbiased_score, + float& selection_score) { + unbiased_score = 0.0F; + selection_score = InvalidScore; + float const input_float = cuda_cast(input); + float const bias = cuda_cast(correction_bias); + if (!is_finite(input_float) || !is_finite(bias)) { + return; + } + + float const unbiased = apply_scoring(input_float); + float const biased = unbiased + bias; + if constexpr (SF == SCORING_NONE) { + if (!is_finite(biased)) { + return; + } + } + unbiased_score = unbiased; + selection_score = biased == 0.0F ? 0.0F : biased; +} + +template +__device__ __forceinline__ void write_outputs( + cg::thread_block_tile const& warp, float lane_selection_score, + float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token, + int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize, + float routed_scaling_factor) { + bool const finite_selection = + lane < topk && lane_selection_score != InvalidScore; + lane_unbiased = finite_selection ? lane_unbiased : 0.0F; + unsigned const finite_mask = __ballot_sync(FullWarpMask, finite_selection); + float const sum = cg::reduce(warp, lane_unbiased, cg::plus{}); + + if (lane < topk) { + float output = 0.0F; + if (finite_mask == 0) { + if (renormalize) { + output = 1.0F / static_cast(topk); + } + } else if (finite_selection) { + float scale = routed_scaling_factor; + if (renormalize) { + scale /= sum + 1e-20F; + } + output = lane_unbiased * scale; + } + + int64_t const output_index = int64_t{token} * topk + lane; + topk_values[output_index] = output; + topk_indices[output_index] = static_cast(lane_expert); + } +} + +template +__global__ void __launch_bounds__(BlockDim) + single_group_topk_block_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_experts, int64_t topk, + bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WorkerValuesPerLane = + UseTunedBlockPath ? 8 : 4; + static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + static constexpr int NumWorkerWarps = + (MaxNumExperts + ExpertsPerWorkerWarp - 1) / ExpertsPerWorkerWarp; + static constexpr int NumIntermediate = NumWorkerWarps * MaxNumTopExperts; + static constexpr int MergeValuesPerLane = + (NumIntermediate + WARP_SIZE - 1) / WARP_SIZE; + static constexpr bool LaneOwnedResourcesFit = + NumWorkerWarps <= BlockDim / WARP_SIZE && MergeValuesPerLane <= 64; + static constexpr bool UseHierarchicalLaneTopK = + LaneOwnedRange::kEnabled && LaneOwnedResourcesFit; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + __shared__ float __attribute((aligned(128))) biased_scores[MaxNumExperts]; + __shared__ float __attribute((aligned(128))) unbiased_scores[MaxNumExperts]; + + int32_t const token = static_cast(blockIdx.x); + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + T const* token_scores = scores + int64_t{token} * num_experts; + + for (int32_t expert = static_cast(threadIdx.x); + expert < num_experts_i32; expert += BlockDim) { + preprocess_score(token_scores[expert], bias[expert], + unbiased_scores[expert], + biased_scores[expert]); + } + __syncthreads(); + + auto warp = cg::tiled_partition(cg::this_thread_block()); + + if constexpr (UseHierarchicalLaneTopK) { + __shared__ float + __attribute((aligned(128))) intermediate_scores[NumIntermediate]; + __shared__ int32_t + __attribute((aligned(128))) intermediate_indices[NumIntermediate]; + + if (warp_id < NumWorkerWarps) { + float local_scores[WorkerValuesPerLane]; + int32_t local_indices[WorkerValuesPerLane]; +#pragma unroll + for (int index = 0; index < WorkerValuesPerLane; ++index) { + int32_t const expert = + warp_id * ExpertsPerWorkerWarp + index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + if (lane < MaxNumTopExperts) { + int32_t const intermediate = warp_id * MaxNumTopExperts + lane; + bool const active = lane < topk_i32; + intermediate_scores[intermediate] = active ? lane_score : InvalidScore; + intermediate_indices[intermediate] = + active ? lane_expert : MaxNumExperts; + } + } + __syncthreads(); + + if (warp_id != 0) { + return; + } + + float merge_scores[MergeValuesPerLane]; + int32_t merge_indices[MergeValuesPerLane]; +#pragma unroll + for (int index = 0; index < MergeValuesPerLane; ++index) { + int32_t const intermediate = index * WARP_SIZE + lane; + bool const active = intermediate < NumIntermediate; + merge_scores[index] = + active ? intermediate_scores[intermediate] : InvalidScore; + merge_indices[index] = + active ? intermediate_indices[intermediate] : MaxNumExperts; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, merge_scores, merge_indices, + InvalidScore, lane); + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } else { + if (warp_id != 0) { + return; + } + + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + float const lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + int32_t const lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct WarpTopKLaunchConfig { + static constexpr int DefaultBlockDim = + MaxNumExperts <= 1024 ? MaxNumExperts : 1024; + static constexpr int BlockDim = DefaultBlockDim > 256 ? 256 : DefaultBlockDim; + static constexpr int NumWarps = BlockDim / WARP_SIZE; + static constexpr int MaxBlockScale = + (DefaultBlockDim + BlockDim - 1) / BlockDim; + static constexpr int MaxBlocks = 1024 * MaxBlockScale; + + static_assert(BlockDim % WARP_SIZE == 0); + + static uint32_t grid_dim(int64_t num_tokens) { + int64_t const token_blocks = (num_tokens + NumWarps - 1) / NumWarps; + int64_t const selected = + token_blocks < MaxBlocks ? token_blocks : MaxBlocks; + return static_cast(selected > 0 ? selected : 1); + } +}; + +template +__global__ void __launch_bounds__(WarpTopKLaunchConfig::BlockDim) + single_group_topk_warp_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WarpBlockDim = + WarpTopKLaunchConfig::BlockDim; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const global_warp = + static_cast(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id; + int32_t const global_warp_stride = + static_cast(gridDim.x) * WarpBlockDim / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + auto warp = cg::tiled_partition(cg::this_thread_block()); + + for (int32_t token = global_warp; token < num_tokens; + token += global_warp_stride) { + T const* token_scores = scores + int64_t{token} * num_experts; + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + float unbiased; + float selection; + if (expert < num_experts_i32) { + preprocess_score(token_scores[expert], bias[expert], + unbiased, selection); + } else { + selection = InvalidScore; + } + local_scores[index] = selection; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + if constexpr (LaneOwnedRange::kEnabled) { + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + } else { + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + } + + float lane_unbiased = 0.0F; + if (lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32) { + lane_unbiased = lane_score - cuda_cast(bias[lane_expert]); + } + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct Tier { + static constexpr int kExperts = Experts; + static constexpr int kTopK = TopK; +}; + +template +struct TierList {}; + +using SigmoidBiasTiers = + TierList, Tier<256, 8>, Tier<384, 8>, Tier<512, 8>, + Tier<512, 22>, Tier<768, 16>, Tier<896, 16>, Tier<1024, 16>>; + +using PrecomputedSoftmaxBiasTiers = + TierList, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, + Tier<256, 16>, Tier<512, 8>, Tier<512, 16>, Tier<512, 22>, + Tier<512, 32>, Tier<576, 8>, Tier<768, 16>, Tier<896, 16>, + Tier<1024, 16>>; + +template +void launch(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + config.dynamicSmemBytes = 0; + bool const use_block_kernel = + UseTunedBlockPath || + MaxNumExperts > 1024 || num_experts >= 1024 || + (num_experts >= 256 && num_tokens <= 1024); + if (use_block_kernel) { + config.gridDim = static_cast(num_tokens); + config.blockDim = BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_block_kernel, + scores, topk_values, topk_indices, bias, num_experts, topk, renormalize, + static_cast(routed_scaling_factor), enable_pdl); + } else { + using WarpConfig = WarpTopKLaunchConfig; + config.gridDim = WarpConfig::grid_dim(num_tokens); + config.blockDim = WarpConfig::BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_warp_kernel, + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, static_cast(routed_scaling_factor), enable_pdl); + } +} + +template +bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t, + int64_t, bool, double, bool, cudaLaunchConfig_t&) { + return false; +} + +template +bool dispatch(TierList*, T* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, int64_t num_tokens, + int64_t num_experts, int64_t topk, bool renormalize, + double routed_scaling_factor, bool enable_pdl, + cudaLaunchConfig_t& config) { + if (num_experts <= First::kExperts && topk <= First::kTopK) { + launch( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + return true; + } + return dispatch( + static_cast*>(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); +} + +} // namespace detail + +template +bool invoke(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID); + if constexpr (SF == SCORING_SIGMOID) { + return detail::dispatch( + static_cast(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); + } else { + return detail::dispatch( + static_cast(nullptr), scores, + topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + } +} + +} // namespace single_group_topk + template void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, int64_t const num_tokens, @@ -905,6 +1335,12 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; config.numAttrs = 1; config.attrs = attrs; + if (n_group == 1 && topk_group == 1 && + single_group_topk::invoke( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, + topk, renormalize, routed_scaling_factor, enable_pdl, config)) { + return; + } // Check if we can use the optimized // grouped_topk_fused_small_expert_count_kernel diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh index 70e21cf8773a..6eadae0def8c 100644 --- a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh +++ b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh @@ -1,6 +1,7 @@ /* * Adapted from * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh * Copyright (c) 2026, The vLLM team. * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights * reserved. SPDX-License-Identifier: Apache-2.0 @@ -23,6 +24,9 @@ #include #include +#include +#include + namespace vllm { namespace moe { namespace reduce_topk { @@ -38,11 +42,10 @@ struct TopKRedType { "Top K reduction only implemented for int, float, float16 and bfloat16"); using TypeCmp = std::conditional_t; - using IdxT = std::conditional_t; static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; static constexpr int kMaxIdx = 65535; - TypeCmp compValIdx; + TypeCmp compVal; static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { auto valueBits = cub::Traits::TwiddleIn( @@ -69,69 +72,175 @@ struct TopKRedType { __host__ __device__ TopKRedType() = default; __host__ __device__ TopKRedType(T val, int32_t idx) - : compValIdx(makeCmpVal(val, idx)) {} + : compVal(makeCmpVal(val, idx)) {} - __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + __host__ __device__ operator TypeCmp() const noexcept { return compVal; } __device__ inline TypeCmp reduce( cg::thread_block_tile const& warp) { - return cg::reduce(warp, compValIdx, cg::greater{}); +#ifdef __CUDA_ARCH__ + static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10; +#else + static constexpr bool kHAS_FAST_REDUX = false; +#endif + if constexpr (!kHAS_FAST_REDUX) { + return cg::reduce(warp, compVal, cg::greater{}); + } else if constexpr (sizeof(TypeCmp) == 8) { + uint32_t hi = static_cast(compVal >> 32); + uint32_t lo = static_cast(compVal & 0xffffffffu); + uint32_t maxHi; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxHi) + : "r"(hi)); + uint32_t loContrib = hi == maxHi ? lo : 0u; + uint32_t maxLo; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxLo) + : "r"(loContrib)); + return (static_cast(maxHi) << 32) | static_cast(maxLo); + } else { + TypeCmp result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(result) + : "r"(compVal)); + return result; + } } }; -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct TopKIdx { - // by default, empty +template +struct IsPowerOf2 { + static constexpr bool value = N > 0 && (N & (N - 1)) == 0; }; -template -struct TopKIdx { - static constexpr int K = K_; - int32_t val[K]; +template +struct NextPow2 { + private: + static constexpr unsigned u = static_cast(N - 1); + static constexpr unsigned s1 = u | (u >> 1); + static constexpr unsigned s2 = s1 | (s1 >> 2); + static constexpr unsigned s3 = s2 | (s2 >> 4); + static constexpr unsigned s4 = s3 | (s3 >> 8); + static constexpr unsigned s5 = s4 | (s4 >> 16); + + public: + static constexpr int value = N <= 1 ? 1 : static_cast(s5 + 1); }; -//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__device__ __forceinline__ void topkCompareSwap(T* a) { + if constexpr (A < Size && B < Size) { + if (a[A] < a[B]) { + T tmp = a[A]; + a[A] = a[B]; + a[B] = tmp; + } + } else { + (void)a; + } +} -#define TOPK_SWAP(I, J) \ - { \ - auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \ - auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \ - topK[I].compValIdx = pairMax; \ - topK[J].compValIdx = pairMin; \ +template +__device__ __forceinline__ void topkMergePairs(T* a) { + if constexpr (I + Step < End) { + topkCompareSwap(a); + topkMergePairs(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkOEM(T* a) { + constexpr int M = R * 2; + if constexpr (M < N) { + topkOEM(a); + topkOEM(a); + topkMergePairs(a); + } else if constexpr (R < N) { + topkCompareSwap(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkSortBatcher(T* a) { + if constexpr (N > 1) { + constexpr int Half = N / 2; + topkSortBatcher(a); + topkSortBatcher(a); + topkOEM(a); + } else { + (void)a; } +} template -struct Sort; +struct Sort { + static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]"); + + static __device__ void run(RedType* topK) { + if constexpr (IsPowerOf2::value) { +#pragma unroll + for (int k = 2; k <= N; k *= 2) { +#pragma unroll + for (int j = k / 2; j > 0; j /= 2) { +#pragma unroll + for (int i = 0; i < N; ++i) { + int ixj = i ^ j; + if (ixj > i) { + if ((i & k) == 0) { + if (topK[i].compVal < topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } else { + if (topK[i].compVal > topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } + } + } + } + } + } else { + constexpr int P = NextPow2::value; + topkSortBatcher<0, P, N, RedType>(topK); + } + } +}; template struct Sort<1, RedType> { - static __device__ void run(RedType* topK) {} + static __device__ void run(RedType*) {} }; template struct Sort<2, RedType> { - static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } + static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); } }; template struct Sort<3, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 1); - TOPK_SWAP(1, 2); - TOPK_SWAP(0, 1); + topkCompareSwap<0, 1, 3>(topK); + topkCompareSwap<1, 2, 3>(topK); + topkCompareSwap<0, 1, 3>(topK); } }; template struct Sort<4, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 2); - TOPK_SWAP(1, 3); - TOPK_SWAP(0, 1); - TOPK_SWAP(2, 3); - TOPK_SWAP(1, 2); + topkCompareSwap<0, 2, 4>(topK); + topkCompareSwap<1, 3, 4>(topK); + topkCompareSwap<0, 1, 4>(topK); + topkCompareSwap<2, 3, 4>(topK); + topkCompareSwap<1, 2, 4>(topK); } }; @@ -147,24 +256,23 @@ __forceinline__ __device__ void reduceTopK( typename RedType::TypeCmp packedMax{}; #pragma unroll for (int kk = 0; kk < actualK; ++kk) { - topK = - kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK; - // get the next largest value + topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK; packedMax = topK.reduce(warp); RedType::unpack(out[kk], outIdx[kk], packedMax); } }; -template -__device__ void reduceTopKFunc(cg::thread_block_tile const& warp, - Type (&out)[K], int32_t (&outIdx)[K], - Type (&value)[N], int32_t (&idx)[N], - Type minValue, int actualK = K) { +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, Type (&out)[K], + int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], + Type const minValue, int actualK = K) { static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); static_assert(N > 0, "Top K must have N > 0"); - static_assert(N < 5, - "Only support candidates number less than or equal to 128"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); using RedType = TopKRedType; RedType topK[N]; #pragma unroll @@ -172,85 +280,88 @@ __device__ void reduceTopKFunc(cg::thread_block_tile const& warp, topK[nn] = RedType{value[nn], idx[nn]}; } - if constexpr (!IsSorted) { - Sort::run(topK); - } + Sort::run(topK); + typename RedType::TypeCmp packedMax{}; -#pragma unroll for (int kk = 0; kk < actualK; ++kk) { - bool update = kk > 0 && packedMax == topK[0].compValIdx; + bool update = kk > 0 && packedMax == topK[0].compVal; #pragma unroll for (int nn = 0; nn < N; ++nn) { topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} : update ? topK[nn + 1] : topK[nn]; } - // get the next largest value packedMax = topK[0].reduce(warp); RedType::unpack(out[kk], outIdx[kk], packedMax); } }; +template +struct LaneOwnedTopKRange { + static_assert(MinExperts > 0 && MinExperts <= MaxExperts); + static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts); + static constexpr bool kEnabled = + NumExperts >= MinExperts && NumExperts <= MaxExperts && + NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts; +}; + +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16; + +template +using HighExpertLaneOwnedTopKRange = + LaneOwnedTopKRange; + template -__forceinline__ __device__ void reduceTopK( - cg::thread_block_tile const& warp, Type (&out)[K], - int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], - Type const minValue, int actualK = K) { +__forceinline__ __device__ void reduceTopKForLane( + cg::thread_block_tile const& warp, Type& out, int32_t& outIdx, + Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) { static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); static_assert(N > 0, "Top K must have N > 0"); - static_assert( - N <= 16, - "Only support candidates number less than or equal to 16*32=512"); - static_assert(N <= 4 || N % 4 == 0, - "Only support candidates number is a multiple of 4*32=128 or " - "less than or equal to 4"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } - if constexpr (N <= 4) { - reduceTopKFunc(warp, out, outIdx, value, idx, minValue, - actualK); - } else { - constexpr int numLoops = N / 4; - constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; - - Type topKBufferValue[numResults]; - int32_t topKBufferIdx[numResults]; - int32_t laneIdx = threadIdx.x % kWARP_SIZE; + Sort::run(topK); - for (int ii = 0; ii < numResults; ++ii) { - topKBufferValue[ii] = minValue; - topKBufferIdx[ii] = ii * kWARP_SIZE - 1; + typename RedType::TypeCmp packedMax{}; + typename RedType::TypeCmp lanePacked{}; +#pragma unroll + for (int kk = 0; kk < K; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compVal; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; } - for (int loop = 0; loop < numLoops; ++loop) { - int start = loop * 4; - Type topKValue[K]; - int32_t topKIdx[K]; - Type inValue[4]; - int32_t inIdx[4]; - for (int i = 0; i < 4; ++i) { - inValue[i] = value[start + i]; - inIdx[i] = idx[start + i]; - } - reduceTopKFunc(warp, topKValue, topKIdx, inValue, inIdx, - minValue, actualK); - int inOffset = laneIdx % K; - if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) { - topKBufferValue[0] = topKValue[inOffset]; - topKBufferIdx[0] = topKIdx[inOffset]; - } - if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) { - topKBufferValue[1] = topKValue[inOffset]; - topKBufferIdx[1] = topKIdx[inOffset]; - } + packedMax = topK[0].reduce(warp); + if (laneIdx == kk) { + lanePacked = packedMax; } - - reduceTopKFunc(warp, out, outIdx, topKBufferValue, - topKBufferIdx, minValue, actualK); } -}; -#undef TOPK_SWAP + if (laneIdx < K) { + RedType::unpack(out, outIdx, lanePacked); + } else { + out = minValue; + outIdx = -1; + } +} } // namespace reduce_topk } // namespace moe diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index 86ee5397fbc3..78aed496bba4 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -1086,4 +1086,4 @@ void moe_lora_align_block_size( has_expert_map); } }); -} \ No newline at end of file +} diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 9ea1dae7cae1..03525a05f9df 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, torch::stable::Tensor& gating_output, bool renormalize, - std::optional bias); + std::optional bias, + std::optional is_padding); void topk_sigmoid(torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, torch::stable::Tensor& gating_output, bool renormalize, std::optional bias, - double routed_scaling_factor); + double routed_scaling_factor, + std::optional is_padding); void topk_softplus_sqrt( torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, @@ -25,7 +27,8 @@ void topk_softplus_sqrt( double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid); + const std::optional& tid2eid, + const std::optional& is_padding); void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index b4bcd9479e9d..098644a1b901 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int end_expert, const bool renormalize, const float* bias, - const double routed_scaling_factor) + const double routed_scaling_factor, + const bool* is_padding) { using cub_kvp = cub::KeyValuePair; @@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int expert = result_kvp.key; const bool node_uses_expert = expert >= start_expert && expert < end_expert; const bool should_process_row = row_is_active && node_uses_expert; + const bool is_pad_row = is_padding != nullptr && is_padding[block_row]; const int idx = k * block_row + k_idx; // Return the unbiased scores for output weights output[idx] = inputs_after_softmax[thread_read_offset + expert]; - indices[idx] = should_process_row ? (expert - start_expert) : num_experts; - assert(indices[idx] >= 0); + indices[idx] = is_pad_row ? static_cast(-1) + : (should_process_row ? (expert - start_expert) : num_experts); + assert(is_pad_row || indices[idx] >= 0); source_rows[idx] = k_idx * num_rows + block_row; if (renormalize) { selected_sum += inputs_after_softmax[thread_read_offset + expert]; @@ -277,7 +280,7 @@ template || std::is_same_v || std::is_same_v, @@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Add a guard to ignore experts not included by this node const bool node_uses_expert = expert >= start_expert && expert < end_expert; const bool should_process_row = row_is_active && node_uses_expert; + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; // The lead thread from each sub-group will write out the final results to global memory. (This will be a // single) thread per row of the input/output matrices. const int idx = k * thread_row + k_idx; output[idx] = max_val; - indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + indices[idx] = is_pad_row ? static_cast(-1) + : (should_process_row ? (expert - start_expert) : NUM_EXPERTS); source_rows[idx] = k_idx * num_rows + thread_row; if (renormalize) { selected_sum += max_val; @@ -605,7 +610,7 @@ struct TopkConstants template void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, - const float* bias, const double routed_scaling_factor, cudaStream_t stream) + const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = detail::TopkConstants; @@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); topkGating<<>>( - input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor); + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding); } #ifndef USE_ROCM @@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); #else #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); \ + bias, routed_scaling_factor, stream, is_padding); \ } else if (WARP_SIZE == 32) { \ topkGatingLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); \ + bias, routed_scaling_factor, stream, is_padding); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -661,7 +666,8 @@ void topkGatingKernelLauncher( const bool renormalize, const float* bias, const double routed_scaling_factor, - cudaStream_t stream) { + cudaStream_t stream, + const bool* is_padding) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; #ifndef USE_ROCM @@ -736,7 +742,7 @@ void topkGatingKernelLauncher( } moeTopK<<>>( workspace, nullptr, topk_weights, topk_indices, token_expert_indices, - num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor); + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding); } } } @@ -755,7 +761,8 @@ void dispatch_topk_launch( int num_tokens, int num_experts, int topk, bool renormalize, std::optional bias, double routed_scaling_factor, - cudaStream_t stream) + cudaStream_t stream, + std::optional is_padding) { const float* bias_ptr = nullptr; if (bias.has_value()) { @@ -769,6 +776,18 @@ void dispatch_topk_launch( bias_ptr = bias_tensor.const_data_ptr(); } + const bool* is_padding_ptr = nullptr; + if (is_padding.has_value()) { + const torch::stable::Tensor& is_padding_tensor = is_padding.value(); + STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool, + "is_padding tensor must be bool"); + STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D"); + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens, + "is_padding size mismatch, expected: ", num_tokens); + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous"); + is_padding_ptr = is_padding_tensor.const_data_ptr(); + } + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { vllm::moe::topkGatingKernelLauncher( reinterpret_cast(gating_output.const_data_ptr()), @@ -777,7 +796,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { vllm::moe::topkGatingKernelLauncher( reinterpret_cast(gating_output.const_data_ptr()), @@ -786,7 +805,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); vllm::moe::topkGatingKernelLauncher( @@ -796,7 +815,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } } @@ -806,7 +825,8 @@ void topk_softmax( torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, - std::optional bias) + std::optional bias, + std::optional is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -825,15 +845,15 @@ void topk_softmax( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } @@ -846,7 +866,8 @@ void topk_sigmoid( torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, std::optional bias, - double routed_scaling_factor) + double routed_scaling_factor, + std::optional is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -865,15 +886,15 @@ void topk_sigmoid( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index b6878eb2d2fd..976eebac3ddd 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -80,22 +80,27 @@ __launch_bounds__(128) __global__ OutIndType* indices, int num_rows, int num_experts, float routed_scaling_factor, const HashIndType* input_ids, - const HashIndType* tid2eid) { + const HashIndType* tid2eid, + const bool* is_padding) { const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane = threadIdx.x % 32; if (warp >= num_rows) return; const int64_t token_id = load_index_as_int64(input_ids, warp); + const bool is_pad_row = is_padding != nullptr && is_padding[warp]; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); #endif int expert = 0; float weight = 0.f; - if (lane < 6) { + if (lane < 6 && !is_pad_row) { // only load and calculate for 6 experts expert = static_cast(tid2eid[token_id * 6 + lane]); const float x = input[warp * num_experts + expert]; weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); + if (isnan(weight)) { + weight = 0.f; + } } float weight_sum = weight; #pragma unroll @@ -111,7 +116,8 @@ __launch_bounds__(128) __global__ const int offset = warp * 6 + lane; output[offset] = weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f); - indices[offset] = static_cast(expert); + indices[offset] = !is_pad_row ? static_cast(expert) + : static_cast(-1); } } @@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, int num_rows, int num_experts, double routed_scaling_factor, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream) { + const HashIndType* tid2eid, cudaStream_t stream, + const bool* is_padding) { if (num_rows == 0) return; auto* kernel = &dsv4HashTopkSoftplusSqrt; cudaLaunchConfig_t config = {}; @@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, config.numAttrs = 1; const float scale = static_cast(routed_scaling_factor); cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows, - num_experts, scale, input_ids, tid2eid); + num_experts, scale, input_ids, tid2eid, is_padding); } #endif @@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ const int num_rows, IndType* indices, int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const HashIndType* input_ids, const HashIndType* tid2eid) { + const HashIndType* input_ids, const HashIndType* tid2eid, + const bool* is_padding) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ return; } const bool row_is_active = finished ? !finished[thread_row] : true; + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; // We finally start setting up the read pointers for each thread. First, each // thread jumps to the start of the row it will read. @@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ cudaGridDependencySynchronize(); #endif - // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert - // to float - if constexpr (std::is_same_v) { + if (is_pad_row) { +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] = 0.f; + } + } else if constexpr (std::is_same_v) { using VecType = AlignedArray; VecType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk); const VecType* vec_thread_read_ptr = @@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ if constexpr (USE_HASH) { const int64_t token_id = load_index_as_int64(input_ids, thread_row); const int64_t token_expert_offset = token_id * static_cast(k); + if (!is_pad_row) { #pragma unroll - for (int ii = 0; ii < VPT; ++ii) { - float val = row_chunk[ii]; - float val_b = val * beta; - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; - row_chunk[ii] = sqrtf(val); + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + val = sqrtf(val); + + // Dummy/padding tokens can result in NaN values, so + // clamp them to 0.0. Note: this clamp could likely be removed if + // 'is_padding' is made mandatory + if (isnan(val)) { + val = 0.f; + } + row_chunk[ii] = val; + } } float selected_sum = 0.f; #pragma unroll @@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; if (expert == expert_idx) { - indices[idx] = static_cast(expert); + indices[idx] = !is_pad_row ? static_cast(expert) + : static_cast(-1); selected_sum += row_chunk[ii]; break; } @@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ #endif return; } else { + if (!is_pad_row) { #pragma unroll - for (int ii = 0; ii < VPT; ++ii) { - float val = row_chunk[ii]; - float val_b = val * beta; - // Compute softplus: log(1 + exp(val)) with numerical stability - // When val > threshold, softplus(x) ≈ x to avoid exp overflow - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; - val = sqrtf(val); - if (correction_bias) { - const int group_id = ii / ELTS_PER_LDG; - const int local_id = ii % ELTS_PER_LDG; - const int expert_idx = first_elt_read_by_thread + - group_id * THREADS_PER_ROW * ELTS_PER_LDG + - local_id; - val = val + correction_bias[expert_idx]; + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + // Compute softplus: log(1 + exp(val)) with numerical stability + // When val > threshold, softplus(x) ≈ x to avoid exp overflow + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + val = sqrtf(val); + // Dummy/padding tokens can result in NaN values, so + // clamp them to 0.0. Note: this clamp could likely be removed if + // 'is_padding' is made mandatory + if (isnan(val)) { + val = 0.f; + } + if (correction_bias) { + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + val = val + correction_bias[expert_idx]; + } + row_chunk[ii] = val; } - row_chunk[ii] = val; } // Original TopK path: find top-k experts by score @@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Add a guard to ignore experts not included by this node const bool node_uses_expert = expert >= start_expert && expert < end_expert; - const bool should_process_row = row_is_active && node_uses_expert; + const bool should_process_row = + row_is_active && node_uses_expert && !is_pad_row; // The lead thread from each sub-group will write out the final results // to global memory. (This will be a single) thread per row of the // input/output matrices. const int idx = k * thread_row + k_idx; - if (correction_bias != nullptr) { + if (correction_bias != nullptr && should_process_row) { max_val -= correction_bias[expert]; } output[idx] = max_val; indices[idx] = - should_process_row ? (expert - start_expert) : NUM_EXPERTS; + !is_pad_row ? expert - start_expert : static_cast(-1); source_rows[idx] = k_idx * num_rows + thread_row; if (renormalize) { selected_sum += max_val; @@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper( const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream) { + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper( cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid); + input_ids, tid2eid, is_padding); #else kernel<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid); + input_ids, tid2eid, is_padding); #endif }) } @@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ - stream); + stream, is_padding); #else #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream); \ + tid2eid, stream, is_padding); \ } else if (WARP_SIZE == 32) { \ topkGatingSoftplusSqrtLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream); \ + tid2eid, stream, is_padding); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher( const int topk, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, const HashIndType* tid2eid, - cudaStream_t stream) { + cudaStream_t stream, const bool* is_padding) { #ifndef USE_ROCM if constexpr (std::is_same_v) { if (use_hash && topk == 6 && renormalize && (num_experts == 256 || num_experts == 384)) { launchDsv4HashTopk( gating_output, topk_weights, topk_indices, num_tokens, num_experts, - routed_scaling_factor, input_ids, tid2eid, stream); + routed_scaling_factor, input_ids, tid2eid, stream, is_padding); return; } } @@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch( int num_experts, int topk, bool renormalize, double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid, cudaStream_t stream) { + const std::optional& tid2eid, cudaStream_t stream, + const std::optional& is_padding) { const float* bias_ptr = nullptr; if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); @@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch( auto launch = [&](auto* topk_indices_ptr) { using OutIndType = typename std::remove_pointer::type; + + const bool* is_padding_ptr = nullptr; + if (is_padding.has_value()) { + const torch::stable::Tensor& is_padding_tensor = is_padding.value(); + STD_TORCH_CHECK(is_padding_tensor.scalar_type() == + torch::headeronly::ScalarType::Bool, + "is_padding tensor must be bool"); + STD_TORCH_CHECK(is_padding_tensor.dim() == 1, + "is_padding tensor must be 1D"); + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens, + "is_padding size mismatch, expected: ", num_tokens); + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), + "is_padding tensor must be contiguous"); + is_padding_ptr = is_padding_tensor.const_data_ptr(); + } + if (tid2eid.has_value()) { STD_TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE"); @@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr); } else { STD_TORCH_CHECK(tid2eid.value().scalar_type() == torch::headeronly::ScalarType::Int); @@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr); } } else { vllm::moe::topkGatingSoftplusSqrtKernelLauncher( @@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, false, static_cast(nullptr), - static_cast(nullptr), stream); + static_cast(nullptr), stream, is_padding_ptr); } }; @@ -793,7 +842,8 @@ void topk_softplus_sqrt( bool renormalize, double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid) { + const std::optional& tid2eid, + const std::optional& is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); @@ -806,21 +856,22 @@ void topk_softplus_sqrt( dispatch_topk_softplus_sqrt_launch( gating_output.const_data_ptr(), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, - routed_scaling_factor, correction_bias, input_ids, tid2eid, stream); + routed_scaling_factor, correction_bias, input_ids, tid2eid, stream, + is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_softplus_sqrt_launch<__half>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream); + input_ids, tid2eid, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream); + input_ids, tid2eid, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index a0296adca78b..c6467f1c1ff1 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { m.def( "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? " - "bias) -> ()"); + "bias, Tensor? is_padding) -> ()"); // Apply topk sigmoid to the gating outputs. m.def( "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, " - "Tensor? bias, float routed_scaling_factor) -> ()"); + "Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()"); m.def( "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, float " "routed_scaling_factor, Tensor? " - "bias, Tensor? input_ids, Tensor? tid2eid) -> ()"); + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()"); // Calculate the result of moe by summing up the partial results // from all selected experts. topk_ids/expert_map are optional and, when diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 3834bea58576..7dfe0da78cb4 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -269,6 +269,14 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps, int64_t cache_block_size); +void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& q_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size); + void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, @@ -276,6 +284,61 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size); +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( torch::stable::Tensor const& q, torch::stable::Tensor const& kv, torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, @@ -313,7 +376,32 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional index_cache, int64_t block_size, std::optional q_out, std::optional index_q_out, - const std::string& kv_cache_dtype, bool skip_index_branch); + const std::string& kv_cache_dtype, bool skip_index_branch, + std::optional q_fp8_out, double q_fp8_scale); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps); +#endif // Sampler kernels (shared CUDA/ROCm) void apply_repetition_penalties_( @@ -372,6 +460,20 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, torch::stable::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); void dispose(fptr_t _fa); int64_t meta_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); @@ -410,6 +512,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double threshold); void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); @@ -486,6 +594,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, const std::string& kv_cache_dtype, torch::stable::Tensor& scale); +void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, + torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, + int64_t entry_stride); + // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, diff --git a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu b/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu index c3702c52efcb..43301f0fe206 100644 --- a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu +++ b/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu @@ -528,5 +528,9 @@ torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros, num_in_feats, num_in_channels, num_out_channels, out_feats); } - return torch::stable::sum(_out_feats, 0); + // Pass the reduction dim via a named, lifetime-stable local rather than a + // bare `0`. + constexpr int64_t sum_dim = 0; + return torch::stable::sum( + _out_feats, torch::headeronly::IntHeaderOnlyArrayRef(&sum_dim, 1)); } diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 2152e64dc962..56dd47038727 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,6 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" +#include "../../../core/batch_invariant.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" @@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch( auto num_tokens = input.numel() / hidden_size; dim3 grid(num_tokens); - const int max_block_size = (num_tokens <= 256) ? 512 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu index f8ef6b12a01f..b48601435e3c 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq( // AWQ zeros: (size_k // group_size, size_n // 8) const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, int32_t group_size) { - int32_t val = - qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; - int32_t zero = - qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + - blockIdx.y]; + // Thread mapping: threadIdx.x -> column dim (coalesced read within a row), + // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the + // same row (stride 1) instead of striding across rows (stride size_n/8). + int col = blockIdx.y * 32 + threadIdx.x; + if (col >= size_n / 8) return; + (void)size_k; + + int32_t val = qweight[blockIdx.x * (size_n / 8) + col]; + int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col]; int32_t new_val = 0; #pragma unroll @@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq( zero >>= 4; } - output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; + output[blockIdx.x * (size_n / 8) + col] = new_val; } torch::stable::Tensor marlin_int4_fp8_preprocess( @@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess( "qweight.size(0) % qzeros.size(0) != 0"); STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); - dim3 blocks(size_k / 32, size_n / 8); + dim3 blocks(size_k, (size_n / 8 + 31) / 32); marlin_int4_fp8_preprocess_kernel_awq<<>>( reinterpret_cast(qweight.const_data_ptr()), reinterpret_cast(output.mutable_data_ptr()), diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 0a475d02c6fe..fd4e8a848ccf 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -324,7 +324,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). // conditionally compiled so impl registration is in source file ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). // conditionally compiled so impl registration is in source file @@ -432,6 +433,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size) -> Tensor"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(" + "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size) -> ()"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. @@ -447,6 +453,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " "int cache_block_size) -> ()"); + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + #ifndef USE_ROCM ops.def( "minimax_allreduce_rms_qk(" @@ -466,7 +514,26 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? slot_mapping, Tensor? index_slot_mapping, " "Tensor!? kv_cache, Tensor!? index_cache, " "int block_size, Tensor!? q_out, Tensor!? index_q_out, " - "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); + "str kv_cache_dtype, bool skip_index_branch=False, " + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0) -> ()"); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " + "int num_blocks, float eps, float output_norm_eps) -> ()"); +#endif // Apply repetition penalties to logits in-place. ops.def( @@ -530,6 +597,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "limit=7.0) " "-> ()"); + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + // GELU implementation used in GPT-2. ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); @@ -682,17 +757,38 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out)); ops.impl( "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); ops.impl( "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); #ifndef USE_ROCM ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif // Sampler kernels (shared CUDA/ROCm) ops.impl("apply_repetition_penalties_", @@ -715,6 +811,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); @@ -812,6 +910,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " str kv_cache_dtype," " Tensor scale) -> ()"); + // Grouped concat_and_cache_mla across all layers (bf16 only). Each + // layer's cache base pointer is read from kv_cache_ptrs. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride) -> ()"); + // Rotate Q and K, then write to kv cache for MLA ops.def( "concat_and_cache_mla_rope_fused(" @@ -910,6 +1017,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); ops.impl("concat_and_cache_mla_rope_fused", TORCH_BOX(&concat_and_cache_mla_rope_fused)); ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 6e95d72ed5e9..296978f384dc 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -13,10 +13,18 @@ namespace vllm { namespace fp8 { #ifdef ENABLE_FP8 +// Unspecialized conversions are a compile error: the old passthrough +// (`return x;`) silently skipped fp8 encoding for any (Tout, Tin) pair +// without a specialization below (e.g. the torch stable-ABI scalar types), +// corrupting quantized data with no runtime signal. +template +inline constexpr bool _no_conversion_specialization = false; + template __inline__ __device__ Tout vec_conversion( const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) { - return x; + static_assert(_no_conversion_specialization, + "no vec_conversion specialization for this (Tout, Tin) pair"); } // float -> c10::Float8_e4m3fn @@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion( template __inline__ __device__ Tout scaled_vec_conversion( const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) { - return x; + static_assert( + _no_conversion_specialization, + "no scaled_vec_conversion specialization for this (Tout, Tin) pair"); } // fp8 -> half @@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion( __builtin_unreachable(); // Suppress missing return statement warning } +// torch stable-ABI (headeronly) scalar types delegate to the CUDA-native +// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// c10::Half quantize correctly without manual casts. +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::BFloat16& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::Half& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + // float -> fp8 template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index 6c3456d06f20..f330ed401e8b 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -79,6 +79,7 @@ union BufferResource { }; }; +#if !defined(__gfx1250__) __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); @@ -86,6 +87,16 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); +#else +__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( + int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {} + +__quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data, + int32x4_t srsrc, + int32_t voffset, + int32_t soffset, + int32_t aux) {} +#endif __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/quickreduce/quick_reduce.h index 7506329972ba..449721b57131 100644 --- a/csrc/quickreduce/quick_reduce.h +++ b/csrc/quickreduce/quick_reduce.h @@ -22,17 +22,25 @@ template __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, int rank, uint8_t** dbuffer_list, - uint32_t data_offset, uint32_t flag_color, + uint32_t data_offset, uint32_t* d_flag_counters, int64_t data_size_per_phase) { int block = blockIdx.x; int grid = gridDim.x; + // Load this block's counter from device memory and advance it on-device, + // so the color keeps changing across graph replays instead of being frozen. + uint32_t flag_color = d_flag_counters[blockIdx.x]; + while (block < num_blocks) { AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase); block += grid; flag_color++; } + // All threads compute the same final value; one writer per block is enough. + if (threadIdx.x == 0 && threadIdx.y == 0) { + d_flag_counters[blockIdx.x] = flag_color; + } } #define TWOSHOT_DISPATCH(__codec) \ @@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, hipLaunchKernelGGL((allreduce_prototype_twoshot), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else if (world_size == 4) { \ using LineCodec = __codec; \ using AllReduceKernel = AllReduceTwoshot; \ hipLaunchKernelGGL((allreduce_prototype_twoshot), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else if (world_size == 8) { \ using LineCodec = __codec; \ using AllReduceKernel = AllReduceTwoshot; \ hipLaunchKernelGGL((allreduce_prototype_twoshot), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } // INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8 @@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, hipLaunchKernelGGL((allreduce_prototype_twoshot), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else { \ throw std::runtime_error( \ "INT3 quick all-reduce is only supported for world_size == 2 " \ @@ -94,7 +102,7 @@ struct DeviceComms { static int constexpr kMaxWorldSize = 8; bool initialized = false; - uint32_t flag_color = 1; + uint32_t* d_flag_counters = nullptr; int world_size; int rank; @@ -128,6 +136,16 @@ struct DeviceComms { // Clear the flags buffer. HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size)); + // One flag-color counter per block, advanced by the kernel. Start at 1 + // to stay clear of the flags buffer we just zeroed. + HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t))); + { + std::vector init_color(kMaxNumBlocks, 1u); + HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(), + kMaxNumBlocks * sizeof(uint32_t), + hipMemcpyHostToDevice)); + } + // Device-side list of IPC buffers. buffer_list.resize(world_size); HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*))); @@ -144,6 +162,12 @@ struct DeviceComms { hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; } void destroy() { + // Allocated before `initialized` flips true, so free it on its own guard + // to avoid a leak if init fails partway through. + if (d_flag_counters) { + HIP_CHECK(hipFree(d_flag_counters)); + d_flag_counters = nullptr; + } if (initialized) { for (int i = 0; i < world_size; i++) { if (i != rank) { @@ -211,8 +235,6 @@ struct DeviceComms { break; } HIP_CHECK(cudaGetLastError()); - // Rotate the flag color. - flag_color += divceil(N, grid); } }; diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 4ac255d0a75f..5ca55fe3b9ea 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -2405,6 +2405,16 @@ template __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, const bit16x8& inpB, const floatx8& inpC) { + #if defined(__gfx1250__) + // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by + // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is + // unsupported on gfx1250; trap if ever launched (fail loud, not + // silent-wrong). + (void)inpA; + (void)inpB; + __builtin_trap(); + return inpC; + #else if constexpr (std::is_same::value) { return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); } else if constexpr (std::is_same::value) { @@ -2412,6 +2422,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, } else { static_assert(false, "unsupported 16b dtype"); } + #endif } template diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 021359cb6321..92dbed7f021c 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -14,6 +14,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // vLLM custom ops for rocm +// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 +// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid +// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. +#ifndef VLLM_SKIP_SKINNY_GEMMS // Custom gemm op for matrix-vector multiplication rocm_ops.def( "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " @@ -46,6 +50,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor scale_a, " " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). diff --git a/csrc/spinloop.cpp b/csrc/spinloop.cpp index c29e48a5f0ec..3285b9a3ded5 100644 --- a/csrc/spinloop.cpp +++ b/csrc/spinloop.cpp @@ -7,7 +7,7 @@ extern "C" { #if defined(__i386__) || defined(__x86_64__) #include - #include + #include #endif #if defined(CLOCK_MONOTONIC_RAW) diff --git a/docker/Dockerfile b/docker/Dockerfile index f008fd29e15f..b3fd7c097ded 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -22,9 +22,13 @@ # docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json # ============================================================================= -ARG CUDA_VERSION=13.0.2 +ARG CUDA_VERSION=13.0.3 ARG PYTHON_VERSION=3.12 ARG UBUNTU_VERSION=22.04 +# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend). +# This version is used for CUDA 12+ builds to keep the DeepEP build and +# runtime NCCL versions compatible. +ARG NCCL_VERSION=2.30.7 # By parameterizing the base images, we allow third-party to use their own # base images. One use case is hermetic builds with base images stored in @@ -477,10 +481,17 @@ WORKDIR /workspace # Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management -ARG DEEPEP_COMMIT_HASH=73b6ea4 +ARG DEEPEP_COMMIT_HASH=d4f41e4e93 ARG NVSHMEM_VER +ARG NCCL_VERSION RUN --mount=type=cache,target=/opt/uv/cache \ mkdir -p /tmp/ep_kernels_workspace/dist && \ + CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ + if [ -n "$NCCL_VERSION" ]; then \ + echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ + > /tmp/nccl-override.txt && \ + export UV_OVERRIDE=/tmp/nccl-override.txt; \ + fi && \ export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ @@ -574,13 +585,14 @@ COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_work RUN sha256sum /tmp/ep_kernels_workspace/dist/*.whl \ > /tmp/ep_kernels_workspace/dist/wheels.sha256 -# Check the size of the wheel if RUN_WHEEL_CHECK is true +# Check the size of the CUDA 13 wheel uploaded to PyPI COPY .buildkite/check-wheel-size.py check-wheel-size.py # sync the default value with .buildkite/check-wheel-size.py ARG VLLM_MAX_SIZE_MB=500 ENV VLLM_MAX_SIZE_MB=$VLLM_MAX_SIZE_MB ARG RUN_WHEEL_CHECK=true -RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \ +RUN if [ "$RUN_WHEEL_CHECK" = "true" ] && \ + [ "${CUDA_VERSION%%.*}" = "13" ]; then \ python3 check-wheel-size.py dist; \ else \ echo "Skipping wheel size check."; \ @@ -644,6 +656,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base ARG CUDA_VERSION ARG PYTHON_VERSION +ARG NCCL_VERSION ARG DEADSNAKES_MIRROR_URL ARG DEADSNAKES_GPGKEY_URL ARG GET_PIP_URL @@ -696,7 +709,6 @@ RUN apt-get update -y \ # Install CUDA development tools for runtime JIT compilation # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime) RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \ - CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \ apt-get update -y && \ apt-get install -y --no-install-recommends --allow-change-held-packages \ cuda-nvcc-${CUDA_VERSION_DASH} \ @@ -709,12 +721,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \ libnuma-dev \ # numactl CLI for NUMA binding at runtime numactl && \ - # Fixes nccl_allocator requiring nccl.h at runtime - # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22 - # NCCL packages don't use the cuda-MAJOR-MINOR naming convention, - # so we pin the version to match our CUDA version - NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \ - apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \ rm -rf /var/lib/apt/lists/* # Install uv for faster pip installs @@ -734,6 +740,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \ && chgrp -R 0 /opt/uv \ && chmod -R g+rwX,a+rX /opt/uv +# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins +# an older version as a transitive dep; this override forces uv to use our +# pinned version whenever nvidia-nccl-cu* is resolved. +RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ + if [ -n "$NCCL_VERSION" ]; then \ + echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ + > /etc/uv-overrides.txt; \ + else \ + touch /etc/uv-overrides.txt; \ + fi +ENV UV_OVERRIDE=/etc/uv-overrides.txt + # ---------------------------------------------------------------------- # Non-root support (opt-in) # ---------------------------------------------------------------------- @@ -793,7 +811,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.14 +ARG FLASHINFER_VERSION=0.6.16.post3 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') @@ -931,7 +949,8 @@ COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt -RUN --mount=type=cache,target=/opt/uv/cache \ +# Source-built wheels in this cache must match the Ubuntu 22.04 ABI. +RUN --mount=type=cache,id=uv-v027-ubuntu2204,target=/opt/uv/cache \ CUDA_MAJOR="${CUDA_VERSION%%.*}"; \ if [ "$CUDA_MAJOR" -ge 12 ]; then \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ @@ -953,6 +972,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ && uv pip install --system -r requirements/dev.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi \ + && if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \ + python3 -c "from arctic_inference.suffix_decoding import _C"; \ + fi; \ fi # install development dependencies (for testing) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index a7cc8c0870bc..a32ecb620f1d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -339,18 +339,17 @@ 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 -FROM base AS build_rixl -ARG RIXL_BRANCH="39be1de8" -ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" -ARG UCX_BRANCH="bfb51733" +# NIXL/UCX build stages +FROM base AS build_nixl +ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2" +ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git" +ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194" ARG UCX_REPO="https://github.com/openucx/ucx.git" ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx -ENV RIXL_HOME=/usr/local/rixl -ENV RIXL_BENCH_HOME=/usr/local/rixl_bench +ENV NIXL_HOME=/usr/local/nixl -# RIXL build system dependences and RDMA support +# NIXL build system dependencies and RDMA support RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ libgrpc-dev \ libgrpc++-dev \ @@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ && rm -rf /var/lib/apt/lists/* RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system meson auditwheel patchelf tomlkit + uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \ + auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0" RUN --mount=type=cache,target=/root/.cache/ccache \ cd /usr/local/src && \ @@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} RUN --mount=type=cache,target=/root/.cache/ccache \ - git clone ${RIXL_REPO} /opt/rixl && \ - cd /opt/rixl && \ - git checkout ${RIXL_BRANCH} && \ + git clone ${NIXL_REPO} /opt/nixl && \ + cd /opt/nixl && \ + git checkout ${NIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ - meson setup build --prefix=${RIXL_HOME} \ + meson setup build --prefix=${NIXL_HOME} \ -Ducx_path=${UCX_HOME} \ - -Drocm_path=${ROCM_PATH} && \ + -Dwheel_variant=rocm \ + -Dbuild_tests=false \ + -Dbuild_examples=false && \ cd build && \ ninja -j$(nproc) && \ - ninja install - -# Generate RIXL wheel + ninja install && \ + echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \ + > /etc/ld.so.conf.d/nixl.conf && \ + echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \ + >> /etc/ld.so.conf.d/nixl.conf && \ + ldconfig + +# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA, +# so configure the ROCm wheel variant directly through Meson. # Exclude libcore and libpull from auditwheel: transitive dependencies # that are not shipped in the wheel and vary across base images. -RUN cd /opt/rixl && \ - sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ - contrib/build-wheel.sh && \ - mkdir -p /app/install && \ - _ucx_install_dir=${UCX_HOME} \ - ./contrib/build-wheel.sh \ - --output-dir /app/install \ - --rocm-dir ${ROCM_PATH} \ +RUN cd /opt/nixl && \ + ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \ + CC="ccache gcc" CXX="ccache g++" \ + uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \ + --python ${PYTHON_VERSION} \ + -Csetup-args=-Ducx_path=${UCX_HOME} \ + -Csetup-args=-Dwheel_variant=rocm \ + -Csetup-args=-Dbuild_tests=false \ + -Csetup-args=-Dbuild_examples=false && \ + mkdir -p /tmp/nixl_wheels/repaired /app/install && \ + auditwheel repair \ + --exclude 'libamdhip64*' \ + --exclude 'libcore*' \ + --exclude 'libpull*' \ + /tmp/nixl_wheels/nixl_rocm*.whl \ + --plat manylinux_2_34_$(uname -m) \ + --wheel-dir /tmp/nixl_wheels/repaired && \ + ./contrib/wheel_add_ucx_plugins.py \ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ - --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \ + /tmp/nixl_wheels/repaired/*.whl && \ + cp /tmp/nixl_wheels/repaired/*.whl /app/install # ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not # invalidate the slow ROCShmem build. @@ -447,18 +467,21 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ # DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. FROM build_rocshmem AS build_deepep -ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_BRANCH="0f63d3e9" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" # Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. # DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +# We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8 +# quantization bound with vLLM's own clamp +# (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ RUN --mount=type=cache,target=/root/.cache/ccache \ export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ && git clone ${DEEPEP_REPO} \ && cd DeepEP \ && git checkout ${DEEPEP_BRANCH} \ - && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # not force users to rebuild the long-lived Dockerfile.rocm_base image. @@ -660,10 +683,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ fi -# Install RIXL + DeepEP wheels. -RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ +# Install NIXL + DeepEP wheels. +RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /rixl_install/*.whl /deep_install/*.whl + uv pip install --system /nixl_install/*.whl /deep_install/*.whl # Copy ROCShmem runtime libraries. COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem @@ -724,6 +747,8 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0 # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. # See: https://github.com/ROCm/rocm-libraries/issues/6266 ENV HSA_ENABLE_IPC_MODE_LEGACY=1 +ENV UCX_RMA_PPLN_ENABLE=y +ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf # ROCm profiler limits workaround. RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf @@ -796,9 +821,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ && pip uninstall -y vllm \ && uv pip install --system *.whl -# Install RIXL wheel -RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ - uv pip install --system /rixl_install/*.whl +# Install NIXL ROCm wheel +RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ + uv pip install --system /nixl_install/*.whl ARG COMMON_WORKDIR ARG BASE_IMAGE @@ -813,6 +838,8 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc # See: https://github.com/ROCm/rocm-libraries/issues/6266 ENV HSA_ENABLE_IPC_MODE_LEGACY=1 +ENV UCX_RMA_PPLN_ENABLE=y +ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf ENV TOKENIZERS_PARALLELISM=false diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 2faaf774cf62..d923693d112e 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.16.post3" +ARG AITER_BRANCH="v0.1.19" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV AITER_ROCM_ARCH=gfx942;gfx950 -ENV MORI_GPU_ARCHS=gfx942;gfx950 +# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 @@ -121,8 +121,6 @@ RUN cd triton \ && if [ ! -f setup.py ]; then cd python; fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ && mkdir -p /app/install && cp dist/*.whl /app/install -RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \ - && python3 -m build --wheel && cp dist/*.whl /app/install; fi ### diff --git a/docker/Dockerfile.rocm_base_gfx1250 b/docker/Dockerfile.rocm_base_gfx1250 new file mode 100644 index 000000000000..3fad6800ab59 --- /dev/null +++ b/docker/Dockerfile.rocm_base_gfx1250 @@ -0,0 +1,325 @@ +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=7.14.0a20260623 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 +ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 + + +ARG FA_BRANCH="jpvillam/gfx1250_wip" +ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git" +ARG AITER_BRANCH="main" +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG MORI_BRANCH="v1.1.0" +ARG MORI_REPO="https://github.com/ROCm/mori.git" + +# Sccache configuration (only used in release pipeline) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG PYTORCH_ROCM_ARCH=gfx1250 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV FA_GPU_ARCHS=gfx942;gfx950;gfx1250 + +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 +ENV SKIP_CK_BUILD="TRUE" +ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" +ARG PREBUILD_KERNELS=0 +ARG TRITON_KERNEL_BRANCH="padroute" +ARG TRITON_KERNEL_REPO="https://github.com/jpvillam-amd/triton.git" + +# Required for RCCL in ROCm7.1 +ENV HSA_NO_SCRATCH_RECLAIM=1 + +ARG PYTHON_VERSION=3.12 +ENV PYTHON_VERSION=${PYTHON_VERSION} + +RUN mkdir -p /app +WORKDIR /app +ENV DEBIAN_FRONTEND=noninteractive + +# Install Python and other dependencies +RUN apt-get update -y \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython +RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +## +## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +## +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +ARG APEX_VERSION +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + --extra-index-url https://pypi.org/simple \ + "torch[device-all]==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + + + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + +# The ROCm SDK wheel ships a broken CMake export for hsakmt +# This patch includes the right paths for the numa build target +RUN <<'EOF' +set -eu +TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" +[ -f "$TARGETS" ] || exit 0 # nothing to patch +grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched + +# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. +sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" + +# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). +sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" + +# 3. Define the numa::numa target the export references but forgot to create. +cat >> "$TARGETS" <<'CMAKE' + +if(NOT TARGET numa::numa) + find_library(NUMA_LIBRARY NAMES numa REQUIRED) + add_library(numa::numa UNKNOWN IMPORTED) + set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") +endif() +CMAKE +EOF + +# Clone custom triton_kernels for install later +RUN mkdir -p /app/patched_triton_kernels; \ + cd /app/patched_triton_kernels \ + && git clone ${TRITON_KERNEL_REPO} \ + && cd triton \ + && git checkout ${TRITON_KERNEL_BRANCH} \ + && git submodule update --init --recursive +ENV TRITON_KERNELS_SRC_DIR="/app/patched_triton_kernels/triton/python/triton_kernels/triton_kernels/" + + +# Setup sccache for HIP compilation via HIP_CLANG_PATH +# This creates wrapper scripts in a separate directory and points HIP to use them +# This avoids modifying the original ROCm binaries which can break detection +# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) +# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + + +### +### AMD SMI Build +### +FROM base AS build_amdsmi +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ + && pip wheel . --wheel-dir=dist +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install + + +### +### MORI Build TODO(Build needs fixing) +### +FROM base AS build_mori +ARG MORI_BRANCH +ARG MORI_REPO +ARG MORI_GPU_ARCHS +RUN mkdir -p /app/install; \ + git clone ${MORI_REPO} \ + && cd mori \ + && git checkout ${MORI_BRANCH} \ + && git submodule update --init --recursive \ + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; + + +### +### FlashAttention Build +### +# Remove && git submodule update --init \ for CK +FROM base AS build_fa +ARG FA_BRANCH +ARG FA_REPO +ARG USE_SCCACHE +RUN mkdir -p /app/install; \ + git clone ${FA_REPO} \ + && cd flash-attention \ + && git checkout ${FA_BRANCH} \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; + + +### +### AITER Build +### +FROM base AS build_aiter +ARG AITER_BRANCH +ARG AITER_REPO +ARG USE_SCCACHE +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} +RUN cd aiter \ + && git submodule update --init --recursive \ + && pip install -r requirements.txt +RUN pip install pyyaml && cd aiter \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && ls /app/aiter/dist/*.whl +RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install + + +### +### Final Build +### + +# Wheel release stage - +# only includes dependencies used by wheel release pipeline +FROM base AS debs_wheel_release +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs + +# Full debs stage - includes Mori (used by Docker releases) +FROM base AS debs +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + +FROM base AS final +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ + pip install /install/*.whl + +ARG BASE_IMAGE +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG FA_BRANCH +ARG FA_REPO +ARG AITER_BRANCH +ARG AITER_REPO +ARG MORI_BRANCH +ARG MORI_REPO +RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ + && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ + && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ + && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ + && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 new file mode 100644 index 000000000000..4c0008f6ecb7 --- /dev/null +++ b/docker/Dockerfile.rocm_gfx1250 @@ -0,0 +1,729 @@ +# default base image +ARG REMOTE_VLLM="0" +ARG COMMON_WORKDIR=/app +ARG BASE_IMAGE=rocm/vllm-dev:base +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +# NIC backend for MoRI RDMA support. +# By default (all), drivers and userspace libraries for all supported NIC types +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only +# --build-arg NIC_BACKEND=none # Install nothing. +ARG NIC_BACKEND=all +# AMD AINIC apt repo settings +# Users can specify a custom version compatible with their host drivers. +# The default version has been tested with ioinic-dkms=25.11.1.001 +ARG AINIC_VERSION=1.117.3-hydra +ARG UBUNTU_CODENAME=jammy + +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 +ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} + +# Install build dependencies and utilities +RUN apt-get update -q -y && apt-get install -q -y \ + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + apt-transport-https ca-certificates wget curl \ + build-essential libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". +# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) +ARG USE_SCCACHE +RUN if [ "$USE_SCCACHE" != "1" ]; then \ + apt-get purge -y sccache || true; \ + python3 -m pip uninstall -y sccache || true; \ + rm -f "$(which sccache)" || true; \ + fi + +# Install UV — download first, then run, so a curl failure is not masked by the pipe +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ + && rm -f /tmp/uv-install.sh \ + && uv --version + +# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out +# Reference: https://github.com/astral-sh/uv/pull/1694 +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +# Use copy mode to avoid hardlink failures with Docker cache mounts +ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + if command -v sccache >/dev/null 2>&1; then \ + echo "sccache already installed, skipping installation"; \ + sccache --version; \ + else \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR} + + +# ----------------------- +# vLLM fetch stages +FROM base AS fetch_vllm_0 +ONBUILD COPY ./ vllm/ +FROM base AS fetch_vllm_1 +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" +ENV VLLM_REPO=${VLLM_REPO} +ENV VLLM_BRANCH=${VLLM_BRANCH} +ONBUILD RUN git clone ${VLLM_REPO} \ + && cd vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git checkout FETCH_HEAD \ + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ + git remote add upstream "https://github.com/vllm-project/vllm.git" \ + && git fetch upstream ; fi +FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm + +# ----------------------- +# Rust build stage +# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages +# don't need the rust toolchain or protoc. +FROM fetch_vllm AS rust-build +ARG COMMON_WORKDIR + +# protoc is used by tonic-build/prost-build. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + ca-certificates curl unzip \ + && rm -rf /var/lib/apt/lists/* + +COPY tools/install_protoc.sh /tmp/install_protoc.sh +RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh + +# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit +# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt + +# Build the release binary. Cargo's registry/git caches can be written by +# concurrent BuildKit jobs on shared workers, so lock those cache mounts while +# keeping the cache benefit. Do not cache target/, because stale target metadata +# can outlive source updates across BuildKit cache reuse. +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + cd ${COMMON_WORKDIR}/vllm \ + && bash build_rust.sh \ + && test -x vllm/vllm-rs + +# ----------------------- +# vLLM native build stages +# +# csrc-build intentionally copies only files that affect ROCm native extension +# compilation. That keeps unrelated CI/test/docs edits from invalidating the +# expensive HIP/C++ build layer. +FROM base AS csrc-build +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY requirements/rocm.txt requirements/rocm.txt +COPY requirements/common.txt requirements/common.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/rocm.txt + +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do +# not invalidate the expensive native build layer. +COPY setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/envs.py +COPY vllm/__init__.py vllm/__init__.py + +ENV VLLM_TARGET_DEVICE=rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" + +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + export CCACHE_BASEDIR="$PWD" \ + && echo "=== ccache stats before ROCm native build ===" \ + && (ccache --show-stats || true) \ + && (ccache --zero-stats || true) \ + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null \ + && echo "=== ccache stats after ROCm native build ===" \ + && (ccache --show-stats || true) + +# Build the full vLLM ROCm wheel by reusing the native extension wheel from +# csrc-build. This stage still rebuilds for Python/package changes, but skips +# the expensive HIP/C++ compile when native inputs are unchanged. +FROM fetch_vllm AS build_vllm +ARG COMMON_WORKDIR +ENV VLLM_TARGET_DEVICE=rocm + +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system -r requirements/rocm.txt \ + && export VLLM_USE_PRECOMPILED=1 \ + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ + && export VLLM_DOCKER_BUILD_CONTEXT=1 \ + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null +FROM scratch AS export_vllm +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +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/vllm/v1 /vllm_v1 + +# RIXL/UCX build stages +FROM base AS build_rixl +ARG RIXL_BRANCH="39be1de8" +ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" +ARG UCX_BRANCH="bfb51733" +ARG UCX_REPO="https://github.com/openucx/ucx.git" +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image +ENV UCX_HOME=/usr/local/ucx +ENV RIXL_HOME=/usr/local/rixl +ENV RIXL_BENCH_HOME=/usr/local/rixl_bench + +# RIXL build system dependences and RDMA support +RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ + libgrpc-dev \ + libgrpc++-dev \ + libprotobuf-dev \ + protobuf-compiler-grpc \ + libcpprest-dev \ + libaio-dev \ + librdmacm1 \ + librdmacm-dev \ + libibverbs1 \ + libibverbs-dev \ + ibverbs-utils \ + rdmacm-utils \ + ibverbs-providers \ + && rm -rf /var/lib/apt/lists/* + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit + +RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ + git clone ${UCX_REPO} && \ + cd ucx && \ + git checkout ${UCX_BRANCH} && \ + ./autogen.sh && \ + mkdir build && cd build && \ + CC="ccache gcc" CXX="ccache g++" \ + ../configure \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ + make -j$(nproc) && \ + make install + +ENV PATH=/usr/local/ucx/bin:$PATH +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ + cd /opt/rixl && \ + git checkout ${RIXL_BRANCH} && \ + CC="ccache gcc" CXX="ccache g++" \ + meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ + cd build && \ + ninja -j$(nproc) && \ + ninja install + +# Generate RIXL wheel +# Exclude libcore and libpull from auditwheel: transitive dependencies +# that are not shipped in the wheel and vary across base images. +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ + mkdir -p /app/install && \ + _ucx_install_dir=${UCX_HOME} \ + ./contrib/build-wheel.sh \ + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem +ARG ROCSHMEM_BRANCH="f0acb0c6" +ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image +ENV ROCSHMEM_DIR=/opt/rocshmem + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF + +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep +ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" +ARG DEEPEP_NIC="cx7" + +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +RUN --mount=type=cache,target=/root/.cache/ccache \ + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ + && git clone ${DEEPEP_REPO} \ + && cd DeepEP \ + && git checkout ${DEEPEP_BRANCH} \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + +# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do +# not force users to rebuild the long-lived Dockerfile.rocm_base image. +FROM base AS mori_base +ARG NIC_BACKEND +ARG AINIC_VERSION +ARG UBUNTU_CODENAME +RUN /bin/bash -lc 'set -euo pipefail; \ + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' + +# ----------------------- +# vLLM wheel release build stage (for building distributable wheels) +# This stage pins dependencies to custom ROCm wheel versions and handles version detection +FROM fetch_vllm AS build_vllm_wheel_release + +ARG COMMON_WORKDIR + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +# Create /install directory for custom wheels +RUN mkdir -p /install + +# Copy custom ROCm wheels from docker/context if they exist +# COPY ensures Docker cache is invalidated when wheels change +# .keep file ensures directory always exists for COPY to work +COPY docker/context/base-wheels/ /tmp/base-wheels/ +# This is how we know if we are building for a wheel release or not. +# If there are not wheels found there, we are not building for a wheel release. +# So we exit with an error. To skip this stage. +RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ + echo "Found custom wheels - copying to /install"; \ + cp /tmp/base-wheels/*.whl /install/ && \ + echo "Copied custom wheels:"; \ + ls -lh /install/; \ + else \ + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ + echo "Wheel releases require pre-built ROCm wheels."; \ + exit 1; \ + fi + +# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) +# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm +ARG GIT_REPO_CHECK=0 +RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ + echo "Running repository checks..."; \ + cd vllm && bash tools/check_repo.sh; \ + fi + +# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) +# This ensures setuptools_scm sees clean repo state for version detection +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system setuptools_scm regex \ + && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ + && echo "Detected vLLM version: ${VLLM_VERSION}" \ + && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + +# Fail if git-based package dependencies are found in requirements files +# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) +# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. +RUN echo "Checking for git-based packages in requirements files..." \ + && echo "Checking common.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ + echo "ERROR: Git-based packages found in common.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in common.txt"; \ + fi \ + && echo "Checking rocm.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ + echo "ERROR: Git-based packages found in rocm.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in rocm.txt"; \ + fi \ + && echo "All requirements files are clean - no git-based packages found" + +# Pin vLLM dependencies to exact versions of custom ROCm wheels +# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi +COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py +RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ + && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt + +# Install dependencies using custom wheels from /install +RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && echo "Building vLLM with custom wheels from /install" \ + && uv pip install --system --find-links /install -r requirements/rocm.txt + +# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# (setup.py auto-detects ccache/sccache in PATH) +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + cd vllm \ + && export CCACHE_BASEDIR="$PWD" \ + && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist + +FROM scratch AS export_vllm_wheel_release +ARG COMMON_WORKDIR +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +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/vllm/v1 /vllm_v1 + +# ----------------------- +# CI base image (Tier 1) - stable, rarely changing CI dependencies. +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers +# in the vLLM artifacts for the current commit. +FROM mori_base AS ci_base +ARG COMMON_WORKDIR + +# Update rdma-core to support latest rocshmem. +ARG DEEPEP_NIC +RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ + git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && \ + mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ +fi + +# Install RIXL + DeepEP wheels. +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl + +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem + +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + librdmacm1 \ + libibverbs1 \ + ibverbs-providers \ + ibverbs-utils \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ + && rm -rf /var/lib/apt/lists/* + +# Install torchcodec from source for ROCm/torch ABI compatibility. +COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ + && rm /tmp/install_torchcodec.sh \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Pre-install shared ROCm runtime dependencies. +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + && rm -rf /tmp/ci-base-requirements + +# Enable fast and less brittle model downloads in tests. +ENV HF_XET_HIGH_PERFORMANCE=1 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60 + +# Pre-install vLLM test dependencies. +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/rocm-test-reqs.txt + +# Rebuild fastsafetensors from source so its C++ extension is compiled with +# USE_ROCM and can detect libamdhip64.so at runtime. +RUN --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ + && test -n "${FASTSAFETENSORS_REQ}" \ + && python3 -m pip install --force-reinstall --no-deps \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + && rm /tmp/rocm-test-reqs.txt + +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. +# See: https://github.com/pytorch/pytorch/issues/169857 +ENV MIOPEN_DEBUG_CONV_DIRECT=0 +ENV MIOPEN_DEBUG_CONV_GEMM=0 + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" + +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + +# Hide source under src/ so it won't shadow the installed package in tests. +RUN mkdir src && mv vllm src/vllm + +# ----------------------- +# Final vLLM image +FROM mori_base AS final + +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* + +# Clean up sccache from release image (not needed at runtime) +# This removes the binary and wrappers that may have been installed during build +RUN rm -f /usr/bin/sccache || true \ + && rm -rf /opt/sccache-wrappers || true + +# Unset sccache environment variables for the release image +# This prevents S3 bucket config from leaking into production images +ENV SCCACHE_BUCKET= +ENV SCCACHE_REGION= +ENV SCCACHE_ENDPOINT= +ENV SCCACHE_S3_NO_CREDENTIALS= +ENV SCCACHE_IDLE_TIMEOUT= + +# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. +# Manually remove it so that later steps of numpy upgrade can continue +RUN case "$(which python3)" in \ + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system --upgrade huggingface-hub[cli] + +# Install vLLM using uv (inherited from base stage) +# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system -r requirements/rocm.txt \ + && pip uninstall -y vllm \ + && uv pip install --system *.whl + +# Install RIXL wheel +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + uv pip install --system /rixl_install/*.whl + +ARG COMMON_WORKDIR +ARG BASE_IMAGE +ARG NIC_BACKEND +ARG AINIC_VERSION + +# Copy over the benchmark scripts as well +COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks +COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples +COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +ENV TOKENIZERS_PARALLELISM=false + +# ENV that can improve safe tensor loading, and end-to-end time +ENV SAFETENSORS_FAST_GPU=1 + +# Performance environment variable. +ENV HIP_FORCE_DEV_KERNARG=1 + +# Workaround for ROCm profiler limits +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt + +# Download bench scripts for gfx1250 +RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh + + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + +CMD ["/bin/bash"] + +#Set entrypoint for vllm-openai official images +FROM final AS vllm-openai +ENTRYPOINT ["vllm", "serve"] diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index c2adfafdee27..d645e37e78c3 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -61,13 +61,13 @@ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" FROM python-install AS torch-vision # Install torchvision -ARG TORCH_VISION_VERSION=v0.26.0 +ARG TORCH_VISION_VERSION=v0.28.0 WORKDIR /tmp RUN --mount=type=cache,target=/root/.cache/uv \ git clone https://github.com/pytorch/vision.git && \ cd vision && \ git checkout $TORCH_VISION_VERSION && \ - uv pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu && \ + uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu && \ python setup.py bdist_wheel FROM python-install AS hf-xet-builder @@ -86,6 +86,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p /tmp/hf-xet/dist && \ cp dist/*.whl /tmp/hf-xet/dist/ +# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which +# llvmlite v0.47 does not support; only SystemZ target is needed). +FROM base AS llvm20-build +ARG LLVM_VERSION=20.1.8 +WORKDIR /tmp +RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \ + curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \ + tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \ + cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=/opt/llvm20 \ + -DLLVM_TARGETS_TO_BUILD="SystemZ" \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_BUILD_TOOLS=OFF \ + -DLLVM_BUILD_UTILS=ON \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_BUILD_TESTS=OFF \ + -DLLVM_INCLUDE_TESTS=OFF \ + -DLLVM_INCLUDE_EXAMPLES=OFF \ + -DLLVM_INCLUDE_BENCHMARKS=OFF && \ + ninja -C build install && \ + rm -rf build llvm-project-${LLVM_VERSION}.src* + # Build numba FROM python-install AS numba-builder @@ -96,11 +119,13 @@ WORKDIR /tmp # Clone all required dependencies RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \ microdnf install ninja-build gcc gcc-c++ -y && \ git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \ git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ cd llvmlite && \ uv pip install 'cmake<4' 'setuptools<70' numpy && \ + CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \ python setup.py bdist_wheel && \ cd ../numba && \ if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \ @@ -158,7 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \ OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \ uv pip install -v \ - $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ $LLVM_WHL_FILE \ diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 3bd16e8629ba..c3b72d1af6dd 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -63,21 +63,6 @@ RUN apt-get update -y && \ python3-pip && \ rm -rf /var/lib/apt/lists/* -# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer. -RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ - echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \ - printf '%s\n' \ - 'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \ - 'Pin: version 2025.3*' \ - 'Pin-Priority: 1001' \ - > /etc/apt/preferences.d/oneapi-2025.3.pref && \ - apt-get update -y && \ - apt-get install -y --no-install-recommends \ - intel-oneapi-compiler-dpcpp-cpp-2025.3 \ - intel-oneapi-mkl-devel-2025.3 \ - intel-oneapi-dnnl-devel-2025.3 && \ - rm -rf /var/lib/apt/lists/* - # Install UMD RUN mkdir neo && \ cd neo && \ @@ -100,22 +85,6 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh \ && uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} ENV PATH="$VIRTUAL_ENV/bin:$PATH" -# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3. -ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" -RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \ - bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \ - rm "${ONECCL_INSTALLER}" && \ - echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ - echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \ - rm -f /opt/intel/oneapi/ccl/latest && \ - ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \ - printf '%s\n' \ - '/opt/intel/oneapi/ccl/2021.15/lib' \ - '/opt/intel/oneapi/mpi/2021.15/lib' \ - '/opt/intel/oneapi/compiler/2025.3/lib' \ - > /etc/ld.so.conf.d/oneapi-ccl.conf && \ - ldconfig - SHELL ["bash", "-c"] CMD ["bash", "-c", "source /root/.bashrc && exec bash"] @@ -135,8 +104,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --upgrade pip - -ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib +ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib CMD ["/bin/bash"] ######################### UCX + NIXL BUILD STAGE ######################### @@ -216,8 +184,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r /workspace/vllm/requirements/xpu.txt && \ uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \ uv pip uninstall triton triton-xpu && \ - uv pip install triton-xpu==3.7.1 && \ - uv pip uninstall oneccl oneccl-devel + uv pip install triton-xpu==3.7.2 # Keep source-dependent layers near the end so frequent code-only changes # don't invalidate heavy dependency and UCX/NIXL layers. diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index 0ae991bd52d8..7062a0be9346 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" { } # Pre-built CI base image (Tier 1). Per-PR builds pull this instead of -# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in +# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in # Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test # stage and is irrelevant when building --target ci_base itself. variable "CI_BASE_IMAGE" { @@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" { # Upstream dependency commit pins -- extracted from Dockerfile.rocm by # ci-bake-rocm.sh at build time. Empty defaults are safe: the cache # functions produce no entries when the variable is empty. -variable "RIXL_BRANCH" { +variable "NIXL_BRANCH" { default = "" } @@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" { default = "" } -variable "RIXL_CACHE_KEY" { +variable "NIXL_CACHE_KEY" { default = "" } @@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" { ]) } -# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP). +# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP). # These stages are pinned to specific upstream commit hashes, so cache keys use # those hashes rather than the Buildkite commit. This means the cache persists # across all vLLM commits as long as the upstream dependency pins don't change. @@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" { function "get_cache_from_rocm_deps" { params = [] result = compact([ - RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), + NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""), DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""), ]) } -function "get_cache_to_rocm_rixl" { +function "get_cache_to_rocm_nixl" { params = [] result = compact([ - RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""), + NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""), ]) } @@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" { # in the registry cache keyed by its upstream commit hash. When ci_base rebuilds # (e.g., requirements change), these stages are cache hits if their upstream # pins haven't changed -- saving ~35min of compilation. -target "rixl-rocm-ci" { +target "nixl-rocm-ci" { inherits = ["_common-rocm", "_ci-rocm"] - target = "build_rixl" + target = "build_nixl" cache-from = get_cache_from_rocm_deps() - cache-to = get_cache_to_rocm_rixl() + cache-to = get_cache_to_rocm_nixl() output = ["type=cacheonly"] } @@ -396,7 +396,7 @@ target "deepep-rocm-ci" { output = ["type=cacheonly"] } -# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.) +# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.) # Invoked by the ensure-ci-base step when the content hash of ci_base-affecting # files drifts from the remote image label. Per-PR builds then pull the result # as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit. @@ -412,7 +412,7 @@ target "ci-base-rocm-ci" { CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "", CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "", ]), - # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages + # Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages # are cache hits even when ci_base itself needs rebuilding. get_cache_from_rocm_deps(), ) @@ -424,5 +424,5 @@ target "ci-base-rocm-ci" { # Group for ci_base builds -- exports dependency stage caches alongside the # ci_base image so future rebuilds can reuse them independently. group "ci-base-rocm-ci-with-deps" { - targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] + targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] } diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl index 6b51781834b2..9d73a6505eab 100644 --- a/docker/docker-bake-rocm.hcl +++ b/docker/docker-bake-rocm.hcl @@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" { # Upstream dependency commit pins. Plain local bake builds use the Dockerfile # ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env # overrides) and writes a small HCL override before invoking CI targets. -variable "RIXL_BRANCH" { +variable "NIXL_BRANCH" { default = "" } @@ -106,7 +106,7 @@ target "test-rocm" { output = ["type=docker"] } -# CI base image target - builds only the ci_base stage (RIXL, DeepEP, +# CI base image target - builds only the ci_base stage (NIXL, DeepEP, # torchcodec, requirements, etc.). Used by the weekly scheduled build and # the auto-rebuild trigger when requirements change in a PR. target "ci-base-rocm" { diff --git a/docker/versions.json b/docker/versions.json index e6839bbb05cf..009fa9c31cae 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -2,7 +2,7 @@ "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py", "variable": { "CUDA_VERSION": { - "default": "13.0.2" + "default": "13.0.3" }, "PYTHON_VERSION": { "default": "3.12" @@ -10,11 +10,14 @@ "UBUNTU_VERSION": { "default": "22.04" }, + "NCCL_VERSION": { + "default": "2.30.7" + }, "BUILD_BASE_IMAGE": { - "default": "nvidia/cuda:13.0.2-devel-ubuntu22.04" + "default": "nvidia/cuda:13.0.3-devel-ubuntu22.04" }, "FINAL_BASE_IMAGE": { - "default": "nvidia/cuda:13.0.2-base-ubuntu22.04" + "default": "nvidia/cuda:13.0.3-base-ubuntu22.04" }, "BUILD_OS": { "default": "ubuntu" @@ -56,7 +59,7 @@ "default": "cuda" }, "DEEPEP_COMMIT_HASH": { - "default": "73b6ea4" + "default": "d4f41e4e93" }, "GIT_REPO_CHECK": { "default": "0" @@ -68,7 +71,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.14" + "default": "0.6.16.post3" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/.nav.yml b/docs/.nav.yml index 7d985fdeb58c..a2fae5d655b9 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -56,7 +56,9 @@ nav: - API Reference: - api/README.md - api/vllm - - CLI Reference: cli + - CLI Reference: + - cli/README.md + - vllm: cli - Community: - community/* - Governance: governance diff --git a/docs/assets/benchmarking/latency-metrics-speculative-decoding-dark.svg b/docs/assets/benchmarking/latency-metrics-speculative-decoding-dark.svg new file mode 100644 index 000000000000..5ff1742ebf9c --- /dev/null +++ b/docs/assets/benchmarking/latency-metrics-speculative-decoding-dark.svg @@ -0,0 +1,77 @@ + + Latency metrics with token bundling + The client receives token 1 at 100 milliseconds, tokens 2 through 4 together at 140 milliseconds, and token 5 at 180 milliseconds. The two observed ITL gaps average 40 milliseconds. TPOT divides the 80 millisecond generation span by the four tokens after the first, giving 20 milliseconds per token. + + + + + + + + + + + Latency with token bundling + 5 tokens · 3 stream outputs + + + + Request + + + 0 ms + + Output 1 + + + T1 + + + 100 ms + + Output 2 · bundled + + + + + T2 + T3 + T4 + + + 140 ms + + Output 3 + + + T5 + + + 180 ms + + + + TTFT 100 ms + + + ITL 40 ms + + ITL 40 ms + + + 80 ms generation span + + + + Mean ITL + (40 ms + 40 ms) ÷ 2 + 40 ms + 2 observed gaps between stream outputs + + + + TPOT + (180 − 100) ms ÷ (5 − 1) + 20 ms/token + 4 output tokens after the first + diff --git a/docs/assets/benchmarking/latency-metrics-speculative-decoding-light.svg b/docs/assets/benchmarking/latency-metrics-speculative-decoding-light.svg new file mode 100644 index 000000000000..06cb45a03f04 --- /dev/null +++ b/docs/assets/benchmarking/latency-metrics-speculative-decoding-light.svg @@ -0,0 +1,77 @@ + + Latency metrics with token bundling + The client receives token 1 at 100 milliseconds, tokens 2 through 4 together at 140 milliseconds, and token 5 at 180 milliseconds. The two observed ITL gaps average 40 milliseconds. TPOT divides the 80 millisecond generation span by the four tokens after the first, giving 20 milliseconds per token. + + + + + + + + + + + Latency with token bundling + 5 tokens · 3 stream outputs + + + + Request + + + 0 ms + + Output 1 + + + T1 + + + 100 ms + + Output 2 · bundled + + + + + T2 + T3 + T4 + + + 140 ms + + Output 3 + + + T5 + + + 180 ms + + + + TTFT 100 ms + + + ITL 40 ms + + ITL 40 ms + + + 80 ms generation span + + + + Mean ITL + (40 ms + 40 ms) ÷ 2 + 40 ms + 2 observed gaps between stream outputs + + + + TPOT + (180 − 100) ms ÷ (5 − 1) + 20 ms/token + 4 output tokens after the first + diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 8cb98a8f4e45..7d069d0ed45f 100644 Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 692cb4918ba6..476171f73226 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -111,6 +111,45 @@ P99 ITL (ms): 8.39 ================================================== ``` +#### Understanding the Latency Metrics + +`vllm bench serve` measures latency at the benchmark client: + +!!! note + Metric terminology is not standardized across benchmarking tools. When + comparing results, use the measurement points and formulas rather than the + metric names alone. This section explains how we refer to them in vLLM. + +- **Time to first token (TTFT)** is the time from sending a request to receiving + its first streamed output. +- **Inter-token latency (ITL)** records the time between consecutive streamed + outputs. The reported ITL statistics aggregate these individual gaps across + all successful requests. +- **Time per output token (TPOT)** is calculated once per request, excluding the + first token, and then aggregated across requests: + + $$ + \text{TPOT} = + \frac{\text{end-to-end latency} - \text{TTFT}} + {\text{number of output tokens} - 1} + $$ + +With standard decoding, each streamed output usually contains one token, so ITL +and TPOT are typically similar. + +With speculative decoding, one streamed output can contain multiple tokens, +such as several accepted draft tokens within a single engine tstep. ITL records +only the gaps between streamed outputs; it does not add zero-duration gaps for +tokens in the same output. TPOT instead amortizes the request's decoding time +over every output token. + +![Latency metrics with bundled tokens (light theme)](../assets/benchmarking/latency-metrics-speculative-decoding-light.svg#only-light) +![Latency metrics with bundled tokens (dark theme)](../assets/benchmarking/latency-metrics-speculative-decoding-dark.svg#only-dark) + +In this example, the benchmark observes two 40 ms ITL samples. The three tokens +in the second streamed output do not create additional ITL samples, so mean ITL +is 40 ms. TPOT is `(180 ms - 100 ms) / (5 - 1) = 20 ms/token`. + #### Results Visualization The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis. @@ -583,6 +622,31 @@ The following arguments can be used to control the ramp-up: - `--ramp-up-start-rps`: The request rate at the beginning of the benchmark. - `--ramp-up-end-rps`: The request rate at the end of the benchmark. +#### Probe Requests + +The benchmark tool also supports sending probe requests alongside the main +workload. This can be useful for measuring how the main workload affects +unrelated traffic sharing the server, e.g. a few requests with large images +stalling a concurrent lightweight request while their multimodal preprocessing +occupies the frontend. + +Setting `--probe-request-rate` to a positive value sends single-token text-only +probe requests at that rate (requests per second) alongside the main workload. +Probes bypass `--max-concurrency` and their latency is reported separately, so +the probe percentiles directly measure the interference that the main workload +inflicts on unrelated requests. + +```bash +vllm bench serve \ + --model Qwen/Qwen2.5-VL-3B-Instruct \ + --backend openai-chat \ + --endpoint /v1/chat/completions \ + --dataset-name random-mm \ + --random-mm-bucket-config '{(2048, 2048, 1): 1.0}' \ + --request-rate 4 \ + --probe-request-rate 20 +``` + #### Load Pattern Configuration vLLM's benchmark serving script provides sophisticated load pattern simulation capabilities through three key parameters that control request generation and concurrency behavior: diff --git a/docs/cli/.nav.yml b/docs/cli/.nav.yml index 586685c5a10a..1a6a5f693865 100644 --- a/docs/cli/.nav.yml +++ b/docs/cli/.nav.yml @@ -1,10 +1,8 @@ nav: - - README.md - - serve.md - - chat.md - - complete.md - - run-batch.md - - vllm bench: - - bench/**/*.md - - vllm launch: - - launch/**/*.md + - "*.md" + - bench: + - bench/*.md + - sweep: + - bench/sweep/*.md + - launch: + - launch/*.md diff --git a/docs/cli/bench/latency.md b/docs/cli/bench/latency.md deleted file mode 100644 index 9e1b90533975..000000000000 --- a/docs/cli/bench/latency.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench latency - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_latency.inc.md" diff --git a/docs/cli/bench/mm_processor.md b/docs/cli/bench/mm_processor.md deleted file mode 100644 index 26746ce12d8e..000000000000 --- a/docs/cli/bench/mm_processor.md +++ /dev/null @@ -1,55 +0,0 @@ -# vllm bench mm-processor - -## Overview - -`vllm bench mm-processor` profiles the multimodal input processor pipeline of -vision-language models. It measures per-stage latency from the HuggingFace -processor through to the encoder forward pass, helping you identify -preprocessing bottlenecks and understand how different image resolutions or -item counts affect end-to-end request time. - -The benchmark supports two data sources: synthetic random multimodal inputs -(`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before -measurement to ensure stable results. - -## Quick Start - -```bash -vllm bench mm-processor \ - --model Qwen/Qwen2-VL-7B-Instruct \ - --dataset-name random-mm \ - --num-prompts 50 \ - --random-input-len 300 \ - --random-output-len 40 \ - --random-mm-base-items-per-request 2 \ - --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \ - --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' -``` - -## Measured Stages - -| Stage | Description | -| ----- | ----------- | -| `get_mm_hashes_secs` | Time spent hashing multimodal inputs | -| `get_cache_missing_items_secs` | Time spent looking up the processor cache | -| `apply_hf_processor_secs` | Time spent in the HuggingFace processor | -| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs | -| `apply_prompt_updates_secs` | Time spent updating prompt tokens | -| `preprocessor_total_secs` | Total preprocessing time | -| `encoder_forward_secs` | Time spent in the encoder model forward pass | -| `num_encoder_calls` | Number of encoder invocations per request | - -The benchmark also reports end-to-end latency (TTFT + decode time) per -request. Use `--metric-percentiles` to select which percentiles to report -(default: p99) and `--output-json` to save results. - -For more examples (HF datasets, warmup, JSON output), see -[Benchmarking CLI — Multimodal Processor Benchmark](../../benchmarking/cli.md#multimodal-processor-benchmark). - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_mm_processor.inc.md" diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md deleted file mode 100644 index 792c6e094b35..000000000000 --- a/docs/cli/bench/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_serve.inc.md" diff --git a/docs/cli/bench/sweep/plot.md b/docs/cli/bench/sweep/plot.md deleted file mode 100644 index d7dc65e6df62..000000000000 --- a/docs/cli/bench/sweep/plot.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep plot - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_plot.inc.md" diff --git a/docs/cli/bench/sweep/plot_pareto.md b/docs/cli/bench/sweep/plot_pareto.md deleted file mode 100644 index 13dffd7f2b5c..000000000000 --- a/docs/cli/bench/sweep/plot_pareto.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep plot_pareto - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_plot_pareto.inc.md" diff --git a/docs/cli/bench/sweep/serve.md b/docs/cli/bench/sweep/serve.md deleted file mode 100644 index 6a8182feb406..000000000000 --- a/docs/cli/bench/sweep/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_serve.inc.md" diff --git a/docs/cli/bench/sweep/serve_workload.md b/docs/cli/bench/sweep/serve_workload.md deleted file mode 100644 index 8c21788e8d93..000000000000 --- a/docs/cli/bench/sweep/serve_workload.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep serve_workload - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_serve_workload.inc.md" diff --git a/docs/cli/bench/throughput.md b/docs/cli/bench/throughput.md deleted file mode 100644 index 66434c87819f..000000000000 --- a/docs/cli/bench/throughput.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench throughput - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_throughput.inc.md" diff --git a/docs/cli/chat.md b/docs/cli/chat.md deleted file mode 100644 index 7b8e718f625f..000000000000 --- a/docs/cli/chat.md +++ /dev/null @@ -1,5 +0,0 @@ -# vllm chat - -## Arguments - ---8<-- "docs/generated/argparse/chat.inc.md" diff --git a/docs/cli/complete.md b/docs/cli/complete.md deleted file mode 100644 index 65d953a7c046..000000000000 --- a/docs/cli/complete.md +++ /dev/null @@ -1,5 +0,0 @@ -# vllm complete - -## Arguments - ---8<-- "docs/generated/argparse/complete.inc.md" diff --git a/docs/cli/json_tip.inc.md b/docs/cli/json_tip.inc.md deleted file mode 100644 index 56c9cb2cc8e2..000000000000 --- a/docs/cli/json_tip.inc.md +++ /dev/null @@ -1,10 +0,0 @@ - -When passing JSON CLI arguments, the following sets of arguments are equivalent: - -- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'` -- `--json-arg.key1 value1 --json-arg.key2.key3 value2` - -Additionally, list elements can be passed individually using `+`: - -- `--json-arg '{"key4": ["value3", "value4", "value5"]}'` -- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'` diff --git a/docs/cli/launch/render.md b/docs/cli/launch/render.md deleted file mode 100644 index 4d15e5f1162d..000000000000 --- a/docs/cli/launch/render.md +++ /dev/null @@ -1,22 +0,0 @@ -# vllm launch render - -## Overview - -`vllm launch render` starts a GPU-less rendering server for preprocessing and -postprocessing only. - -```bash -vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 -``` - -This command reuses the standard serving parser, so model, frontend, -networking, and related CLI options follow the same conventions as -[`vllm serve`](../serve.md). - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/launch_render.inc.md" diff --git a/docs/cli/run-batch.md b/docs/cli/run-batch.md deleted file mode 100644 index f2255e66373d..000000000000 --- a/docs/cli/run-batch.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm run-batch - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/run-batch.inc.md" diff --git a/docs/cli/serve.md b/docs/cli/serve.md deleted file mode 100644 index 0326fe29ec7f..000000000000 --- a/docs/cli/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/serve.inc.md" diff --git a/docs/configuration/engine_args.md b/docs/configuration/engine_args.md index b619cbf3db02..11886cca7be9 100644 --- a/docs/configuration/engine_args.md +++ b/docs/configuration/engine_args.md @@ -11,12 +11,4 @@ Engine arguments control the behavior of the vLLM engine. The engine argument classes, [EngineArgs][vllm.engine.arg_utils.EngineArgs] and [AsyncEngineArgs][vllm.engine.arg_utils.AsyncEngineArgs], are a combination of the configuration classes defined in [vllm.config][]. Therefore, if you are interested in developer documentation, we recommend looking at these configuration classes as they are the source of truth for types, defaults and docstrings. ---8<-- "docs/cli/json_tip.inc.md" - -## `EngineArgs` - ---8<-- "docs/generated/argparse/engine_args.inc.md" - -## `AsyncEngineArgs` - ---8<-- "docs/generated/argparse/async_engine_args.inc.md" +--8<-- "gen:engine-args" diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 34dc385db78d..89acc6b7f5ae 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -301,8 +301,10 @@ review process: isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion. - Note that not all CI checks will be executed due to limited computational - resources. The reviewer will add `ready` label to the PR when the PR is - ready to merge or a full CI run is needed. + resources. Reviewers with write access and configured trusted contributors + can comment `/ci run` when CI signals are needed before a PR is ready. After + the PR is approved or has the `ready` label, the PR author can use `/ci run` + or `/ci retry`. New commits do not start CI automatically. ### Pull Request Limits and Escalation diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index b076ef84a46c..f82bafa44d52 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -195,7 +195,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics: The API server takes care of basic audio I/O and optional chunking before building prompts: - Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`. -- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`. +- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point. - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words. Relevant server logic: diff --git a/docs/deployment/docker.md b/docs/deployment/docker.md index a8debf2cdb3a..b5f3a8b049b6 100644 --- a/docs/deployment/docker.md +++ b/docs/deployment/docker.md @@ -8,6 +8,26 @@ toc_depth: 2 --8<-- "docs/getting_started/installation/gpu.md:pre-built-images" +## Persist the compile cache across containers + +Mounting the Hugging Face cache keeps model weights across containers, but each +new container still starts with an empty `VLLM_CACHE_ROOT` (default +`~/.cache/vllm`) and recompiles the model's `torch.compile` artifacts. Mount a +named volume at that path to reuse the inductor, Triton, and AOT artifacts from +the second container onward: + +```bash +docker run --rm --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -v vllm-cache:/root/.cache/vllm \ + -p 8000:8000 \ + vllm/vllm-openai:latest \ + meta-llama/Llama-3.1-8B-Instruct +``` + +See [Faster Startup](../configuration/optimization.md#faster-startup) for the +mechanism and for what invalidates the cache. + ## Run as a non-root user The CUDA `vllm/vllm-openai` image runs as root by default for backward diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md index 6060b98f6421..7d261eb910b6 100644 --- a/docs/deployment/integrations/llm-d.md +++ b/docs/deployment/integrations/llm-d.md @@ -1,5 +1,37 @@ # llm-d -vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers. +[llm-d](https://llm-d.ai/) is a Kubernetes-native distributed inference framework for serving large language models at scale, with vLLM as its primary inference engine. llm-d coordinates a fleet of vLLM instances across a cluster so that performance holds up under real production traffic, achieving the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators. -You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). +It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d-to-the-cncf-evolving-kubernetes-into-sota-ai-infrastructure/) founded by Red Hat, Google Cloud, IBM Research, CoreWeave, and NVIDIA. + +## What llm-d adds to vLLM + +A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own: + +- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. +- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. +- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. +- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. +- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. + +These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear. + +## Performance + +Representative benchmarked results across accelerators: + +- **3x higher output throughput** and **2x faster TTFT** from prefix-aware routing vs round-robin (Llama 3.1 70B, AMD MI300X) +- **Up to 70% higher tokens/sec** from prefill/decode disaggregation (GPT-OSS, NVIDIA B200) +- **13.9x throughput** from hierarchical KV offloading at high concurrency vs GPU-only (NVIDIA H100) + +See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and reproducible benchmarks on [Prism](https://prism.llm-d.ai/). + +## Get started + +1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. +2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. +3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. + +You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). + +Questions and contributions are welcome on [GitHub](https://github.com/llm-d/llm-d) and [Slack](https://llm-d.ai/slack). diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 8062ce237db2..c4cedd9e4e14 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -1,15 +1,9 @@ # Attention Backend Feature Support -This document is auto-generated by `tools/pre_commit/generate_attention_backend_docs.py`. -It shows the feature support for each registered attention backend -based on the checks in `AttentionBackend.validate_configuration()`. - -**Do not edit this file manually.** Run the following command to -regenerate it: - -```bash -python tools/pre_commit/generate_attention_backend_docs.py -``` +The priority and feature tables on this page are auto-generated from the +attention backend registry by +`docs/mkdocs/gen_files/generate_attention_backends.py`, based on the checks in +`AttentionBackend.validate_configuration()`. ## Setting the Attention Backend @@ -98,40 +92,11 @@ Priority is **1 = highest** (tried first). ### Standard Attention (MHA, MQA, GQA) -**Blackwell (SM 10.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASHINFER` | -| 2 | `FLASH_ATTN` | -| 3 | `TRITON_ATTN` | -| 4 | `FLEX_ATTENTION` | -| 5 | `TURBOQUANT` | - -**Ampere/Hopper (SM 8.x-9.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASH_ATTN` | -| 2 | `FLASHINFER` | -| 3 | `TRITON_ATTN` | -| 4 | `FLEX_ATTENTION` | -| 5 | `TURBOQUANT` | +--8<-- "gen:priority-standard" ### MLA Attention (DeepSeek-style) -**Blackwell (SM 10.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASHINFER_MLA` | -| 2 | `TOKENSPEED_MLA` | -| 3 | `CUTLASS_MLA` | -| 4 | `FLASH_ATTN_MLA` | -| 5 | `FLASHMLA` | -| 6 | `TRITON_MLA` | -| 7 | `FLASHINFER_MLA_SPARSE`**\*** | -| 8 | `FLASHMLA_SPARSE` | +--8<-- "gen:priority-mla" > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > @@ -157,24 +122,7 @@ Priority is **1 = highest** (tried first). ## Standard Attention (MHA, MQA, GQA) Backends -| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A | -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x | -| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 | -| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | -| `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`, `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` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %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 | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any | -| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | -| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-standard" > **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`. > @@ -188,9 +136,7 @@ automatic priority lists above. A lightning indexer scores KV blocks, the top-k blocks (plus fixed init/local blocks) are selected, and attention attends only to those blocks; index keys live in a separate side cache. -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-minimax" ## MLA (Multi-head Latent Attention) Backends @@ -203,38 +149,20 @@ To explicitly select a prefill backend, use Otherwise, the prefill backend is selected automatically at runtime based on hardware and configuration. -| Backend | Description | Dtypes | Compute Cap. | Notes | -| ------- | ----------- | ------ | ------------ | ----- | -| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) | -| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only | -| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | -| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | +--8<-- "gen:table-mla-prefill" > **‡** Automatic selection tries FlashAttention first. On Blackwell > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then -> TokenSpeed MLA. On other GPUs, only FlashAttention is considered. +> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64, +> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other +> GPUs, only FlashAttention is considered. ### Decode Backends MLA decode backends are selected using the standard `-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`). -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x | -| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | -| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-mla-decode" ### DeepSeek V4 Decode Backends @@ -245,8 +173,4 @@ pipeline (compressor + SWA + indexer, 256-token blocks, head 512); default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and `FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures. -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x | -| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +--8<-- "gen:table-mla-v4-decode" diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 7eab425d6e21..e7a847f4c742 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -41,13 +41,13 @@ class BudgetGraphMetadata: Budgets are auto-generated as power-of-2 levels from a model-provided range via `get_encoder_cudagraph_budget_range()`, with the maximum budget always included even if it does not fall on a power-of-2 boundary. Budgets can also be explicitly specified by the user via `encoder_cudagraph_token_budgets` in `CompilationConfig`. -When `EncoderCudaGraphConfig.enable_dual_path_graph` is `True`, the manager generates two independent budget lists — `global_token_budgets` (multiples of `global_token_per_image`) and `local_token_budgets` (multiples of `local_token_per_patch`) — and stores captured graphs under `budget_graphs["global"]` and `budget_graphs["local"]` respectively. +Each entry in `EncoderCudaGraphConfig.paths` defines an independently captured encoder path. A path can provide its own minimum token budget and opt into zero-token batches; the manager generates and stores a separate budget graph set for every configured path. ### Greedy bin-packing at runtime When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. -For dual-path models, the manager routes to `_execute_local_dual_path()`, which constrains both global and local token budgets simultaneously during packing (see [Dual-Path graph capture](#dual-path-graph-capture)). +For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)). For each graph replay: @@ -56,38 +56,29 @@ For each graph replay: 3. Replay the CUDA Graph. 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). -### Dual-Path graph capture +### Multi-Path graph capture -For two-tower vision encoders (e.g., DeepSeek-OCR), the `EncoderCudaGraphConfig` sets `enable_dual_path_graph=True` and provides `global_token_per_image` / `local_token_per_patch`. The manager captures two independent sets of CUDA graphs — one for the **global** image path and one for the **local** patch path — stored under `budget_graphs["global"]` and `budget_graphs["local"]` respectively. +`EncoderCudaGraphConfig.paths` maps path names to `EncoderCudaGraphPathConfig` capture policies. For example, DeepSeek-OCR configures a **global** image path and a **local** patch path, which are captured independently under `budget_graphs["global"]` and `budget_graphs["local"]`. -**Budget generation.** Two separate budget lists are generated: +**Budget generation.** Each path gets a separate budget list. For DeepSeek-OCR: -* `global_token_budgets` — power-of-2 multiples of `global_token_per_image` (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR). -* `local_token_budgets` — power-of-2 multiples of `local_token_per_patch` (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800]` for DeepSeek-OCR). A budget of `0` is always included to handle images with no local patches (images ≤ 640×640 that produce only global features). +* the `global` path — power-of-2 budgets starting at the global path minimum (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR). +* the `local` path — power-of-2 budgets starting at the local path minimum (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800, 13824]` for DeepSeek-OCR). A budget of `0` is included when `allow_zero_tokens=True` to handle images with no local patches (images ≤ 640×640 that produce only global features). Both lists are capped at the same `max_budget`. -**Dual-path greedy packing.** Each `EncoderItemSpec` provides both `global_output_tokens` (constant per image) and `local_output_tokens` (proportional to the patch count). The dual-path packing algorithm constrains both budgets simultaneously: +**Multi-path greedy packing.** Each `EncoderItemSpec` provides `path_output_tokens`, mapping each path to its contribution for that item. The packing algorithm constrains every path simultaneously: * Sort images by total output tokens (global + local), smallest first. -* Greedily pack images: an image is added to the current sub-batch only if both the accumulated global tokens ≤ `max_global_budget` **and** the accumulated local tokens ≤ `max_local_budget`, with the image count ≤ `max_batch_size`. -* Once either constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path. +* Greedily pack images: an image is added to the current sub-batch only if every accumulated path token count is within that path's maximum budget, with the image count ≤ `max_batch_size`. +* Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path. * Repeat until all images are packed. -**Partial graph fallback.** After packing, each sub-batch falls into one of four execution scenarios: - -| Global budget | Local budget | Execution | -| :---: | :---: | --- | -| Found | Found | Both paths use CUDA graph replay | -| Found | `None` | Global graph replay + local path skipped (no patches) | -| `None` | Found | Global eager fallback + local graph replay | -| `None` | `None` | Both paths fall back to eager execution | - -Note that the `0`-budget graph is never actually replayed for local — it signals that local patch processing should be skipped entirely. +**Partial graph fallback.** For each non-empty path, the manager replays the smallest fitting graph or runs only that path eagerly when no graph fits. Paths with zero tokens are skipped; a `0`-budget graph is never captured or replayed. **Buffer keys per path.** Global and local paths use different buffer keys. For DeepSeek-OCR, the global path uses `pixel_values` (full images, shape `[B, 3, 1280, 1280]`) while the local path uses `images_crop` (patches, shape `[P, 3, 1024, 1024]`). The manager iterates over each captured graph's own `input_buffers.keys()` rather than a shared `buffer_keys` list, so both paths can use different buffers. -**Post-processing.** The `postprocess_encoder_output` method receives a `local_output` parameter (a tensor or `None`) containing the local-path encoder output. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, this means reshaping the global output into `[B, 272, n_embed]`, the local output into `[P, 100, n_embed]`, assembling patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image. +**Post-processing.** The `postprocess_encoder_output` method receives an `outputs` dictionary keyed by path name. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, it reshapes the global output into `[B, 272, n_embed]`, reshapes the local output into `[P, 100, n_embed]`, assembles patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image. !!! note The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image. @@ -101,7 +92,7 @@ When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. !!! note - Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture. + Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture. Mixed inputs (image+video) per prompt are also supported now. @@ -111,13 +102,13 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra * `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video). * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets. -* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and optionally per-path token counts (`global_output_tokens`, `local_output_tokens`) for dual-path models. +* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and per-path token counts (`path_output_tokens`) for multi-path models. * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding. * `prepare_encoder_cudagraph_capture_inputs(..., path="default")` — creates dummy inputs for graph capture. The `path` parameter (`"global"` or `"local"`) tells the model which path to generate dummy inputs for. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph. * `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch, path="default")` — computes buffer values from actual batch inputs. The `path` parameter selects which modality keys to extract from `mm_kwargs`. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match the captured graph's `input_buffers.keys()`. * `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor], path="default")` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `path` parameter dispatches to the correct encoder sub-module (e.g., global vs. local path for DeepSeek-OCR). * `encoder_eager_forward(mm_kwargs, path="default")` — fallback eager forward when no graph fits. When `path` is `"global"` or `"local"`, runs only that encoder path without graph capture. -* `postprocess_encoder_output(..., local_output=None)` — post-process encoder output. The `local_output` parameter receives the local-path encoder output tensor (or `None`), enabling dual-path models to assemble global and local features into the final per-image embedding. +* `postprocess_encoder_output(outputs, ...)` — post-process encoder outputs keyed by path name, enabling multi-path models to assemble their path-specific features into the final per-item embedding. !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. @@ -129,6 +120,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | @@ -145,14 +137,14 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra ## Configuration -Three fields in `CompilationConfig` control encoder CUDA Graphs: +Four fields in `CompilationConfig` control encoder CUDA Graphs: * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level. * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference. * `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. * `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode). -Dual-path mode is configured at the model level via `EncoderCudaGraphConfig` fields (`enable_dual_path_graph`, `global_token_per_image`, `local_token_per_patch`) — no additional user configuration is required. The manager automatically generates separate budget lists and routes to dual-path execution when the model opts in. +Multi-path mode is configured at the model level through `EncoderCudaGraphConfig.paths`. Each `EncoderCudaGraphPathConfig` can set a path-specific minimum budget and whether zero-token batches are allowed. The manager automatically generates separate budget lists and uses the same execution loop for single- and multi-path models. ## Usage guide diff --git a/docs/design/custom_op.md b/docs/design/custom_op.md index d2557a2281cf..f99c8043fa04 100644 --- a/docs/design/custom_op.md +++ b/docs/design/custom_op.md @@ -122,8 +122,6 @@ For example: --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm" ---8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer" - --8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv" ``` diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 2654b323ff06..4be5ff1de63e 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -4,7 +4,7 @@ FusedMoEModularKernel is implemented [here](../../vllm/model_executor/layers/fused_moe/modular_kernel.py) -Based on the format of the input activations, FusedMoE implementations are broadly classified into 2 types. +Based on the format of the input activations, fused MoE implementations are broadly classified into 2 types. * Contiguous / Standard / Non-Batched, and * Batched @@ -17,24 +17,24 @@ The input activation format completely depends on the All2All Dispatch being use * In the Contiguous variant, the All2All Dispatch returns the activations as a contiguous tensor of shape (M, K) along with TopK Ids and TopK weights of shape (M, num_topk). Look at `DeepEPHTPrepareAndFinalize` for an example. * In the Batched variant, the All2All Dispatch returns the activations as a tensor of shape (num_experts, max_tokens, K). Here, the activations/tokens that subscribe to the same expert are batched together. Note that not all entries of the tensor are valid. The activations tensor is typically accompanied by an `expert_num_tokens` tensor of size `num_experts`, where `expert_num_tokens[i]` indicates the number of valid tokens that subscribe to the ith expert. Look at `DeepEPLLPrepareAndFinalize` for an example. -The FusedMoE operation is generally made of multiple operations, in both the Contiguous and Batched variants, as described in the diagrams below +The fused MoE operation is generally made of multiple operations, in both the Contiguous and Batched variants, as described in the diagrams below -![FusedMoE Non-Batched](../assets/design/fused_moe_modular_kernel/fused_moe_non_batched.png) +![Fused MoE Non-Batched](../assets/design/fused_moe_modular_kernel/fused_moe_non_batched.png) -![FusedMoE Batched](../assets/design/fused_moe_modular_kernel/fused_moe_batched.png) +![Fused MoE Batched](../assets/design/fused_moe_modular_kernel/fused_moe_batched.png) !!! note The main difference, in terms of operations, between the Batched and Non-Batched cases is the Permute / Unpermute operations. All other operations remain. ## Motivation -As can be seen from the diagrams, there are a lot of operations and there can be a variety of implementations for each operation. The set of ways the operations can be put together to make a valid FusedMoE implementation quickly becomes intractable. The Modular Kernel framework addresses this issue, by grouping the operations into logical components. This broad categorization makes the combinations manageable and prevents code-duplication. This also decouples the All2All Dispatch & Combine implementations from the FusedMoE implementations and allows for their independent development and testing. Furthermore, the Modular Kernel framework introduces Abstract classes for the different components thus providing a well-defined skeleton for future implementations. +As can be seen from the diagrams, there are a lot of operations and there can be a variety of implementations for each operation. The set of ways the operations can be put together to make a valid fused MoE implementation quickly becomes intractable. The Modular Kernel framework addresses this issue, by grouping the operations into logical components. This broad categorization makes the combinations manageable and prevents code-duplication. This also decouples the All2All Dispatch & Combine implementations from the fused MoE implementations and allows for their independent development and testing. Furthermore, the Modular Kernel framework introduces Abstract classes for the different components thus providing a well-defined skeleton for future implementations. The rest of the document will focus on the Contiguous / Non-Batched case. Extrapolating to the Batched case should be straight-forward. ## ModularKernel Components -FusedMoEModularKernel splits the FusedMoE operation into 3 parts, +FusedMoEModularKernel splits the fused MoE operation into 3 parts, 1. TopKWeightAndReduce 2. FusedMoEPrepareAndFinalizeModular @@ -81,7 +81,7 @@ The `apply` method is where the implementations perform #### workspace_shapes() -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. +The core fused MoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the fused MoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the fused MoE implementation. #### finalize_weight_and_reduce_impl() @@ -163,7 +163,7 @@ We suggest picking an already existing `FusedMoEPrepareAndFinalizeModular` imple ### How To Add a FusedMoEExpertsModular Type -FusedMoEExpertsModular performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows, +FusedMoEExpertsModular performs the core of the fused MoE operations. The various functions exposed by the abstract class and their significance is as follows, `FusedMoEExpertsModular::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format. @@ -205,7 +205,7 @@ derived classes. Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). -**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. +**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what fused MoE implementation is used. ### How To Unit Test @@ -242,4 +242,4 @@ See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2a ## FusedMoEExpertsModular -See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts. +See [Fused MoE Kernel features](./moe_kernel_features.md#fused-experts-kernels) for a list of all the available modular experts. diff --git a/docs/design/fusions.md b/docs/design/fusions.md index c9991f75cdb8..002984c113a1 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) -- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu) +- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu) ### SiLU+Mul + Quantization (`fuse_act_quant`) @@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) -- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) +- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu) ### RMSNorm + Padding (`fuse_act_padding`) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 07d2a5398013..59271c481283 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -4,7 +4,7 @@ The purpose of this document is to provide an overview of the various MoE kernel ## Fused MoE Modular All2All backends -There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalizeModular` subclasses provide an interface for each all2all backend. +There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `MoERunner` layer. The different `FusedMoEPrepareAndFinalizeModular` subclasses provide an interface for each all2all backend. The following table describes the relevant features of each backend, i.e. activation format, supported quantization schemes and async support. @@ -32,7 +32,7 @@ th { | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass | | ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- | -| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] | +| naive | standard | all1 | G,A,T | N | 6 | [`MoERunner`][vllm.model_executor.layers.fused_moe.runner.moe_runner.MoERunner] | | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index b99ba6659f74..6f66be5eb9f5 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -47,7 +47,8 @@ sequenceDiagram else only one side present PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched end - PWriter->>PWriter: ensure D handshake (one-time) + PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE) + PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif note over DWorker,DWriter: D side - completion accounting @@ -100,15 +101,20 @@ event: D, completion notifs after a WRITE, late-arriving ``PUSH_REG``) even when there is no new metadata to act on. 3. **Handshake-completion callback** (background handshake executor - thread) — when a deferred D→P handshake finishes successfully, the - future's done-callback re-enqueues the registration onto - ``_reg_send_inbox`` and sets the wake so the corresponding - ``send_notif`` runs on the writer (we never call ``send_notif`` from - the executor thread). On this second pass ``_ensure_handshake`` - returns ``None`` (the agent is now connected), so the writer sends - the ``PUSH_REG`` directly. If the handshake *failed*, the callback - fails the request instead of re-enqueuing, so there is no retry - loop. + thread) — both handshakes run on the executor and never block the + writer; their done-callbacks re-enqueue the deferred op and set the + wake, since neither ``send_notif`` nor the NIXL WRITE may run off the + writer thread: + * the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the + registration onto ``_reg_send_inbox``; + * the **P→D** handshake (before a WRITE) re-enqueues the matched + ``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``. + + On this second pass ``_ensure_handshake`` returns ``None`` (the agent + is now connected), so the writer sends the ``PUSH_REG`` / issues the + WRITE directly. If a handshake *failed*, the callback fails or drops + the request instead of re-enqueuing, so there is no retry loop (see + Failure handling). In addition to event-driven wakes, the writer self-polls at ``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side @@ -227,6 +233,13 @@ Two per-request timers are armed on the scheduler: * **D-side ``send_notif`` failure when shipping the PUSH_REG to P** — identical handling: ``_handle_failed_transfer`` marks the recv as failed. +* **P-side handshake failure (P→D handshake before a WRITE)** — the + future's done-callback logs ``push_handshake_failed`` and drops the + request without re-queuing. It deliberately does *not* call + ``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to + invalidate on the producer side, same reasoning as the WRITE-submission + failure below). P's blocks are reclaimed by the ``_kv_lease_duration`` + lease and D's stale registration by its watchdog. * **P-side WRITE submission failure** — the WRITE handle (if any) is released and ``xfer_stats.record_failed_transfer()`` bumps the failure counter. We deliberately do not call @@ -246,10 +259,11 @@ existing NIXL connector: class — all subclasses of the existing base classes; * one dedicated background thread per worker; * a few cross-thread queues, each with a single consumer (the writer); - most have one producer, except ``_reg_send_inbox``, which is fed both - by the engine main thread (new registrations) and by the - handshake-completion callback (registrations replayed after their - D→P handshake finishes); + most have one producer, except the two replay queues fed by both the + engine main thread and a handshake-completion callback: + ``_reg_send_inbox`` (registrations replayed after their D→P handshake) + and ``_deferred_push_inbox`` (matched pushes replayed after their P→D + handshake); * one new notification type (`PUSH_REG:`). Behavior on the engine main thread is otherwise unchanged. The writer diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md index 578343096df7..4c5265449c97 100644 --- a/docs/features/disagg_prefill.md +++ b/docs/features/disagg_prefill.md @@ -49,6 +49,34 @@ Now supports 9 types of connectors: --kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}' ``` +## Reusing prefill token ids on decode + +!!! note + This applies to disaggregated prefill and decode serving on the `/v1/chat/completions` endpoint, using a KV connector configured as in the Usage example above. It is experimental and subject to change. + +In disaggregated serving, the prefill and decode stages both render the chat prompt from `messages` and tokenize it. Because the prefill stage has already produced the token ids, the decode stage can reuse them and skip its own templating and tokenization. The output is otherwise identical to a normal chat completion: it is detokenized to text, and tool and reasoning parsing, streaming, and structured output constraints all still apply. + +The token ids are passed to the decode stage through `kv_transfer_params`, the dict already attached to the decode request to coordinate the transfer: + +1. Send the prefill request with `return_token_ids` enabled, and read `prompt_token_ids` from the response. +2. Set `kv_transfer_params["prompt_token_ids"]` to those ids on the decode request. `messages` is still required, but its content is not tokenized when the ids are present. + +```python +prefill = client.chat.completions.create( + model=model, + messages=messages, + extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}}, +) +ids = prefill.prompt_token_ids + +decode = client.chat.completions.create( + model=model, + messages=messages, + stream=True, + extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}}, +) +``` + ## Development We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance. diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 6349b6329154..72f838c7d2ca 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -68,15 +68,47 @@ vllm serve \ | --- | --- | --- | --- | --- | | `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. | | `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). | -| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. | -| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. | +| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. | +| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. | +| `eviction_policy` | no | `lru` | both | Primary tier policy: built-in `lru`/`arc`, or a custom `CachePolicy` name (see [Custom Eviction Policies](#custom-eviction-policies)). | +| `cache_policy_module_path` | no | — | both | Python import path for a custom `CachePolicy` not in the built-in registry. Required only when `eviction_policy` is not built-in and wasn't pre-registered via `CachePolicyFactory` (advanced). | | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. | | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | -| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | +| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | +## Custom Eviction Policies + +`eviction_policy` resolves through `CachePolicyFactory` (`vllm/v1/kv_offload/cpu/policies/factory.py`), which pre-registers the built-in `lru` and `arc` policies. + +### Out-of-tree (recommended) + +Implement `CachePolicy` (`vllm/v1/kv_offload/cpu/policies/base.py`) in your own package — no vLLM fork or patch required — and point `kv_connector_extra_config` at it directly: + +```json +{ + "cpu_bytes_to_use": 10737418240, + "eviction_policy": "MyCachePolicy", + "cache_policy_module_path": "my_package.my_module" +} +``` + +`eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts. + +### Registering a friendly short name (in-process only) + +If you control the process that constructs the vLLM engine (e.g. an embedding application), you can register a short name once at startup instead of repeating the module path in every config: + +```python +from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory + +CachePolicyFactory.register_cache_policy("my_policy", "my_package.my_module", "MyCachePolicy") +``` + +Then set `"eviction_policy": "my_policy"` in `kv_connector_extra_config`, the same as `"lru"`/`"arc"`. This only takes effect within the process that ran the `register_cache_policy` call — it does not help when the server is launched as a separate process (e.g. via the `vllm serve` CLI), where the out-of-tree `cache_policy_module_path` config above is the only option. + ## Secondary Tiers Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields. @@ -156,7 +188,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. -PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes. +The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected. | Key | Required | Default | Notes | | --- | --- | --- | --- | @@ -175,11 +207,76 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on - `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide. - `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision. +#### Orchestration-Layer Protocol + +The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations. + +Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes. + +Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache. + +Each key names the **remote counterpart** this peer transfers with (not this +peer's own role), so the name reads as "the remote ___ I transfer with". + +| Key | Set on | Value fields | Meaning | +| --- | --- | --- | --- | +| `remote_decoder` | prefill producer request | `kv_request_id` | Peer computes KV and keeps it available in CPU cache for the remote decoder to pull. | +| `remote_prefiller` | decode consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer pulls KV from the remote prefiller at the given address (classic P/D disaggregation). | +| `remote_kv_source` | P2P consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer looks up and pulls whatever blocks the remote source currently holds in CPU cache. | + +Field semantics: + +- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID. +- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)). +- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank. + +Allowed and forbidden combinations: + +- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder. +- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together. + +Minimal examples (values that would appear in the request's `kv_transfer_params`): + +```python +# Prefill producer — compute and keep KV for a remote decoder to pull +kv_transfer_params = {"remote_decoder": {"kv_request_id": ""}} + +# Decode consumer — pull KV from a specific prefiller (classic P/D) +kv_transfer_params = { + "remote_prefiller": { + "kv_request_id": "", + "remote_host": "", + "remote_port": 5710, + } +} + +# P2P consumer — pull whatever the source already has cached +kv_transfer_params = { + "remote_kv_source": { + "kv_request_id": "", + "remote_host": "", + "remote_port": 5710, + } +} +``` + +Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above: + +1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)). +2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step. +3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes. +4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits. +5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes). +6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status. +7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine. + +In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached. + ## Tuning Tips - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload. - For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate. -- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size. +- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups. - FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high. - Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk. diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index df33cea05426..387049f7cc71 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -350,6 +350,31 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py) +#### Video Token Pruning + +For supported models, vLLM can prune video tokens after the vision encoder to +reduce prefill time and KV cache usage, at some cost in accuracy. Set +`--video-pruning-rate ` to prune the fraction `q` of video tokens from each +video, and `--video-pruning-method` to choose the training-free algorithm: + +- **`evs`** (Efficient Video Sampling, default): drops the tokens with the + lowest temporal dissimilarity to the previous frame. The first frame is + always fully retained. +- **`vidcom2`** (Video Compression Commander): scores tokens by similarity to + video-level and frame-level feature centers and gives distinctive frames a + larger share of the budget. At least one token per frame is retained. + +```bash +vllm serve Qwen/Qwen3-VL-8B-Instruct \ + --video-pruning-rate 0.75 --video-pruning-method vidcom2 +``` + +!!! note + `evs` is supported by all models implementing multimodal pruning; + `vidcom2` is currently supported by Qwen3-VL only. Unsupported combinations + are rejected at startup. Enabling video pruning also disables encoder CUDA + graphs, since the retained token count becomes data-dependent. + ### Audio Inputs You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary. @@ -396,6 +421,7 @@ full_transcription = " ".join(transcriptions) The `split_audio` function: +- Expects 1D mono audio (`load_audio` downmixes by default) - Splits audio at quiet points to avoid cutting through speech - Uses RMS energy to find low-amplitude regions within the overlap window - Preserves all audio samples (no data loss) @@ -818,16 +844,20 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu #### Video Decoding Backend -vLLM decodes video bytes into frames using a selectable decoding backend. Three +vLLM decodes video bytes into frames using a selectable decoding backend. Five backends are supported: -- `opencv` (default): OpenCV-based decoder. -- `pyav`: PyAV decoder. -- `torchcodec`: TorchCodec (PyTorch-native) decoder. +| Backend | Device | Description | +| --- | --- | --- | +| `opencv` (default) | CPU | OpenCV-based decoder | +| `pyav` | CPU | PyAV decoder | +| `torchcodec` | CPU | TorchCodec (PyTorch-native) decoder | +| `pynvvideocodec` | GPU | NVIDIA PyNvVideoCodec decoder | +| `deepstream` | GPU | NVIDIA DeepStream 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. +The three CPU 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`: @@ -854,6 +884,21 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' ``` +**PyNvVideoCodec-specific parameters:** + +- `hw_decoders`: Maximum number of concurrent hardware decoder slots retained + by each API server process. It must be a positive integer and defaults to `2`, + which is the recommended starting point for concurrent video workloads. + Because vLLM reserves GPU memory for these slots at startup, this value cannot + be overridden per request. Benchmark before increasing it because each + additional slot increases the GPU memory reservation. + +```bash +# Example: explicitly use the recommended 2 hardware decoders +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "pynvvideocodec", "hw_decoders": 2}}' +``` + #### 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. @@ -879,11 +924,55 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ Works with common video formats like MP4 when using OpenCV backends. +#### GPU Video Decoding with PyNvVideoCodec (NVDEC) + +The `pynvvideocodec` backend uses NVIDIA NVDEC to decode the sampled video +frames on the GPU before copying them into host memory for multimodal +preprocessing. For workloads with large videos and relatively light inference, +such as video tagging, this can alleviate bottlenecks in CPU-based video +decoders. + +!!! warning + [CUDA Multi-Process Service (MPS)](https://docs.nvidia.com/deploy/mps/quick-start.html) + is required when using this backend. Video decoding runs in the API server + process while model serving runs in the engine process, so multiple CUDA + processes share the same GPU. Configure and start MPS before starting vLLM. + +You must also set a positive `--mm-ipc-gpu-memory-gb` value to reserve VRAM for +video decoding. vLLM carves this budget out of the memory available to the KV +cache and uses it to bound concurrent frontend decode allocations. If the +budget is exhausted, decode work waits instead of consuming the engine's VRAM +headroom and potentially causing an out-of-memory error while serving requests. + +Select the backend with an environment variable and specify a workload-appropriate +VRAM budget. For example, to reserve 1 GiB: + +```bash +export VLLM_VIDEO_LOADER_BACKEND=pynvvideocodec +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --mm-ipc-gpu-memory-gb 1 +``` + +Alternatively, select it with `--media-io-kwargs`: + +```bash +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "pynvvideocodec"}}' \ + --mm-ipc-gpu-memory-gb 1 +``` + +Choose a budget large enough for the largest sampled video that a single API +server process must decode. When using multiple API server processes, vLLM +divides the configured budget evenly among them. + +For streaming video sources, use the DeepStream backend instead. + #### GPU Video Decoding with DeepStream (NVDEC) By default vLLM decodes video on the CPU. On NVIDIA GPUs you can instead decode directly on the hardware video engine (NVDEC) with the DeepStream backend, which -keeps decoding off the CPU and can significantly increase video throughput. +keeps decoding off the CPU and can significantly increase video throughput. It +is the recommended GPU backend for streaming video sources. Install the backend (Linux x86-64 only): diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 03b05751c14d..a84e3de15b1a 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat - Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions - The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files -For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already. - -- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information -- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt) -- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages +For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source. For non-cuda platform, please install nixl with ucx build from source, instructed as below. diff --git a/docs/features/quantization/fp8_vit_attn.md b/docs/features/quantization/fp8_vit_attn.md index bf628cd8a72a..18f42caa9089 100644 --- a/docs/features/quantization/fp8_vit_attn.md +++ b/docs/features/quantization/fp8_vit_attn.md @@ -3,9 +3,9 @@ For visual understanding workloads with large images (e.g. QHD, 4K) and relatively short text prompts/generation, the ViT encoder attention can become a significant bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM -supports optional FP8 quantization for the ViT encoder attention via the -FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the -cuDNN attention call. +supports optional FP8 quantization for ViT encoder attention via FlashInfer +cuDNN on NVIDIA GPUs and AITER on AMD GPUs. Q/K/V are quantized on-the-fly to +FP8 before the attention call. !!! note - Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`, @@ -15,20 +15,29 @@ cuDNN attention call. requests. Smaller images may see no speedup due to quantization overhead (3 quantization kernel launches + un-padding). - FP8 tensor-core speedup is more pronounced on GB300 than GB200. + - On ROCm, packed variable-length image and video batches are supported by + AITER's native varlen FP8 attention. ## Requirements -- FlashInfer cuDNN backend with cuDNN >= 9.17.1. +- NVIDIA: FlashInfer cuDNN backend with cuDNN >= 9.17.1. +- AMD: AITER with `flash_attn_varlen_fp8_pertensor_func` support on gfx942 + (MI300 series) or gfx950 (MI350 series). ## Usage -Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together -with `--mm-encoder-attn-backend FLASHINFER`: +Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` and selecting +the backend for the current platform: ```bash vllm serve $MODEL \ --mm-encoder-attn-backend FLASHINFER \ --mm-encoder-attn-dtype fp8 + +# AMD ROCm +vllm serve $MODEL \ + --mm-encoder-attn-backend ROCM_AITER_FA \ + --mm-encoder-attn-dtype fp8 ``` By default (no scale file), **dynamic scaling** is used: a 16-entry circular @@ -45,7 +54,7 @@ reuse them to avoid the dynamic overhead: # Step 1: calibrate and save scales (runs dynamic scaling for 16 passes, # then dumps the learned scales to JSON). vllm bench mm-processor \ - --model $MODEL --mm-encoder-attn-backend FLASHINFER \ + --model $MODEL --mm-encoder-attn-backend $MM_ATTN_BACKEND \ --mm-encoder-attn-dtype fp8 \ --mm-encoder-fp8-scale-save-path /path/to/scales.json \ --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \ @@ -53,7 +62,7 @@ vllm bench mm-processor \ # Step 2: serve with static scales (no dynamic overhead). vllm serve $MODEL \ - --mm-encoder-attn-backend FLASHINFER \ + --mm-encoder-attn-backend $MM_ATTN_BACKEND \ --mm-encoder-attn-dtype fp8 \ --mm-encoder-fp8-scale-path /path/to/scales.json ``` @@ -94,6 +103,16 @@ Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases. Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins. +**Complete AITER attention call on MI300X** (BF16 input, 16 heads, +head_dim=72; FP8 includes Q/K/V quantization): + +| Sequence length | AITER BF16 | AITER FP8 | Speedup | +| --------------- | ---------- | --------- | ------- | +| 2304 | 0.467 ms | 0.337 ms | **1.38x** | +| 4096 | 0.812 ms | 0.764 ms | **1.06x** | +| 8192 | 2.555 ms | 2.364 ms | **1.08x** | +| 16384 | 9.769 ms | 8.655 ms | **1.13x** | + ## Accuracy ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index ad417bcb30ae..7850dc75ccca 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -19,6 +19,20 @@ following `quantization.quant_algo` values: - `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`). - `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`). +!!! note + For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load + time from the backends available on the current platform (CUTLASS, + FlashInfer, Marlin, and others). On GPUs without a supported native FP4 + GEMM kernel, vLLM falls back to weight-only (W4A16) execution via Marlin + and logs a warning; this may reduce throughput for compute-heavy + workloads. Use `--linear-backend` to override the automatic selection + (this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment + variable). Values relevant to NVFP4 include `cutlass`, + `flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and + `marlin`; the full list is documented under `KernelConfig` on the + [Engine Arguments](../../configuration/engine_args.md) page and shown by + `vllm serve --help=KernelConfig`. + ## Quantizing HuggingFace Models with PTQ You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory. diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index 3950adc0251f..e719d65dd95e 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -28,20 +28,6 @@ uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VE pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` -!!! warning "set `LD_PRELOAD`" - Before use vLLM CPU installed via wheels, make sure TCMalloc is installed and added to `LD_PRELOAD`: - ```bash - # install TCMalloc - sudo apt-get install -y --no-install-recommends libtcmalloc-minimal4 - - # manually find the path - sudo find / -iname *libtcmalloc_minimal.so.4 - TC_PATH=... - - # add them to LD_PRELOAD - export LD_PRELOAD="$TC_PATH:$LD_PRELOAD" - ``` - The `uv` approach works for vLLM `v0.6.6` and later. A unique feature of `uv` is that packages in `--extra-index-url` have [higher priority than the default index](https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes). If the latest public release is `v0.6.6.post1`, `uv`'s behavior allows installing a commit before `v0.6.6.post1` by specifying the `--extra-index-url`. In contrast, `pip` combines packages from `--extra-index-url` and the default index, choosing only the latest version, which makes it difficult to install a development version prior to the released version. #### Install the latest code diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 8b3605e8557d..925f54dcec7a 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel - vLLM CPU supports quantizations: - AWQ (x86 only) - GPTQ (x86 only) - - compressed-tensor INT8 W8A8 (x86, s390x) + - compressed-tensor INT8 W8A8 (x86, s390x only) ### Why do I see `get_mempolicy: Operation not permitted` when running in Docker? diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 15baa487c2a0..0c07c255e2f2 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -10,8 +10,8 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and - OS: `Linux` - SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools -- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above. -- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless` +- Instruction Set Architecture (ISA): VXE support is required. Works with Z15 and above. +- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet` --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -28,13 +28,24 @@ Install the following packages from the package manager before building the vLLM ```bash dnf install -y \ - which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \ + which procps findutils tar vim git patch xz ninja-build \ + gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \ libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \ openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \ clang llvm-devel llvm-static clang-devel ``` -Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation. +Build and install `numactl` from source: + +```bash +curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz +tar -xvzf v2.0.19.tar.gz +cd numactl-2.0.19 +./autogen.sh && ./configure && make && make install +cd .. +``` + +Install rust>=1.80 which is needed for `outlines-core`, `uvloop`, and `hf-xet` python packages installation. ```bash curl https://sh.rustup.rs -sSf | sh -s -- -y && \ @@ -44,26 +55,79 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \ Execute the following commands to build and install vLLM from source. !!! tip - Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM. + Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`. + See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build. + +!!! note "LLVM 20 required for llvmlite" + `llvmlite v0.47` requires LLVM 20, but UBI 9.6 repos ship LLVM 21 which is + not compatible. You must build LLVM 20 from source before building `llvmlite`: + + ```bash + curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-20.1.8/llvm-project-20.1.8.src.tar.xz + tar -xf llvm-project-20.1.8.src.tar.xz + cmake -G Ninja -S llvm-project-20.1.8.src/llvm -B llvm-build \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=/opt/llvm20 \ + -DLLVM_TARGETS_TO_BUILD="SystemZ" \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_BUILD_TOOLS=OFF \ + -DLLVM_BUILD_UTILS=ON \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_BUILD_TESTS=OFF \ + -DLLVM_INCLUDE_TESTS=OFF \ + -DLLVM_INCLUDE_EXAMPLES=OFF \ + -DLLVM_INCLUDE_BENCHMARKS=OFF + ninja -C llvm-build install + ``` + + Then build `llvmlite` pointing to LLVM 20: + + ```bash + CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \ + python setup.py bdist_wheel + ``` ```bash - uv pip install -v \ - -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 +uv pip install -v \ + /path/to/torchvision.whl \ + /path/to/llvmlite.whl \ + /path/to/numba.whl \ + /path/to/opencv_python_headless.whl \ + /path/to/hf_xet.whl \ + -r requirements/build/cpu.txt \ + -r requirements/cpu.txt \ + --torch-backend cpu \ + --index-strategy unsafe-best-match && \ +VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ + uv pip install dist/*.whl ``` ??? console "pip" ```bash - pip install -v \ - --extra-index-url https://download.pytorch.org/whl/cpu \ - -r requirements/build/cpu.txt \ - -r requirements/cpu.txt \ - VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ - pip install dist/*.whl + pip install -v \ + --extra-index-url https://download.pytorch.org/whl/cpu \ + /path/to/torchvision.whl \ + /path/to/llvmlite.whl \ + /path/to/numba.whl \ + /path/to/opencv_python_headless.whl \ + /path/to/hf_xet.whl \ + -r requirements/build/cpu.txt \ + -r requirements/cpu.txt && \ + VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ + pip install dist/*.whl + ``` + +!!! warning "Protobuf workaround for s390x" + The C++ protobuf extension crashes on s390x. After installation, set the + following environment variable and remove the C++ extensions: + + ```bash + export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python + + # Remove C++ protobuf extensions that crash on s390x + SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])") + rm -rf "$SITE_PKGS/google/_upb/"*.so \ + "$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true ``` --8<-- [end:build-wheel-from-source] @@ -80,19 +144,20 @@ docker build -f docker/Dockerfile.s390x \ # Launch OpenAI server docker run --rm \ - --privileged true \ + --security-opt seccomp=unconfined \ + --cap-add SYS_NICE \ --shm-size 4g \ -p 8000:8000 \ -e VLLM_CPU_KVCACHE_SPACE= \ -e VLLM_CPU_OMP_THREADS_BIND= \ vllm-cpu-env \ --model meta-llama/Llama-3.2-1B-Instruct \ - --dtype float \ + --dtype bfloat16 \ other vLLM OpenAI server arguments ``` !!! tip - An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`. + Alternatively, `--privileged=true` also works but is broader and not generally recommended. --8<-- [end:build-image-from-source] --8<-- [start:extra-information] diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index 6ded3b508321..b2ef384bd93b 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -33,19 +33,14 @@ uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VE pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" - Before use vLLM CPU installed via wheels, make sure TCMalloc and Intel OpenMP are installed and added to `LD_PRELOAD`: + Before use vLLM CPU installed via wheels, make Intel OpenMP is added to `LD_PRELOAD`: ```bash - # install TCMalloc, Intel OpenMP is installed with vLLM CPU - sudo apt-get install -y --no-install-recommends libtcmalloc-minimal4 - # manually find the path - sudo find / -iname *libtcmalloc_minimal.so.4 sudo find / -iname *libiomp5.so - TC_PATH=... IOMP_PATH=... - # add them to LD_PRELOAD - export LD_PRELOAD="$TC_PATH:$IOMP_PATH:$LD_PRELOAD" + # add it to LD_PRELOAD + export LD_PRELOAD="$IOMP_PATH:$LD_PRELOAD" ``` #### Install the latest code @@ -131,7 +126,7 @@ uv pip install dist/*.whl ``` !!! warning "set `LD_PRELOAD`" - Before use vLLM CPU installed via wheels, make sure TCMalloc and Intel OpenMP are installed and added to `LD_PRELOAD`: + Before using vLLM CPU installed via wheels, make sure TCMalloc and Intel OpenMP are installed and added to `LD_PRELOAD`: ```bash # install TCMalloc, Intel OpenMP is installed with vLLM CPU sudo apt-get install -y --no-install-recommends libtcmalloc-minimal4 diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 0e86c0e6049c..7755ad6278da 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -384,8 +384,7 @@ A docker container can be built for aarch64 systems such as the Nvidia Grace-Hop -t vllm/vllm-gh200-openai:latest \ --build-arg max_jobs=66 \ --build-arg nvcc_threads=2 \ - --build-arg torch_cuda_arch_list="9.0 10.0+PTX" \ - --build-arg RUN_WHEEL_CHECK=false + --build-arg torch_cuda_arch_list="9.0 10.0+PTX" ``` For (G)B300, we recommend using CUDA 13, as shown in the following command. @@ -398,7 +397,6 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command. --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ --build-arg max_jobs=256 \ --build-arg nvcc_threads=2 \ - --build-arg RUN_WHEEL_CHECK=false \ --build-arg torch_cuda_arch_list='9.0 10.0+PTX' \ --platform "linux/arm64" \ --tag vllm/vllm-gb300-openai:latest \ diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index 8564f2a7265b..ef207c8d83d0 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels. - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers). - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)): -- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue. +- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue. ```bash git clone https://github.com/vllm-project/vllm.git @@ -42,12 +42,12 @@ pip install -v -r requirements/xpu.txt ```bash pip uninstall -y triton triton-xpu - pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu + pip install triton-xpu==3.7.2 --extra-index-url https://download.pytorch.org/whl/xpu ``` !!! note - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues. - - For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu). + - For torch 2.13 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.2`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu). - Finally, build and install vLLM XPU backend: @@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v --8<-- [end:build-wheel-from-source] --8<-- [start:pre-built-images] -Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm). +vLLM offers official Docker images for deployment. +The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags). + +- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0 +- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes + +```bash +docker run --rm \ + --network=host \ + --device /dev/dri:/dev/dri \ + -v /dev/dri/by-path:/dev/dri/by-path \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + --privileged \ + vllm/vllm-openai-xpu: \ + --model Qwen/Qwen3-0.6B +``` + +To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint. + +???+ console "Commands" + ```bash + docker run --rm -it \ + --network=host \ + --device /dev/dri:/dev/dri \ + -v /dev/dri/by-path:/dev/dri/by-path \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + --privileged \ + --entrypoint /bin/bash \ + vllm/vllm-openai-xpu: + ``` --8<-- [end:pre-built-images] --8<-- [start:build-image-from-source] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index f674853f81e5..8a15a669c475 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform: !!! tip A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds. +=== "Intel GPU" + + vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon. + + Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags). + + !!! tip + For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab. + === "Google TPU" To run vLLM on Google TPUs, you need to install the `vllm-tpu` package. diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py similarity index 51% rename from docs/mkdocs/hooks/generate_argparse.py rename to docs/mkdocs/gen_files/generate_argparse.py index 4548e33f8814..c0fa707a1cda 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.metadata import importlib.util +import inspect import logging import sys import textwrap @@ -10,17 +11,21 @@ from collections.abc import Callable, Iterable from importlib.machinery import ModuleSpec from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch +import mkdocs_gen_files +import regex as re from pydantic_core import core_schema logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse" sys.path.insert(0, str(ROOT_DIR)) +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 def mock_if_no_torch(mock_module: str, mock: MagicMock): @@ -132,8 +137,8 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): bench_latency = auto_mock("vllm.benchmarks", "latency") -bench_mm_processor = auto_mock("vllm.benchmarks", "mm_processor") bench_serve = auto_mock("vllm.benchmarks", "serve") +bench_startup = auto_mock("vllm.benchmarks", "startup") bench_sweep_plot = auto_mock("vllm.benchmarks.sweep.plot", "SweepPlotArgs") bench_sweep_plot_pareto = auto_mock( "vllm.benchmarks.sweep.plot_pareto", "SweepPlotParetoArgs" @@ -142,12 +147,28 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): bench_sweep_serve_workload = auto_mock( "vllm.benchmarks.sweep.serve_workload", "SweepServeWorkloadArgs" ) +bench_sweep_startup = auto_mock("vllm.benchmarks.sweep.startup", "SweepStartupArgs") bench_throughput = auto_mock("vllm.benchmarks", "throughput") AsyncEngineArgs = auto_mock("vllm.engine.arg_utils", "AsyncEngineArgs") EngineArgs = auto_mock("vllm.engine.arg_utils", "EngineArgs") ChatCommand = auto_mock("vllm.entrypoints.cli.openai", "ChatCommand") CompleteCommand = auto_mock("vllm.entrypoints.cli.openai", "CompleteCommand") +BenchmarkSubcommand = auto_mock( + "vllm.entrypoints.cli.benchmark.main", "BenchmarkSubcommand" +) +import_bench_subcommands = auto_mock( + "vllm.entrypoints.cli.benchmark.main", "_import_bench_subcommand_modules" +) +BenchmarkSubcommandBase = auto_mock( + "vllm.entrypoints.cli.benchmark.base", "BenchmarkSubcommandBase" +) +BenchmarkMMProcessorSubcommand = auto_mock( + "vllm.entrypoints.cli.benchmark.mm_processor", "BenchmarkMMProcessorSubcommand" +) +LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcommandBase") +launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION") RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand") +sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS") openai_cli_args = auto_mock("vllm.entrypoints.openai", "cli_args") openai_run_batch = auto_mock("vllm.entrypoints.openai", "run_batch") @@ -179,7 +200,7 @@ def end_section(self): def add_text(self, text: str): if text: - self._markdown_output.append(f"{text.strip()}\n\n") + self._markdown_output.append(f"{inspect.cleandoc(text)}\n\n") def add_usage(self, usage, actions, groups, prefix=None): pass @@ -241,49 +262,163 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser: return _parser or parser -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - logger.info("Generating argparse documentation") - logger.debug("Root directory: %s", ROOT_DIR.resolve()) - logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve()) - - # Create the ARGPARSE_DOC_DIR if it doesn't exist - if not ARGPARSE_DOC_DIR.exists(): - ARGPARSE_DOC_DIR.mkdir(parents=True) - - # Create parsers to document - parsers = { - # Engine args - "engine_args": create_parser(EngineArgs.add_cli_args), - "async_engine_args": create_parser( - AsyncEngineArgs.add_cli_args, async_args_only=True - ), - # CLI - "serve": create_parser(openai_cli_args.make_arg_parser), - "chat": create_parser(ChatCommand.add_cli_args), - "complete": create_parser(CompleteCommand.add_cli_args), - "launch_render": create_parser(RenderSubcommand.add_cli_args), - "run-batch": create_parser(openai_run_batch.make_arg_parser), - # Benchmark CLI - "bench_latency": create_parser(bench_latency.add_cli_args), - "bench_mm_processor": create_parser(bench_mm_processor.add_cli_args), - "bench_serve": create_parser(bench_serve.add_cli_args), - "bench_sweep_plot": create_parser(bench_sweep_plot.add_cli_args), - "bench_sweep_plot_pareto": create_parser(bench_sweep_plot_pareto.add_cli_args), - "bench_sweep_serve": create_parser(bench_sweep_serve.add_cli_args), - "bench_sweep_serve_workload": create_parser( - bench_sweep_serve_workload.add_cli_args - ), - "bench_throughput": create_parser(bench_throughput.add_cli_args), - } - - # Generate documentation for each parser - for stem, parser in parsers.items(): - doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md" - # Specify encoding for building on Windows - with open(doc_path, "w", encoding="utf-8") as f: - f.write(super(type(parser), parser).format_help()) - logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR)) - - -if __name__ == "__main__": - on_startup("build", False) +def format_help(parser: FlexibleArgumentParser) -> str: + """Format a parser's help as markdown using `MarkdownFormatter`.""" + return super(type(parser), parser).format_help() + + +# Absolute docs URLs are kept in the help text because they are useful in the +# terminal. Wrap them as markdown links so the `url_schemes` hook can rewrite +# them into doc-relative links / cross-references at render time. +_DOCS_URL = re.compile(r"https://docs\.vllm\.ai/en/[^/\s]+/[^\s)>]+") + + +def linkify_docs_urls(text: str) -> str: + """Wrap bare docs.vllm.ai URLs in help text as markdown links.""" + return _DOCS_URL.sub(lambda m: f"[{m.group()}]({m.group()})", text) + + +logger.info("Generating argparse documentation") +logger.debug("Root directory: %s", ROOT_DIR.resolve()) + +# The JSON tip is always rendered immediately before generated argument content, +# and the generator is its only consumer, so it lives here rather than in a +# separate snippet file. (The runtime terminal equivalent is +# `FlexibleArgumentParser._json_tip` in vllm/utils/argparse_utils.py.) +JSON_TIP = """## JSON CLI Arguments + +When passing JSON CLI arguments, the following sets of arguments are equivalent: + +- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'` +- `--json-arg.key1 value1 --json-arg.key2.key3 value2` + +Additionally, list elements can be passed individually using `+`: + +- `--json-arg '{"key4": ["value3", "value4", "value5"]}'` +- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'` + +""" + +# Argument sections filled into `gen:` markers on handwritten pages +engine_args = create_parser(EngineArgs.add_cli_args) +async_engine_args = create_parser(AsyncEngineArgs.add_cli_args, async_args_only=True) +fill_markers( + "configuration/engine_args.md", + { + "engine-args": ( + f"{JSON_TIP}## `EngineArgs`\n\n" + f"{linkify_docs_urls(format_help(engine_args))}" + f"## `AsyncEngineArgs`\n\n" + f"{linkify_docs_urls(format_help(async_engine_args))}" + ) + }, +) + +# CLI reference pages generated entirely from their parser: page -> (parser, JSON tip) +pages = { + "cli/serve.md": (create_parser(openai_cli_args.make_arg_parser), True), + "cli/chat.md": (create_parser(ChatCommand.add_cli_args), False), + "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False), + "cli/run-batch.md": (create_parser(openai_run_batch.make_arg_parser), True), + "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True), + "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True), + # URL kept as `mm_processor` for back-compat; command name is `mm-processor` + "cli/bench/mm_processor.md": ( + create_parser(BenchmarkMMProcessorSubcommand.add_cli_args), + True, + ), + "cli/bench/serve.md": (create_parser(bench_serve.add_cli_args), True), + "cli/bench/startup.md": (create_parser(bench_startup.add_cli_args), True), + "cli/bench/throughput.md": (create_parser(bench_throughput.add_cli_args), True), + "cli/bench/sweep/plot.md": (create_parser(bench_sweep_plot.add_cli_args), True), + "cli/bench/sweep/plot_pareto.md": ( + create_parser(bench_sweep_plot_pareto.add_cli_args), + True, + ), + "cli/bench/sweep/serve.md": (create_parser(bench_sweep_serve.add_cli_args), True), + "cli/bench/sweep/serve_workload.md": ( + create_parser(bench_sweep_serve_workload.add_cli_args), + True, + ), + "cli/bench/sweep/startup.md": ( + create_parser(bench_sweep_startup.add_cli_args), + True, + ), +} + +# Command name for pages whose file stem differs (URL kept for back-compat). +COMMAND_NAMES = {"cli/bench/mm_processor.md": "mm-processor"} + +for doc_path, (parser, json_tip) in pages.items(): + segments = Path(doc_path).relative_to("cli").with_suffix("").parts + label = COMMAND_NAMES.get(doc_path, segments[-1]) + command = " ".join([*segments[:-1], label]) + # `title` frontmatter keeps the nav label to just this command's segment, + # while the H1 stays the full `vllm ...` command for the page heading. + content = f"---\ntitle: {label}\n---\n\n" + content += f"# vllm {command}\n\n" + if parser.description: + content += f"## Overview\n\n{parser.description}\n\n" + # Rendered above instead of at the top of the Arguments section + parser.description = None + if json_tip: + content += JSON_TIP + content += f"## Arguments\n\n{linkify_docs_urls(format_help(parser))}" + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(content) + logger.debug("CLI reference generated: %s", doc_path) + +logger.info("Total argparse docs generated: %d", len(pages) + 2) + + +# --- Bare subcommand (group) pages ------------------------------------------- +# Mirror `vllm --help`: an overview plus a table of child subcommands, +# each linked to its reference page. Children are read from the CLI registries +# so the listing can never drift from the actual subcommands. Each page is the +# `README.md` of its command directory so it becomes that section's index and is +# picked up by the existing nav globs. +import_bench_subcommands() # populate BenchmarkSubcommandBase.__subclasses__() +bench_subcommands = BenchmarkSubcommandBase.__subclasses__() +bench_children = [(cmd.name, cmd.help) for cmd in bench_subcommands] + +groups = { + "cli/bench/README.md": (BenchmarkSubcommand.help, bench_children), + "cli/launch/README.md": ( + launch_description, + [(cmd.name, cmd.help) for cmd in LaunchSubcommandBase.__subclasses__()], + ), + "cli/bench/sweep/README.md": ( + dict(bench_children).get("sweep"), + [(args.parser_name, args.parser_help) for args, _ in sweep_subcommands], + ), +} + +# Doc paths that exist, so we only link a child that has a reference page. +existing_pages = set(pages) | set(groups) + + +def child_link(group_doc: str, name: str) -> str | None: + group_dir = Path(group_doc).parent # cli/bench/README.md -> cli/bench + for stem in (name, name.replace("-", "_")): + # A leaf page (bench/latency.md) or a nested group index (sweep/README.md) + for candidate in (group_dir / f"{stem}.md", group_dir / stem / "README.md"): + if candidate.as_posix() in existing_pages: + return candidate.relative_to(group_dir).as_posix() + return None + + +for doc_path, (overview, children) in groups.items(): + title = "vllm " + Path(doc_path).parent.relative_to("cli").as_posix() + lines = [f"# {title.replace('/', ' ')}", ""] + if overview: + lines += ["## Overview", "", overview.strip(), ""] + lines += ["## Subcommands", "", "| Command | Description |", "| --- | --- |"] + for name, summary in children: + link = child_link(doc_path, name) + command = f"[`{name}`]({link})" if link else f"`{name}`" + lines.append(f"| {command} | {(summary or '').strip()} |") + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write("\n".join(lines) + "\n") + logger.debug("CLI group reference generated: %s", doc_path) + +logger.info("CLI group reference pages generated: %d", len(groups)) diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/docs/mkdocs/gen_files/generate_attention_backends.py similarity index 80% rename from tools/pre_commit/generate_attention_backend_docs.py rename to docs/mkdocs/gen_files/generate_attention_backends.py index d44456530d4b..942165eaa940 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -9,33 +9,28 @@ This approach avoids requiring CUDA/ROCm/GPU libraries to be installed. -When used as a pre-commit hook, this script receives filenames as arguments -and only runs the check if any of the relevant files were modified. +It runs as an mkdocs-gen-files script, so the page is generated at docs build +time rather than being committed to the repository. """ -import argparse import ast -import fnmatch +import logging import sys from collections.abc import Callable from pathlib import Path from typing import Any +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 + +logger = logging.getLogger("mkdocs") + # --------------------------------------------------------------------------- # Constants and file paths # --------------------------------------------------------------------------- -REPO_ROOT = Path(__file__).parent.parent.parent - -RELEVANT_PATTERNS = [ - "vllm/v1/attention/backends/*.py", - "vllm/v1/attention/backends/**/*.py", - "vllm/models/minimax_m3/common/sparse_attention.py", - "vllm/model_executor/layers/attention/mla_attention.py", - "vllm/platforms/cuda.py", - "tools/pre_commit/generate_attention_backend_docs.py", - "docs/design/attention_backends.md", -] +REPO_ROOT = Path(__file__).parent.parent.parent.parent BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends" REGISTRY_FILE = BACKENDS_DIR / "registry.py" @@ -55,19 +50,6 @@ } -def is_relevant_file(filepath: str) -> bool: - """Check if a file matches any of the relevant patterns.""" - path = Path(filepath) - if path.is_absolute(): - try: - path = path.relative_to(REPO_ROOT) - except ValueError: - return False - path_str = str(path) - - return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS) - - MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" @@ -960,7 +942,7 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None try: tree = ast.parse(file_path.read_text()) except Exception as e: - print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr) + logger.warning("Could not parse %s: %s", file_path, e) return None class_name = class_path.rsplit(".", 1)[1] @@ -1083,7 +1065,8 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: return {} # Analyze the functions to determine FA3/FA4-specific features - fa3_supports_fp8 = True + fa3_supports_fp8 = False + fa4_supports_fp8 = False fa3_supports_sinks = False fa4_supports_sinks = False fa3_compute_cap: str | None = None @@ -1093,6 +1076,44 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: if not isinstance(node, ast.FunctionDef): continue + # Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version. + # Accept both equality checks and membership checks such as `in (3, 4)`. + if node.name == "flash_attn_supports_kv_cache_dtype": + for n in ast.walk(node): + if not ( + isinstance(n, ast.Compare) + and len(n.ops) == 1 + and len(n.comparators) == 1 + ): + continue + is_version_compare = ( + isinstance(n.left, ast.Name) and n.left.id == "fa_version" + ) or ( + isinstance(n.left, ast.Call) + and isinstance(n.left.func, ast.Name) + and n.left.func.id == "get_flash_attn_version" + ) + if not is_version_compare: + continue + + versions: list[Any] = [] + comparator = n.comparators[0] + if isinstance(n.ops[0], ast.Eq) and isinstance( + comparator, ast.Constant + ): + versions = [comparator.value] + elif isinstance(n.ops[0], ast.In) and isinstance( + comparator, (ast.Tuple, ast.List, ast.Set) + ): + versions = [ + elt.value + for elt in comparator.elts + if isinstance(elt, ast.Constant) + ] + + fa3_supports_fp8 |= 3 in versions + fa4_supports_fp8 |= 4 in versions + # Check flash_attn_supports_sinks - looks for `fa_version == 3/4` # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`) if node.name == "flash_attn_supports_sinks": @@ -1199,7 +1220,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: }, "fa4": { "compute_capability": fa4_compute_cap, - "supports_fp8": False, + "supports_fp8": fa4_supports_fp8, "supports_sink": fa4_supports_sinks, }, } @@ -1255,6 +1276,22 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]: # --------------------------------------------------------------------------- +def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str: + """Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string. + + The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``, + but actual support varies by FA version (e.g. FA2 has no fp8 support), so each + variant's row is normalized to match its own capability. + """ + fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"} + dtypes = kv_cache_dtypes.split(", ") + base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes] + if supports_fp8: + advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes] + return ", ".join(base_dtypes + advertised_fp8_dtypes) + return ", ".join(base_dtypes) + + def _expand_flash_attn_variants( all_backends: list[dict[str, Any]], fa_features: dict[str, dict[str, Any]], @@ -1275,6 +1312,9 @@ def _expand_flash_attn_variants( fa2["_sort_key"] = "FLASH_ATTN" fa2["_sort_order"] = 0 fa2["supports_sink"] = fa_features["fa2"]["supports_sink"] + fa2["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"] + ) # Create FA3 entry (uses parsed compute_capability from fa_utils) fa3 = backend.copy() @@ -1284,11 +1324,9 @@ def _expand_flash_attn_variants( if fa_features["fa3"]["compute_capability"]: fa3["compute_capability"] = fa_features["fa3"]["compute_capability"] fa3["supports_sink"] = fa_features["fa3"]["supports_sink"] - if fa_features["fa3"]["supports_fp8"]: - base_dtypes = backend["kv_cache_dtypes"].split(", ") - fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"] - new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes] - fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes) + fa3["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"] + ) expanded.append(fa2) expanded.append(fa3) @@ -1302,6 +1340,9 @@ def _expand_flash_attn_variants( if fa_features["fa4"].get("compute_capability"): fa4["compute_capability"] = fa_features["fa4"]["compute_capability"] fa4["supports_sink"] = fa_features["fa4"]["supports_sink"] + fa4["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"] + ) expanded.append(fa4) return expanded @@ -1598,113 +1639,12 @@ def _render_table( return lines -def generate_markdown_table( - backends: list[dict[str, Any]], title: str, is_mla_table: bool = False -) -> str: - """Generate a titled markdown table from backend info.""" - if not backends: - return f"## {title}\n\nNo backends found.\n" - has_versions = any(b.get("version") for b in backends) - columns = _build_columns(is_mla_table, has_versions) - lines = [f"## {title}", ""] - lines.extend(_render_table(columns, backends)) - lines.append("") - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Markdown section generators (usage, priority, legend, MLA) -# --------------------------------------------------------------------------- - - -def generate_usage_section() -> str: - """Generate the usage documentation section.""" - return """## Setting the Attention Backend - -### Command Line - -There are two ways to specify the backend from the command line: - -**Option 1: Using `--attention-backend` (simple)** - -```bash -vllm serve --attention-backend FLASH_ATTN -``` - -**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)** - -```bash -# Dot notation -vllm serve --attention-config.backend FLASH_ATTN -vllm serve -ac.backend FLASH_ATTN - -# JSON format -vllm serve --attention-config '{"backend": "FLASH_ATTN"}' -vllm serve -ac '{"backend": "FLASH_ATTN"}' -``` - -> **Note:** `--attention-backend` and `--attention-config.backend` are mutually -> exclusive. Use one or the other, not both. - -### Python API - -Use `AttentionConfig` with the `LLM` class: - -```python -from vllm import LLM -from vllm.config import AttentionConfig -from vllm.v1.attention.backends.registry import AttentionBackendEnum - -# Method 1: Using AttentionConfig with enum -llm = LLM( - model="Qwen/Qwen3-0.6B", - attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN), -) - -# Method 2: Using attention_backend parameter with string -llm = LLM( - model="Qwen/Qwen3-0.6B", - attention_backend="FLASH_ATTN", -) -``` - -## Backend Selection Behavior - -### Manual Selection - -When you explicitly set a backend via `--attention-backend` or `AttentionConfig`: - -1. The backend is **validated** against your configuration (model dtype, head - size, compute capability, etc.) -2. If the backend **doesn't support** your configuration, an error is raised - with the specific reason -3. If valid, the backend is used - -Example error when selecting an incompatible backend: - -```text -ValueError: Selected backend FLASHMLA is not valid for this configuration. -Reason: ['compute capability not supported'] -``` - -### Automatic Selection - -When no backend is specified (the default): - -1. vLLM iterates through backends in **priority order** (see tables below) -2. Each backend is validated against your configuration -3. The **first compatible backend** is selected -4. If no backend is compatible, an error is raised listing all backends and - their incompatibility reasons -""" - - def _priority_table( title: str, backends: list[str], annotations: dict[str, str] | None = None, ) -> list[str]: - """Generate a priority table for a list of backends.""" + """Render a priority table for a list of backends.""" def _fmt(b: str) -> str: suffix = annotations.get(b, "") if annotations else "" @@ -1720,102 +1660,38 @@ def _fmt(b: str) -> str: ] -def generate_priority_section(priorities: dict[str, list[str]]) -> str: - """Generate the priority ranking section.""" - lines = [ - "## Backend Priority (CUDA)", - "", - "When no backend is explicitly selected, vLLM chooses the first", - "compatible backend from these priority-ordered lists.", - "", - "Priority is **1 = highest** (tried first).", - "", - "### Standard Attention (MHA, MQA, GQA)", - "", - ] - - sm100 = "Blackwell (SM 10.x)" - ampere = "Ampere/Hopper (SM 8.x-9.x)" - - if "standard_sm100" in priorities: - lines.extend(_priority_table(sm100, priorities["standard_sm100"])) - if "standard_default" in priorities: - lines.extend(_priority_table(ampere, priorities["standard_default"])) +_SM100 = "Blackwell (SM 10.x)" +_AMPERE = "Ampere/Hopper (SM 8.x-9.x)" - lines.extend(["### MLA Attention (DeepSeek-style)", ""]) - mla_sm100_annotations = { - "FLASHINFER_MLA_SPARSE": "**\\***", - } - if "mla_sm100" in priorities: - lines.extend( - _priority_table(sm100, priorities["mla_sm100"], mla_sm100_annotations) - ) - if "mla_default" in priorities: - lines.extend(_priority_table(ampere, priorities["mla_default"])) - - if "mla_sm100" in priorities: - lines.append( - "> **\\*** For sparse MLA, FP8 KV cache always prefers " - "`FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` " - "is preferred for low query-head counts (<= 16), while " - "`FLASHMLA_SPARSE` is preferred otherwise." - ) - lines.append(">") - - lines.append( - "> **Note:** ROCm and CPU platforms have their own selection logic. " - "See the platform-specific documentation for details." - ) - lines.append("") - - return "\n".join(lines) +def _priority_block( + priorities: dict[str, list[str]], + sm100_key: str, + default_key: str, + sm100_annotations: dict[str, str] | None = None, +) -> str: + """Render whichever priority tables exist for one attention category.""" + lines: list[str] = [] + if sm100_key in priorities: + lines += _priority_table(_SM100, priorities[sm100_key], sm100_annotations) + if default_key in priorities: + lines += _priority_table(_AMPERE, priorities[default_key]) + return "\n".join(lines).strip() -def generate_legend() -> str: - """Generate a legend explaining the table columns.""" - return """## Legend - -| Column | Description | -| ------ | ----------- | -| **Dtypes** | Supported model data types (fp16, bf16, fp32) | -| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) | -| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | -| **Head Sizes** | Supported attention head sizes | -| **Sink** | Attention sink support (for StreamingLLM) | -| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | -| **Sparse** | Sparse attention support (MLA only) | -| **MM Prefix** | Multimodal prefix full attention support | -| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | -| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) | -| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) | - -**Symbols:** ✅ = Supported, ❌ = Not supported -""" +def _feature_table(backends: list[dict[str, Any]], is_mla: bool) -> str: + """Render a backend feature table (header, separator, one row per backend).""" + has_versions = any(b.get("version") for b in backends) + columns = _build_columns(is_mla, has_versions) + return "\n".join(_render_table(columns, backends)) -def generate_mla_section( - prefill_backends: list[dict[str, Any]], - decode_backends: list[dict[str, Any]], - v4_decode_backends: list[dict[str, Any]] | None = None, -) -> str: - """Generate the complete MLA section with prefill and decode tables.""" +def _mla_prefill_table(prefill_backends: list[dict[str, Any]]) -> str: + """Render the MLA prefill backend table.""" lines = [ - "## MLA (Multi-head Latent Attention) Backends", - "", - "MLA uses separate backends for prefill and decode phases.", - "", - "### Prefill Backends", - "", - "To explicitly select a prefill backend, use", - "`-ac.mla_prefill_backend=` (e.g., `FLASH_ATTN`, `FLASHINFER`).", - "Otherwise, the prefill backend is selected automatically at runtime based on", - "hardware and configuration.", - "", "| Backend | Description | Dtypes | Compute Cap. | Notes |", "| ------- | ----------- | ------ | ------------ | ----- |", ] - for backend in prefill_backends: row = "| `{}`{} | {} | {} | {} | {} |".format( backend["name"], @@ -1826,87 +1702,21 @@ def generate_mla_section( backend.get("notes", ""), ) lines.append(row.replace(" ", " ")) - - lines.extend( - [ - "", - "> **‡** Automatic selection tries FlashAttention first. On Blackwell", - "> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then", - "> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.", - "", - "### Decode Backends", - "", - "MLA decode backends are selected using the standard", - "`-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`).", - "", - ] - ) - - # Reuse data-driven table rendering for decode backends - columns = _build_columns(is_mla=True, has_versions=False) - lines.extend(_render_table(columns, decode_backends)) - - if v4_decode_backends: - lines.extend( - [ - "", - "### DeepSeek V4 Decode Backends", - "", - "DeepSeek V4 sparse MLA uses its own decode backends, selected via", - "`--attention-backend=` (e.g., `FLASHMLA_SPARSE_DSV4`,", - "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index", - "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);", - "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and", - "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.", - "", - ] - ) - lines.extend(_render_table(columns, v4_decode_backends)) - - lines.append("") - return "\n".join(lines) - - -def generate_minimax_section(backends: list[dict[str, Any]]) -> str: - """Generate the MiniMax M3 sparse attention section.""" - lines = [ - "## MiniMax M3 Sparse Attention Backends", - "", - 'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")', - "layers. It is wired in directly by the model and is not part of the", - "automatic priority lists above. A lightning indexer scores KV blocks, the", - "top-k blocks (plus fixed init/local blocks) are selected, and attention", - "attends only to those blocks; index keys live in a separate side cache.", - "", - ] - columns = _build_columns(is_mla=False, has_versions=False) - lines.extend(_render_table(columns, backends)) - lines.append("") return "\n".join(lines) -# --------------------------------------------------------------------------- -# Top-level orchestration -# --------------------------------------------------------------------------- +def build_blocks() -> dict[str, str]: + """Build the generated table blocks keyed by their `gen:` marker name. - -def generate_docs() -> str: - """Generate the complete documentation.""" + Only the tables are generated here; the surrounding prose lives in the + handwritten ``docs/design/attention_backends.md`` page. + """ attention_backends_map = parse_registry() - - # Parse priority lists from cuda.py priorities = parse_cuda_priority_lists() - - # Parse FlashAttention FA2/FA3 feature differences fa_features = parse_flash_attn_features() - - # Parse FlashInfer TRTLLM feature differences (native vs TRTLLM on Blackwell) fi_features = parse_flashinfer_trtllm_features() - - # Parse MLA prefill backends mla_prefill_backends = parse_mla_prefill_backends() - # Collect backend info all_backends = [] for backend_name, class_path in attention_backends_map.items(): if backend_name in SKIP_BACKENDS: @@ -1914,17 +1724,14 @@ def generate_docs() -> str: info = analyze_backend(backend_name, class_path) if info: all_backends.append(info) - - # Expand backends into version variants if fa_features: all_backends = _expand_flash_attn_variants(all_backends, fa_features) if fi_features: all_backends = _expand_flashinfer_variants(all_backends, fi_features) - # DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each - # get their own subsection rather than mixing into the main MLA / standard - # tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so - # filter purely on the name). + # DeepSeek V4 (*_DSV4) and MiniMax M3 sparse backends get their own tables + # rather than mixing into the main MLA / standard tables (the ROCm V4 backend + # isn't flagged is_mla by the AST heuristic, so filter purely on the name). def _is_v4(b: dict[str, Any]) -> bool: return b["name"].endswith("_DSV4") @@ -1940,112 +1747,21 @@ def _is_minimax(b: dict[str, Any]) -> bool: if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b) ] - # Generate documentation - script_path = "tools/pre_commit/generate_attention_backend_docs.py" - doc_lines = [ - "# Attention Backend Feature Support", - "", - f"This document is auto-generated by `{script_path}`.", - "It shows the feature support for each registered attention backend", - "based on the checks in `AttentionBackend.validate_configuration()`.", - "", - "**Do not edit this file manually.** Run the following command to", - "regenerate it:", - "", - "```bash", - f"python {script_path}", - "```", - "", - ] - - # Add usage documentation - doc_lines.append(generate_usage_section()) - - # Add priority section - doc_lines.append(generate_priority_section(priorities)) - - # Add legend and feature tables - doc_lines.append(generate_legend()) - standard_title = "Standard Attention (MHA, MQA, GQA) Backends" - doc_lines.append( - generate_markdown_table(non_mla_backends, standard_title, is_mla_table=False) - ) - # Add footnotes for version/variant distinctions (in table order) - footnotes = [] - if fi_features: - footnotes.append( - "> **†** FlashInfer Native is the regular FlashInfer path. XQA is the " - "SM90 decode path exposed through FlashInfer's TRTLLM decode API. " - "trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen " - "via `--attention-config.use_trtllm_attention=0`." - ) - if fa_features: - footnotes.append( - "> **\\*** Specify the FlashAttention version via " - "`--attention-config.flash_attn_version=2`, `3`, or `4`. " - "Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), " - "FA2 otherwise." - ) - if footnotes: - doc_lines.append("\n>\n".join(footnotes) + "\n") - - # Add MiniMax M3 sparse section (separate category after standard GQA) - if minimax_backends: - doc_lines.append(generate_minimax_section(minimax_backends)) - - # Add MLA section with prefill and decode backends - doc_lines.append( - generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends) - ) - - return "\n".join(doc_lines) - - -def main(): - parser = argparse.ArgumentParser( - description="Generate attention backend documentation table" - ) - parser.add_argument( - "--output", - "-o", - type=str, - default=str(REPO_ROOT / "docs" / "design" / "attention_backends.md"), - help="Output file path (default: docs/design/attention_backends.md)", - ) - parser.add_argument( - "--check", - action="store_true", - help="Check if the documentation is up to date (for pre-commit)", - ) - parser.add_argument( - "files", - nargs="*", - help="Files to check (passed by pre-commit). If none are relevant, skip.", - ) - args = parser.parse_args() - - if args.files and not any(is_relevant_file(f) for f in args.files): - sys.exit(0) - - output_path = Path(args.output) - new_content = generate_docs() - - if args.check: - needs_update = ( - not output_path.exists() or output_path.read_text() != new_content - ) - if needs_update: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(new_content) - print(f"🔄 Regenerated: {output_path}") - sys.exit(1) - print(f"✅ Up to date: {output_path}") - sys.exit(0) - - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(new_content) - print(f"Generated: {output_path}") + mla_sm100_annotations = {"FLASHINFER_MLA_SPARSE": "**\\***"} + return { + "priority-standard": _priority_block( + priorities, "standard_sm100", "standard_default" + ), + "priority-mla": _priority_block( + priorities, "mla_sm100", "mla_default", mla_sm100_annotations + ), + "table-standard": _feature_table(non_mla_backends, is_mla=False), + "table-minimax": _feature_table(minimax_backends, is_mla=False), + "table-mla-prefill": _mla_prefill_table(mla_prefill_backends), + "table-mla-decode": _feature_table(mla_backends, is_mla=True), + "table-mla-v4-decode": _feature_table(v4_decode_backends, is_mla=True), + } -if __name__ == "__main__": - main() +logger.info("Generating attention backend documentation") +fill_markers("design/attention_backends.md", build_blocks()) diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py similarity index 76% rename from docs/mkdocs/hooks/generate_examples.py rename to docs/mkdocs/gen_files/generate_examples.py index 07fbd7e4d555..c9c10018ecfb 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -5,16 +5,15 @@ from dataclasses import dataclass from functools import cached_property from pathlib import Path -from typing import Literal +import mkdocs_awesome_nav.nav.directory as _nav_dir +import mkdocs_gen_files import regex as re logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -ROOT_DIR_RELATIVE = "../../../../.." EXAMPLE_DIR = ROOT_DIR / "examples" -EXAMPLE_DOC_DIR = ROOT_DIR / "docs/examples" def title(text: str) -> str: @@ -197,44 +196,38 @@ def generate(self) -> str: return content -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - # Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are - # title-cased (e.g. "Offline Inference" instead of "Offline inference"). - import mkdocs_awesome_nav.nav.directory as _nav_dir - - _nav_dir.dirname_to_title = title - logger.info("Generating example documentation") - logger.debug("Root directory: %s", ROOT_DIR.resolve()) - logger.debug("Example directory: %s", EXAMPLE_DIR.resolve()) - logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve()) - - # Create the EXAMPLE_DOC_DIR if it doesn't exist - if not EXAMPLE_DOC_DIR.exists(): - EXAMPLE_DOC_DIR.mkdir(parents=True) - - categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir()) - - examples = [] - glob_patterns = ["*.py", "*.md", "*.sh"] - # Find categorised examples - for category in categories: - logger.info("Processing category: %s", category.stem) - globs = [category.glob(pattern) for pattern in glob_patterns] - for path in itertools.chain(*globs): - examples.append(Example(path, category.stem)) - # Find examples in subdirectories - globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns] - for path in itertools.chain(*globs): - examples.append(Example(path.parent, category.stem)) - - # Generate the example documentation - for example in sorted(examples, key=lambda e: e.path.stem): - example_name = f"{example.path.stem}.md" - doc_path = EXAMPLE_DOC_DIR / example.category / example_name - if not doc_path.parent.exists(): - doc_path.parent.mkdir(parents=True) - # Specify encoding for building on Windows - with open(doc_path, "w+", encoding="utf-8") as f: - f.write(example.generate()) - logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR)) - logger.info("Total examples generated: %d", len(examples)) +# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are +# title-cased (e.g. "Offline Inference" instead of "Offline inference"). +_nav_dir.dirname_to_title = title +logger.info("Generating example documentation") +logger.debug("Root directory: %s", ROOT_DIR.resolve()) +logger.debug("Example directory: %s", EXAMPLE_DIR.resolve()) + +categories = sorted( + p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".") +) + +examples = [] +glob_patterns = ["*.py", "*.md", "*.sh"] +# Find categorised examples +for category in categories: + logger.info("Processing category: %s", category.stem) + globs = [category.glob(pattern) for pattern in glob_patterns] + for path in itertools.chain(*globs): + examples.append(Example(path, category.stem)) + # Find examples in subdirectories + globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns] + for path in itertools.chain(*globs): + examples.append(Example(path.parent, category.stem)) + +# Generate the example documentation +for example in sorted(examples, key=lambda e: e.path.stem): + doc_path = f"examples/{example.category}/{example.path.stem}.md" + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(example.generate()) + if example.main_file is not None: + # Point the edit button at the example's source file + edit_path = Path("..") / example.main_file.relative_to(ROOT_DIR) + mkdocs_gen_files.set_edit_path(doc_path, str(edit_path)) + logger.debug("Example generated: %s", doc_path) +logger.info("Total examples generated: %d", len(examples)) diff --git a/docs/mkdocs/hooks/generate_metrics.py b/docs/mkdocs/gen_files/generate_metrics.py similarity index 63% rename from docs/mkdocs/hooks/generate_metrics.py rename to docs/mkdocs/gen_files/generate_metrics.py index 97282aaee7d0..b952f53977a3 100644 --- a/docs/mkdocs/hooks/generate_metrics.py +++ b/docs/mkdocs/gen_files/generate_metrics.py @@ -2,27 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import ast import logging +import sys from pathlib import Path -from typing import Literal + +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -DOCS_DIR = ROOT_DIR / "docs" -GENERATED_METRICS_DIR = DOCS_DIR / "generated" / "metrics" -# Files to scan for metric definitions - each will generate a separate table +# Files to scan for metric definitions - each fills a `gen:` marker in +# docs/usage/metrics.md with its table (the section heading and any preamble +# live in the tracked page next to the marker). METRIC_SOURCE_FILES = [ - {"path": "vllm/v1/metrics/loggers.py", "output": "general.inc.md"}, - { - "path": "vllm/v1/spec_decode/metrics.py", - "output": "spec_decode.inc.md", - }, + {"path": "vllm/v1/metrics/loggers.py", "key": "metrics-general"}, + {"path": "vllm/v1/spec_decode/metrics.py", "key": "metrics-spec-decode"}, { "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py", - "output": "nixl_connector.inc.md", + "key": "metrics-nixl", }, - {"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"}, + {"path": "vllm/v1/metrics/perf.py", "key": "metrics-mfu"}, ] @@ -110,41 +111,27 @@ def generate_markdown_table(metrics: list[dict[str, str]]) -> str: return "\n".join(lines) + "\n" -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - """Generate metrics documentation tables from source files.""" - logger.info("Generating metrics documentation") - - # Create generated directory if it doesn't exist - GENERATED_METRICS_DIR.mkdir(parents=True, exist_ok=True) - - total_metrics = 0 - for source_config in METRIC_SOURCE_FILES: - source_path = source_config["path"] - output_file = source_config["output"] - - filepath = ROOT_DIR / source_path - if not filepath.exists(): - raise FileNotFoundError(f"Metrics source file not found: {filepath}") - - logger.debug("Extracting metrics from: %s", source_path) - metrics = extract_metrics_from_file(filepath) - logger.debug("Found %d metrics in %s", len(metrics), source_path) - - # Generate and write the markdown table for this source - table_content = generate_markdown_table(metrics) - output_path = GENERATED_METRICS_DIR / output_file - with open(output_path, "w", encoding="utf-8") as f: - f.write(table_content) - - total_metrics += len(metrics) - logger.info( - "Generated metrics table: %s (%d metrics)", - output_path.relative_to(ROOT_DIR), - len(metrics), - ) - - logger.info( - "Total metrics generated: %d across %d files", - total_metrics, - len(METRIC_SOURCE_FILES), - ) +logger.info("Generating metrics documentation") + +blocks = {} +total_metrics = 0 +for source_config in METRIC_SOURCE_FILES: + source_path = source_config["path"] + + filepath = ROOT_DIR / source_path + if not filepath.exists(): + raise FileNotFoundError(f"Metrics source file not found: {filepath}") + + logger.debug("Extracting metrics from: %s", source_path) + metrics = extract_metrics_from_file(filepath) + logger.debug("Found %d metrics in %s", len(metrics), source_path) + + blocks[source_config["key"]] = generate_markdown_table(metrics).strip() + total_metrics += len(metrics) + +fill_markers("usage/metrics.md", blocks) +logger.info( + "Total metrics generated: %d across %d files", + total_metrics, + len(METRIC_SOURCE_FILES), +) diff --git a/docs/mkdocs/gen_files/generated_content.py b/docs/mkdocs/gen_files/generated_content.py new file mode 100644 index 000000000000..fa3f31b294bf --- /dev/null +++ b/docs/mkdocs/gen_files/generated_content.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inline build-time generated content into existing docs pages. + +Source pages mark where generated content goes with a snippet-style marker, +`--8<-- "gen:"`, so the insertion point is explicit and readable. The +substitution happens here (at gen-files time, before mkdocs-gen-files shadows +the page), not via pymdownx.snippets, so the content can be generated at build +time without living in a real file on disk. + +The `gen:` prefix keeps these markers distinct from real pymdownx.snippets +includes, and `fill_markers` fails loudly if a marker is missing or left behind +(pymdownx.snippets would otherwise silently drop an unsubstituted marker). +""" + +from pathlib import Path + +import mkdocs_gen_files +import regex as re + +DOCS_DIR = Path(__file__).parent.parent.parent + +_MARKER = '--8<-- "gen:{key}"' +_ANY_MARKER = re.compile(r'--8<-- "gen:[^"]*"') + + +def fill_markers(doc_path: str, blocks: dict[str, str]) -> None: + """Replace `--8<-- "gen:"` markers in a docs page with generated content. + + Args: + doc_path: Docs-relative path of the source page to fill. + blocks: Mapping of marker key to the markdown to insert in its place. + + Raises: + FileNotFoundError: If the source page does not exist. + ValueError: If an expected marker is missing, or any `gen:` marker is + left unsubstituted after filling. + """ + source = DOCS_DIR / doc_path + if not source.exists(): + raise FileNotFoundError(f"Cannot fill markers in missing page: {doc_path}") + + text = source.read_text() + for key, content in blocks.items(): + marker = _MARKER.format(key=key) + if marker not in text: + raise ValueError(f"{doc_path}: missing marker {marker}") + text = text.replace(marker, content) + + if leftover := _ANY_MARKER.search(text): + raise ValueError(f"{doc_path}: unsubstituted marker {leftover.group()}") + + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(text) + # Keep the edit button pointing at the real source page + mkdocs_gen_files.set_edit_path(doc_path, doc_path) diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py index e6faf95cd469..d208b9d6779e 100644 --- a/docs/mkdocs/hooks/url_schemes.py +++ b/docs/mkdocs/hooks/url_schemes.py @@ -19,6 +19,7 @@ each page is converted. """ +import posixpath from pathlib import Path import regex as re @@ -38,18 +39,22 @@ REPO = r"(?P.+?/.+?)" TYPE = r"(?Pissues|pull|projects)" NUMBER = r"(?P\d+)" +VERSION = r"[^/\s]+" PATH = r"(?P[^\s]+?)" FRAGMENT = r"(?P#[^\s]+)?" -URL = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}" +URL_GITHUB = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}" RELATIVE = rf"(?!(https?|ftp)://|#){PATH}{FRAGMENT}" +URL_DOCS = f"https://docs.vllm.ai/en/{VERSION}/{PATH}{FRAGMENT}" # Common titles to use for GitHub links when none is provided in the link. TITLES = {"issues": "Issue ", "pull": "Pull Request ", "projects": "Project "} # Regex to match GitHub issue, PR, and project links with optional titles. -github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL}(\)|>)") +github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL_GITHUB}(\)|>)") # Regex to match relative file links with optional titles. relative_link = re.compile(rf"\[{TITLE}\]\({RELATIVE}\)") +# Regex to match absolute docs.vllm.ai links (should only exist in CLI). +docs_link = re.compile(rf"\[{TITLE}\]\({URL_DOCS}\)") class UrlSchemesPreprocessor(Preprocessor): @@ -61,7 +66,8 @@ def __init__(self, md, ext): def run(self, lines): page = self.ext.page - if page is None or getattr(page.file, "abs_src_path", None) is None: + files = self.ext.files + if page is None: return lines def replace_relative_link(match: re.Match) -> str: @@ -70,7 +76,7 @@ def replace_relative_link(match: re.Match) -> str: """ title = match.group("title") path = match.group("path") - path = (Path(page.file.abs_src_path).parent / path).resolve() + path = ((DOC_DIR / page.file.src_uri).parent / path).resolve() fragment = match.group("fragment") or "" # Check if the path exists and is outside the docs dir @@ -105,9 +111,36 @@ def replace_github_link(match: re.Match) -> str: url = f"https://github.com/{repo}/{type}/{number}{fragment}" return f"[{gh_icon} {title}]({url})" + def replace_docs_link(match: re.Match) -> str: + """Rewrite absolute docs.vllm.ai links as doc-relative links.""" + title = match.group("title") + path = match.group("path").rstrip("/") + fragment = match.group("fragment") or "" + + # vllm.config. API reference -> mkdocstrings cross-reference + if path == "api/vllm/config" and re.fullmatch( + r"#vllm\.config\.\w+", fragment + ): + ident = fragment[1:] + return f"[`{ident}`][{ident}]" + + # Other docs pages -> link relative to the current page, but only + # when the target is a known docs page (real or generated); leave + # unknown/external URLs untouched. This is correct even when the same + # docstring is also rendered on its API reference page. + src = f"{path.removesuffix('.html')}.md" + if files.get_file_from_path(src) is None: + return match.group(0) + rel = posixpath.relpath(src, posixpath.dirname(page.file.src_uri)) + # Auto-wrapped bare URLs use the URL as their title; make it readable. + if title.startswith("http"): + title = path.removesuffix(".html") + return f"[{title}]({rel}{fragment})" + markdown = "\n".join(lines) markdown = github_link.sub(replace_github_link, markdown) markdown = relative_link.sub(replace_relative_link, markdown) + markdown = docs_link.sub(replace_docs_link, markdown) return markdown.split("\n") @@ -116,6 +149,7 @@ class UrlSchemesExtension(Extension): def __init__(self, **kwargs): self.page = None + self.files = None super().__init__(**kwargs) def extendMarkdown(self, md): @@ -138,4 +172,5 @@ def on_page_markdown( ) -> str: """Pass the current page context to the preprocessor.""" _ext.page = page + _ext.files = files return markdown diff --git a/docs/mkdocs/javascript/reo.js b/docs/mkdocs/javascript/reo.js new file mode 100644 index 000000000000..cb7430978b1d --- /dev/null +++ b/docs/mkdocs/javascript/reo.js @@ -0,0 +1,3 @@ +// Reo.Dev documentation tracking +// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation +!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}(); diff --git a/docs/models/hardware_supported_models/xpu.md b/docs/models/hardware_supported_models/xpu.md index d065b4b68903..49636d0d5c98 100644 --- a/docs/models/hardware_supported_models/xpu.md +++ b/docs/models/hardware_supported_models/xpu.md @@ -31,10 +31,8 @@ | THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | | | 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | | -| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | | | meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | | -| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | | | Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | | | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | | diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index d69a4dc616eb..4bad45d7a8fa 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -184,7 +184,7 @@ Our online Server provides endpoints that correspond to the offline APIs: - [Classification API](classify.md#online-serving)(`/classify`) - Corresponding to `LLM.score`: - [Score API](scoring.md#score-api) (`/score`, `/v1/score`) - - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) + - [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) - Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models. The following introduces the Pooling API. For other APIs, please refer to the link above. diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 1b9d14d7a0a4..c548ed80d201 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -39,12 +39,13 @@ You can compute pairwise similarity scores to build a similarity matrix using th | ------------ | ------ | ----------------- | ------------------------------ | ------------------------------------------ | | `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | | `BertSpladeSparseEmbeddingModel` | SPLADE | `naver/splade-v3` | | | +| `BgeM3EmbeddingModel` | BGE-M3 | `BAAI/bge-m3` | | | | `Gemma2Model`C | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | | `Gemma3TextModel`C | Gemma 3-based | `google/embeddinggemma-300m`, etc. | ✅︎ | ✅︎ | | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | | `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | | `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | -| `JinaEmbeddingsV5Model`C | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ | +| `JinaEmbeddingsV5Model`C | Qwen3-decoder or EuroBERT-encoder backbone with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small`, `jinaai/jina-embeddings-v5-text-nano` (see note) | ✅︎ | ✅︎ | | `LlamaBidirectionalModel`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ | | `LlamaModel`C, `LlamaForCausalLM`C, `MistralModel`C, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | @@ -74,7 +75,9 @@ You can compute pairwise similarity scores to build a similarity matrix using th `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights. !!! note - `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters + `jinaai/jina-embeddings-v5-text-small` (Qwen3 decoder) and + `jinaai/jina-embeddings-v5-text-nano` (bidirectional EuroBERT encoder, + `is_decoder=false`) ship with four task-specific LoRA adapters (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the selected adapter into the base weights at load time. Choose the task with `--hf-overrides '{"jina_task": ""}'`; the default is `retrieval`. diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index e3b54b020751..58083763cb41 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -20,7 +20,7 @@ The score models is designed to compute similarity scores between two input prom - `LLM.score` - Online APIs: - [Score API](scoring.md#score-api) (`/score`, `/v1/score`) - - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) + - [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) !!! note Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 98f9cb5b65d9..de937ec520cd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -19,7 +19,7 @@ vLLM also supports model implementations that are available in Transformers. We Currently, the Transformers modeling backend works for the following: -- Modalities: embedding models, language models and vision-language models* +- Modalities: embedding models, language models, vision-language models* and audio-language models - Architectures: encoder-only, decoder-only, mixture-of-experts - Attention types: full attention and/or sliding attention @@ -364,7 +364,7 @@ th { | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ | | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ | -| `ExaoneMoEForCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B`, etc. | | | +| `ExaoneMoeForCausalLM` | K-EXAONE, K-EXAONE-2 | `LGAI-EXAONE/K-EXAONE-236B-A23B`, `LGAI-EXAONE/K-EXAONE-2.0-750B-A37B`, etc. | | | | `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ | | `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ | | `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ | @@ -427,7 +427,6 @@ th { | `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ | | `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ | -| `OuroForCausalLM` | ouro | `ByteDance/Ouro-1.4B`, `ByteDance/Ouro-2.6B`, etc. | ✅︎ | | | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ | | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ | | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ | @@ -435,7 +434,6 @@ th { | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | -| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | | `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | @@ -466,6 +464,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | +| `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ | !!! note Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096. @@ -548,6 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ | | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A+ | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I+ | `moonshotai/Kimi-K2.5` | | ✅︎ | +| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I+ | `moonshotai/Kimi-K3` | | ✅︎ | | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I+ | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ | | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I+ | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ | | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I+ | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ | @@ -608,7 +608,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ------ | ----------------- | --------------------------- | --------------------------------------- | -| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ | +| `Emu3ForConditionalGeneration` | Emu3 | T + I+ | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ | +| `VibeVoiceAsrForConditionalGeneration` | VibeVoice-ASR | T + A+ | `microsoft/VibeVoice-ASR-HF` | ✅︎ | ✅︎ | ^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
E Pre-computed embeddings can be inputted for this modality.
diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh index d55f8c8db12d..d611e616071e 100644 --- a/docs/pre_run_check.sh +++ b/docs/pre_run_check.sh @@ -28,7 +28,7 @@ DOCS_PATHS=( docs/ # Actual docs content examples/ # Examples are rendered in docs vllm/ # API & CLI reference - requirements/test/cuda.txt # CLI reference (see docs/mkdocs/hooks/generate_argparse.py) + requirements/test/cuda.txt # CLI reference (see docs/mkdocs/gen_files/generate_argparse.py) mkdocs.yaml # Affects build process .readthedocs.yaml # Affects build process requirements/docs.txt # Affects build process diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 9a8bb68bc11f..d1dff32fc946 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -12,6 +12,11 @@ Before using EP, you need to install the necessary dependencies. We are actively 2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation). 3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). +!!! note "NCCL version (CUDA 13+)" + The `deepep_v2` backend requires NCCL >= 2.30.4. PyTorch ships an older + NCCL, so you must upgrade it before building or running DeepEP. See the + [EP kernels guide](../../tools/ep_kernels) for instructions. + ### Backend Selection Guide vLLM provides multiple communication backends for EP. Use `--all2all-backend` to select one: @@ -26,9 +31,6 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to ## Single Node Deployment -!!! warning - EP is an experimental feature. Argument names and default values may change in the future. - ### Configuration Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as: diff --git a/docs/serving/offline_inference.md b/docs/serving/offline_inference.md index 4512f4a07201..9a71612f262b 100644 --- a/docs/serving/offline_inference.md +++ b/docs/serving/offline_inference.md @@ -65,6 +65,8 @@ For further details on Weight Transfer, please refer to [this page](../training/ - `LLM.start_weight_update` - Starts a new weight update cycle. - `LLM.update_weights` - Updates the model weights. - `LLM.finish_weight_update` - Finishes the current weight update cycle. +- `LLM.update_weight_version` - Sets the weight version without updating model weights. +- `LLM.get_weight_version` - Returns the latest committed weight version. ## Additional APIs diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 5df1821f3b38..90ff7a3e3d8b 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -10,7 +10,7 @@ We currently support the following OpenAI APIs: - Only applicable to [text generation models](../../models/generative_models.md). - *Note: `suffix` parameter is not supported.* - [Chat Completions API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions`) - - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](./openai_compatible_server.md#chat-template). + - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](#chat-template). - *Note: `user` parameter is ignored.* - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls. - [Chat Completions batch API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions/batch`) @@ -32,7 +32,7 @@ We currently support the following OpenAI APIs: - [Cohere Embed API](../../models/pooling_models/embed.md#cohere-embed-api) (`/v2/embed`) - Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed) - Works with any [embedding model](../../models/pooling_models/embed.md#supported-models), including multimodal models. -- [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) +- [Cohere Rerank API](../../models/pooling_models/scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) - Implements [Jina AI's v1 rerank API](https://jina.ai/reranker/) - compatible with [Cohere's v1 & v2 rerank APIs](https://docs.cohere.com/v2/reference/rerank) @@ -49,7 +49,7 @@ For further details on pooling models, please refer to [this page](../../models/ - Only applicable to [embedding models](../../models/pooling_models/embed.md). - [Scoring Usages](../../models/pooling_models/scoring.md) - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`, `/v1/score`) - - [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) + - [Cohere Rerank API](../../models/pooling_models/scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction). - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`) - Applicable to all [pooling models](../../models/pooling_models/README.md). @@ -73,7 +73,7 @@ For further details on speech to text, please refer to [this page](speech_to_tex - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction). - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`) - Applicable to all [pooling models](../../models/pooling_models/README.md). -- [Generative Scoring API](generative_scoring.md#generative-scoring-api) (`/generative_scoring`) +- [Generative Scoring API](generative_scoring.md) (`/generative_scoring`) - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`). - Computes next-token probabilities for specified `label_token_ids`. @@ -179,6 +179,8 @@ For further details on Weight Transfer, please refer to [this page](../../traini - `/start_weight_update` - Prepares the inference engine for a weight update. - `/update_weights` - Update model weights (can alter model behavior) - `/finish_weight_update` - Finalizes the weight update +- `/update_weight_version` - Set the weight version without updating model weights +- `/weight_info` - Get the latest committed weight version - `/get_world_size` - Get distributed world size ### Collective RPC diff --git a/docs/serving/online_serving/speech_to_text.md b/docs/serving/online_serving/speech_to_text.md index d503923c4f9d..a956367cc30b 100644 --- a/docs/serving/online_serving/speech_to_text.md +++ b/docs/serving/online_serving/speech_to_text.md @@ -64,10 +64,11 @@ The Transcriptions API supports uploading audio files in various formats includi - `model`: The model to use for transcription (required) - `language`: The language code (e.g., "en", "zh") (optional) - `prompt`: Optional text to guide the transcription style (optional) -- `response_format`: Format of the response ("json", "text") (optional) +- `response_format`: Format of the response ("json", "text", "verbose_json", or + "diarized_json") (optional) - `temperature`: Sampling temperature between 0 and 1 (optional) -For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182). +For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py). **Response Format:** @@ -98,6 +99,29 @@ For `verbose_json` response format: ``` Currently “verbose_json” response format doesn’t support no_speech_prob. +For models with diarization support, `diarized_json` returns OpenAI-compatible +speaker segments. Currently, this is supported by +`OpenMOSS-Team/MOSS-Transcribe-Diarize`. + +```json +{ + "task": "transcribe", + "duration": 6.1, + "text": "Hello. Hi, how are you?", + "segments": [ + { + "type": "transcript.text.segment", + "id": "segment_0", + "start": 0.0, + "end": 2.8, + "text": "Hello.", + "speaker": "S01" + } + ], + "usage": {"type": "duration", "seconds": 7} +} +``` + ### Extra Parameters The following [sampling parameters](../../api/README.md#inference-parameters) are supported. diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index e655f9c39ffe..9e75a24eaa1a 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -38,11 +38,12 @@ Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will ### HTTP Endpoints -When using the vLLM HTTP server, the same functionality is available via: +With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionality via: - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) +- `GET /weight_info` - Return the latest committed `weight_version` !!! note "Data Parallelism" When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 7579e5fd4d02..b8d39763181d 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -53,7 +53,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info | | `/start_weight_update` | POST | Start a weight update | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | -| `/finish_weight_update` | POST | Finish the weight update and run post-processing | +| `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | +| `/update_weight_version` | POST | Update `weight_version` without changing model weights | +| `/weight_info` | GET | Get the latest committed weight version | | `/pause` | POST | Pause generation before weight sync to handle inflight requests | | `/resume` | POST | Resume generation after weight sync | | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) | @@ -79,7 +81,7 @@ EngineClass.trainer_send_weights( ) # 4. Finish weight update on inference side -llm.finish_weight_update() +llm.finish_weight_update(weight_version="step-42") ``` See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples. diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 44c9c7cbfe50..305c21478fe9 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -35,21 +35,21 @@ The following metrics are exposed: ## General Metrics ---8<-- "docs/generated/metrics/general.inc.md" +--8<-- "gen:metrics-general" ## Speculative Decoding Metrics ---8<-- "docs/generated/metrics/spec_decode.inc.md" +--8<-- "gen:metrics-spec-decode" ## NIXL KV Connector Metrics ---8<-- "docs/generated/metrics/nixl_connector.inc.md" +--8<-- "gen:metrics-nixl" ## Model Flops Utilization (MFU) Performance Metrics These metrics are available via `--enable-mfu-metrics`: ---8<-- "docs/generated/metrics/perf.inc.md" +--8<-- "gen:metrics-mfu" ## Deprecation Policy diff --git a/docs/usage/security.md b/docs/usage/security.md index 2a2e1e886d64..7eb57c9aa24e 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -396,7 +396,7 @@ FIPS compliance depends on many factors, so a vLLM deployment is not automatical Operators running vLLM on FIPS-enabled hosts should select FIPS-approved algorithms via the following knobs: -- **Multimodal input hashing** — `VLLM_MM_HASHER_ALGORITHM` defaults to `blake3`, which is not FIPS-approved. Set it to `sha256` or `sha512` in FIPS-enabled environments. +- **Multimodal input hashing** — `--mm-hasher-algorithm` (config field `mm_hasher_algorithm`) defaults to `blake3`, which is not FIPS-approved. Set it to `sha256` or `sha512` in FIPS-enabled environments. - **Prefix-cache hashing** — set `--prefix-caching-hash-algo` (config field `prefix_caching_hash_algo`) to `sha256` or `sha256_cbor`. The `xxhash` and `xxhash_cbor` options are not FIPS-approved. - **TLS ciphers** — use `--ssl-ciphers` to restrict the API server's TLS handshake to FIPS-approved cipher suites that match your environment's policy. @@ -408,7 +408,13 @@ vLLM uses MD5 in a few places to derive non-security cache keys (for example, co Some dependencies expose hash implementations that are not FIPS-approved. vLLM only invokes them when the corresponding algorithm is selected, but operators with strict cryptographic controls may want to ensure the code paths are not exercised — and, where policy requires, that the packages themselves are absent: -- `blake3` — currently listed in `requirements/common.txt`, so a standard install pulls it in. It is imported lazily and only used when `VLLM_MM_HASHER_ALGORITHM=blake3` (the default). Setting `VLLM_MM_HASHER_ALGORITHM` to `sha256` or `sha512` is sufficient to keep the non-FIPS code path dormant. If your policy additionally forbids the package being present, uninstall it after `pip install` (`pip uninstall blake3`); vLLM will continue to function as long as `VLLM_MM_HASHER_ALGORITHM` is set to a non-blake3 value. +- `blake3` — currently listed in `requirements/common.txt`, so a standard + install pulls it in. It is imported lazily and only used when + `mm_hasher_algorithm=blake3` (the default). Setting + `--mm-hasher-algorithm sha256` or `--mm-hasher-algorithm sha512` is sufficient + to keep the non-FIPS code path dormant. If your policy additionally forbids + the package being present, uninstall it after installation; vLLM will + continue to function as long as a non-blake3 algorithm is selected. - `xxhash` — a true optional dependency (not in `requirements/common.txt`). It is only imported when an `xxhash`-based prefix-cache algorithm is selected. Leave it uninstalled and select a `sha256`-based prefix-cache algorithm. ### Beyond hashing: other FIPS considerations @@ -425,6 +431,104 @@ Hashing is the area where vLLM has explicit FIPS-aware code, but a FIPS-complian In short: the configuration knobs above let vLLM avoid non-approved algorithms, and the automatic fallbacks let it run without crashing on FIPS-enabled hosts. End-to-end FIPS compliance, however, is a property of the full deployment — host OS, crypto provider, transitive dependencies, and network architecture — not of vLLM alone. +## Ray Cluster Trust Model and Environment Variable Propagation + +### Trust Assumption + +vLLM treats the entire Ray cluster as a single trust domain. Any principal +with the ability to execute code within the Ray cluster (e.g. submit actors +or tasks) is considered to have the same level of trust as the driver/API +server process. This means that vLLM does **not** attempt to isolate +driver-side credentials from worker-side processes within the same Ray +cluster. + +This assumption is consistent with +[Ray's own security model](https://docs.ray.io/en/latest/ray-core/security.html), +which states that any user who can connect to a Ray cluster can run arbitrary +code on any node in that cluster. In other words, Ray cluster access already +implies full code execution on worker nodes, so restricting environment +variable propagation alone would not constitute a meaningful security +boundary. + +### Driver-to-Worker Environment Variable Propagation + +When using `RayExecutorV2` in multi-node deployments, vLLM propagates +environment variables from the driver process to remote Ray workers so that +workers have the configuration they need to function correctly (e.g. vLLM +settings, NCCL tuning, Hugging Face tokens for gated model downloads). + +The propagation uses a **copy-all-except-denylist** policy via +`get_driver_env_vars()` in `vllm/v1/executor/ray_env_utils.py`: every +environment variable present in the driver's `os.environ` is sent to +workers, except for a small set of worker-specific variables and any +names the operator has explicitly excluded. + +On the worker side, propagated variables are applied with `setdefault` +semantics — they fill in missing variables but never overwrite values +already present in the worker's environment. + +### When This Matters + +In deployments where operators intentionally scope credentials (such as +`HF_TOKEN`, cloud storage keys, registry tokens, or internal service +tokens) to the driver/API server alone — for example, when GPU workers +run in a different node, pod, or trust domain — the default propagation +behavior will copy those credentials into worker environments. A process +running on a worker node under the same OS user may then be able to read +those credentials (e.g. via `/proc//environ` on Linux). + +If your deployment treats Ray workers as less trusted than the driver, be +aware that environment-variable isolation is **not** enforced by default. + +### Hardening Recommendations + +Operators who want to limit which environment variables are propagated to +Ray workers can use the following mechanisms: + +#### 1. Denylist via Configuration File + +Create a JSON file at `$VLLM_CONFIG_ROOT/ray_non_carry_over_env_vars.json` +(default: `~/.config/vllm/ray_non_carry_over_env_vars.json`) containing an +array of environment variable names to exclude from propagation: + +```json +[ + "HF_TOKEN", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "GOOGLE_APPLICATION_CREDENTIALS", + "AZURE_CLIENT_SECRET", + "REGISTRY_TOKEN", + "MY_INTERNAL_SERVICE_KEY" +] +``` + +Any variable listed here will **not** be copied from the driver to workers. + +#### 2. Minimize Driver Environment + +Rather than setting credentials in the driver's shell environment, inject +them through a secrets manager, a mounted file, or a short-lived +subprocess so that they are not present in `os.environ` when vLLM starts. + +#### 3. Network and Process Isolation + +- Restrict `procfs` visibility on worker nodes (e.g. mount `/proc` with + `hidepid=2` or use a container runtime that isolates `/proc` between + pods) so that same-UID processes cannot read each other's + `/proc//environ`. +- Run the driver and workers under different OS users or in separate + containers with non-overlapping UIDs. + +#### 4. Limit Ray Cluster Access + +Because Ray cluster access is equivalent to arbitrary code execution, +ensure that only trusted principals can submit work to the cluster: + +- Use Ray's TLS authentication to restrict cluster membership. +- Place the Ray cluster on an isolated network segment. +- Do not expose the Ray client port or dashboard to untrusted networks. + ## Reporting Security Vulnerabilities If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md). diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 5613d5ba4e85..4ef2b0b5b6a7 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -126,7 +126,7 @@ Models using selective state-space mechanisms instead of standard transformer at Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported. Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., -`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). +`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`). Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 661401046bd4..d8e088351975 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -503,6 +503,45 @@ def run_gemma3n(questions: list[str], modality: str) -> ModelRequestData: ) +# Gemma 4 +def run_gemma4(questions: list[str], modality: str) -> ModelRequestData: + assert modality in ("image", "video") + model_name = "google/gemma-4-31B-it" + + # NOTE: Gemma-4-31B is a large model. Users running into Out-Of-Memory (OOM) + # errors might need to set `tensor_parallel_size` to > 1. + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=2, + limit_mm_per_prompt={modality: 1}, + ) + + if modality == "image": + prompts = [ + ( + "user\n" + f"<|image|>\n{question}\n" + "model\n" + ) + for question in questions + ] + else: # video + prompts = [ + ( + "user\n" + f"<|video|>\n{question}\n" + "model\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # GLM-4v def run_glm4v(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2303,6 +2342,7 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData: "exaone4_5": run_exaone4_5, "gemma3": run_gemma3, "gemma3n": run_gemma3n, + "gemma4": run_gemma4, "glm4v": run_glm4v, "glm4_1v": run_glm4_1v, "glm4_5v": run_glm4_5v, @@ -2374,6 +2414,7 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData: MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "llama4", + "gemma4", "qwen2_vl", "qwen2_5_vl", "qwen3_vl", diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py index 0a0efcbee361..8c9526e16e6a 100644 --- a/examples/rl/rlhf_http_ipc.py +++ b/examples/rl/rlhf_http_ipc.py @@ -46,10 +46,12 @@ from openai import OpenAI from transformers import AutoModelForCausalLM -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo BASE_URL = "http://localhost:8000" MODEL_NAME = "facebook/opt-125m" @@ -72,28 +74,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list return results -def init_weight_transfer_engine(base_url: str) -> None: - """Initialize weight transfer via HTTP endpoint (no-op for IPC).""" - url = f"{base_url}/init_weight_transfer_engine" - payload = {"init_info": dict()} - response = requests.post(url, json=payload, timeout=60) - response.raise_for_status() - - -def start_weight_update(base_url: str) -> None: - """Start a weight update via HTTP endpoint.""" - url = f"{base_url}/start_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - -def finish_weight_update(base_url: str) -> None: - """Finish a weight update via HTTP endpoint.""" - url = f"{base_url}/finish_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - def pause_generation(base_url: str) -> None: """Pause generation via HTTP endpoint.""" url = f"{base_url}/pause" @@ -159,23 +139,20 @@ def main(): print("Initializing weight transfer (IPC backend)...") - # Initialize weight transfer on vLLM server (no-op for IPC, but still required) - init_weight_transfer_engine(BASE_URL) + # The trainer engine drives the inference side over HTTP. init for IPC is a + # no-op rendezvous; the same client carries start/update/finish. + engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender + client=HTTPVLLMWeightSyncClient(BASE_URL), + source=ModuleSource(train_model), + ) # Pause generation before weight sync pause_generation(BASE_URL) - # Start weight update, broadcast via IPC, then finish - start_weight_update(BASE_URL) - print("Broadcasting weights via CUDA IPC (HTTP)...") - trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL) - IPCWeightTransferEngine.trainer_send_weights( - iterator=train_model.named_parameters(), - trainer_args=trainer_args, - ) - - finish_weight_update(BASE_URL) + # One call drives start_weight_update / update_weights / finish_weight_update. + engine.send_weights() # Resume generation after weight sync resume_generation(BASE_URL) diff --git a/examples/rl/rlhf_ipc.py b/examples/rl/rlhf_ipc.py index cb8542898792..6ed0343ade77 100644 --- a/examples/rl/rlhf_ipc.py +++ b/examples/rl/rlhf_ipc.py @@ -29,10 +29,12 @@ from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo class MyLLM(LLM): @@ -65,23 +67,19 @@ def __init__(self, llm_handle: ray.actor.ActorHandle): self.llm_handle = llm_handle def init_weight_transfer(self): - # IPC backend doesn't need initialization info - ray.get( - self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict())) + """Build the trainer-side IPC engine (no rendezvous needed for IPC).""" + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender + client=RayVLLMWeightSyncClient(self.llm_handle), + source=ModuleSource(self.train_model), ) - def broadcast_weights( - self, llm_handle: ray.actor.ActorHandle, packed: bool = False - ): - """Broadcast weights to the inference engine using IPC.""" - self.llm_handle = llm_handle - trainer_args = IPCTrainerSendWeightsArgs( - send_mode="ray", llm_handle=llm_handle, packed=packed - ) - IPCWeightTransferEngine.trainer_send_weights( - iterator=self.train_model.named_parameters(), - trainer_args=trainer_args, - ) + def broadcast_weights(self): + """Broadcast weights to the inference engine using IPC. + + Drives start/update/finish on the inference side internally. + """ + self.engine.send_weights() ray.init() @@ -138,10 +136,8 @@ def broadcast_weights( ray.get(llm.sleep.remote(level=0)) ray.get(train_model.init_weight_transfer.remote()) -# Start weight update, sync weights, then finish -ray.get(llm.start_weight_update.remote()) -ray.get(train_model.broadcast_weights.remote(llm)) -ray.get(llm.finish_weight_update.remote()) +# One call drives start_weight_update / update_weights / finish_weight_update. +ray.get(train_model.broadcast_weights.remote()) ray.get(llm.wake_up.remote(tags=["scheduling"])) diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py index 77ac6b4cfca0..c4c355948989 100644 --- a/examples/rl/rlhf_ipc_fsdp_ep.py +++ b/examples/rl/rlhf_ipc_fsdp_ep.py @@ -14,8 +14,10 @@ * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors, waits for initialization, and orchestrates generation / weight-sync. -Uses the built-in ``ray`` send_mode: each FSDP worker calls -``trainer_send_weights`` targeting its colocated LLM actor. +Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``) +and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only +rank 0 (the sender) ships the merged handles and drives the DP LLM actors through +its ``RayVLLMWeightSyncClient``. This example was run on 4xH100. """ @@ -23,7 +25,6 @@ from __future__ import annotations import os -from dataclasses import asdict import ray import torch @@ -31,19 +32,26 @@ from huggingface_hub import snapshot_download from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from torch.distributed._tensor import DTensor from torch.distributed.fsdp import fully_shard from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, - IPCWeightTransferInitInfo, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo from vllm.utils.network_utils import get_ip, get_open_port +# The inference side only needs the backend; the packed wire params are +# trainer-side and get propagated to the workers at the init handshake. +WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc") +# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size). +PACKED = True +PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 + TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42")) VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42")) @@ -125,72 +133,34 @@ def get_rank(self): def get_weight_metadata(self): return self.weight_names, self.weight_dtype_names, self.weight_shapes - def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True): - """All-gather full params; all ranks create IPC handles, rank 0 sends. + def setup_engine(self, llm_handles): + """Build the trainer IPC engine on every FSDP rank. - All ranks must call trainer_send_weights so they participate in the - all_gather_object collective inside _all_gather_and_merge_handles. - Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero). + Called on all ranks: rank 0 becomes the sender (drives the inference + side and performs the no-op IPC init handshake against all DP LLM + actors); the other ranks hold a null-client engine and only join the + IPC handle all-gather during send_weights. """ - - def _full_param_iter(): - # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs - # all experts into two fused 3-D tensors per layer: - # experts.gate_up_proj shape (E, 2*I, H) - # experts.down_proj shape (E, H, I) - # vLLM's Qwen3MoE load_weights still expects the older - # per-expert HF layout (experts..gate_proj.weight, - # experts..up_proj.weight, experts..down_proj.weight), - # so we un-fuse on the fly. Split order matches HF's forward: - # gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1) - # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up. - params = self.model.state_dict() - for name in list(params.keys()): - param = params.pop(name) - if isinstance(param, DTensor): - tensor = param.full_tensor().detach().contiguous() - else: - tensor = param.detach().contiguous() - del param - - if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3: - prefix = name[: -len(".gate_up_proj")] - num_experts, two_inter, _ = tensor.shape - inter = two_inter // 2 - for i in range(num_experts): - expert = tensor[i] - yield ( - f"{prefix}.{i}.gate_proj.weight", - expert[:inter].contiguous(), - ) - yield ( - f"{prefix}.{i}.up_proj.weight", - expert[inter:].contiguous(), - ) - del tensor - elif name.endswith(".experts.down_proj") and tensor.dim() == 3: - prefix = name[: -len(".down_proj")] - num_experts = tensor.shape[0] - for i in range(num_experts): - yield ( - f"{prefix}.{i}.down_proj.weight", - tensor[i].contiguous(), - ) - del tensor - else: - yield name, tensor - - trainer_args = IPCTrainerSendWeightsArgs( - send_mode="ray", - llm_handle=llm_handle, - packed=packed, - packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB - ) - IPCWeightTransferEngine.trainer_send_weights( - iterator=_full_param_iter(), - trainer_args=trainer_args, + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo( + rank=self.rank, # FSDP rank; sender is 0 + packed=PACKED, + packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES, + ), + client=RayVLLMWeightSyncClient(llm_handles), + source=ModuleSource(self.model), ) + def gather_and_broadcast_weights_ipc(self): + """All-gather full params across FSDP ranks; rank 0 sends to vLLM. + + Called on all ranks concurrently. `send_weights` gathers each param and + contributes to the IPC handle all-gather on every rank; only rank 0 (the + sender) drives start/update/finish on the inference side (the other + ranks' client RPCs no-op). + """ + self.engine.send_weights() + @ray.remote(num_cpus=1) class DataParallelInferenceEngine: @@ -224,7 +194,7 @@ def __init__( distributed_executor_backend="ray", enable_expert_parallel=True, gpu_memory_utilization=0.35, - weight_transfer_config=WeightTransferConfig(backend="ipc"), + weight_transfer_config=WEIGHT_TRANSFER_CONFIG, enable_sleep_mode=True, load_format="dummy", dp_rank=r, @@ -267,22 +237,6 @@ def generate(self, prompts: list[str], sampling_params): rank_idx += 1 return ordered - def init_weight_transfer(self): - ray.get( - [ - actor.init_weight_transfer_engine.remote( - dict(init_info=asdict(IPCWeightTransferInitInfo())) - ) - for actor in self.llm_actors - ] - ) - - def start_weight_update(self): - ray.get([actor.start_weight_update.remote() for actor in self.llm_actors]) - - def finish_weight_update(self): - ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors]) - def sleep(self, level: int = 0): ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors]) @@ -370,8 +324,10 @@ def main(): print("-" * 60) # --- Weight transfer --- - print("[transfer] Initializing IPC weight transfer...") - ray.get(inference_engine.init_weight_transfer.remote()) + # The rank-0 FSDP worker owns the trainer engine and drives the inference + # side (init/start/update/finish) for all DP actors via the Ray client. + print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...") + ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers]) # Two-phase sleep/wake pattern: # 1. sleep(level=1) — offload weights to CPU, discard KV cache @@ -384,18 +340,11 @@ def main(): print("[sync] Waking weights (KV cache stays free)...") ray.get(inference_engine.wake_up.remote(tags=["weights"])) - print("[sync] Starting weight update...") - ray.get(inference_engine.start_weight_update.remote()) - + # All FSDP ranks participate in the IPC handle all-gather; rank 0's engine + # additionally drives start_weight_update / update_weights / + # finish_weight_update on the inference side. print("[sync] Packed IPC transfer FSDP → vLLM...") - ray.get( - [ - w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True) - for w in fsdp_workers - ] - ) - - ray.get(inference_engine.finish_weight_update.remote()) + ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers]) print("[sync] Weight transfer complete.") print("[sync] Waking KV cache + scheduling...") diff --git a/mkdocs.yaml b/mkdocs.yaml index a32cea618068..f3b9deab7871 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -3,7 +3,6 @@ site_url: !ENV READTHEDOCS_CANONICAL_URL repo_url: https://github.com/vllm-project/vllm edit_uri: edit/main/docs/ exclude_docs: | - argparse *.inc.md *.template.md theme: @@ -50,24 +49,22 @@ theme: hooks: - docs/mkdocs/hooks/remove_announcement.py - - docs/mkdocs/hooks/generate_examples.py - - docs/mkdocs/hooks/generate_argparse.py - - docs/mkdocs/hooks/generate_metrics.py - docs/mkdocs/hooks/url_schemes.py - docs/mkdocs/hooks/autoref_code.py plugins: - meta - search + - gen-files: + scripts: + - docs/mkdocs/gen_files/generate_examples.py + - docs/mkdocs/gen_files/generate_argparse.py + - docs/mkdocs/gen_files/generate_metrics.py + - docs/mkdocs/gen_files/generate_attention_backends.py - autorefs - awesome-nav - glightbox - - git-revision-date-localized: - # exclude autogenerated files - exclude: - - api/* - - examples/* - - generated/* + - git-revision-date-localized - minify: minify_html: true minify_js: true @@ -160,3 +157,4 @@ extra_javascript: - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js - mkdocs/javascript/edit_and_feedback.js - mkdocs/javascript/slack_and_forum.js + - mkdocs/javascript/reo.js diff --git a/pyproject.toml b/pyproject.toml index 04f1df204ba8..0766645fc748 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ requires = [ "setuptools>=77.0.3,<81.0.0", "setuptools-scm>=8.0", "setuptools-rust>=1.9.0", - "torch == 2.11.0", + "torch == 2.13.0", "wheel", "jinja2", ] @@ -129,10 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", - "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] + "rust/src/tokenizer/src/incremental.rs"] ignore-hidden = false [tool.typos.default] @@ -156,6 +155,8 @@ view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" [tool.typos.default.extend-words] +Hel = "Hel" +wether = "wether" iy = "iy" indx = "indx" # intel cpu features diff --git a/requirements/build/cpu.txt b/requirements/build/cpu.txt index 27a3ac65c986..5a03960e5608 100644 --- a/requirements/build/cpu.txt +++ b/requirements/build/cpu.txt @@ -4,8 +4,8 @@ packaging>=24.2 setuptools==77.0.3 # this version can reuse CMake build dir setuptools-scm>=8 setuptools-rust>=1.9.0 -torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64" -torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64" +torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64" +torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64" wheel jinja2>=3.1.6 regex diff --git a/requirements/build/cuda.txt b/requirements/build/cuda.txt index 70da484a4133..d9d68022b307 100644 --- a/requirements/build/cuda.txt +++ b/requirements/build/cuda.txt @@ -5,7 +5,7 @@ packaging>=24.2 setuptools>=77.0.3,<81.0.0 setuptools-scm>=8 setuptools-rust>=1.9.0 -torch==2.11.0 +torch==2.13.0 wheel jinja2>=3.1.6 regex diff --git a/requirements/build/tpu.txt b/requirements/build/tpu.txt index 56348e757ecd..1a791a633f46 100644 --- a/requirements/build/tpu.txt +++ b/requirements/build/tpu.txt @@ -4,5 +4,5 @@ ninja setuptools>=77.0.3,<81.0.0 setuptools-scm>=8 setuptools-rust>=1.9.0 -torch==2.11.0+cpu +torch==2.13.0+cpu wheel diff --git a/requirements/common.txt b/requirements/common.txt index ec4386baa936..748ffe88f45d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -31,7 +31,7 @@ partial-json-parser # used for parsing partial JSON outputs jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation pyzmq >= 25.0.0 msgspec -mistral_common[image] >= 1.11.5 +mistral_common[image] >= 1.11.6 opencv-python-headless >= 4.13.0 # required for video IO pyyaml six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12 diff --git a/requirements/cpu.txt b/requirements/cpu.txt index c0b98d22c9b9..30e47a8cab1a 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding # Dependencies for CPUs -torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64" -torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64" +torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64" +torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64" # required for the image processor of minicpm-o-2_6, this must be updated alongside torch torchaudio; platform_machine != "s390x" and platform_machine != "riscv64" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 3b8ff816f180..af1f5f61ca6a 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -4,20 +4,20 @@ numba == 0.65.0 # Required for N-gram speculative decoding # Dependencies for NVIDIA GPUs -torch==2.11.0 +torch==2.13.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 +torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version torchcodec >= 0.14 PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from # install_requires so the published wheel does not carry an unresolvable pin --extra-index-url https://flashinfer.ai/whl/ -flashinfer-python==0.6.14 -flashinfer-cubin==0.6.14 -apache-tvm-ffi==0.1.10 -tilelang==0.1.9 +flashinfer-python==0.6.16.post3 +flashinfer-cubin==0.6.16.post3 +apache-tvm-ffi==0.1.11 +tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1 nvtx==0.2.15 @@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) nvidia-cutlass-dsl[cu13]==4.6.0 -quack-kernels>=0.4.0 # Required for tml-fa4 +quack-kernels==0.6.1 # Required for CUTLASS DSL 4.6 by MSA # Tokenspeed_MLA for faster mla with spec decode tokenspeed-mla==0.1.8; platform_system == "Linux" diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index d7cec7766972..aed7dd09a9c3 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -124,6 +126,8 @@ click==8.4.2 # uvicorn cloudpickle==3.1.2 # via -r requirements/test/../common.txt +cohere==7.0.8 + # via -r requirements/test/cuda.in cohere-melody==0.9.0 # via -r requirements/test/cuda.in colorama==0.4.6 @@ -224,6 +228,8 @@ fastar==0.11.0 # via # fastapi # fastapi-cloud-cli +fastavro==1.12.2 + # via cohere fastparquet==2024.11.0 # via genai-perf fastrlock==0.8.2 @@ -324,6 +330,7 @@ httpx==0.27.2 # via # -r requirements/test/cuda.in # anthropic + # cohere # fastapi # fastapi-cloud-cli # huggingface-hub @@ -474,7 +481,7 @@ mcp==1.28.1 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.5 +mistral-common==1.11.6 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in @@ -763,11 +770,14 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -r requirements/test/../common.txt # albumentations # anthropic + # cohere # compressed-tensors # datamodel-code-generator # fastapi @@ -785,7 +795,9 @@ pydantic==2.12.0 # ray # xgrammar pydantic-core==2.41.1 - # via pydantic + # via + # cohere + # pydantic pydantic-extra-types==2.10.5 # via # fastapi @@ -907,6 +919,7 @@ requests==2.32.3 # -r requirements/test/../common.txt # azure-core # buildkite-test-collector + # cohere # datasets # docker # evaluate @@ -1106,8 +1119,9 @@ tokenizers==0.22.2 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in + # cohere # transformers -torch==2.11.0+cpu +torch==2.13.0+cpu # via # -r requirements/test/cuda.in # accelerate @@ -1127,6 +1141,8 @@ torch==2.11.0+cpu # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cpu # via # -r requirements/test/cuda.in @@ -1134,7 +1150,7 @@ torchaudio==2.11.0+cpu # vocos torchcodec==0.14.0+cpu # via -r requirements/test/cuda.in -torchvision==0.26.0+cpu +torchvision==0.28.0+cpu # via # -r requirements/test/cuda.in # open-clip-torch @@ -1157,7 +1173,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.13.1 +transformers==5.14.1 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in @@ -1185,6 +1201,8 @@ typer==0.26.8 # fastsafetensors # perceptron # transformers +types-requests==2.33.0.20260712 + # via cohere typing-extensions==4.15.0 # via # -r requirements/test/../common.txt @@ -1198,6 +1216,7 @@ typing-extensions==4.15.0 # azure-identity # azure-storage-blob # chz + # cohere # fastapi # grpcio # huggingface-hub @@ -1242,6 +1261,7 @@ urllib3==2.2.3 # responses # sentry-sdk # tritonclient + # types-requests uvicorn==0.35.0 # via # fastapi diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index d91cba0601ac..1a652a035f82 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -28,24 +28,24 @@ soundfile # required for audio tests jiwer # required for audio tests tblib # for pickling test exceptions timm >=1.0.17 # required for internvl and gemma3n-mm test -torch==2.11.0 +torch==2.13.0 torchaudio==2.11.0 -torchvision==0.26.0 +torchvision==0.28.0 transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.5 # required for voxtral test +mistral_common[image,audio] >= 1.11.6 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.13.1 +transformers==5.14.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 - +torch-abi-audit # CI check for PyTorch stable ABI compliance genai_perf>=0.0.8 tritonclient>=2.51.0 @@ -73,6 +73,7 @@ gpt-oss >= 0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test cohere_melody>=0.9.0 # required for cohere command reasoning parser test +cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with teerratorch requirements. diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index c4a9ea1ae2db..07d18e4f8a7b 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -43,7 +45,7 @@ anyio==4.14.1 # sse-starlette # starlette # watchfiles -apache-tvm-ffi==0.1.10 +apache-tvm-ffi==0.1.11 # via # -c requirements/cuda.txt # xgrammar @@ -129,6 +131,8 @@ click==8.4.2 # uvicorn cloudpickle==3.1.2 # via -r requirements/test/../common.txt +cohere==7.0.8 + # via -r requirements/test/cuda.in cohere-melody==0.9.0 # via -r requirements/test/cuda.in colorama==0.4.6 @@ -159,7 +163,7 @@ cuda-bindings==13.0.3 # via torch cuda-pathfinder==1.3.3 # via cuda-bindings -cuda-toolkit==13.0.2 +cuda-toolkit==13.0.3.0 # via torch cupy-cuda12x==13.6.0 # via ray @@ -240,6 +244,8 @@ fastar==0.11.0 # via # fastapi # fastapi-cloud-cli +fastavro==1.12.2 + # via cohere fastparquet==2024.11.0 # via genai-perf fastrlock==0.8.2 @@ -343,6 +349,7 @@ httpx==0.27.2 # via # -r requirements/test/cuda.in # anthropic + # cohere # fastapi # fastapi-cloud-cli # huggingface-hub @@ -500,7 +507,7 @@ mcp==1.28.1 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.5 +mistral-common==1.11.6 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -599,7 +606,7 @@ numpy==2.2.6 # tritonclient # vocos # xgrammar -nvidia-cublas==13.1.0.3 +nvidia-cublas==13.1.1.3 # via # cuda-toolkit # nvidia-cudnn-cu13 @@ -607,10 +614,12 @@ nvidia-cublas==13.1.0.3 nvidia-cuda-cupti==13.0.85 # via cuda-toolkit nvidia-cuda-nvrtc==13.0.88 - # via cuda-toolkit + # via + # cuda-toolkit + # nvidia-cublas nvidia-cuda-runtime==13.0.96 # via cuda-toolkit -nvidia-cudnn-cu13==9.19.0.56 +nvidia-cudnn-cu13==9.20.0.48 # via torch nvidia-cufft==12.0.0.61 # via cuda-toolkit @@ -624,9 +633,9 @@ nvidia-cusparse==12.6.3.3 # via # cuda-toolkit # nvidia-cusolver -nvidia-cusparselt-cu13==0.8.0 +nvidia-cusparselt-cu13==0.8.1 # via torch -nvidia-nccl-cu13==2.28.9 +nvidia-nccl-cu13==2.29.7 # via torch nvidia-nvjitlink==13.0.88 # via @@ -848,12 +857,15 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt # albumentations # anthropic + # cohere # compressed-tensors # datamodel-code-generator # fastapi @@ -871,7 +883,9 @@ pydantic==2.12.0 # ray # xgrammar pydantic-core==2.41.1 - # via pydantic + # via + # cohere + # pydantic pydantic-extra-types==2.10.5 # via # fastapi @@ -996,6 +1010,7 @@ requests==2.32.3 # -r requirements/test/../common.txt # azure-core # buildkite-test-collector + # cohere # datasets # docker # evaluate @@ -1201,8 +1216,9 @@ tokenizers==0.22.2 # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/cuda.in + # cohere # transformers -torch==2.11.0+cu130 +torch==2.13.0+cu130 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in @@ -1223,6 +1239,8 @@ torch==2.11.0+cu130 # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cu130 # via # -c requirements/cuda.txt @@ -1233,7 +1251,7 @@ torchcodec==0.14.0+cu130 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in -torchvision==0.26.0+cu130 +torchvision==0.28.0+cu130 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in @@ -1257,7 +1275,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.13.1 +transformers==5.14.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1270,7 +1288,7 @@ transformers==5.13.1 # xgrammar transformers-stream-generator==0.0.5 # via -r requirements/test/cuda.in -triton==3.6.0 +triton==3.7.1 # via # torch # xgrammar @@ -1288,6 +1306,8 @@ typer==0.26.8 # fastsafetensors # perceptron # transformers +types-requests==2.33.0.20260712 + # via cohere typing-extensions==4.15.0 # via # -c requirements/common.txt @@ -1302,6 +1322,7 @@ typing-extensions==4.15.0 # azure-identity # azure-storage-blob # chz + # cohere # fastapi # grpcio # huggingface-hub @@ -1346,6 +1367,7 @@ urllib3==2.2.3 # responses # sentry-sdk # tritonclient + # types-requests uvicorn==0.35.0 # via # fastapi diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index b308a4401399..1beb4294d0fb 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -23,13 +23,13 @@ jiwer # required for audio tests timm # required for internvl test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.5 # required for voxtral test +mistral_common[image,audio] >= 1.11.6 # required for voxtral test num2words # required for smolvlm test opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.13.1 +transformers==5.14.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 7cb3b91f7bd7..9edb08dc4611 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -21,7 +21,7 @@ vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test peft>=0.19.1 # required for phi-4-mm test pqdm -ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests +ray[cgraph,default]==2.55.1 # Includes worker startup getenv/setenv race fix sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests @@ -29,13 +29,13 @@ tblib # for pickling test exceptions timm>=1.0.17 # required for internvl and gemma3n-mm test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio]>=1.11.5 # required for voxtral test +mistral_common[image,audio]>=1.11.6 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.13.1 +transformers==5.14.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test # quantization @@ -71,6 +71,7 @@ gpt-oss>=0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank>=1.18.7 # required for fireredasr2 test cohere_melody>=0.9.0 # required for cohere command reasoning parser test +cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with terratorch requirements. diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index f274a93362ce..c7979e4efaf9 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -131,6 +131,8 @@ cloudpickle==3.1.2 # via # -r requirements/test/../common.txt # tilelang +cohere==7.0.8 + # via -r requirements/test/rocm.in cohere-melody==0.9.0 # via -r requirements/test/rocm.in colorama==0.4.6 @@ -231,6 +233,8 @@ fastapi-cloud-cli==0.16.1 # via fastapi-cli fastar==0.10.0 # via fastapi-cloud-cli +fastavro==1.12.2 + # via cohere fastparquet==2026.3.0 # via genai-perf fastsafetensors==0.3.2 @@ -336,6 +340,7 @@ httpx==0.27.2 # via # -r requirements/test/rocm.in # anthropic + # cohere # fastapi # fastapi-cloud-cli # huggingface-hub @@ -497,7 +502,7 @@ mcp==1.27.0 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.5 +mistral-common==1.11.6 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -832,6 +837,7 @@ pydantic==2.12.5 # -r requirements/test/../common.txt # albumentations # anthropic + # cohere # compressed-tensors # datamodel-code-generator # fastapi @@ -849,7 +855,9 @@ pydantic==2.12.5 # ray # xgrammar pydantic-core==2.41.5 - # via pydantic + # via + # cohere + # pydantic pydantic-extra-types==2.11.1 # via # fastapi @@ -956,7 +964,7 @@ rapidfuzz==3.12.1 # via # -r requirements/test/rocm.in # jiwer -ray==2.54.0 +ray==2.55.1 # via -r requirements/test/rocm.in redis==7.3.0 # via tensorizer @@ -978,6 +986,7 @@ requests==2.32.5 # -r requirements/test/../common.txt # azure-core # buildkite-test-collector + # cohere # datasets # docker # evaluate @@ -1193,6 +1202,7 @@ tokenizers==0.22.2 # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/rocm.in + # cohere # transformers torch-c-dlpack-ext==0.1.5 # via tilelang @@ -1214,7 +1224,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.13.1 +transformers==5.14.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1243,6 +1253,8 @@ typer==0.24.1 # fastsafetensors # perceptron # transformers +types-requests==2.33.0.20260712 + # via cohere typing-extensions==4.15.0 # via # -c requirements/common.txt @@ -1257,6 +1269,7 @@ typing-extensions==4.15.0 # azure-identity # azure-storage-blob # chz + # cohere # fastapi # grpcio # huggingface-hub @@ -1302,6 +1315,7 @@ urllib3==2.6.3 # responses # sentry-sdk # tritonclient + # types-requests uvicorn==0.42.0 # via # fastapi diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index d2380a751319..a786740fceb9 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -17,7 +17,7 @@ accelerate arctic-inference lm_eval[api]>=0.4.12 modelscope<1.38 -transformers==5.13.1 +transformers==5.14.1 # --- Audio Processing --- librosa diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index d5a33b8a1fc1..4ecf0bede9ee 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -140,7 +140,7 @@ docopt==0.6.2 # via num2words docstring-parser==0.18.0 # via anthropic -dpcpp-cpp-rt==2025.3.2 +dpcpp-cpp-rt==2026.0.0 # via # onemkl-sycl-blas # onemkl-sycl-dft @@ -253,27 +253,27 @@ ijson==3.5.0 # via -r requirements/test/../common.txt imageio==2.37.3 # via scikit-image -impi-rt==2021.17.2 +impi-rt==2021.18.0 # via # oneccl # torch iniconfig==2.3.0 # via pytest -intel-cmplr-lib-rt==2025.3.2 +intel-cmplr-lib-rt==2026.0.0 # via # intel-sycl-rt # torch -intel-cmplr-lib-ur==2025.3.2 +intel-cmplr-lib-ur==2026.0.0 # via # intel-openmp # intel-sycl-rt # torch -intel-cmplr-lic-rt==2025.3.2 +intel-cmplr-lic-rt==2026.0.0 # via # intel-opencl-rt # intel-sycl-rt # torch -intel-opencl-rt==2025.3.2 +intel-opencl-rt==2026.0.0 # via # dpcpp-cpp-rt # onemkl-sycl-blas @@ -282,14 +282,14 @@ intel-opencl-rt==2025.3.2 # onemkl-sycl-rng # onemkl-sycl-sparse # torch -intel-openmp==2025.3.2 +intel-openmp==2026.0.0 # via # dpcpp-cpp-rt # mkl # torch -intel-pti==0.16.0 +intel-pti==0.17.0 # via torch -intel-sycl-rt==2025.3.2 +intel-sycl-rt==2026.0.0 # via # dpcpp-cpp-rt # oneccl @@ -373,12 +373,12 @@ mcp==1.28.1 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.5 +mistral-common==1.11.6 # via # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/xpu.in -mkl==2025.3.1 +mkl==2026.0.0 # via # onemkl-sycl-blas # onemkl-sycl-dft @@ -453,28 +453,28 @@ numpy==2.2.6 # torchvision # transformers # xgrammar -oneccl==2021.17.2 +oneccl==2022.0.0 # via # oneccl-devel # torch -oneccl-devel==2021.17.2 +oneccl-devel==2022.0.0 # via torch -onemkl-license==2025.3.1 +onemkl-license==2026.0.0 # via # mkl # torch -onemkl-sycl-blas==2025.3.1 +onemkl-sycl-blas==2026.0.0 # via # onemkl-sycl-lapack # onemkl-sycl-sparse # torch -onemkl-sycl-dft==2025.3.1 +onemkl-sycl-dft==2026.0.0 # via torch -onemkl-sycl-lapack==2025.3.1 +onemkl-sycl-lapack==2026.0.0 # via torch -onemkl-sycl-rng==2025.3.1 +onemkl-sycl-rng==2026.0.0 # via torch -onemkl-sycl-sparse==2025.3.1 +onemkl-sycl-sparse==2026.0.0 # via torch openai==2.44.0 # via @@ -719,6 +719,8 @@ pyyaml==6.0.3 # timm # transformers # uvicorn +pyzes==0.1.1 + # via torch pyzmq==27.1.0 # via # -c requirements/common.txt @@ -871,14 +873,14 @@ tabledata==1.3.4 # via pytablewriter tabulate==0.10.0 # via sacrebleu -tbb==2022.3.1 +tbb==2023.0.0 # via # intel-opencl-rt # mkl # torch tblib==3.1.0 # via -r requirements/test/xpu.in -tcmlib==1.4.1 +tcmlib==1.5.0 # via # tbb # torch @@ -910,7 +912,7 @@ tokenizers==0.22.2 # -c requirements/common.txt # -r requirements/test/../common.txt # transformers -torch==2.12.0+xpu +torch==2.13.0+xpu # via # -c requirements/xpu.txt # accelerate @@ -920,7 +922,7 @@ torch==2.12.0+xpu # timm # torchvision # xgrammar -torchvision==0.27.0+xpu +torchvision==0.28.0+xpu # via timm tqdm==4.67.3 # via @@ -936,7 +938,7 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==5.13.1 +transformers==5.14.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -946,7 +948,7 @@ transformers==5.13.1 # xgrammar triton==3.7.1 # via xgrammar -triton-xpu==3.7.1 +triton-xpu==3.7.2 # via torch typepy==1.3.4 # via @@ -1001,7 +1003,7 @@ typing-inspection==0.4.2 # mcp # pydantic # pydantic-settings -umf==1.0.3 +umf==1.1.0 # via # intel-cmplr-lib-ur # torch diff --git a/requirements/tpu.txt b/requirements/tpu.txt index da477a68461c..bc25cd3c23ab 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.24.0 +tpu-inference==0.26.0 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 8ae5649b20bc..dc81cce64f44 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -12,10 +12,10 @@ jinja2>=3.1.6 datasets # for benchmark scripts numba == 0.65.0 # Required for N-gram speculative decoding --extra-index-url=https://download.pytorch.org/whl/xpu -torch==2.12.0 +torch==2.13.0 torchaudio torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib==0.14.1 -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl +auto_round_lib==0.14.2 +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.12/vllm_xpu_kernels-0.1.12-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 0709ef8fc3ac..38f1c257756f 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5" +source = "git+https://github.com/smg-project/llm-multimodal?rev=15adba5e025d8636ba4a334fb379b1371f6196a1#15adba5e025d8636ba4a334fb379b1371f6196a1" dependencies = [ "anyhow", "base64 0.22.1", @@ -3446,7 +3446,6 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", - "futures-channel", "futures-core", "futures-util", "h2", @@ -4876,6 +4875,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] @@ -4937,9 +4937,9 @@ dependencies = [ [[package]] name = "tonic" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", @@ -4966,9 +4966,9 @@ dependencies = [ [[package]] name = "tonic-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734" +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322" dependencies = [ "prettyplease", "proc-macro2", @@ -4976,11 +4976,24 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "tonic-health" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082" +dependencies = [ + "prost", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", +] + [[package]] name = "tonic-prost" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", "prost", @@ -4989,9 +5002,9 @@ dependencies = [ [[package]] name = "tonic-prost-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a" +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" dependencies = [ "prettyplease", "proc-macro2", @@ -5484,12 +5497,15 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.18", + "thiserror-ext", "tiktoken-rs 0.9.1", "tokenizers", "tokio", "tokio-stream", + "tracing", "url", "uuid", + "vllm-tracing", ] [[package]] @@ -5554,16 +5570,16 @@ dependencies = [ "serde_json", "serde_with", "thiserror-ext", - "time", "tokio", "tokio-util", "tracing", - "tracing-subscriber", "uuid", + "vllm-bench", "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", + "vllm-tracing", ] [[package]] @@ -5729,6 +5745,7 @@ dependencies = [ "tokio-stream", "tokio-util", "tonic", + "tonic-health", "tonic-prost", "tonic-prost-build", "tower", @@ -5809,6 +5826,15 @@ dependencies = [ "vllm-parser", ] +[[package]] +name = "vllm-tracing" +version = "0.1.0" +dependencies = [ + "time", + "tracing", + "tracing-subscriber", +] + [[package]] name = "walkdir" version = "2.5.0" @@ -6320,9 +6346,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" [[package]] name = "xgrammar-structural-tag" -version = "0.1.0+xgrammar.0.2.2.4d145cc" +version = "0.2.0+xgrammar.0.2.4.dd729e7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6" +checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a" dependencies = [ "auto_impl", "serde", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 492b1c5e2607..09f55cf07cdf 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,6 +13,7 @@ members = [ "src/server", "src/text", "src/tokenizer", + "src/tracing", ] resolver = "3" @@ -58,7 +59,7 @@ indexmap = "2.13.0" indicatif = "0.18.4" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } @@ -118,10 +119,11 @@ tokio = { version = "1.47.1", features = [ tokio-openssl = "0.6" tokio-stream = "0.1" tokio-util = { version = "0.7.18", features = ["rt"] } -tonic = "0.14.5" -tonic-build = "0.14.5" -tonic-prost = "0.14.5" -tonic-prost-build = "0.14.5" +tonic = "0.14.6" +tonic-build = "0.14.6" +tonic-health = "0.14.6" +tonic-prost = "0.14.6" +tonic-prost-build = "0.14.6" tool-parser = "1.2.0" tower = { version = "0.5.3", features = ["util"] } tower-http = { version = "0.6.8", features = ["cors", "trace"] } @@ -132,6 +134,7 @@ trait-set = "0.3.0" url = "2.5.7" uuid = { version = "1.22.0", features = ["v4"] } validator = { version = "0.20.0", features = ["derive"] } +vllm-bench = { path = "src/bench" } vllm-chat = { path = "src/chat" } vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } @@ -141,8 +144,9 @@ vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } +vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } -xgrammar-structural-tag = "0.1.0" +xgrammar-structural-tag = "0.2.0" zeromq = { version = "0.6.0", default-features = false, features = [ "tokio-runtime", "all-transport", diff --git a/rust/proto/control.proto b/rust/proto/control.proto new file mode 100644 index 000000000000..0e25aea26474 --- /dev/null +++ b/rust/proto/control.proto @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +syntax = "proto3"; +package vllm; + +service Control { + rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} + rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} + rpc Abort (AbortRequest) returns (AbortResponse) {} + rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} +} + +message GetServerInfoRequest {} + +message ServerInfo { + string engine_version = 1; + string api_version = 2; + string instance_id = 3; + ParallelismInfo parallelism = 4; + uint32 max_model_len = 5; + uint32 kv_block_size = 6; + uint64 total_kv_blocks = 7; + uint64 max_running_requests = 8; + uint64 max_batched_tokens = 9; +} + +message ParallelismInfo { + uint32 tensor_parallel_size = 1; + uint32 pipeline_parallel_size = 2; + uint32 data_parallel_size = 3; + uint32 data_parallel_rank = 4; + uint32 decode_context_parallel_size = 5; +} + +message GetModelInfoRequest {} + +message ModelInfo { + string model_id = 1; + string served_model_name = 2; + repeated string served_model_aliases = 3; + + bool supports_text_input = 20; + bool supports_token_ids_input = 21; + bool supports_multimodal = 23; + string reasoning_parser = 24; + string tool_call_parser = 25; +} + +message AbortRequest { + repeated string request_ids = 1; +} + +message AbortResponse {} + +// ====================================================================================== +// KV discovery +// ====================================================================================== + +message GetKvEventSourcesRequest {} +message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } + +message KvEventSource { + string transport = 1; + string endpoint = 2; + string topic = 3; + string replay_endpoint = 4; + optional uint32 data_parallel_rank = 5; + string encoding = 6; + uint32 schema_version = 7; + uint32 buffer_steps = 8; + uint32 hwm = 9; + uint32 max_queue_size = 10; +} diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/inference.proto similarity index 99% rename from rust/proto/vllm_grpc.proto rename to rust/proto/inference.proto index 2509d5071b63..b1c08ae76e60 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/inference.proto @@ -7,7 +7,7 @@ package vllm; import "google/protobuf/struct.proto"; -service Generate { +service Inference { // Generates text given a prompt rpc Generate (GenerateRequest) returns (GenerateResponse) {} // Generates text given a prompt, streaming the outputs @@ -200,4 +200,3 @@ message CandidateTokenInfo { message TokenIds { repeated uint32 ids = 1; } - diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 22c8f84418b2..960e7a62f7dd 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -20,18 +20,21 @@ mimalloc.workspace = true rand.workspace = true rand_distr.workspace = true rayon.workspace = true -reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] } +reqwest = { workspace = true, features = ["json", "stream", "http2"] } rlimit.workspace = true rustc-hash.workspace = true serde = { workspace = true, features = ["rc"] } serde_json = { workspace = true, features = ["raw_value"] } thiserror.workspace = true +thiserror-ext.workspace = true tiktoken-rs.workspace = true tokenizers.workspace = true tokio.workspace = true tokio-stream.workspace = true +tracing.workspace = true url.workspace = true uuid.workspace = true +vllm-tracing.workspace = true [lints] workspace = true diff --git a/rust/src/bench/src/backends/pooling.rs b/rust/src/bench/src/backends/pooling.rs index 6a1cdf4fe160..8f9d4deab996 100644 --- a/rust/src/bench/src/backends/pooling.rs +++ b/rust/src/bench/src/backends/pooling.rs @@ -158,9 +158,10 @@ impl PoolingBackend { // (mirrors Python async_request_vllm_rerank). if let Some(ref list) = input.prompt_list { if list.len() < 2 { - eprintln!( - "WARNING: vllm-rerank request has no documents \ - (prompt_list needs [query, doc, ...])" + tracing::warn!( + backend = "vllm-rerank", + inputs = list.len(), + "rerank request has no documents" ); } let query = list.first().map(|s| s.as_ref()).unwrap_or(""); @@ -175,10 +176,10 @@ impl PoolingBackend { // Legacy path: text prompt as query, documents via --extra-body. let query = input.prompt.as_ref(); if query.is_empty() && input.prompt_token_ids.is_some() { - eprintln!( - "WARNING: vllm-rerank received empty query (random dataset uses \ - token IDs only). Use --dataset-name random-rerank for meaningful \ - rerank benchmarks." + tracing::warn!( + backend = "vllm-rerank", + dataset = "random", + "rerank request has an empty query; use the random-rerank dataset" ); } serde_json::json!({ diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 8c834641acb0..cc4acc4f0614 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::time::Instant; use indicatif::{ProgressBar, ProgressStyle}; +use thiserror_ext::AsReport as _; use tokio::sync::Semaphore; use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend}; @@ -72,12 +73,12 @@ pub fn pre_resolve_dns( v4.extend(v6); if !v4.is_empty() { let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect(); - println!("Pre-resolved {host} -> {ips:?}"); + tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS"); builder = builder.resolve_to_addrs(host, &v4); } } Err(e) => { - eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}"); + tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS"); } } @@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let (model_id, model_name) = if let Some(ref m) = config.model { (m.clone(), config.model_name.clone()) } else { - println!("Model not specified, fetching first model from server..."); + tracing::info!(base_url = %config.base_url, "fetching first model from server"); let (name, id) = get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?; - println!("First model name: {name}, first model id: {id}"); + tracing::info!( + model_name = name, + model_id = id, + "selected first model from server" + ); (id, Some(name)) }; @@ -358,10 +363,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { None } else { let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); - println!("Loading tokenizer: {tid}"); + tracing::info!(tokenizer = tid, "loading tokenizer"); let server_info = Some((config.base_url.as_str(), model_id.as_str())); - let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?; - println!("Tokenizer loaded successfully."); + let t = + crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?; Some(t) }; let has_tokenizer = tokenizer.is_some(); @@ -421,7 +426,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.num_prompts, config.random_batch_size, config.is_reranker, ), }; - println!("Generating {dataset_label}..."); + tracing::info!( + dataset = ?config.dataset_name, + prompts = config.num_prompts, + description = %dataset_label, + "generating benchmark dataset" + ); let gen_start = Instant::now(); let mut input_requests = match config.dataset_name { @@ -472,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let path = match config.dataset_path.as_deref() { Some(p) => p, None => { - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?; downloaded.as_str() } }; @@ -512,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { None => { downloaded = crate::datasets::speed_bench::download_speed_bench( config.speed_bench_config, - )?; + ) + .await?; downloaded.as_str() } }; @@ -543,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.hf_subset.as_deref(), config.hf_split.as_deref(), config.num_prompts, - )?; + ) + .await?; crate::datasets::hf_dataset::load_hf_dataset( tok, &downloaded_path, @@ -608,18 +620,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { }; let gen_elapsed = gen_start.elapsed(); - println!( - "Generated {} prompts in {:.2}s", - input_requests.len(), - gen_elapsed.as_secs_f64() + tracing::info!( + prompts = input_requests.len(), + elapsed_seconds = gen_elapsed.as_secs_f64(), + "generated benchmark dataset" ); let filtered_count = filter_requests_by_max_model_len(&mut input_requests, config.max_model_len); if filtered_count > 0 { - println!( - "Filtered {filtered_count} prompt(s) above --max-model-len {}.", - config.max_model_len.unwrap() + tracing::info!( + filtered_prompts = filtered_count, + max_model_len = config.max_model_len.unwrap(), + "filtered prompts above maximum model length" ); } if input_requests.is_empty() { @@ -670,7 +683,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { // Ready check if config.ready_check_timeout_sec > 0 { - println!("Starting initial single prompt test run..."); + tracing::info!("starting initial single-prompt test run"); let test_output = wait_for_endpoint( config.backend, &client, @@ -685,7 +698,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { test_output.error ))); } - println!("Initial test run completed."); + tracing::info!("initial single-prompt test run completed"); } // Verify and fix prompt token lengths against the server's /tokenize endpoint. @@ -703,12 +716,15 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { DatasetName::Random | DatasetName::PrefixRepetition ); if verifiable_dataset && has_token_ids && !config.backend.is_pooling() { - println!("Using prompt_token_ids, skipping server-side tokenizer verification."); + tracing::info!( + reason = "prompt_token_ids", + "skipping server tokenizer verification" + ); } if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() { let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id); if is_tokenizer_verified(&cache_key) { - println!("Tokenizer verified in previous run (cached), skipping verification."); + tracing::info!(reason = "cached", "skipping server tokenizer verification"); } else { let num_special = tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0); @@ -723,14 +739,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { .await? { SampleVerifyOutcome::Passed => { - println!("Sample verification passed, skipping full verification."); + tracing::info!("tokenizer sample verification passed"); mark_tokenizer_verified(&cache_key); } SampleVerifyOutcome::Skipped(reason) => { - println!("Server /tokenize unavailable ({reason}), skipping verification."); + tracing::warn!( + reason = %reason, + "server tokenizer unavailable; skipping prompt verification" + ); } SampleVerifyOutcome::Mismatch => { - println!("Sample verification found mismatch, running full verify+fix..."); + tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts"); match verify_and_fix_prompt_lengths( &client, &config.base_url, @@ -742,16 +761,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { .await { Ok(()) => { - println!( - "All {} prompts verified: exact token length match.", - input_requests.len() + tracing::info!( + prompts = input_requests.len(), + "verified exact prompt token lengths" ); mark_tokenizer_verified(&cache_key); } Err(BenchError::TokenizeUnavailable(reason)) => { - println!( - "Server /tokenize became unavailable during verification \ - ({reason}); proceeding with client-side token counts." + tracing::warn!( + reason = %reason, + "server tokenizer became unavailable; using client token counts" ); } Err(e) => return Err(e), @@ -763,7 +782,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { // Warmup if config.num_warmups > 0 { - println!("Warming up with {} requests...", config.num_warmups); + tracing::info!(requests = config.num_warmups, "starting benchmark warmup"); run_warmup( config.backend, &client, @@ -776,7 +795,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.disable_tqdm, ) .await; - println!("Warmup run completed."); + tracing::info!(requests = config.num_warmups, "benchmark warmup completed"); } // Start profiler if requested (immediate mode — no batch threshold) @@ -814,28 +833,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let spec_decode_before = fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; if spec_decode_before.is_some() { - println!("Speculative decoding detected, will collect metrics."); + tracing::info!("detected speculative decoding; collecting metrics"); } // Main benchmark - println!("Starting main benchmark run..."); let distribution = if config.burstiness == 1.0 { "Poisson process" } else { "Gamma distribution" }; - println!( - "Traffic request rate: {}", - if config.request_rate.is_infinite() { - "inf".to_string() - } else { - format!("{}", config.request_rate) - } - ); - println!("Burstiness factor: {} ({distribution})", config.burstiness); - println!( - "Maximum request concurrency: {}", - config.max_concurrency.unwrap_or(config.num_prompts) + tracing::info!( + request_rate = config.request_rate, + burstiness = config.burstiness, + distribution, + max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), + prompts = config.num_prompts, + "starting main benchmark run" ); // Pre-assign LoRA adapters to each request (None when --lora-modules not set). @@ -847,11 +860,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { ); if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) { let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect(); - println!( - "LoRA adapters ({}): {:?} [assignment={:?}]", - modules.len(), - names, - config.lora_assignment + tracing::info!( + adapters = modules.len(), + names = ?names, + assignment = ?config.lora_assignment, + "assigned LoRA adapters" ); } @@ -1125,7 +1138,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { if let Some((cancel_tx, task)) = profile_task { let _ = cancel_tx.send(()); if let Err(e) = task.await { - eprintln!("WARNING: Profile background task failed: {e}"); + tracing::error!(error = %e.as_report(), "profiler background task failed"); } } @@ -1289,12 +1302,14 @@ pub(crate) async fn start_profiler_immediate( base_url: &str, extra_headers: &Option>, ) { - println!("Starting profiler..."); let profile_url = format!("{base_url}/start_profile"); + tracing::info!(url = %profile_url, "starting profiler"); match send_profile_request(client, &profile_url, extra_headers).await { - Ok(true) => println!("Profiler started"), - Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"), - Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"), + Ok(true) => tracing::info!(url = %profile_url, "profiler started"), + Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"), + Err(e) => { + tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler") + } } } @@ -1304,12 +1319,14 @@ pub(crate) async fn stop_profiler_immediate( base_url: &str, extra_headers: &Option>, ) { - println!("Stopping profiler..."); let profile_url = format!("{base_url}/stop_profile"); + tracing::info!(url = %profile_url, "stopping profiler"); match send_profile_request(client, &profile_url, extra_headers).await { - Ok(true) => println!("Profiler stopped"), - Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), - Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), + Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"), + Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"), + Err(e) => { + tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler") + } } } @@ -1371,25 +1388,30 @@ pub(crate) async fn profile_on_batch_threshold( duration_secs: f64, mut cancel_rx: tokio::sync::oneshot::Receiver<()>, ) { - println!( - "Waiting for batch size >= {threshold} before starting profiler \ - (will capture {duration_secs}s)..." + tracing::info!( + threshold, + duration_seconds = duration_secs, + "waiting for profiler batch threshold" ); loop { if let Some(running) = fetch_num_requests_running(client, base_url).await && running >= threshold { - println!("Batch size {running} >= {threshold}, starting profiler..."); + tracing::info!( + running_requests = running, + threshold, + "profiler batch threshold reached" + ); break; } // Wait 500ms or until the benchmark signals cancellation tokio::select! { _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} _ = &mut cancel_rx => { - eprintln!( - "NOTE: Benchmark finished before batch threshold {threshold} was reached; \ - profiling skipped." + tracing::warn!( + threshold, + "benchmark finished before profiler batch threshold; skipping profiling" ); return; } @@ -1398,13 +1420,13 @@ pub(crate) async fn profile_on_batch_threshold( let start_url = format!("{base_url}/start_profile"); match send_profile_request(client, &start_url, extra_headers).await { - Ok(true) => println!("Profiler started"), + Ok(true) => tracing::info!(url = %start_url, "profiler started"), Ok(false) => { - eprintln!("WARNING: Profiler start request returned non-success"); + tracing::warn!(url = %start_url, "profiler start request was unsuccessful"); return; } Err(e) => { - eprintln!("WARNING: Failed to start profiler: {e}"); + tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler"); return; } } @@ -1413,15 +1435,17 @@ pub(crate) async fn profile_on_batch_threshold( tokio::select! { _ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {} _ = &mut cancel_rx => { - println!("Benchmark finished, stopping profiler early..."); + tracing::info!("benchmark finished; stopping profiler early"); } } let stop_url = format!("{base_url}/stop_profile"); match send_profile_request(client, &stop_url, extra_headers).await { - Ok(true) => println!("Profiler stopped after capturing"), - Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), - Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), + Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"), + Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"), + Err(e) => { + tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler") + } } } @@ -1502,10 +1526,11 @@ async fn verify_and_fix_prompt_lengths( let excess = tokens.len().saturating_sub(expected_input_len); let compensate = if excess > 0 && last_excess == Some(excess) { if _iter == 1 { - eprintln!( - "Prompt {i}: server consistently adds {excess} extra token(s) \ - (likely BOS), compensating target to {}.", - expected_input_len.saturating_sub(excess), + tracing::warn!( + prompt_index = i, + extra_tokens = excess, + adjusted_target = expected_input_len.saturating_sub(excess), + "server consistently adds prompt tokens; compensating verification target" ); } excess @@ -1563,7 +1588,10 @@ async fn verify_and_fix_prompt_lengths( let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed); if fc > 0 { - println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence."); + tracing::info!( + fixed_prompts = fc, + "fixed prompt lengths using server tokenizer" + ); } Ok(()) @@ -1818,7 +1846,7 @@ async fn sample_verify_prompts( let tokenize_url = format!("{base_url}/tokenize"); let api_key = std::env::var("OPENAI_API_KEY").ok(); - println!("Sampling {sample_size} prompts for verification..."); + tracing::info!(sample_size, "sampling prompts for tokenizer verification"); for (i, request) in requests.iter().enumerate().take(sample_size) { let tokens = match server_tokenize( @@ -1841,9 +1869,11 @@ async fn sample_verify_prompts( let expected = request.prompt_len + num_special; if tokens.len() != expected { - println!( - "Prompt {i}: expected {expected} tokens, server returned {}", - tokens.len() + tracing::warn!( + prompt_index = i, + expected_tokens = expected, + actual_tokens = tokens.len(), + "tokenizer verification sample mismatch" ); return Ok(SampleVerifyOutcome::Mismatch); } diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 9aa8f388633b..06e32e14ce3b 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -3,8 +3,6 @@ use std::fmt; -use clap::Parser; - /// Backend type for the benchmark endpoint. #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] pub enum BackendKind { @@ -77,7 +75,7 @@ pub enum DatasetName { ShareGpt, #[value(name = "sonnet")] Sonnet, - #[value(name = "speed-bench")] + #[value(name = "speed-bench", alias = "speed_bench")] SpeedBench, #[value(name = "hf")] Hf, @@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig { } /// High-performance benchmark client for vLLM serving endpoints. -#[derive(Parser, Debug, Clone)] -#[command( - name = "vllm-bench", - about = "Benchmark online serving throughput", - version -)] -pub struct Cli { +#[derive(clap::Args, Debug, Clone)] +pub struct BenchServeArgs { /// The type of backend or endpoint to use for the benchmark. #[arg(long, default_value = "openai")] pub backend: BackendKind, @@ -659,7 +652,7 @@ pub struct Cli { pub lora_assignment: LoraAssignment, } -impl Cli { +impl BenchServeArgs { /// Resolve the base URL from explicit --base-url or from --host/--port. pub fn resolve_base_url(&self) -> String { if let Some(ref base) = self.base_url { diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 9696612c30fd..0700cf1fad3e 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -4,7 +4,9 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig}; +use crate::cli::{ + BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig, +}; use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt}; use crate::error::{BenchError, Result}; @@ -215,63 +217,63 @@ pub struct BenchConfig { } impl BenchConfig { - pub fn from_cli(cli: &Cli) -> Result { - if cli.burstiness <= 0.0 { + pub fn from_args(args: &BenchServeArgs) -> Result { + if args.burstiness <= 0.0 { return Err(BenchError::Config("Burstiness must be positive".into())); } - if cli.num_prompts == 0 { + if args.num_prompts == 0 { return Err(BenchError::Config( "--num-prompts must be at least 1".into(), )); } - if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() { + if args.request_rate <= 0.0 && !args.request_rate.is_infinite() { return Err(BenchError::Config( "--request-rate must be positive (or inf)".into(), )); } - if cli.max_model_len == Some(0) { + if args.max_model_len == Some(0) { return Err(BenchError::Config( "--max-model-len must be at least 1".into(), )); } - let base_url = cli.resolve_base_url(); - let api_url = cli.resolve_api_url(); + let base_url = args.resolve_base_url(); + let api_url = args.resolve_api_url(); - let extra_headers = cli.parse_headers()?; - let mut extra_body = cli.parse_extra_body()?; + let extra_headers = args.parse_headers()?; + let mut extra_body = args.parse_extra_body()?; // Merge sampling parameters into extra_body (matches Python behavior). // Python collects non-None sampling params and merges them UNDER extra_body, // meaning extra_body keys take precedence over sampling params. { let mut sampling_params = serde_json::Map::new(); - if let Some(v) = cli.top_p { + if let Some(v) = args.top_p { sampling_params.insert("top_p".into(), serde_json::json!(v)); } - if let Some(v) = cli.top_k { + if let Some(v) = args.top_k { sampling_params.insert("top_k".into(), serde_json::json!(v)); } - if let Some(v) = cli.min_p { + if let Some(v) = args.min_p { sampling_params.insert("min_p".into(), serde_json::json!(v)); } - if let Some(v) = cli.temperature { + if let Some(v) = args.temperature { sampling_params.insert("temperature".into(), serde_json::json!(v)); } - if let Some(v) = cli.frequency_penalty { + if let Some(v) = args.frequency_penalty { sampling_params.insert("frequency_penalty".into(), serde_json::json!(v)); } - if let Some(v) = cli.presence_penalty { + if let Some(v) = args.presence_penalty { sampling_params.insert("presence_penalty".into(), serde_json::json!(v)); } - if let Some(v) = cli.repetition_penalty { + if let Some(v) = args.repetition_penalty { sampling_params.insert("repetition_penalty".into(), serde_json::json!(v)); } if !sampling_params.is_empty() { - if !cli.backend.is_openai_compatible() { + if !args.backend.is_openai_compatible() { return Err(BenchError::Config( "Sampling parameters are only supported by openai-compatible backends." .into(), @@ -286,10 +288,18 @@ impl BenchConfig { } Some(other) => { // extra_body was not an object — just use sampling params - eprintln!( - "Warning: --extra-body is not a JSON object, sampling params may be lost" + let value_type = match &other { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "boolean", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => unreachable!(), + }; + tracing::warn!( + value_type, + "sampling parameters may be lost because --extra-body is not a JSON object" ); - let _ = other; sampling_params } None => sampling_params, @@ -299,7 +309,7 @@ impl BenchConfig { } // Parse metadata - let metadata = match &cli.metadata { + let metadata = match &args.metadata { None => None, Some(items) => { let mut pairs = Vec::new(); @@ -314,24 +324,24 @@ impl BenchConfig { }; // Parse goodput SLOs - let goodput = parse_goodput(&cli.goodput)?; + let goodput = parse_goodput(&args.goodput)?; // Parse ramp-up config - let ramp_up = parse_ramp_up(cli)?; + let ramp_up = parse_ramp_up(args)?; // Default percentile metrics based on backend type - let default_percentile_metrics = if cli.backend.is_pooling() { + let default_percentile_metrics = if args.backend.is_pooling() { "e2el" } else { "ttft,tpot,itl,e2el" }; let percentile_metrics_str = - cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics); + args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics); let selected_percentile_metrics: Vec = percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect(); - let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?; - let sweep_summary_percentiles = cli + let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?; + let sweep_summary_percentiles = args .sweep_summary_percentiles .as_deref() .map(|raw| parse_percentiles(raw, true)) @@ -344,38 +354,38 @@ impl BenchConfig { selected_percentiles.push(90.0); } - let tokenizer_id = if cli.skip_tokenizer_init { + let tokenizer_id = if args.skip_tokenizer_init { None } else { - Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default()) + args.tokenizer.clone().or_else(|| args.model.clone()) }; // Resolve input/output lengths - let random_input_len = cli.resolved_random_input_len(); - let random_output_len = cli.resolved_random_output_len(); - let per_turn_input_len = cli.resolved_per_turn_input_len(); + let random_input_len = args.resolved_random_input_len(); + let random_output_len = args.resolved_random_output_len(); + let per_turn_input_len = args.resolved_per_turn_input_len(); // Normalized multi-turn turn counts (computed in validation block below, defaults // to num_turns if multi-turn mode is not active) - let mut multi_turn_min_turns = cli.multi_turn_num_turns; - let mut multi_turn_max_turns = cli.multi_turn_num_turns; + let mut multi_turn_min_turns = args.multi_turn_num_turns; + let mut multi_turn_max_turns = args.multi_turn_num_turns; // For random datasets with openai-compatible backends, default to ignore_eos. // Exception: multi-turn mode, where ignore_eos causes unbounded context growth // across turns. Multi-turn uses min_tokens instead for output length control. // Pooling backends don't generate tokens, so ignore_eos is irrelevant. - let ignore_eos = if cli.backend.is_pooling() { + let ignore_eos = if args.backend.is_pooling() { false } else { - cli.ignore_eos - || ((cli.dataset_name == DatasetName::Random - || cli.dataset_name == DatasetName::RandomMm) - && cli.backend.is_openai_compatible() - && !cli.multi_turn) + args.ignore_eos + || ((args.dataset_name == DatasetName::Random + || args.dataset_name == DatasetName::RandomMm) + && args.backend.is_openai_compatible() + && !args.multi_turn) }; // Pooling backends don't support multi-turn - if cli.backend.is_pooling() && cli.multi_turn { + if args.backend.is_pooling() && args.multi_turn { return Err(BenchError::Config( "Pooling/embedding backends do not support --multi-turn".into(), )); @@ -383,7 +393,7 @@ impl BenchConfig { // LoRA validation. Adapter names must be non-empty after trim; pooling // backends are out of scope (vLLM LoRA routing is for generative paths). - let lora_modules = match cli.lora_modules.as_ref() { + let lora_modules = match args.lora_modules.as_ref() { None => None, Some(names) => { if names.is_empty() { @@ -391,7 +401,7 @@ impl BenchConfig { "--lora-modules requires at least one adapter name".into(), )); } - if cli.backend.is_pooling() { + if args.backend.is_pooling() { return Err(BenchError::Config( "--lora-modules is not supported for pooling/embedding backends".into(), )); @@ -411,18 +421,18 @@ impl BenchConfig { }; // Random-MM validation and config parsing - let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm { - if cli.backend != BackendKind::OpenaiChat { + let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm { + if args.backend != BackendKind::OpenaiChat { return Err(BenchError::Config( "Multi-modal content (images) is only supported on 'openai-chat' backend." .into(), )); } let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt( - &cli.random_mm_limit_mm_per_prompt, + &args.random_mm_limit_mm_per_prompt, )?; let buckets = - crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?; + crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?; (limit, buckets) } else { (MmLimitPerPrompt::default(), Vec::new()) @@ -432,18 +442,18 @@ impl BenchConfig { // sonnet (uses built-in Shakespeare's sonnets). // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1)) - let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?; + let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?; // Batched inputs only make sense for pooling backends (the generation // backends send one prompt per request). - if cli.random_batch_size == 0 { + if args.random_batch_size == 0 { return Err(BenchError::Config( "--random-batch-size must be at least 1".into(), )); } - if cli.random_batch_size > 1 - && !cli.backend.is_pooling() - && cli.dataset_name != DatasetName::RandomRerank + if args.random_batch_size > 1 + && !args.backend.is_pooling() + && args.dataset_name != DatasetName::RandomRerank { return Err(BenchError::Config( "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(), @@ -451,16 +461,16 @@ impl BenchConfig { } // random-rerank validation (mirrors Python RandomDatasetForReranking) - let is_reranker = !cli.no_reranker; - if cli.dataset_name == DatasetName::RandomRerank { - if !cli.backend.is_pooling() { + let is_reranker = !args.no_reranker; + if args.dataset_name == DatasetName::RandomRerank { + if !args.backend.is_pooling() { return Err(BenchError::Config( "--dataset-name random-rerank requires an embeddings/pooling backend \ (e.g. --backend vllm-rerank)" .into(), )); } - if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) { + if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) { return Err(BenchError::Config( "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \ (the query is folded into the first batch slot)" @@ -470,8 +480,8 @@ impl BenchConfig { } // Custom dataset validation - if cli.dataset_name == DatasetName::Custom { - match cli.dataset_path.as_deref() { + if args.dataset_name == DatasetName::Custom { + match args.dataset_path.as_deref() { None => { return Err(BenchError::Config( "--dataset-path is required for --dataset-name custom \ @@ -486,38 +496,38 @@ impl BenchConfig { } _ => {} } - if !cli.skip_chat_template { - eprintln!( - "NOTE: client-side chat template rendering is not supported; custom \ - dataset prompts are sent raw (equivalent to --skip-chat-template)." + if !args.skip_chat_template { + tracing::warn!( + dataset = "custom", + "client-side chat template rendering is unsupported; sending prompts raw" ); } } // Prefix repetition validation - if cli.dataset_name == DatasetName::PrefixRepetition { - if cli.prefix_repetition_num_prefixes == 0 { + if args.dataset_name == DatasetName::PrefixRepetition { + if args.prefix_repetition_num_prefixes == 0 { return Err(BenchError::Config( "--prefix-repetition-num-prefixes must be at least 1".into(), )); } - if cli.num_prompts < cli.prefix_repetition_num_prefixes { + if args.num_prompts < args.prefix_repetition_num_prefixes { return Err(BenchError::Config(format!( "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})", - cli.num_prompts, cli.prefix_repetition_num_prefixes + args.num_prompts, args.prefix_repetition_num_prefixes ))); } } // HF dataset validation - if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() { + if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() { return Err(BenchError::Config( "--dataset-path is required for --dataset-name hf \ (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')" .into(), )); } - if let Some(len) = cli.hf_output_len + if let Some(len) = args.hf_output_len && len == 0 { return Err(BenchError::Config( @@ -526,13 +536,13 @@ impl BenchConfig { } // Multi-turn validation - if cli.multi_turn { - if cli.backend != BackendKind::OpenaiChat { + if args.multi_turn { + if args.backend != BackendKind::OpenaiChat { return Err(BenchError::Config( "--multi-turn requires --backend openai-chat".into(), )); } - if cli.multi_turn_num_turns == 0 { + if args.multi_turn_num_turns == 0 { return Err(BenchError::Config( "--multi-turn-num-turns must be at least 1".into(), )); @@ -541,18 +551,18 @@ impl BenchConfig { // Normalize and validate min/max turns. ShareGPT only consumes max_turns // (the loader walks all available turns up to the cap), so the // min/num/max coupling used for synthetic generation does not apply. - if cli.dataset_name == DatasetName::ShareGpt { - if cli.multi_turn_max_turns == 1 { + if args.dataset_name == DatasetName::ShareGpt { + if args.multi_turn_max_turns == 1 { return Err(BenchError::Config( "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), )); } } else { (multi_turn_min_turns, multi_turn_max_turns) = - match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) { - (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns), - (m, 0) => (m, cli.multi_turn_num_turns), - (0, x) => (cli.multi_turn_num_turns, x), + match (args.multi_turn_min_turns, args.multi_turn_max_turns) { + (0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns), + (m, 0) => (m, args.multi_turn_num_turns), + (0, x) => (args.multi_turn_num_turns, x), (m, x) => (m, x), }; if multi_turn_min_turns < 1 { @@ -568,15 +578,16 @@ impl BenchConfig { } if ignore_eos { - eprintln!( - "WARNING: --ignore-eos is set with --multi-turn. The server may not \ - respect output length limits, causing unbounded context growth." + tracing::warn!( + ignore_eos, + multi_turn = true, + "output length limits may be ignored, causing unbounded context growth" ); } // Validate prefix sharing ratios - let pg = cli.multi_turn_prefix_global_ratio; - let pc = cli.multi_turn_prefix_conversation_ratio; + let pg = args.multi_turn_prefix_global_ratio; + let pc = args.multi_turn_prefix_conversation_ratio; if !(0.0..=1.0).contains(&pg) { return Err(BenchError::Config( "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(), @@ -592,20 +603,20 @@ impl BenchConfig { "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(), )); } - if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random { + if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random { return Err(BenchError::Config( "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(), )); } } - if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) { + if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) { return Err(BenchError::Config(format!( "--steady-state-threshold must be in (0.0, 1.0], got {}", - cli.steady_state_threshold + args.steady_state_threshold ))); } - if let Some(mw) = cli.steady_state_min_window + if let Some(mw) = args.steady_state_min_window && mw < 0.0 { return Err(BenchError::Config(format!( @@ -613,122 +624,122 @@ impl BenchConfig { ))); } - if cli.profile_batch_threshold.is_some() && !cli.profile { + if args.profile_batch_threshold.is_some() && !args.profile { return Err(BenchError::Config( "--profile-batch-threshold requires --profile".into(), )); } - if cli.profile_duration <= 0.0 { + if args.profile_duration <= 0.0 { return Err(BenchError::Config( "--profile-duration must be positive".into(), )); } - if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 { + if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 { return Err(BenchError::Config( "--profile-duration requires --profile-batch-threshold".into(), )); } Ok(BenchConfig { - backend: cli.backend, + backend: args.backend, base_url, api_url, - model: cli.model.clone(), - model_name: cli.served_model_name.clone(), + model: args.model.clone(), + model_name: args.served_model_name.clone(), tokenizer_id, - tokenizer_mode: cli.tokenizer_mode.clone(), - trust_remote_code: cli.trust_remote_code, - skip_tokenizer_init: cli.skip_tokenizer_init, - dataset_name: cli.dataset_name, - dataset_path: cli.dataset_path.clone(), - max_model_len: cli.max_model_len, + tokenizer_mode: args.tokenizer_mode.clone(), + trust_remote_code: args.trust_remote_code, + skip_tokenizer_init: args.skip_tokenizer_init, + dataset_name: args.dataset_name, + dataset_path: args.dataset_path.clone(), + max_model_len: args.max_model_len, random_input_len, random_output_len, - random_prefix_len: cli.random_prefix_len, + random_prefix_len: args.random_prefix_len, random_range_ratio, - random_batch_size: cli.random_batch_size, + random_batch_size: args.random_batch_size, is_reranker, - custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len), - prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len, - prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len, - prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes, - prefix_repetition_output_len: cli + custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len), + prefix_repetition_prefix_len: args.prefix_repetition_prefix_len, + prefix_repetition_suffix_len: args.prefix_repetition_suffix_len, + prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes, + prefix_repetition_output_len: args .output_len - .unwrap_or(cli.prefix_repetition_output_len), - random_cache_hit_fraction: cli.random_cache_hit_fraction, - random_cache_ratio: cli.random_cache_ratio, - sharegpt_output_len: cli.sharegpt_output_len, - sonnet_input_len: cli.sonnet_input_len, - sonnet_output_len: cli.sonnet_output_len, - sonnet_prefix_len: cli.sonnet_prefix_len, - no_oversample: cli.no_oversample, - disable_shuffle: cli.disable_shuffle, - num_prompts: cli.num_prompts, - request_rate: cli.request_rate, - burstiness: cli.burstiness, - max_concurrency: cli.max_concurrency, - steady_state_threshold: cli.steady_state_threshold, - steady_state_min_window: cli.steady_state_min_window, - no_steady_state: cli.no_steady_state, - disable_tqdm: cli.disable_tqdm, - num_warmups: cli.num_warmups, - profile: cli.profile, - profile_batch_threshold: cli.profile_batch_threshold, - profile_duration: cli.profile_duration, - save_result: cli.save_result, - save_detailed: cli.save_detailed, - append_result: cli.append_result, - result_dir: cli.result_dir.clone(), - result_filename: cli.result_filename.clone(), - seed: cli.seed, + .unwrap_or(args.prefix_repetition_output_len), + random_cache_hit_fraction: args.random_cache_hit_fraction, + random_cache_ratio: args.random_cache_ratio, + sharegpt_output_len: args.sharegpt_output_len, + sonnet_input_len: args.sonnet_input_len, + sonnet_output_len: args.sonnet_output_len, + sonnet_prefix_len: args.sonnet_prefix_len, + no_oversample: args.no_oversample, + disable_shuffle: args.disable_shuffle, + num_prompts: args.num_prompts, + request_rate: args.request_rate, + burstiness: args.burstiness, + max_concurrency: args.max_concurrency, + steady_state_threshold: args.steady_state_threshold, + steady_state_min_window: args.steady_state_min_window, + no_steady_state: args.no_steady_state, + disable_tqdm: args.disable_tqdm, + num_warmups: args.num_warmups, + profile: args.profile, + profile_batch_threshold: args.profile_batch_threshold, + profile_duration: args.profile_duration, + save_result: args.save_result, + save_detailed: args.save_detailed, + append_result: args.append_result, + result_dir: args.result_dir.clone(), + result_filename: args.result_filename.clone(), + seed: args.seed, ignore_eos, - insecure: cli.insecure, + insecure: args.insecure, selected_percentile_metrics, selected_percentiles, sweep_summary_percentiles, - label: cli.label.clone(), - logprobs: cli.logprobs, - request_id_prefix: cli.get_request_id_prefix(), - ready_check_timeout_sec: cli.ready_check_timeout_sec, + label: args.label.clone(), + logprobs: args.logprobs, + request_id_prefix: args.get_request_id_prefix(), + ready_check_timeout_sec: args.ready_check_timeout_sec, extra_headers, extra_body, metadata, - dry_run: cli.dry_run, + dry_run: args.dry_run, goodput, ramp_up, - multi_turn: cli.multi_turn, - multi_turn_num_turns: cli.multi_turn_num_turns, + multi_turn: args.multi_turn, + multi_turn_num_turns: args.multi_turn_num_turns, multi_turn_min_turns, multi_turn_max_turns, - sharegpt_multi_turn_max_turns: if cli.multi_turn - && cli.dataset_name == DatasetName::ShareGpt - && cli.multi_turn_max_turns != 0 + sharegpt_multi_turn_max_turns: if args.multi_turn + && args.dataset_name == DatasetName::ShareGpt + && args.multi_turn_max_turns != 0 { - Some(cli.multi_turn_max_turns) + Some(args.multi_turn_max_turns) } else { None }, per_turn_input_len, - multi_turn_concurrency: cli.multi_turn_concurrency, - multi_turn_delay_ms: cli.multi_turn_delay_ms, - multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio, - multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio, - speed_bench_config: cli.speed_bench_config, - speed_bench_category: cli.speed_bench_category.clone(), - speed_bench_max_input_len: cli.speed_bench_max_input_len, - hf_split: cli.hf_split.clone(), - hf_subset: cli.hf_subset.clone(), - hf_output_len: cli.hf_output_len, - hf_text_column: cli.hf_text_column.clone(), - reset_prefix_cache: cli.reset_prefix_cache, - prompt_token_ids: cli.prompt_token_ids, - random_mm_base_items_per_request: cli.random_mm_base_items_per_request, - random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio, + multi_turn_concurrency: args.multi_turn_concurrency, + multi_turn_delay_ms: args.multi_turn_delay_ms, + multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio, + multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio, + speed_bench_config: args.speed_bench_config, + speed_bench_category: args.speed_bench_category.clone(), + speed_bench_max_input_len: args.speed_bench_max_input_len, + hf_split: args.hf_split.clone(), + hf_subset: args.hf_subset.clone(), + hf_output_len: args.hf_output_len, + hf_text_column: args.hf_text_column.clone(), + reset_prefix_cache: args.reset_prefix_cache, + prompt_token_ids: args.prompt_token_ids, + random_mm_base_items_per_request: args.random_mm_base_items_per_request, + random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio, random_mm_limit, random_mm_buckets, - enable_multimodal_chat: cli.enable_multimodal_chat, + enable_multimodal_chat: args.enable_multimodal_chat, lora_modules, - lora_assignment: cli.lora_assignment, + lora_assignment: args.lora_assignment, }) } } @@ -811,17 +822,17 @@ fn parse_goodput(goodput_args: &Option>) -> Result { Ok(config) } -fn parse_ramp_up(cli: &Cli) -> Result> { - let strategy = match cli.ramp_up_strategy { +fn parse_ramp_up(args: &BenchServeArgs) -> Result> { + let strategy = match args.ramp_up_strategy { None => return Ok(None), Some(s) => s, }; - let start_rps = cli.ramp_up_start_rps.ok_or_else(|| { + let start_rps = args.ramp_up_start_rps.ok_or_else(|| { BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into()) })?; - let end_rps = cli.ramp_up_end_rps.ok_or_else(|| { + let end_rps = args.ramp_up_end_rps.ok_or_else(|| { BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into()) })?; @@ -843,7 +854,21 @@ mod tests { use clap::Parser; use super::*; - use crate::cli::Cli; + use crate::cli::BenchServeArgs; + + #[derive(Parser)] + struct TestCli { + #[command(flatten)] + args: BenchServeArgs, + } + + fn parse_args(args: I) -> BenchServeArgs + where + I: IntoIterator, + T: Into + Clone, + { + TestCli::parse_from(args).args + } fn base_multi_turn_args() -> Vec<&'static str> { vec![ @@ -859,8 +884,8 @@ mod tests { #[test] fn test_prefix_sharing_defaults_to_zero() { let args = base_multi_turn_args(); - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.multi_turn_prefix_global_ratio, 0.0); assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0); } @@ -874,8 +899,8 @@ mod tests { "--multi-turn-prefix-conversation-ratio", "0.8", ]); - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10); assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10); } @@ -889,8 +914,8 @@ mod tests { "--multi-turn-prefix-conversation-ratio", "0.6", ]); - let cli = Cli::parse_from(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + let args = parse_args(args); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] @@ -902,16 +927,16 @@ mod tests { "--multi-turn-prefix-conversation-ratio", "0.5", ]); - let cli = Cli::parse_from(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + let args = parse_args(args); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] fn test_prefix_sharing_out_of_range_fails() { let mut args = base_multi_turn_args(); args.extend(["--multi-turn-prefix-global-ratio", "1.5"]); - let cli = Cli::parse_from(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + let args = parse_args(args); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] @@ -928,8 +953,8 @@ mod tests { "--multi-turn-prefix-global-ratio", "0.1", ]; - let cli = Cli::parse_from(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + let args = parse_args(args); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] @@ -944,8 +969,8 @@ mod tests { "--dataset-name", "sharegpt", ]; - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.multi_turn_max_turns, 3); assert_eq!(config.sharegpt_multi_turn_max_turns, None); @@ -968,8 +993,8 @@ mod tests { "--multi-turn-max-turns", "2", ]; - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); } @@ -987,8 +1012,8 @@ mod tests { "--multi-turn-max-turns", "1", ]; - let cli = Cli::parse_from(args); - let err = BenchConfig::from_cli(&cli).unwrap_err().to_string(); + let args = parse_args(args); + let err = BenchConfig::from_args(&args).unwrap_err().to_string(); assert!( err.contains("at least 2 for ShareGPT"), "expected ShareGPT-specific error, got: {err}" @@ -1009,8 +1034,8 @@ mod tests { "--multi-turn-max-turns", "20", ]; - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); } @@ -1018,8 +1043,8 @@ mod tests { #[test] fn test_sweep_summary_percentiles_default_empty() { let args = base_multi_turn_args(); - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert!(config.sweep_summary_percentiles.is_empty()); assert_eq!(config.selected_percentiles, vec![99.0, 90.0]); @@ -1034,8 +1059,8 @@ mod tests { "--sweep-summary-percentiles", "90,95,90", ]); - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]); assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]); @@ -1045,8 +1070,8 @@ mod tests { fn test_invalid_sweep_summary_percentile_fails() { let mut args = base_multi_turn_args(); args.extend(["--sweep-summary-percentiles", "101"]); - let cli = Cli::parse_from(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + let args = parse_args(args); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] @@ -1058,12 +1083,20 @@ mod tests { "--max-model-len", "4096", ]; - let cli = Cli::parse_from(args); - let config = BenchConfig::from_cli(&cli).unwrap(); + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); assert_eq!(config.max_model_len, Some(4096)); } + #[test] + fn test_tokenizer_id_deferred_when_model_is_unspecified() { + let args = parse_args(["vllm-bench"]); + let config = BenchConfig::from_args(&args).unwrap(); + + assert_eq!(config.tokenizer_id, None); + } + #[test] fn test_zero_max_model_len_fails() { let args = vec![ @@ -1073,9 +1106,9 @@ mod tests { "--max-model-len", "0", ]; - let cli = Cli::parse_from(args); + let args = parse_args(args); - assert!(BenchConfig::from_cli(&cli).is_err()); + assert!(BenchConfig::from_args(&args).is_err()); } #[test] fn test_range_ratio_parse_float() { diff --git a/rust/src/bench/src/datasets/custom.rs b/rust/src/bench/src/datasets/custom.rs index 04bbfcfdbac2..11ddee02b17d 100644 --- a/rust/src/bench/src/datasets/custom.rs +++ b/rust/src/bench/src/datasets/custom.rs @@ -128,8 +128,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } #[test] diff --git a/rust/src/bench/src/datasets/hf_dataset.rs b/rust/src/bench/src/datasets/hf_dataset.rs index 35932a39833f..93c9f4cdbd2a 100644 --- a/rust/src/bench/src/datasets/hf_dataset.rs +++ b/rust/src/bench/src/datasets/hf_dataset.rs @@ -8,6 +8,7 @@ use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use super::SampleRequest; +use super::progress::RowDownloadReporter; use crate::error::{BenchError, Result}; use crate::tokenizer::TokenizerKind; @@ -50,18 +51,19 @@ enum ColumnFormat { /// Make a GET request with retry logic (3 retries with exponential backoff). /// Returns the parsed JSON response. -fn get_with_retry( - client: &reqwest::blocking::Client, +async fn get_with_retry( + client: &reqwest::Client, url: &str, label: &str, ) -> Result { let max_retries = 3; for attempt in 0..=max_retries { - let resp = match client.get(url).send() { + let resp = match client.get(url).send().await { Ok(r) => r, Err(e) => { if attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))) + .await; continue; } return Err(BenchError::Config(format!( @@ -80,7 +82,7 @@ fn get_with_retry( } if status.is_server_error() && attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await; continue; } @@ -92,6 +94,7 @@ fn get_with_retry( let data: serde_json::Value = resp .json() + .await .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?; return Ok(data); } @@ -105,7 +108,7 @@ fn get_with_retry( /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization. /// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected /// or the dataset is exhausted. -pub fn download_hf_dataset( +pub async fn download_hf_dataset( dataset: &str, subset: Option<&str>, split: Option<&str>, @@ -115,7 +118,7 @@ pub fn download_hf_dataset( url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); let mut client_builder = - reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120)); + reqwest::Client::builder().timeout(std::time::Duration::from_secs(120)); // Add HF_TOKEN auth header if available if let Ok(token) = std::env::var("HF_TOKEN") { @@ -138,7 +141,7 @@ pub fn download_hf_dataset( // Call /info to discover available configs and splits let info_url = format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}"); - let info = get_with_retry(&client, &info_url, "HF dataset /info")?; + let info = get_with_retry(&client, &info_url, "HF dataset /info").await?; let dataset_info = info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| { @@ -201,7 +204,12 @@ pub fn download_hf_dataset( (resolved_config, resolved_split) }; - println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})"); + tracing::info!( + dataset, + config = resolved_config, + split = resolved_split, + "resolved Hugging Face dataset" + ); // Check cache let dir = cache_dir(); @@ -215,11 +223,16 @@ pub fn download_hf_dataset( if cache_path.exists() { let path_str = cache_path.to_string_lossy().to_string(); - println!("HF dataset cached: {path_str}"); + tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset"); return Ok((path_str, resolved_config, resolved_split)); } - println!("Downloading HF dataset '{dataset}' from datasets-server..."); + tracing::info!( + dataset, + config = resolved_config, + split = resolved_split, + "downloading Hugging Face dataset" + ); let encoded_config: String = url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect(); @@ -229,6 +242,7 @@ pub fn download_hf_dataset( let mut all_rows: Vec = Vec::new(); let mut offset = 0usize; let page_size = 100usize; + let mut progress = RowDownloadReporter::new(); loop { let url = format!( @@ -240,7 +254,7 @@ pub fn download_hf_dataset( &length={page_size}" ); - let data = get_with_retry(&client, &url, "HF dataset /rows")?; + let data = get_with_retry(&client, &url, "HF dataset /rows").await?; let rows = data["rows"] .as_array() @@ -260,14 +274,14 @@ pub fn download_hf_dataset( offset += fetched; let total = data["num_rows_total"].as_u64().unwrap_or(0); - eprint!("\r Fetched {offset}/{total} rows..."); + progress.update(offset, total); // Stop if we have enough rows or reached end of dataset if all_rows.len() >= num_rows_needed || fetched < page_size { break; } } - eprintln!(); // newline after progress + progress.finish(); if all_rows.is_empty() { return Err(BenchError::Config(format!( @@ -280,7 +294,12 @@ pub fn download_hf_dataset( std::fs::write(&cache_path, &json_str)?; let path_str = cache_path.to_string_lossy().to_string(); - println!("HF dataset: {} rows saved to {path_str}", all_rows.len()); + tracing::info!( + dataset, + rows = all_rows.len(), + path = %path_str, + "saved Hugging Face dataset" + ); Ok((path_str, resolved_config, resolved_split)) } @@ -483,21 +502,31 @@ pub fn load_hf_dataset( // Detect column format from first row let format = detect_column_format(&entries[0], text_column_override)?; - // Print detected format match &format { - ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"), + ColumnFormat::Chat(col) => { + tracing::info!( + format = "chat", + column = col, + "detected Hugging Face dataset format" + ); + } ColumnFormat::Text { prompt_col, output_col, } => { - let out_msg = output_col.as_deref().unwrap_or("none"); - println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}"); + tracing::info!( + format = "text", + prompt_column = prompt_col, + output_column = output_col.as_deref().unwrap_or("none"), + "detected Hugging Face dataset format" + ); } ColumnFormat::Combined { cols, output_col } => { - let out_msg = output_col.as_deref().unwrap_or("none"); - println!( - "HF dataset: detected combined columns {:?}, output column: {out_msg}", - cols + tracing::info!( + format = "combined", + prompt_columns = ?cols, + output_column = output_col.as_deref().unwrap_or("none"), + "detected Hugging Face dataset format" ); } } @@ -608,9 +637,10 @@ pub fn load_hf_dataset( if len == 0 { 128 } else { len } } else { if !warned_no_output { - eprintln!( - "WARNING: No output column detected and --hf-output-len not set. \ - Using default output length of 128 tokens." + tracing::warn!( + path = dataset_path, + default_output_tokens = 128, + "no dataset output column or --hf-output-len; using default output length" ); warned_no_output = true; } @@ -618,9 +648,10 @@ pub fn load_hf_dataset( } } else { if !warned_no_output { - eprintln!( - "WARNING: No output column detected and --hf-output-len not set. \ - Using default output length of 128 tokens." + tracing::warn!( + path = dataset_path, + default_output_tokens = 128, + "no dataset output column or --hf-output-len; using default output length" ); warned_no_output = true; } @@ -640,9 +671,11 @@ pub fn load_hf_dataset( // Oversample if needed if samples.len() < num_requests { if no_oversample { - println!( - "Skipping oversampling. Total samples: {} (requested: {num_requests})", - samples.len() + tracing::info!( + dataset = "hf", + samples = samples.len(), + requested = num_requests, + "skipping dataset oversampling" ); } else if !samples.is_empty() { let original_len = samples.len(); @@ -652,9 +685,11 @@ pub fn load_hf_dataset( req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); samples.push(req); } - println!( - "Oversampled HF dataset from {original_len} to {} total samples.", - samples.len() + tracing::info!( + dataset = "hf", + original_samples = original_len, + samples = samples.len(), + "oversampled dataset" ); } } @@ -1002,8 +1037,10 @@ mod tests { /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required). fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + crate::tokenizer::TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } /// Write JSON data to a unique temp file and return the path string. diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs index 3917b4d833c1..8f81f8d3c75d 100644 --- a/rust/src/bench/src/datasets/mod.rs +++ b/rust/src/bench/src/datasets/mod.rs @@ -5,6 +5,7 @@ pub mod custom; pub mod hf_dataset; pub mod multi_turn; pub mod prefix_repetition; +mod progress; pub mod random; pub mod random_mm; pub mod random_rerank; @@ -90,9 +91,10 @@ pub fn oversample_requests( return; } if no_oversample { - println!( - "Skipping oversampling. Total samples: {} (requested: {num_requests})", - requests.len() + tracing::info!( + samples = requests.len(), + requested = num_requests, + "skipping dataset oversampling" ); return; } @@ -103,9 +105,10 @@ pub fn oversample_requests( req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); requests.push(req); } - println!( - "Oversampled requests from {original_len} to {} total samples.", - requests.len() + tracing::info!( + original_samples = original_len, + samples = requests.len(), + "oversampled dataset" ); } diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 43098b9725e3..97ec522e29ad 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn( conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i); conversations.push(conv); } - println!( - "Oversampled multi-turn conversations from {original_len} to {} total.", - conversations.len() + tracing::info!( + original_conversations = original_len, + conversations = conversations.len(), + "oversampled multi-turn conversations" ); } @@ -525,10 +526,12 @@ mod tests { len } - #[test] + #[tokio::test] #[ignore] - fn test_prefix_sharing_structure() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_prefix_sharing_structure() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 5, @@ -610,10 +613,12 @@ mod tests { println!("All prefix sharing checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_per_turn_input_len_default_mode() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_per_turn_input_len_default_mode() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 4, @@ -650,10 +655,12 @@ mod tests { println!("per_turn_input_len default-mode checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_variable_turns_range() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_variable_turns_range() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 50, @@ -684,10 +691,12 @@ mod tests { println!("variable_turns_range checks passed! counts: {distinct_counts:?}"); } - #[test] + #[tokio::test] #[ignore] - fn test_variable_turns_fixed() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_variable_turns_fixed() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 10, @@ -709,10 +718,12 @@ mod tests { println!("variable_turns_fixed checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_per_turn_input_len_prefix_sharing() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_per_turn_input_len_prefix_sharing() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600 // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100 diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs index 7b06ef6077d1..e4dfcd36a3a9 100644 --- a/rust/src/bench/src/datasets/prefix_repetition.rs +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset( } let total = prompts_per_prefix * num_prefixes; if total != num_requests { - println!( - "prefix_repetition: generating {total} requests \ - ({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \ - {} dropped to divide evenly)", - num_requests - total + tracing::info!( + requested = num_requests, + generated = total, + prefixes = num_prefixes, + prompts_per_prefix, + dropped = num_requests - total, + "adjusted prefix-repetition request count" ); } @@ -109,8 +111,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } #[test] diff --git a/rust/src/bench/src/datasets/progress.rs b/rust/src/bench/src/datasets/progress.rs new file mode 100644 index 000000000000..fb2df327e608 --- /dev/null +++ b/rust/src/bench/src/datasets/progress.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::time::{Duration, Instant}; + +use indicatif::{ProgressBar, ProgressStyle}; + +const REPORT_INTERVAL: Duration = Duration::from_secs(10); + +/// Reports row download progress to an interactive progress bar, or through +/// periodic tracing events when the progress bar is hidden on a non-TTY. +pub(super) struct RowDownloadReporter { + progress: ProgressBar, + next_report: Instant, +} + +impl RowDownloadReporter { + /// Creates a reporter that emits non-TTY updates every 10 seconds. + pub fn new() -> Self { + let progress = ProgressBar::new(0); + progress.set_style( + ProgressStyle::with_template( + "{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}", + ) + .unwrap() + .progress_chars("#>-"), + ); + Self { + progress, + next_report: Instant::now() + REPORT_INTERVAL, + } + } + + /// Updates the current row count and reports progress when due. + pub fn update(&mut self, rows: usize, total: u64) { + let rows = rows as u64; + let total = total.max(rows); + self.progress.set_length(total); + self.progress.set_position(rows); + + if self.should_report(Instant::now()) { + tracing::info!(rows, total, "fetching dataset rows"); + } + } + + /// Clears the interactive progress bar after the download completes. + pub fn finish(self) { + self.progress.finish_and_clear(); + } + + fn should_report(&mut self, now: Instant) -> bool { + if !self.progress.is_hidden() || now < self.next_report { + return false; + } + self.next_report = now + REPORT_INTERVAL; + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hidden_reporter_uses_ten_second_deadline() { + let start = Instant::now(); + let mut reporter = RowDownloadReporter { + progress: ProgressBar::hidden(), + next_report: start + REPORT_INTERVAL, + }; + + assert!(!reporter.should_report(start + Duration::from_secs(9))); + assert!(reporter.should_report(start + Duration::from_secs(10))); + assert!(!reporter.should_report(start + Duration::from_secs(19))); + assert!(reporter.should_report(start + Duration::from_secs(20))); + } +} diff --git a/rust/src/bench/src/datasets/random.rs b/rust/src/bench/src/datasets/random.rs index 7cf5311278f5..132dd8420e96 100644 --- a/rust/src/bench/src/datasets/random.rs +++ b/rust/src/bench/src/datasets/random.rs @@ -49,9 +49,12 @@ pub fn generate_random_dataset( let (input_low, input_high) = range_ratio.input_bounds(real_input_len); let (output_low, output_high) = range_ratio.output_bounds(output_len); if !range_ratio.is_fixed() { - println!( - "Sampling input_len from [{input_low}, {input_high}] and \ - output_len from [{output_low}, {output_high}]" + tracing::info!( + input_low, + input_high, + output_low, + output_high, + "sampling random request lengths" ); } @@ -305,7 +308,8 @@ mod tests { #[test] #[ignore] fn test_generate_random_dataset_token_ids() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let requests = generate_random_dataset( &tokenizer, 10, // num_requests @@ -337,7 +341,8 @@ mod tests { #[test] #[ignore] fn test_generate_random_dataset_text() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let requests = generate_random_dataset( &tokenizer, 10, // num_requests @@ -371,7 +376,8 @@ mod tests { #[test] #[ignore] fn test_token_length_exact_local() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let target_len = 512; let requests = generate_random_dataset( &tokenizer, @@ -405,11 +411,11 @@ mod tests { } /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode). - #[test] + #[tokio::test] #[ignore] - fn test_token_length_exact_tiktoken() { + async fn test_token_length_exact_tiktoken() { // Use Qwen2.5 which has a tiktoken-format tokenizer - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await; let tokenizer = match tokenizer { Ok(t) => t, Err(e) => { @@ -453,10 +459,10 @@ mod tests { /// Test encode/decode roundtrip stability for tiktoken. /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift. - #[test] + #[tokio::test] #[ignore] - fn test_tiktoken_roundtrip_stability() { - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + async fn test_tiktoken_roundtrip_stability() { + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await; let tokenizer = match tokenizer { Ok(t) => t, Err(e) => { diff --git a/rust/src/bench/src/datasets/random_rerank.rs b/rust/src/bench/src/datasets/random_rerank.rs index 9c335e7e44f0..ad3686cf3f21 100644 --- a/rust/src/bench/src/datasets/random_rerank.rs +++ b/rust/src/bench/src/datasets/random_rerank.rs @@ -139,8 +139,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } fn fixed_ratio() -> RangeRatio { diff --git a/rust/src/bench/src/datasets/sharegpt.rs b/rust/src/bench/src/datasets/sharegpt.rs index 2d3901820df6..b8f7543b21d8 100644 --- a/rust/src/bench/src/datasets/sharegpt.rs +++ b/rust/src/bench/src/datasets/sharegpt.rs @@ -22,18 +22,21 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json"; /// Download the default ShareGPT dataset from HuggingFace Hub. /// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly. -pub fn download_sharegpt_dataset() -> Result { - println!( - "Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..." +pub async fn download_sharegpt_dataset() -> Result { + tracing::info!( + repository = DEFAULT_SHAREGPT_REPO, + file = DEFAULT_SHAREGPT_FILE, + "downloading ShareGPT dataset" ); - let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string()); - let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| { + let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string()) + .map_err(BenchError::Config)?; + let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| { BenchError::Config(format!( "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}" )) })?; let path_str = path.to_string_lossy().to_string(); - println!("ShareGPT dataset ready: {path_str}"); + tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready"); Ok(path_str) } @@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset( // Oversample if dataset is smaller than requested if samples.len() < num_requests { if no_oversample { - println!( - "Skipping oversampling. Total samples: {} (requested: {num_requests})", - samples.len() + tracing::info!( + dataset = "sharegpt", + samples = samples.len(), + requested = num_requests, + "skipping dataset oversampling" ); } else if !samples.is_empty() { let needed = num_requests - samples.len(); @@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset( req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); samples.push(req); } - println!( - "Oversampled requests from {original_len} to {} total samples.", - samples.len() + tracing::info!( + dataset = "sharegpt", + original_samples = original_len, + samples = samples.len(), + "oversampled dataset" ); } } diff --git a/rust/src/bench/src/datasets/speed_bench.rs b/rust/src/bench/src/datasets/speed_bench.rs index a6490b76e851..d615bb49039a 100644 --- a/rust/src/bench/src/datasets/speed_bench.rs +++ b/rust/src/bench/src/datasets/speed_bench.rs @@ -8,6 +8,7 @@ use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use super::SampleRequest; +use super::progress::RowDownloadReporter; use crate::cli::SpeedBenchConfig; use crate::error::{BenchError, Result}; use crate::tokenizer::TokenizerKind; @@ -25,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf { /// Download SPEED-Bench dataset from HuggingFace datasets-server API. /// Results are cached as JSON locally for subsequent runs. -pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { +pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result { let config_name = config.as_str(); let dir = cache_dir(); @@ -35,13 +36,13 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { // Return cached file if it exists if cache_path.exists() { let path_str = cache_path.to_string_lossy().to_string(); - println!("SPEED-Bench ({config_name}) cached: {path_str}"); + tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset"); return Ok(path_str); } - println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server..."); + tracing::info!(config = config_name, "downloading SPEED-Bench dataset"); - let client = reqwest::blocking::Client::builder() + let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(120)) .build() .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?; @@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { let mut all_rows: Vec = Vec::new(); let mut offset = 0usize; let page_size = 100usize; + let mut progress = RowDownloadReporter::new(); loop { let url = format!( @@ -64,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { let max_retries = 3; let mut data: Option = None; for attempt in 0..=max_retries { - let resp = match client.get(&url).send() { + let resp = match client.get(&url).send().await { Ok(r) => r, Err(e) => { if attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs( + tokio::time::sleep(std::time::Duration::from_secs( 2 * (attempt as u64 + 1), - )); + )) + .await; continue; } return Err(BenchError::Config(format!( @@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { }; if resp.status().is_server_error() && attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await; continue; } @@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { ))); } - data = Some(resp.json().map_err(|e| { + data = Some(resp.json().await.map_err(|e| { BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}")) })?); break; @@ -116,15 +119,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { let fetched = rows.len(); offset += fetched; - // Print progress let total = data["num_rows_total"].as_u64().unwrap_or(0); - eprint!("\r Fetched {offset}/{total} rows..."); + progress.update(offset, total); if fetched < page_size { break; } } - eprintln!(); // newline after progress + progress.finish(); if all_rows.is_empty() { return Err(BenchError::Config( @@ -137,9 +139,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { std::fs::write(&cache_path, &json_str)?; let path_str = cache_path.to_string_lossy().to_string(); - println!( - "SPEED-Bench ({config_name}): {} rows saved to {path_str}", - all_rows.len() + tracing::info!( + config = config_name, + rows = all_rows.len(), + path = %path_str, + "saved SPEED-Bench dataset" ); Ok(path_str) } @@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset( // Oversample if needed if samples.len() < num_requests { if no_oversample { - println!( - "Skipping oversampling. Total samples: {} (requested: {num_requests})", - samples.len() + tracing::info!( + dataset = "speed-bench", + samples = samples.len(), + requested = num_requests, + "skipping dataset oversampling" ); } else if !samples.is_empty() { let original_len = samples.len(); @@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset( req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); samples.push(req); } - println!( - "Oversampled SPEED-Bench from {original_len} to {} total samples.", - samples.len() + tracing::info!( + dataset = "speed-bench", + original_samples = original_len, + samples = samples.len(), + "oversampled dataset" ); } } @@ -288,7 +296,6 @@ pub fn load_speed_bench_dataset( )); } - // Print category distribution let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new(); for entry in &filtered[..filtered.len().min(samples.len())] { let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown"); @@ -297,7 +304,7 @@ pub fn load_speed_bench_dataset( let mut cats: Vec<_> = cat_counts.into_iter().collect(); cats.sort_by_key(|b| std::cmp::Reverse(b.1)); let cat_str: Vec = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect(); - println!("SPEED-Bench categories: {}", cat_str.join(", ")); + tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution"); Ok(samples) } diff --git a/rust/src/bench/src/hub.rs b/rust/src/bench/src/hub.rs index 0e1b3e962c8e..0b446450fb2d 100644 --- a/rust/src/bench/src/hub.rs +++ b/rust/src/bench/src/hub.rs @@ -1,51 +1,39 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -//! Sync facade over the async `hf_hub` API. -//! -//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq` -//! backend unconditionally pulls ureq's default rustls feature. So we use the -//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset -//! loaders, tokenizer fallback in rayon threads) by running each download on a -//! dedicated thread with its own single-threaded runtime. - use std::path::PathBuf; +use hf_hub::Repo; +use hf_hub::api::tokio::{ApiBuilder, ApiRepo}; + /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache. pub struct HubRepo { - repo: hf_hub::Repo, + repo: ApiRepo, } impl HubRepo { - pub fn model(model_id: String) -> Self { - Self { - repo: hf_hub::Repo::model(model_id), - } + pub fn model(model_id: String) -> Result { + Self::new(Repo::model(model_id)) } - pub fn dataset(repo_id: String) -> Self { - Self { - repo: hf_hub::Repo::dataset(repo_id), + pub fn dataset(repo_id: String) -> Result { + Self::new(Repo::dataset(repo_id)) + } + + fn new(repo: Repo) -> Result { + let mut builder = ApiBuilder::from_env(); + if let Ok(token) = std::env::var("HF_TOKEN") { + builder = builder.with_token(Some(token)); } + let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?; + Ok(Self { + repo: api.repo(repo), + }) } /// Download (or fetch from cache) a single file from the repo. /// Auth is handled by hf-hub via HF_TOKEN / the cached login token. - pub fn get(&self, filename: &str) -> Result { - let repo = self.repo.clone(); - let filename = filename.to_string(); - std::thread::spawn(move || { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .map_err(|e| format!("Failed to build download runtime: {e}"))?; - rt.block_on(async move { - let api = hf_hub::api::tokio::Api::new() - .map_err(|e| format!("Failed to init HF API: {e}"))?; - api.repo(repo).get(&filename).await.map_err(|e| format!("{e}")) - }) - }) - .join() - .map_err(|_| "HF Hub download thread panicked".to_string())? + pub async fn get(&self, filename: &str) -> Result { + self.repo.get(filename).await.map_err(|e| format!("{e}")) } } diff --git a/rust/src/bench/src/lib.rs b/rust/src/bench/src/lib.rs new file mode 100644 index 000000000000..0a4f19c0d6ee --- /dev/null +++ b/rust/src/bench/src/lib.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +mod backends; +mod benchmark; +mod cli; +mod compare; +mod config; +mod datasets; +mod error; +mod hub; +mod metrics; +mod multi_run; +mod multi_turn; +mod output; +mod rate_control; +mod ready_checker; +mod sweep; +mod tiktoken; +mod tokenizer; + +use anyhow::Context; + +pub use cli::{ + BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig, +}; +use config::BenchConfig; + +/// Prepare process-wide resources for a benchmark run. +pub fn prepare_process() { + // Raise the open-file soft limit to the hard limit. High-concurrency + // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit. + if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX) + && new > 1024 + { + tracing::info!(soft_limit = new, "raised open-file limit"); + } +} + +/// Run the online serving benchmark. +pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> { + // --- Compare mode: no server needed, just diff two JSON files --- + if let Some(ref files) = args.compare { + return compare::compare_results(&files[0], &files[1]).context("Comparison failed"); + } + + let config = BenchConfig::from_args(&args).context("Configuration error")?; + + async { + if config.multi_turn { + if let Some(ref sweep_mc) = args.sweep_max_concurrency { + // --- Sweep over concurrency in multi-turn mode --- + let values = sweep::parse_concurrency_values(sweep_mc) + .context("Invalid --sweep-max-concurrency")?; + sweep::run_multi_turn_concurrency_sweep( + &config, + &values, + args.sweep_num_prompts_factor, + ) + .await?; + } else { + // --- Single multi-turn conversation benchmark --- + multi_turn::run_multi_turn_benchmark(&config).await?; + } + } else if let Some(ref sweep_mc) = args.sweep_max_concurrency { + // --- Sweep over max-concurrency --- + let values = sweep::parse_concurrency_values(sweep_mc) + .context("Invalid --sweep-max-concurrency")?; + sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?; + } else if let Some(ref sweep_rate) = args.sweep_request_rate { + // --- Sweep over request-rate --- + let values = + sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?; + sweep::run_rate_sweep(&config, &values).await?; + } else if args.num_runs > 1 { + // --- Multi-run with statistical aggregation --- + multi_run::run_multi(&config, args.num_runs).await?; + } else { + // --- Normal single benchmark --- + benchmark::run_benchmark(&config).await?; + } + anyhow::Ok(()) + } + .await + .context("Benchmark failed") +} diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index dd37ec56db38..1764983e4487 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -1,92 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -mod backends; -mod benchmark; -mod cli; -mod compare; -mod config; -mod datasets; -mod error; -mod hub; -mod metrics; -mod multi_run; -mod multi_turn; -mod output; -mod rate_control; -mod ready_checker; -mod sweep; -mod tiktoken; -mod tokenizer; - #[cfg(not(target_env = "msvc"))] #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; use anyhow::Context; use clap::Parser; -use cli::Cli; -use config::BenchConfig; + +#[derive(Parser)] +#[command( + name = "vllm-bench", + about = "Benchmark online serving throughput", + version +)] +struct Cli { + #[command(flatten)] + args: vllm_bench::BenchServeArgs, +} fn main() -> anyhow::Result<()> { - // Raise the open-file soft limit to the hard limit. High-concurrency - // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit. - if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX) - && new > 1024 - { - eprintln!("Open-file limit: {new}"); - } + vllm_tracing::init_tracing("Bench"); let cli = Cli::parse(); - - // --- Compare mode: no server needed, just diff two JSON files --- - if let Some(ref files) = cli.compare { - return compare::compare_results(&files[0], &files[1]).context("Comparison failed"); - } - - let config = BenchConfig::from_cli(&cli).context("Configuration error")?; + vllm_bench::prepare_process(); let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build() - .expect("Failed to build tokio runtime"); + .context("Failed to build tokio runtime")?; - runtime - .block_on(async { - if config.multi_turn { - if let Some(ref sweep_mc) = cli.sweep_max_concurrency { - // --- Sweep over concurrency in multi-turn mode --- - let values = sweep::parse_concurrency_values(sweep_mc) - .context("Invalid --sweep-max-concurrency")?; - sweep::run_multi_turn_concurrency_sweep( - &config, - &values, - cli.sweep_num_prompts_factor, - ) - .await?; - } else { - // --- Single multi-turn conversation benchmark --- - multi_turn::run_multi_turn_benchmark(&config).await?; - } - } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency { - // --- Sweep over max-concurrency --- - let values = sweep::parse_concurrency_values(sweep_mc) - .context("Invalid --sweep-max-concurrency")?; - sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor) - .await?; - } else if let Some(ref sweep_rate) = cli.sweep_request_rate { - // --- Sweep over request-rate --- - let values = - sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?; - sweep::run_rate_sweep(&config, &values).await?; - } else if cli.num_runs > 1 { - // --- Multi-run with statistical aggregation --- - multi_run::run_multi(&config, cli.num_runs).await?; - } else { - // --- Normal single benchmark --- - benchmark::run_benchmark(&config).await?; - } - anyhow::Ok(()) - }) - .context("Benchmark failed") + runtime.block_on(vllm_bench::run(cli.args)) } diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs index 195dad93de3e..3bce8dd1a7f4 100644 --- a/rust/src/bench/src/metrics/calculator.rs +++ b/rust/src/bench/src/metrics/calculator.rs @@ -7,6 +7,22 @@ use crate::datasets::SampleRequest; use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; use crate::multi_turn::ConversationOutput; +fn log_failed_requests(outputs: &[RequestFuncOutput]) { + let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect(); + if failed_outputs.is_empty() { + return; + } + + tracing::warn!( + failed_requests = failed_outputs.len(), + displayed_errors = failed_outputs.len().min(10), + "benchmark requests failed" + ); + for (index, output) in failed_outputs.into_iter().take(10).enumerate() { + tracing::warn!(index, error = %output.error, "benchmark request failed"); + } +} + /// Calculate benchmark metrics from request outputs. /// /// Mirrors Python's `calculate_metrics()` from serve.py:392-599. @@ -63,14 +79,7 @@ pub fn calculate_metrics( let failed = outputs.len() - completed; - // Print failed request errors (capped to 10) - let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect(); - if !failed_outputs.is_empty() { - eprintln!("Failed requests during benchmark run detected (capping to 10):"); - for (i, err) in failed_outputs.iter().take(10).enumerate() { - eprintln!("Error {i}: {}", err.error); - } - } + log_failed_requests(outputs); // Calculate max output tokens per second and max concurrent requests let mut max_output_tokens_per_s = 0.0_f64; @@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics( let failed = outputs.len() - completed; - // Print failed request errors (capped to 10) - let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect(); - if !failed_outputs.is_empty() { - eprintln!("Failed requests during benchmark run detected (capping to 10):"); - for (i, err) in failed_outputs.iter().take(10).enumerate() { - eprintln!("Error {i}: {}", err.error); - } - } + log_failed_requests(outputs); // Compute peak concurrent requests from start_time + latency windows let successful_outputs: Vec<&RequestFuncOutput> = diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 8f0aedf31335..4b962b879c01 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::time::Instant; use indicatif::{ProgressBar, ProgressStyle}; +use thiserror_ext::AsReport as _; use tokio::sync::Semaphore; use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend}; @@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Result Result p, None => { - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?; downloaded.as_str() } }; @@ -179,8 +188,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result 0 || filtered_conversations > 0 { - println!( - "Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}." + tracing::info!( + filtered_turns, + filtered_conversations, + max_model_len, + "filtered conversations above maximum model length" ); } if conversations.is_empty() { @@ -192,11 +204,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Result Result Result Result Result Result Result = modules.iter().map(|s| s.as_ref()).collect(); - println!( - "LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]", - modules.len(), - names, - config.lora_assignment + tracing::info!( + adapters = modules.len(), + names = ?names, + assignment = ?config.lora_assignment, + scope = "conversation", + "assigned LoRA adapters" ); } @@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result all_outputs.push(output), Err(e) => { - eprintln!("Conversation task panicked: {e}"); + tracing::error!(error = %e.as_report(), "conversation task panicked"); } } } @@ -453,7 +467,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Result<()> { let content = serde_json::to_string(json)?; std::fs::write(file_path, content)?; - println!("Results saved to {file_path}"); + tracing::info!(path = file_path, "saved benchmark results"); Ok(()) } @@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> { file.write_all(b"\n")?; } file.write_all(content.as_bytes())?; - println!("Results appended to {file_path}"); + tracing::info!(path = file_path, "appended benchmark results"); Ok(()) } diff --git a/rust/src/bench/src/ready_checker.rs b/rust/src/bench/src/ready_checker.rs index 86af8663d115..144a1ebd99ea 100644 --- a/rust/src/bench/src/ready_checker.rs +++ b/rust/src/bench/src/ready_checker.rs @@ -23,7 +23,11 @@ pub async fn wait_for_endpoint( let backend = get_backend(backend)?; let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds); - println!("Waiting for endpoint to become up in {timeout_seconds}s"); + tracing::info!( + timeout_seconds, + retry_interval, + "waiting for endpoint readiness" + ); let pb = ProgressBar::new(timeout_seconds); pb.set_style( @@ -53,7 +57,9 @@ pub async fn wait_for_endpoint( Ok(output) => { let err = output.error.clone(); let err_last_line = err.lines().last().unwrap_or(&err); - eprintln!("Endpoint is not ready. Error='{err_last_line}'"); + pb.suspend(|| { + tracing::warn!(error = err_last_line, "endpoint is not ready"); + }); last_error = err; } Err(e) => { diff --git a/rust/src/bench/src/sweep.rs b/rust/src/bench/src/sweep.rs index 68cee1582530..2d695356dc15 100644 --- a/rust/src/bench/src/sweep.rs +++ b/rust/src/bench/src/sweep.rs @@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> { .await .map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?; if resp.status().is_success() { - println!("Prefix cache reset successfully."); + tracing::info!(url = %url, "reset prefix cache"); } else { let status = resp.status(); let body = resp.text().await.unwrap_or_default(); diff --git a/rust/src/bench/src/tiktoken.rs b/rust/src/bench/src/tiktoken.rs index b0cfbfeae499..d02e72828659 100644 --- a/rust/src/bench/src/tiktoken.rs +++ b/rust/src/bench/src/tiktoken.rs @@ -19,9 +19,9 @@ const NUM_RESERVED_SPECIAL_TOKENS: u32 = 256; pub struct TiktokenTokenizer { bpe: tiktoken_rs::CoreBPE, vocab_size: u32, - #[allow(dead_code)] num_base_tokens: u32, - /// IDs of all special tokens (for filtering from allowed_tokens) + /// Registered special token IDs. File-based tokenizers exclude them from + /// allowed tokens; built-in tokenizers retain the previously allowed IDs. special_token_ids: Vec, /// Reverse mapping: token_id -> byte sequence, for lossy UTF-8 decoding. /// Empty for built-in encodings (use bpe.decode instead). @@ -113,12 +113,17 @@ impl TiktokenTokenizer { } /// Create from a built-in tiktoken encoding (o200k_base, cl100k_base, etc.) - pub fn from_builtin_bpe(bpe: tiktoken_rs::CoreBPE, vocab_size: u32) -> Self { + pub fn from_builtin_bpe( + bpe: tiktoken_rs::CoreBPE, + vocab_size: u32, + num_base_tokens: u32, + special_token_ids: Vec, + ) -> Self { Self { bpe, vocab_size, - num_base_tokens: vocab_size, - special_token_ids: Vec::new(), + num_base_tokens, + special_token_ids, decoder: Vec::new(), is_builtin: true, } @@ -152,14 +157,15 @@ impl TiktokenTokenizer { self.vocab_size } - /// Get non-special token IDs whose byte representation is valid UTF-8. - /// Excludes ALL special tokens (like Python's `set(all_tokens) - set(prohibited_tokens)`). + /// Get token IDs that can be decoded safely. + /// + /// File-based tokenizers exclude special tokens and tokens whose byte + /// representation is not valid UTF-8. pub fn get_allowed_tokens(&self) -> Vec { if self.is_builtin { - // For built-in encodings, return full token range. - // Note: when used with random dataset, these IDs are sent to vLLM as-is; - // ensure the model's tokenizer is compatible (e.g. GPT-4o for o200k_base). - return (0..self.vocab_size).collect(); + return (0..self.num_base_tokens) + .chain(self.special_token_ids.iter().copied()) + .collect(); } let special_set: std::collections::HashSet = self.special_token_ids.iter().copied().collect(); @@ -185,12 +191,22 @@ impl TiktokenTokenizer { /// These encodings are bundled with tiktoken-rs — no network download required. /// Useful for consistent cross-model token counting (e.g. Artificial Analysis methodology). pub fn load_builtin_tiktoken(encoding: &str) -> Result { - let (bpe, vocab_size) = match encoding { - "o200k_base" => (tiktoken_rs::o200k_base(), 200_275u32), - "cl100k_base" => (tiktoken_rs::cl100k_base(), 100_277u32), - "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32), - "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32), - "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32), + let (bpe, vocab_size, num_base_tokens, special_token_ids) = match encoding { + "o200k_base" => ( + tiktoken_rs::o200k_base(), + 200_275u32, + 199_998u32, + vec![199_999, 200_018], + ), + "cl100k_base" => ( + tiktoken_rs::cl100k_base(), + 100_277u32, + 100_256u32, + vec![100_257, 100_258, 100_259, 100_260, 100_276], + ), + "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32, 50_281u32, vec![]), + "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32, 50_281u32, vec![]), + "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32, 50_256u32, vec![50_256]), _ => { return Err(BenchError::Tokenizer(format!( "Unknown built-in tiktoken encoding: '{encoding}'. \ @@ -199,12 +215,22 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result { } }; let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?; - println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})"); - Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size)) + tracing::info!( + encoding, + kind = "built-in-tiktoken", + vocab_size, + "loaded tokenizer" + ); + Ok(TiktokenTokenizer::from_builtin_bpe( + bpe, + vocab_size, + num_base_tokens, + special_token_ids, + )) } /// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo. -pub fn try_load_tiktoken(model_id: &str) -> Result { +pub async fn try_load_tiktoken(model_id: &str) -> Result { // Phase 1: If model_id is a local directory, look for tiktoken files there let local_dir = Path::new(model_id); if local_dir.is_dir() { @@ -212,7 +238,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result { } // Phase 2: Fall back to HuggingFace Hub download - try_load_tiktoken_from_hf(model_id) + try_load_tiktoken_from_hf(model_id).await } /// Common tiktoken model filenames to search for. @@ -247,25 +273,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result Result { - let repo = crate::hub::HubRepo::model(model_id.to_string()); - - let model_path = repo - .get("tiktoken.model") - .or_else(|_| repo.get("qwen.tiktoken")) - .or_else(|_| repo.get("vocab.tiktoken")) - .map_err(|_| { - BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) - })?; +async fn try_load_tiktoken_from_hf(model_id: &str) -> Result { + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?; + + let mut model_path = None; + for filename in TIKTOKEN_MODEL_FILENAMES { + if let Ok(path) = repo.get(filename).await { + model_path = Some(path); + break; + } + } + let model_path = model_path.ok_or_else(|| { + BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) + })?; let num_base_tokens = count_base_tokens(&model_path)?; - let config = match repo.get("tokenizer_config.json") { + let config = match repo.get("tokenizer_config.json").await { Ok(config_path) => read_tokenizer_config(&config_path), Err(_) => None, }; - let pattern = extract_pat_str_from_repo(&repo); + let pattern = extract_pat_str_from_repo(&repo).await; build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens) } @@ -309,15 +338,16 @@ fn build_tiktoken( } } - println!( - "Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...", - num_base_tokens, - all_special_tokens.len(), - if pattern.is_some() { + tracing::info!( + model = model_id, + base_tokens = num_base_tokens, + special_tokens = all_special_tokens.len(), + pattern = if pattern.is_some() { "custom" } else { "default" }, + "loading tiktoken model" ); TiktokenTokenizer::from_file( @@ -397,9 +427,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option { /// Try to download the Python tokenizer source file and extract pat_str via regex. /// Returns None if unavailable or unparsable. -fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option { +async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option { // Try common Python tokenizer filenames - let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?; + let py_path = match repo.get("tokenization_kimi.py").await { + Ok(path) => path, + Err(_) => repo.get("tokenizer.py").await.ok()?, + }; let source = std::fs::read_to_string(&py_path).ok()?; @@ -438,9 +471,9 @@ fn extract_pat_str_from_source(source: &str) -> Option { if !fragments.is_empty() { let pattern = fragments.join("|"); - println!( - "Extracted pat_str from Python source: {} fragments", - fragments.len() + tracing::debug!( + fragments = fragments.len(), + "extracted tiktoken pattern from Python source" ); return Some(pattern); } @@ -540,3 +573,44 @@ fn find_unescaped(s: &str, ch: char) -> Option { } None } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builtin_allowed_tokens_are_decodable() { + let cases = [ + ("o200k_base", 200_000), + ("cl100k_base", 100_261), + ("p50k_base", 50_281), + ("p50k_edit", 50_281), + ("r50k_base", 50_257), + ("gpt2", 50_257), + ]; + + for (encoding, expected_count) in cases { + let tokenizer = load_builtin_tiktoken(encoding).unwrap(); + let allowed_tokens = tokenizer.get_allowed_tokens(); + + assert_eq!(allowed_tokens.len(), expected_count, "{encoding}"); + for token_id in allowed_tokens { + tokenizer + .decode(&[token_id]) + .unwrap_or_else(|error| panic!("{encoding} token {token_id}: {error}")); + } + } + } + + #[test] + fn sparse_builtin_allowed_tokens_skip_unassigned_ids() { + let o200k = load_builtin_tiktoken("o200k_base").unwrap().get_allowed_tokens(); + assert_eq!(&o200k[199_997..], &[199_997, 199_999, 200_018]); + + let cl100k = load_builtin_tiktoken("cl100k_base").unwrap().get_allowed_tokens(); + assert_eq!( + &cl100k[100_255..], + &[100_255, 100_257, 100_258, 100_259, 100_260, 100_276] + ); + } +} diff --git a/rust/src/bench/src/tokenizer.rs b/rust/src/bench/src/tokenizer.rs index 3d0d4c9c3155..add890c39166 100644 --- a/rust/src/bench/src/tokenizer.rs +++ b/rust/src/bench/src/tokenizer.rs @@ -2,8 +2,10 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashSet; +use std::future::Future; use std::path::Path; +use thiserror_ext::AsReport as _; use tokenizers::Tokenizer; use crate::error::{BenchError, Result}; @@ -18,7 +20,8 @@ pub enum TokenizerKind { /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints. pub struct ServerTokenizer { - client: reqwest::blocking::Client, + client: reqwest::Client, + runtime: tokio::runtime::Handle, tokenize_url: String, detokenize_url: String, model: String, @@ -27,8 +30,8 @@ pub struct ServerTokenizer { impl ServerTokenizer { /// Create a new server tokenizer and verify connectivity. - pub fn new(base_url: &str, model: &str) -> Result { - let client = reqwest::blocking::Client::builder() + pub async fn new(base_url: &str, model: &str) -> Result { + let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .build() .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?; @@ -38,6 +41,7 @@ impl ServerTokenizer { let st = Self { client, + runtime: tokio::runtime::Handle::current(), tokenize_url, detokenize_url, model: model.to_string(), @@ -45,7 +49,7 @@ impl ServerTokenizer { }; // Probe the endpoint to verify it works and discover vocab size - let test_tokens = st.encode_inner("test")?; + let test_tokens = st.encode_async("test").await?; let max_id = test_tokens.iter().copied().max().unwrap_or(0); let estimated_vocab = (max_id * 2).max(131072); @@ -56,6 +60,10 @@ impl ServerTokenizer { } fn encode_inner(&self, text: &str) -> Result> { + self.block_on(self.encode_async(text)) + } + + async fn encode_async(&self, text: &str) -> Result> { let payload = serde_json::json!({ "model": self.model, "prompt": text, @@ -66,6 +74,7 @@ impl ServerTokenizer { .post(&self.tokenize_url) .json(&payload) .send() + .await .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?; if !resp.status().is_success() { @@ -75,7 +84,7 @@ impl ServerTokenizer { ))); } - let data: serde_json::Value = resp.json().map_err(|e| { + let data: serde_json::Value = resp.json().await.map_err(|e| { BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}")) })?; @@ -95,6 +104,10 @@ impl ServerTokenizer { } fn decode_inner(&self, ids: &[u32]) -> Result { + self.block_on(self.decode_async(ids)) + } + + async fn decode_async(&self, ids: &[u32]) -> Result { let payload = serde_json::json!({ "model": self.model, "tokens": ids, @@ -105,6 +118,7 @@ impl ServerTokenizer { .post(&self.detokenize_url) .json(&payload) .send() + .await .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?; if !resp.status().is_success() { @@ -114,7 +128,7 @@ impl ServerTokenizer { ))); } - let data: serde_json::Value = resp.json().map_err(|e| { + let data: serde_json::Value = resp.json().await.map_err(|e| { BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}")) })?; @@ -123,6 +137,26 @@ impl ServerTokenizer { .map(|s| s.to_string()) .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into())) } + + fn block_on(&self, future: impl Future>) -> Result { + if matches!( + self.runtime.runtime_flavor(), + tokio::runtime::RuntimeFlavor::CurrentThread + ) { + return Err(BenchError::Tokenizer( + "Server tokenizer fallback requires a multi-thread Tokio runtime".into(), + )); + } + + // Sync tokenizer calls can come from a Tokio worker or a Rayon worker. + // Tokio workers must enter a blocking region before re-entering the runtime; + // Rayon workers can drive the future directly with the saved runtime handle. + if tokio::runtime::Handle::try_current().is_ok() { + tokio::task::block_in_place(|| self.runtime.block_on(future)) + } else { + self.runtime.block_on(future) + } + } } // --- TokenizerKind methods --- @@ -192,7 +226,7 @@ impl TokenizerKind { /// 3. Server-side /tokenize + /detokenize endpoints /// /// `server_info` is `Some((base_url, model))` to enable server-side fallback. -pub fn load_tokenizer( +pub async fn load_tokenizer( model_id: &str, _trust_remote_code: bool, server_info: Option<(&str, &str)>, @@ -212,31 +246,48 @@ pub fn load_tokenizer( } // 1. Try local HuggingFace tokenizer (tokenizer.json) - match try_load_local(model_id) { + match try_load_local(model_id).await { Ok(tok) => { - println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true)); + tracing::info!( + model = model_id, + kind = "local", + vocab_size = tok.get_vocab_size(true), + "loaded tokenizer" + ); Ok(TokenizerKind::Local(Box::new(tok))) } Err(local_err) => { // 2. Try tiktoken format - println!("No tokenizer.json for '{model_id}', trying tiktoken format..."); - match crate::tiktoken::try_load_tiktoken(model_id) { + tracing::info!( + model = model_id, + error = %local_err.as_report(), + "local tokenizer unavailable; trying tiktoken" + ); + match crate::tiktoken::try_load_tiktoken(model_id).await { Ok(tok) => { - println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size()); + tracing::info!( + model = model_id, + kind = "tiktoken", + vocab_size = tok.vocab_size(), + "loaded tokenizer" + ); Ok(TokenizerKind::Tiktoken(tok)) } Err(tiktoken_err) => { // 3. Try server-side fallback if let Some((base_url, model)) = server_info { - println!( - "Tiktoken also not available ({tiktoken_err}), \ - trying server-side tokenization..." + tracing::info!( + model = model_id, + error = %tiktoken_err.as_report(), + "tiktoken unavailable; trying server-side tokenization" ); - match ServerTokenizer::new(base_url, model) { + match ServerTokenizer::new(base_url, model).await { Ok(srv) => { - println!( - "Tokenizer: Server (vocab_size≈{})", - srv.cached_vocab_size + tracing::info!( + model = model_id, + kind = "server", + vocab_size = srv.cached_vocab_size, + "loaded tokenizer" ); return Ok(TokenizerKind::Server(srv)); } @@ -264,7 +315,7 @@ pub fn load_tokenizer( } /// Try loading tokenizer.json from local path or HuggingFace Hub. -fn try_load_local(model_id: &str) -> Result { +async fn try_load_local(model_id: &str) -> Result { // 1. Try local directory with tokenizer.json let local_path = Path::new(model_id).join("tokenizer.json"); if local_path.exists() { @@ -290,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result { } // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token) - let repo = crate::hub::HubRepo::model(model_id.to_string()); + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?; let tokenizer_path = repo .get("tokenizer.json") + .await .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?; Tokenizer::from_file(&tokenizer_path) .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}"))) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_server_tokenizer_sync_bridge() { + let tokenizer = std::sync::Arc::new(ServerTokenizer { + client: reqwest::Client::new(), + runtime: tokio::runtime::Handle::current(), + tokenize_url: String::new(), + detokenize_url: String::new(), + model: String::new(), + cached_vocab_size: 0, + }); + + assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1); + + let (tx, rx) = tokio::sync::oneshot::channel(); + rayon::spawn(move || { + let _ = tx.send(tokenizer.block_on(async { Ok(2) })); + }); + assert_eq!(rx.await.unwrap().unwrap(), 2); + } +} diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index d4999d07cd05..20d2b81cc178 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -20,7 +20,7 @@ use crate::output::{ use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; use crate::renderer::{ DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, - InklingChatRenderer, + InklingChatRenderer, KimiK3ChatRenderer, }; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -57,6 +57,7 @@ impl HfChatBackend { processor_config: files.processor_config_path.as_deref(), }, tokenizer.clone(), + options.limit_mm_per_prompt.clone(), )? }; let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref()); @@ -73,6 +74,7 @@ impl HfChatBackend { RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), + RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())), }; info!( @@ -230,6 +232,7 @@ mod tests { chat_template_content_format: Default::default(), chat_template: None, default_chat_template_kwargs: HashMap::new(), + limit_mm_per_prompt: HashMap::new(), }, test_tokenizer(), ) diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 68f2f6a568ca..0a0f757d654f 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -8,7 +8,7 @@ use serde_json::Value; use vllm_text::{DynTextBackend, TextBackend}; use crate::error::Result; -use crate::multimodal::MultimodalModelInfo; +use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo}; use crate::output::DynChatOutputProcessor; use crate::renderer::DynChatRenderer; use crate::request::ChatRequest; @@ -74,6 +74,9 @@ pub struct LoadModelBackendsOptions { /// Optional server-default keyword arguments merged into every /// chat-template render before request-level `chat_template_kwargs`. pub default_chat_template_kwargs: HashMap, + /// Maximum number of input items allowed per prompt for each modality. + /// Unspecified modalities are unlimited. + pub limit_mm_per_prompt: MmLimitPerPrompt, } /// Shared backends loaded from a model id. diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index e09c3734d988..b032e9bf24cf 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -23,6 +23,8 @@ pub enum Error { UnsupportedMultimodalContent(&'static str), #[error("`{modality}` input is not supported by this model")] UnsupportedModality { modality: String }, + #[error("At most {limit} {modality}(s) may be provided in one prompt.")] + MmLimitExceeded { modality: String, limit: usize }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -87,7 +89,8 @@ impl Error { Self::Text(error) => error.is_request_validation_error(), Self::UnsupportedMultimodalRenderer | Self::UnsupportedMultimodalContent(_) - | Self::UnsupportedModality { .. } => true, + | Self::UnsupportedModality { .. } + | Self::MmLimitExceeded { .. } => true, _ => false, } diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 3188657b9096..0d538a24ba1a 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, - HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection, + HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, + RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -54,6 +55,7 @@ mod stream; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::Llm; use vllm_text::{Prompt, TextLlm, TextRequest}; @@ -88,6 +90,92 @@ pub fn validate_parser_overrides( Ok(()) } +/// Chat request preparation shared by inference and render-only frontends. +pub struct ChatRequestProcessor { + backend: DynChatBackend, + /// Effective model dtype reported by the engine. + /// Absent for text-only frontends without an engine handshake. + model_dtype: Option, +} + +impl ChatRequestProcessor { + /// Create a processor with multimodal support using the effective model + /// dtype reported by the engine. + fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self { + Self { + backend, + model_dtype: Some(model_dtype), + } + } + + /// Create a render-only processor that rejects multimodal requests. + pub fn render_only(backend: DynChatBackend) -> Self { + Self { + backend, + model_dtype: None, + } + } + + async fn finalize_rendered_prompt( + &self, + request: &ChatRequest, + rendered: RenderedPrompt, + ) -> Result<(Prompt, Option)> { + match self.model_dtype { + Some(model_dtype) => { + multimodal::finalize_rendered_prompt( + request, + rendered, + self.backend.multimodal_model_info(), + model_dtype, + ) + .await + } + None if !request.has_multimodal() => Ok((rendered.prompt, None)), + None => Err(Error::UnsupportedMultimodalRenderer), + } + } + + /// Prepare one chat request without submitting it to an engine. + pub async fn prepare( + &self, + mut request: ChatRequest, + options: NewChatOutputProcessorOptions<'_>, + ) -> Result<(TextRequest, DynChatOutputProcessor)> { + 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, options)?; + let rendered = self.backend.chat_renderer().render(&request)?; + let reasoning_parser_kwargs = + request + .sampling_params + .structured_outputs + .is_some() + .then(|| ReasoningParserKwargs { + chat_template_kwargs: rendered.effective_template_kwargs.clone(), + }); + let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?; + let text_request = TextRequest { + request_id: request.request_id, + prompt, + mm_features, + sampling_params: request.sampling_params, + decode_options: request.decode_options, + intermediate: request.intermediate, + priority: request.priority, + cache_salt: request.cache_salt, + add_special_tokens: request.add_special_tokens, + data_parallel_rank: request.data_parallel_rank, + reasoning_parser_kwargs, + lora_request: request.lora_request, + arrival_time: Some(arrival_time), + }; + Ok((text_request, output_processor)) + } +} + /// Structured chat facade above [`TextLlm`]. /// /// This layer stays above raw text semantics: it takes care of chat-template @@ -95,9 +183,7 @@ pub fn validate_parser_overrides( /// request semantics such as tool calls. pub struct ChatLlm { text: TextLlm, - backend: DynChatBackend, - /// Effective model dtype reported by the engine. - model_dtype: ModelDtype, + processor: ChatRequestProcessor, /// Tool-call parser selection. tool_call_parser: ParserSelection, /// Reasoning parser selection. @@ -112,8 +198,7 @@ impl ChatLlm { Self { text, - backend, - model_dtype, + processor: ChatRequestProcessor::new(backend, model_dtype), tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, } @@ -140,7 +225,7 @@ impl ChatLlm { /// Override the effective model dtype used for multimodal tensor encoding. pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self { - self.model_dtype = model_dtype; + self.processor.model_dtype = Some(model_dtype); self } @@ -171,58 +256,51 @@ impl ChatLlm { self.text.engine_core_client() } - /// Render, tokenize, and submit one chat request. - 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(); + /// Whether the loaded backend has a registered multimodal processor. + pub fn supports_multimodal(&self) -> bool { + self.processor.backend.multimodal_model_info().is_some() + } - let output_processor = self.backend.new_chat_output_processor( - &mut request, - NewChatOutputProcessorOptions { - tool_call_parser: &self.tool_call_parser, - reasoning_parser: &self.reasoning_parser, - }, - )?; - let rendered = self.backend.chat_renderer().render(&request)?; - let reasoning_parser_kwargs = - request - .sampling_params - .structured_outputs - .is_some() - .then(|| ReasoningParserKwargs { - chat_template_kwargs: rendered.effective_template_kwargs.clone(), - }); + /// Effective tool-call parser name for this model, if parsing is enabled. + pub fn tool_call_parser_name(&self) -> Option<&str> { + match &self.tool_call_parser { + ParserSelection::Auto => { + ToolParserFactory::global().resolve_name_for_model(self.model_id()) + } + ParserSelection::None => None, + ParserSelection::Explicit(name) => Some(name), + } + } - let (prompt, mm_features) = multimodal::finalize_rendered_prompt( - &request, - rendered, - self.backend.multimodal_model_info(), - self.model_dtype, - ) - .await?; + /// Effective reasoning parser name for this model, if parsing is enabled. + pub fn reasoning_parser_name(&self) -> Option<&str> { + match &self.reasoning_parser { + ParserSelection::Auto => { + ReasoningParserFactory::global().resolve_name_for_model(self.model_id()) + } + ParserSelection::None => None, + ParserSelection::Explicit(name) => Some(name), + } + } - let text_request = TextRequest { - request_id: request.request_id.clone(), - prompt, - mm_features, - sampling_params: request.sampling_params, - decode_options: request.decode_options, - intermediate: request.intermediate, - priority: request.priority, - cache_salt: request.cache_salt, - add_special_tokens: request.add_special_tokens, - data_parallel_rank: request.data_parallel_rank, - reasoning_parser_kwargs, - lora_request: request.lora_request, - arrival_time: Some(arrival_time), - }; + /// Render, tokenize, and submit one chat request. + pub async fn chat(&self, request: ChatRequest) -> Result { + let (text_request, output_processor) = self + .processor + .prepare( + request, + NewChatOutputProcessorOptions { + tool_call_parser: &self.tool_call_parser, + reasoning_parser: &self.reasoning_parser, + }, + ) + .await?; + let request_id = text_request.request_id.clone(); let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); let structured_stream = output_processor.process(decoded_stream)?; - Ok(ChatEventStream::new(request.request_id, structured_stream)) + Ok(ChatEventStream::new(request_id, structured_stream)) } /// Render through the chat template and tokenize, without submitting to the engine. @@ -233,14 +311,9 @@ impl ChatLlm { pub async fn tokenize_chat(&self, request: ChatRequest) -> Result> { request.validate()?; - let rendered = self.backend.chat_renderer().render(&request)?; - let (prompt, _mm_features) = multimodal::finalize_rendered_prompt( - &request, - rendered, - self.backend.multimodal_model_info(), - self.model_dtype, - ) - .await?; + let rendered = self.processor.backend.chat_renderer().render(&request)?; + let (prompt, _mm_features) = + self.processor.finalize_rendered_prompt(&request, rendered).await?; let tokenizer = self.text.tokenizer(); let token_ids = match prompt { @@ -281,6 +354,12 @@ mod tests { .unwrap(); } + #[test] + fn validate_parser_overrides_accepts_explicit_kimi_k3() { + let selection = ParserSelection::Explicit("kimi_k3".to_string()); + validate_parser_overrides(&selection, &selection).unwrap(); + } + #[test] fn validate_parser_overrides_accepts_auto_and_none() { validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap(); @@ -294,7 +373,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -305,6 +384,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 9c4f57c0413e..c0600a6329ba 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -11,7 +11,7 @@ //! Raw media stays above `vllm-text`; this module lowers it into token IDs and //! opaque tensor payloads before the request is handed to text generation. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fs; use std::path::Path; use std::sync::{Arc, LazyLock}; @@ -24,6 +24,7 @@ use llm_multimodal::{ PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor, VisionProcessorRegistry, }; +use serde::{Deserialize, Serialize}; use thiserror_ext::AsReport as _; use tracing::warn; use vllm_engine_core_client::protocol::dtype::ModelDtype; @@ -52,6 +53,71 @@ pub struct MultimodalModelInfo { video: Option, audio: Option, media_connector: Arc, + /// Maximum number of input items allowed per prompt for each modality. + limit_mm_per_prompt: MmLimitPerPrompt, +} + +/// Per-modality item-count limits configured by `--limit-mm-per-prompt`. +/// +/// Modalities absent from the map are unlimited. +pub type MmLimitPerPrompt = HashMap; + +/// Modalities that `--limit-mm-per-prompt` can be keyed by. +/// +/// Closed on purpose: these are exactly the keys Python accepts, per +/// `MultiModalDummyOptionsBuiltins` in `vllm/config/multimodal.py`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MmLimitModality { + Image, + Audio, + Video, +} + +impl MmLimitModality { + /// The wire name, matching Python's modality strings. + pub fn as_str(self) -> &'static str { + match self { + Self::Image => "image", + Self::Audio => "audio", + Self::Video => "video", + } + } +} + +/// One modality's limit, in either of the two shapes Python accepts. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde( + untagged, + expecting = "an item count, or an object with an optional `count` field" +)] +pub enum MmLimitSpec { + /// Legacy form: `"image": 16` + Count(usize), + + /// Configurable form: + /// `"video": {"count": 1, "num_frames": 32}` + Options { + /// Absent means unlimited, matching an absent modality. + #[serde(default, skip_serializing_if = "Option::is_none")] + count: Option, + + /// Preserve Python-owned options for forwarding. Never interpreted + /// here: they size the engine's dummy-profiling encoder cache, which + /// has no Rust counterpart. + #[serde(flatten)] + extra: BTreeMap, + }, +} + +impl MmLimitSpec { + /// The configured item count, or `None` when this modality is unlimited. + pub fn count(&self) -> Option { + match self { + Self::Count(count) => Some(*count), + Self::Options { count, .. } => *count, + } + } } /// Model metadata and tokenizer access shared by all multimodal specs. @@ -287,6 +353,7 @@ impl MultimodalModelInfo { model_type: Option, files: MultimodalConfigFiles<'_>, tokenizer: DynTokenizer, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> Result> { let config = match files.config { Some(path) => { @@ -319,7 +386,12 @@ impl MultimodalModelInfo { tokenizer: TokenizerResolver(tokenizer), }; - Self::from_loaded(context, preprocessor_config, video_preprocessor_config) + Self::from_loaded( + context, + preprocessor_config, + video_preprocessor_config, + limit_mm_per_prompt, + ) } /// Resolve multimodal support from an assembled context and parsed @@ -328,6 +400,7 @@ impl MultimodalModelInfo { context: MultimodalModelContext, preprocessor_config: PreProcessorConfig, video_preprocessor_config: PreProcessorConfig, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> Result> { let (image, video) = Self::resolve_vision_lanes( &context, @@ -356,6 +429,7 @@ impl MultimodalModelInfo { video, audio, media_connector, + limit_mm_per_prompt, })) } @@ -567,7 +641,55 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result { Ok(format!("data:{mime_type};base64,{data}")) } +/// The modality a content part counts against, or `None` for plain text. +/// +/// Embedding inputs share their base modality's budget rather than getting one +/// of their own, matching Python's `modality.replace("_embeds", "")` in +/// `vllm/entrypoints/chat_utils.py`. +fn media_part_limit_modality(part: &MediaContentPart) -> Option { + match part { + MediaContentPart::Text { .. } => None, + MediaContentPart::ImageUrl { .. } + | MediaContentPart::ImageData { .. } + | MediaContentPart::ImageEmbeds { .. } => Some(MmLimitModality::Image), + MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => { + Some(MmLimitModality::Audio) + } + MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => { + Some(MmLimitModality::Video) + } + } +} + impl MultimodalModelInfo { + /// Reject requests exceeding `--limit-mm-per-prompt`'s configured + /// per-modality item count, before any fetch/decode work is spent on them. + /// + /// Modalities without a configured count are unlimited. + fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> { + let mut counts: HashMap = HashMap::new(); + for part in media_parts { + if let Some(modality) = media_part_limit_modality(part) { + *counts.entry(modality).or_default() += 1; + } + } + + for (modality, count) in counts { + let Some(limit) = self.limit_mm_per_prompt.get(&modality).and_then(MmLimitSpec::count) + else { + continue; + }; + if count > limit { + return Err(Error::MmLimitExceeded { + modality: modality.as_str().to_string(), + limit, + }); + } + } + + Ok(()) + } + /// Run media fetch, per-modality preprocessing, prompt expansion, and /// feature build. /// @@ -584,8 +706,7 @@ impl MultimodalModelInfo { if media_parts_len == 0 { return Ok(Vec::new()); } - // TODO: enforce per-modality item-count limits, aligned with the - // engine's `--limit-mm-per-prompt` semantics. + self.validate_mm_limits(&media_parts)?; let fetched = self.fetch_media(media_parts).await?; let mut prepared = Vec::new(); @@ -753,10 +874,11 @@ mod tests { .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID) } - fn test_info( + fn test_info_with_limits( model_type: &str, config: serde_json::Value, tokenizer: TestTokenizer, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> MultimodalModelInfo { let context = MultimodalModelContext { model_id: format!("{model_type}-test"), @@ -769,11 +891,20 @@ mod tests { context, PreProcessorConfig::default(), PreProcessorConfig::default(), + limit_mm_per_prompt, ) .unwrap() .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve")) } + fn test_info( + model_type: &str, + config: serde_json::Value, + tokenizer: TestTokenizer, + ) -> MultimodalModelInfo { + test_info_with_limits(model_type, config, tokenizer, HashMap::new()) + } + fn llama4_info() -> MultimodalModelInfo { let config = serde_json::json!({ "model_type": "llama4", @@ -783,16 +914,19 @@ mod tests { test_info("llama4", config, llama4_tokenizer()) } - pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { - let config = serde_json::json!({ + fn qwen3_vl_config() -> serde_json::Value { + serde_json::json!({ "model_type": "qwen3_vl", "image_token_id": QWEN3_IMAGE_PAD_ID, "video_token_id": QWEN3_VIDEO_PAD_ID, "vision_start_token_id": 151652, "vision_end_token_id": 151653, "vision_config": {"patch_size": 16} - }); - test_info("qwen3_vl", config, qwen3_vl_tokenizer()) + }) + } + + pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { + test_info("qwen3_vl", qwen3_vl_config(), qwen3_vl_tokenizer()) } #[test] @@ -853,4 +987,123 @@ mod tests { ); assert!(input_audio_data_url("AAAA", Some("flac")).is_err()); } + + fn image_url_part() -> MediaContentPart { + MediaContentPart::ImageUrl { + url: "https://example.com/image.png".to_string(), + detail: None, + uuid: None, + } + } + + fn qwen3_vl_info_with_limits(limit_mm_per_prompt: MmLimitPerPrompt) -> MultimodalModelInfo { + test_info_with_limits( + "qwen3_vl", + qwen3_vl_config(), + qwen3_vl_tokenizer(), + limit_mm_per_prompt, + ) + } + + #[test] + fn validate_mm_limits_ignores_text_parts() { + let info = qwen3_vl_info(); + let parts = vec![ + MediaContentPart::Text { + text: "hello".to_string(), + }, + MediaContentPart::Text { + text: "world".to_string(), + }, + ]; + assert!(info.validate_mm_limits(&parts).is_ok()); + } + + #[test] + fn validate_mm_limits_leaves_unconfigured_modalities_unlimited() { + let info = qwen3_vl_info(); + let parts: Vec<_> = std::iter::repeat_with(image_url_part).take(1_000).collect(); + + assert!(info.validate_mm_limits(&parts).is_ok()); + } + + #[test] + fn validate_mm_limits_enforces_configured_limit_at_the_boundary() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Count(1), + )])); + assert!(info.validate_mm_limits(&[image_url_part()]).is_ok()); + + let error = info.validate_mm_limits(&[image_url_part(), image_url_part()]).unwrap_err(); + assert_eq!( + error.to_report_string(), + "At most 1 image(s) may be provided in one prompt." + ); + // Confirms the HTTP-mapping bug found during implementation stays fixed: + // this must map to 400, not 500. + assert!(error.is_request_validation_error()); + } + + #[test] + fn validate_mm_limits_counts_image_embeds_against_the_image_limit() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Count(1), + )])); + let image_embeds_part = MediaContentPart::ImageEmbeds { + payload: serde_json::Value::String("AAAA".to_string()), + uuid: None, + }; + + let error = info.validate_mm_limits(&[image_url_part(), image_embeds_part]).unwrap_err(); + assert_eq!( + error.to_report_string(), + "At most 1 image(s) may be provided in one prompt." + ); + } + + /// An options object without a `count` carries only profiling keys, which + /// say nothing about how many items are allowed. + #[test] + fn validate_mm_limits_treats_a_count_less_options_object_as_unlimited() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Options { + count: None, + extra: BTreeMap::from([("width".to_string(), serde_json::json!(512))]), + }, + )])); + + assert!(info.validate_mm_limits(&[image_url_part(), image_url_part()]).is_ok()); + } + + fn parse_limits(json: &str) -> MmLimitPerPrompt { + serde_json::from_str(json).expect("limit map should parse") + } + + #[test] + fn limit_map_parses_both_shapes_python_accepts() { + let limits = parse_limits(r#"{"image": 16, "video": {"count": 1, "num_frames": 32}}"#); + + assert_eq!(limits[&MmLimitModality::Image].count(), Some(16)); + assert_eq!(limits[&MmLimitModality::Video].count(), Some(1)); + assert_eq!(limits.get(&MmLimitModality::Audio), None); + } + + #[test] + fn limit_map_rejects_keys_python_does_not_accept() { + assert!(serde_json::from_str::(r#"{"image_embeds": 1}"#).is_err()); + } + + /// Managed mode forwards this map back to Python as JSON, where + /// `BaseDummyOptions.count` is a non-optional `int` under `extra="forbid"`. + /// Emitting `"count": null` would make the engine subprocess fail to start. + #[test] + fn limit_map_round_trips_without_emitting_a_null_count() { + let source = r#"{"video":{"num_frames":32}}"#; + let encoded = serde_json::to_string(&parse_limits(source)).expect("map should serialize"); + + assert_eq!(encoded, source); + } } diff --git a/rust/src/chat/src/multimodal/audio.rs b/rust/src/chat/src/multimodal/audio.rs index 4910f235eea7..41ea47b484ff 100644 --- a/rust/src/chat/src/multimodal/audio.rs +++ b/rust/src/chat/src/multimodal/audio.rs @@ -107,6 +107,7 @@ mod tests { context, PreProcessorConfig::default(), PreProcessorConfig::default(), + HashMap::new(), ) .unwrap() .expect("Inkling multimodal support") @@ -126,6 +127,7 @@ mod tests { context, PreProcessorConfig::default(), PreProcessorConfig::default(), + HashMap::new(), ) .unwrap() .expect("Qwen3-ASR multimodal support") diff --git a/rust/src/chat/src/multimodal/expand.rs b/rust/src/chat/src/multimodal/expand.rs index fc4b47f4d208..e449de59094d 100644 --- a/rust/src/chat/src/multimodal/expand.rs +++ b/rust/src/chat/src/multimodal/expand.rs @@ -222,7 +222,9 @@ mod tests { assert_eq!(tensor.shape, vec![expected.len()]); assert_eq!( tensor.data, - WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) + WireArrayData::RawView( + expected.iter().map(|value| u8::from(*value)).collect::>().into(), + ) ); } diff --git a/rust/src/chat/src/multimodal/item.rs b/rust/src/chat/src/multimodal/item.rs index d5600c568803..69ab71af3073 100644 --- a/rust/src/chat/src/multimodal/item.rs +++ b/rust/src/chat/src/multimodal/item.rs @@ -38,7 +38,7 @@ pub(super) fn build_batched_items( let keep_on_cpu = spec.keep_on_cpu_keys.contains(key); let (value, field) = match spec.field_layout_for(key) { Some(FieldLayout::Batched) => ( - tensor.batched_value_at(index)?, + tensor.batched_wire_value_at(index)?, MmField::Batched(MmBatchedField { keep_on_cpu }), ), Some(FieldLayout::Flat { sizes_key }) => { @@ -47,7 +47,7 @@ pub(super) fn build_batched_items( })?; let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; ( - tensor.flat_value_range(start, end)?, + tensor.flat_wire_value_range(start, end)?, MmField::Flat(MmFlatField { slices: vec![MmSlice::Slice(SliceSpec { start: Some(0), @@ -60,7 +60,7 @@ pub(super) fn build_batched_items( ) } None => ( - tensor.clone(), + tensor.try_into()?, MmField::Shared(MmSharedField { batch_size: len, keep_on_cpu, @@ -71,7 +71,7 @@ pub(super) fn build_batched_items( data.insert( key.clone(), MmFieldElem { - data: Some(value.try_into()?), + data: Some(value), field, }, ); diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs index f701e0646c63..41f0347a4117 100644 --- a/rust/src/chat/src/multimodal/tensor.rs +++ b/rust/src/chat/src/multimodal/tensor.rs @@ -2,28 +2,58 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::mem::size_of; use half::{bf16, f16}; use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; use vllm_engine_core_client::protocol::dtype::ModelDtype; use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; -use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor}; +use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor}; use crate::error::{Error, Result, bail_multimodal, multimodal}; +/// Element type retained alongside an encoded tensor during multimodal lowering. +#[derive(Debug, Clone, Copy)] +pub(super) enum TensorKind { + /// 32-bit floating point. + F32, + /// IEEE 16-bit floating point. + F16, + /// Brain floating point. + Bf16, + /// Signed 64-bit integer. + I64, + /// Unsigned 32-bit integer. + U32, +} + +impl TensorKind { + const fn element_size(self) -> usize { + match self { + Self::F32 => size_of::(), + Self::F16 => size_of::(), + Self::Bf16 => size_of::(), + Self::I64 => size_of::(), + Self::U32 => size_of::(), + } + } + + const fn wire_dtype(self) -> &'static str { + match self { + Self::F32 => "float32", + Self::F16 => "float16", + Self::Bf16 => "bfloat16", + Self::I64 => "int64", + Self::U32 => "uint32", + } + } +} + /// Representation for multimodal kwarg values for transformation. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(super) enum KwargValue { - /// Float tensor with row-major flat data and shape. - F32Tensor { data: Vec, shape: Vec }, - /// Float16 tensor with row-major flat data and shape. - F16Tensor { data: Vec, shape: Vec }, - /// BFloat16 tensor with row-major flat data and shape. - Bf16Tensor { data: Vec, shape: Vec }, - /// Signed integer tensor with row-major flat data and shape. - I64Tensor { data: Vec, shape: Vec }, - /// Unsigned integer tensor with row-major flat data and shape. - U32Tensor { data: Vec, shape: Vec }, + /// Tensor with row-major flat data and shape. + Tensor { kind: TensorKind, wire: WireTensor }, /// Non-tensor kwarg value that is shared or copied as-is. Passthrough(ProtocolKwargValue), } @@ -67,8 +97,14 @@ impl KwargValue { ModelSpecificValue::Tensor { data, shape } => { Self::from_f32_tensor(data, shape, float_dtype)? } - ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape }, - ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape }, + ModelSpecificValue::IntTensor { data, shape } => { + let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?; + Self::tensor(TensorKind::I64, wire) + } + ModelSpecificValue::UintTensor { data, shape } => { + let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?; + Self::tensor(TensorKind::U32, wire) + } ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)), ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)), ModelSpecificValue::IntVec(values) => { @@ -93,42 +129,35 @@ impl KwargValue { /// Convert a float tensor to the target float dtype if needed, keeping the /// same shape. fn from_f32_tensor(data: Vec, shape: Vec, float_dtype: ModelDtype) -> Result { - match float_dtype { - ModelDtype::Float16 => Ok(Self::F16Tensor { - data: data.into_iter().map(f16::from_f32).collect(), - shape, - }), - ModelDtype::BFloat16 => Ok(Self::Bf16Tensor { - data: data.into_iter().map(bf16::from_f32).collect(), - shape, - }), - ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }), - } + let (kind, wire) = match float_dtype { + ModelDtype::Float16 => ( + TensorKind::F16, + WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()), + ), + ModelDtype::BFloat16 => ( + TensorKind::Bf16, + WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()), + ), + ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)), + }; + wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal) + } + + fn tensor(kind: TensorKind, wire: WireTensor) -> Self { + debug_assert_eq!(wire.dtype, kind.wire_dtype()); + Self::Tensor { kind, wire } } } -impl TryFrom for ProtocolKwargValue { +impl TryFrom<&KwargValue> for ProtocolKwargValue { type Error = Error; - fn try_from(value: KwargValue) -> Result { - match value { - KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor( - WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?, - )), - KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor( - WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?, - )), - KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor( - WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?, - )), - KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor( - WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?, - )), - KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor( - WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?, - )), - KwargValue::Passthrough(value) => Ok(value), - } + fn try_from(value: &KwargValue) -> Result { + let wire = match value { + KwargValue::Tensor { wire, .. } => wire.clone(), + KwargValue::Passthrough(value) => return Ok(value.clone()), + }; + Ok(ProtocolKwargValue::Tensor(wire)) } } @@ -136,72 +165,43 @@ impl KwargValue { /// First-axis length for tensor values; `None` for passthrough kwargs. pub(super) fn first_dim(&self) -> Option { match self { - Self::F32Tensor { shape, .. } - | Self::F16Tensor { shape, .. } - | Self::Bf16Tensor { shape, .. } - | Self::I64Tensor { shape, .. } - | Self::U32Tensor { shape, .. } => shape.first().copied(), + Self::Tensor { wire, .. } => wire.shape.first().copied(), Self::Passthrough(_) => None, } } - /// Extract one media item from a batched tensor field. + /// Convert one media item from a batched tensor field to wire bytes. /// /// Batched fields use their first axis as media-item index and drop that /// axis in the per-feature value, matching vLLM's batched-field semantics. - pub(super) fn batched_value_at(&self, index: usize) -> Result { - match self { - Self::F32Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?; - Ok(Self::F32Tensor { data, shape }) - } - Self::F16Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?; - Ok(Self::F16Tensor { data, shape }) - } - Self::Bf16Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?; - Ok(Self::Bf16Tensor { data, shape }) - } - Self::I64Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?; - Ok(Self::I64Tensor { data, shape }) - } - Self::U32Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?; - Ok(Self::U32Tensor { data, shape }) - } - Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())), - } + pub(super) fn batched_wire_value_at(&self, index: usize) -> Result { + self.wire_value_range(index, index + 1, true) } - /// Extract one media item's variable-length range from a flat tensor field. + /// Convert one media item's flat tensor range directly to wire bytes. /// /// Flat fields keep the first axis as the sliced length for this item. - pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result { - match self { - Self::F32Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; - Ok(Self::F32Tensor { data, shape }) - } - Self::F16Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; - Ok(Self::F16Tensor { data, shape }) - } - Self::Bf16Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; - Ok(Self::Bf16Tensor { data, shape }) - } - Self::I64Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; - Ok(Self::I64Tensor { data, shape }) - } - Self::U32Tensor { data, shape } => { - let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; - Ok(Self::U32Tensor { data, shape }) + pub(super) fn flat_wire_value_range( + &self, + start: usize, + end: usize, + ) -> Result { + self.wire_value_range(start, end, false) + } + + fn wire_value_range( + &self, + start: usize, + end: usize, + drop_axis: bool, + ) -> Result { + let wire = match self { + Self::Tensor { kind, wire } => { + slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis) } - Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())), - } + Self::Passthrough(value) => return Ok(value.clone()), + }; + wire.map(ProtocolKwargValue::Tensor) } } @@ -225,44 +225,53 @@ pub(super) fn flat_range_for_index( /// Read a tensor value as per-image sizes for flat slicing. fn tensor_as_usize_vec(tensor: &KwargValue) -> Result> { match tensor { - KwargValue::I64Tensor { data, .. } => data - .iter() + KwargValue::Tensor { + kind: TensorKind::I64, + wire, + } => raw_tensor_bytes(wire, size_of::())? + .chunks_exact(size_of::()) + .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk"))) .map(|value| { - usize::try_from(*value) + usize::try_from(value) .map_err(|_| multimodal!("negative flat tensor size `{value}`")) }) .collect(), - KwargValue::U32Tensor { data, .. } => { - Ok(data.iter().map(|value| *value as usize).collect()) - } + KwargValue::Tensor { + kind: TensorKind::U32, + wire, + } => Ok(raw_tensor_bytes(wire, size_of::())? + .chunks_exact(size_of::()) + .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize) + .collect()), _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")), } } /// Slice a flat row-major tensor along its first axis. -fn slice_first_axis_range( - shape: &[usize], - data: &[T], +fn slice_first_axis_range( + tensor: &WireTensor, + element_size: usize, start: usize, end: usize, drop_axis: bool, -) -> Result<(Vec, Vec)> { +) -> Result { + let shape = tensor.shape.as_slice(); + raw_tensor_bytes(tensor, element_size)?; let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?; if start > end || end > first_dim { bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}"); } - let expected_len = shape - .checked_numel() - .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?; - if expected_len != data.len() { - bail_multimodal!( - "tensor shape {shape:?} expects {expected_len} elements, got {}", - data.len() - ); - } - let stride = shape[1..].iter().product::(); - let data_start = start * stride; - let data_end = end * stride; + let stride = shape[1..] + .iter() + .try_fold(1usize, |acc, dim| acc.checked_mul(*dim)) + .and_then(|stride| stride.checked_mul(element_size)) + .ok_or_else(|| multimodal!("tensor shape {shape:?} byte stride overflowed usize"))?; + let data_start = start + .checked_mul(stride) + .ok_or_else(|| multimodal!("tensor slice start byte offset overflowed usize"))?; + let data_end = end + .checked_mul(stride) + .ok_or_else(|| multimodal!("tensor slice end byte offset overflowed usize"))?; let out_shape = if drop_axis { shape[1..].to_vec() } else { @@ -270,7 +279,38 @@ fn slice_first_axis_range( shape[0] = end - start; shape }; - Ok((out_shape, data[data_start..data_end].to_vec())) + let WireArrayData::RawView(data) = &tensor.data else { + return Err(multimodal!("cannot slice an aux tensor buffer")); + }; + Ok(WireTensor::from_raw_bytes( + tensor.dtype.clone(), + out_shape, + data.slice(data_start..data_end), + )) +} + +fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> { + let WireArrayData::RawView(data) = &tensor.data else { + return Err(multimodal!("expected an inline tensor buffer")); + }; + let expected_bytes = tensor + .shape + .checked_numel() + .and_then(|numel| numel.checked_mul(element_size)) + .ok_or_else(|| { + multimodal!( + "tensor shape {:?} byte length overflowed usize", + tensor.shape + ) + })?; + if expected_bytes != data.len() { + bail_multimodal!( + "tensor shape {:?} expects {expected_bytes} bytes, got {}", + tensor.shape, + data.len() + ); + } + Ok(data) } #[cfg(test)] @@ -278,43 +318,51 @@ mod tests { use super::*; #[test] - fn batched_value_at_drops_first_axis() { - let value = KwargValue::F32Tensor { - data: vec![1.0, 2.0, 3.0, 4.0], - shape: vec![2, 2], - }; + fn batched_wire_value_at_drops_first_axis() { + let data = vec![1.0_f32, 2.0, 3.0, 4.0]; + let expected_ptr = data.as_ptr().cast::().wrapping_add(2 * size_of::()); + let value = KwargValue::tensor( + TensorKind::F32, + WireTensor::from_f32(vec![2, 2], data).unwrap(), + ); - let value = value.batched_value_at(1).unwrap(); + let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else { + panic!("expected tensor"); + }; - assert!(matches!( - value, - KwargValue::F32Tensor { data, shape } - if shape == vec![2] && data == vec![3.0, 4.0] - )); + assert_eq!(tensor.shape, vec![2]); + let raw_view = tensor.data.into_raw_view().unwrap(); + assert_eq!(raw_view.as_ptr(), expected_ptr); + assert_eq!( + raw_view, + [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::>() + ); } #[test] - fn flat_value_range_keeps_first_axis() { - let value = KwargValue::U32Tensor { - data: (0..10).collect(), - shape: vec![5, 2], - }; + fn flat_wire_value_range_keeps_first_axis() { + let value = KwargValue::tensor( + TensorKind::U32, + WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(), + ); - let value = value.flat_value_range(1, 3).unwrap(); + let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else { + panic!("expected tensor"); + }; - assert!(matches!( - value, - KwargValue::U32Tensor { data, shape } - if shape == vec![2, 2] && data == vec![2, 3, 4, 5] - )); + assert_eq!(tensor.shape, vec![2, 2]); + assert_eq!( + tensor.data.into_raw_view().unwrap(), + [2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::>() + ); } #[test] fn flat_range_for_index_uses_size_tensor() { - let sizes = KwargValue::I64Tensor { - data: vec![2, 3, 4], - shape: vec![3], - }; + let sizes = KwargValue::tensor( + TensorKind::I64, + WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(), + ); assert_eq!( flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(), @@ -324,10 +372,11 @@ mod tests { #[test] fn slice_first_axis_range_errors_on_shape_data_mismatch() { - let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err(); + let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::()]); + let error = slice_first_axis_range(&tensor, size_of::(), 0, 1, true).unwrap_err(); assert!( - matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements")) + matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes")) ); } @@ -336,7 +385,7 @@ mod tests { let value = KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap(); - let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap() + let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap() else { panic!("expected tensor"); }; @@ -351,7 +400,7 @@ mod tests { let value = KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap(); - let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap() + let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap() else { panic!("expected tensor"); }; diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs index 9cc3238ac602..c2a25b27bd59 100644 --- a/rust/src/chat/src/multimodal/video.rs +++ b/rust/src/chat/src/multimodal/video.rs @@ -130,7 +130,7 @@ fn build_video_item( let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key); let (value, field) = match support.spec.field_layout_for(&key) { Some(FieldLayout::Batched) => ( - tensor.batched_value_at(0)?, + tensor.batched_wire_value_at(0)?, MmField::Batched(MmBatchedField { keep_on_cpu }), ), Some(FieldLayout::Flat { .. }) => { @@ -138,7 +138,7 @@ fn build_video_item( .first_dim() .ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?; ( - tensor, + (&tensor).try_into()?, MmField::Flat(MmFlatField { slices: vec![MmSlice::Slice(SliceSpec { start: Some(0), @@ -151,7 +151,7 @@ fn build_video_item( ) } None => ( - tensor, + (&tensor).try_into()?, MmField::Shared(MmSharedField { batch_size: 1, keep_on_cpu, @@ -162,7 +162,7 @@ fn build_video_item( data.insert( key, MmFieldElem { - data: Some(value.try_into()?), + data: Some(value), field, }, ); @@ -207,6 +207,7 @@ mod tests { Some("qwen3_vl".to_string()), files, Arc::new(qwen3_vl_tokenizer()), + std::collections::HashMap::new(), ) }; diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index ab288336f4c7..4317e48fce2d 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor { Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box }; - apply_structural_tag_constraint(request, parser.structural_tag_model())?; + apply_structural_tag_constraint(request, parser.structural_tag_builder())?; if parser.preserve_special_tokens() { request.decode_options.skip_special_tokens = false; diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index 79aa1a1f1b52..c0b31bd74564 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -7,7 +7,8 @@ use thiserror_ext::AsReport; use vllm_engine_core_client::protocol::structured_outputs::{ StructuredOutputBackend, StructuredOutputsParams, }; -use vllm_parser::tool::StructuralTagModel; +use vllm_parser::tool::StructuralTagBuilder; +use xgrammar_structural_tag::builders::StructuralTagOptions; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, build_structural_tag, @@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult}; /// support and the request's tool choice. pub(super) fn apply_structural_tag_constraint( request: &mut ChatRequest, - model: Option, + builder: Option<&dyn StructuralTagBuilder>, ) -> ChatResult<()> { - let Some(model) = model else { + let Some(builder) = builder else { return Ok(()); }; let Some(tool_choice) = structural_tag_tool_choice(request) else { @@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint( }) .collect::>(); - let structural_tag = build_structural_tag(model, &tools, tool_choice, false) - .and_then(|tag| tag.to_json_string()) - .map_err(|error| Error::StructuralTag { - message: error.to_report_string(), - })?; + let structural_tag = build_structural_tag( + builder, + &tools, + tool_choice, + StructuralTagOptions::default().with_reasoning(false), + ) + .and_then(|tag| tag.to_json_string()) + .map_err(|error| Error::StructuralTag { + message: error.to_report_string(), + })?; // Overwrite any existing structured output settings with the structural tag constraint. request.sampling_params.structured_outputs = Some(StructuredOutputsParams { @@ -141,7 +147,7 @@ mod tests { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); let tag = structural_tag_value(&request); @@ -154,7 +160,7 @@ mod tests { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); assert!(request.sampling_params.structured_outputs.is_none()); @@ -169,7 +175,7 @@ mod tests { }); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); let params = structured_outputs(&request); @@ -184,7 +190,7 @@ mod tests { let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); let tag = structural_tag_value(&request); @@ -201,7 +207,7 @@ mod tests { }); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); let params = structured_outputs(&request); @@ -221,7 +227,7 @@ mod tests { ); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); let tag = structural_tag_value(&request).to_string(); @@ -234,7 +240,7 @@ mod tests { let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); assert!(request.sampling_params.structured_outputs.is_none()); @@ -249,7 +255,7 @@ mod tests { }); let parser = qwen3_coder_parser(&request.tools); - apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); let params = structured_outputs(&request); diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index ca025268b042..9a1e37533b45 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -27,6 +27,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const MINIMAX_M2: &str = "minimax_m2"; pub const MINIMAX_M3: &str = "minimax_m3"; pub const NEMOTRON_V3: &str = "nemotron_v3"; @@ -68,6 +69,7 @@ impl ReasoningParserFactory { .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::MINIMAX_M2) .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 78ccde8a4b34..e42b45a48628 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -33,6 +33,7 @@ pub mod names { // also routes to `Internlm2ToolParser` despite the version-agnostic name. pub const INTERNLM: &str = "internlm"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; pub const MINIMAX_M2: &str = "minimax_m2"; @@ -78,6 +79,7 @@ impl ToolParserFactory { .register_parser::(names::HY_V3) .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 2536bf78c252..6246733db281 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -5,7 +5,9 @@ use std::sync::LazyLock; -pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser}; +pub use vllm_parser::unified::{ + Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser, +}; use vllm_tokenizer::DynTokenizer; use crate::parser::ParserFactory; @@ -15,6 +17,7 @@ use crate::request::ChatTool; pub mod names { pub const GEMMA4: &str = "gemma4"; pub const INKLING: &str = "inkling"; + pub const KIMI_K3: &str = "kimi_k3"; } /// Constructor signature for one registered unified parser implementation. @@ -39,11 +42,14 @@ impl UnifiedParserFactory { factory.register_parser::(names::GEMMA4); factory.register_parser::(names::INKLING); + factory.register_parser::(names::KIMI_K3); factory .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) - .register_pattern("inkling", names::INKLING); + .register_pattern("inkling", names::INKLING) + .register_pattern("kimi-k3", names::KIMI_K3) + .register_pattern("kimi_k3", names::KIMI_K3); factory } @@ -121,4 +127,20 @@ mod tests { ); factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap(); } + + #[test] + fn factory_registers_kimi_k3() { + let factory = UnifiedParserFactory::new(); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|open|>", 1001) + .with_regular_token("<|close|>", 1002) + .with_regular_token("<|sep|>", 1003); + + assert!(factory.contains(names::KIMI_K3)); + assert_eq!( + factory.resolve_name_for_model("moonshotai/Kimi-K3"), + Some(names::KIMI_K3) + ); + factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap(); + } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 76796edd0965..7960927219b3 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -59,7 +59,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 73380cef260e..9068d460ee61 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -27,7 +27,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs index 301d9707e716..4849c932233b 100644 --- a/rust/src/chat/src/renderer/harmony/tests.rs +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -22,7 +22,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fixture_chat_request( &fixture_path(input_name), FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, }, ) diff --git a/rust/src/chat/src/renderer/hf/format.rs b/rust/src/chat/src/renderer/hf/format.rs index f419b9e3a0ca..e8e9f6412bb0 100644 --- a/rust/src/chat/src/renderer/hf/format.rs +++ b/rust/src/chat/src/renderer/hf/format.rs @@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool { loops.into_iter().any(|loop_ast| { matches!(loop_ast.target, Expr::Var(_)) - && message_varnames - .iter() - .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))) + && (is_var_access(&loop_ast.iter, "content") + || message_varnames.iter().any(|varname| { + is_var_or_elems_access(&loop_ast.iter, varname, Some("content")) + })) }) } @@ -315,6 +316,16 @@ mod tests { ); } + #[test] + fn detects_openai_template_with_content_parameter_loop() { + assert_eq!( + detect( + "{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}" + ), + ChatTemplateContentFormat::OpenAi + ); + } + #[test] fn detects_openai_template_with_messages_alias() { assert_eq!( diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 7ba2a2ccc030..cfbc3f924a34 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -1309,6 +1309,26 @@ mod tests { .assert_eq(&rendered); } + #[test] + fn qwen35_template_auto_detects_openai_multimodal_content() { + let mut request = image_request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_mm( + QWEN3_5_0_8B_TEMPLATE, + &request, + ChatTemplateContentFormatOption::Auto, + ) + .unwrap(); + + expect![[r#" + Text( + "<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n", + ) + "#]] + .assert_debug_eq(&rendered.prompt); + } + #[test] fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() { let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index d0d55e1be2d5..1c4cd7236dbd 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -31,6 +31,10 @@ impl Tokenizer for FixtureTokenizer { Ok(text.bytes().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], @@ -101,7 +105,7 @@ fn fixture_request(name: &str) -> ChatRequest { fn inkling_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, } } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs new file mode 100644 index 000000000000..242c1e84c24b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -0,0 +1,608 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML prompt renderer. +//! +//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`. + +use std::collections::HashMap; + +use serde_json::{Map, Value, json}; +use vllm_tokenizer::Tokenizer; + +use crate::error::{Error, Result}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, +}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +pub(super) const OPEN: &str = "<|open|>"; +pub(super) const CLOSE: &str = "<|close|>"; +pub(super) const SEP: &str = "<|sep|>"; +pub(super) const END_OF_MSG: &str = "<|end_of_msg|>"; +pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>"; + +const DEFAULT_THINKING_EFFORT: &str = "max"; +const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; + +/// K3 prompt encoder preserving Python's per-segment tokenization boundaries. +pub(super) struct K3TokenWriter<'a> { + tokenizer: &'a dyn Tokenizer, + token_ids: Vec, +} + +impl<'a> K3TokenWriter<'a> { + pub(super) fn new(tokenizer: &'a dyn Tokenizer) -> Self { + Self { + tokenizer, + token_ids: Vec::new(), + } + } + + /// Encode one trusted segment with normal added-token recognition. + pub(super) fn control(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode(text, false)?); + } + Ok(()) + } + + /// Encode one literal segment while bypassing every added-token matcher. + pub(super) fn ordinary(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode_ordinary(text)?); + } + Ok(()) + } + + pub(super) fn finish(self) -> Vec { + self.token_ids + } +} + +/// Render and tokenize one chat request using K3's segment-aware contract. +pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { + let thinking = thinking_enabled(request)?; + let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; + let tools = request_tools(request); + let mut out = K3TokenWriter::new(tokenizer); + + if !tools.is_empty() { + write_tool_declare(&mut out, tools, false)?; + } + + if let Some(effort) = thinking_effort { + // Preserve the checkpoint's literal guidance text: it still names + // `medium`, although the validator above no longer accepts it. + write_internal_system( + &mut out, + "thinking-effort", + &format!( + "`thinking_effort` guides on how much to think in your \ + thinking channel (not including the response channel), \ + supported values include `low`, `medium`, `high`, and `max`.\n\ + Now the system is invoked with `thinking_effort={effort}`." + ), + )?; + } + + // Track prior assistant tool-call ids for tool-result reordering / naming. + let mut tool_call_id_index: HashMap = HashMap::new(); + let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new(); + + let flush_tool_run = |out: &mut K3TokenWriter<'_>, + run: &mut Vec<(usize, ChatMessage)>, + id_index: &HashMap| + -> Result<()> { + if run.is_empty() { + return Ok(()); + } + + let mut resolved = Vec::with_capacity(run.len()); + let mut unresolved = false; + for (offset, message) in run.drain(..) { + let ChatMessage::ToolResponse { + content, + tool_call_id, + } = message + else { + unreachable!("pending tool run only holds tool responses"); + }; + match id_index.get(&tool_call_id) { + Some(&(position, ref name)) => { + resolved.push((position, offset, content, Some(name.clone()))); + } + None => { + unresolved = true; + resolved.push((usize::MAX, offset, content, None)); + } + } + } + + if unresolved { + // Preserve original order when the run cannot be fully matched. + resolved.sort_by_key(|item| item.1); + } else { + resolved.sort_by_key(|item| (item.0, item.1)); + } + + for (position, _, content, name) in resolved { + let tool_name = name.as_deref().ok_or_else(|| { + Error::ChatTemplate( + "Kimi K3 tool messages need a resolvable tool name: \ + carry a matching tool_call_id against a preceding \ + assistant tool_call" + .to_string(), + ) + })?; + write_tool_message(out, tool_name, position, &content)?; + } + Ok(()) + }; + + for (message_index, message) in request.messages.iter().enumerate() { + match message { + ChatMessage::ToolResponse { .. } => { + pending_tool_run.push((message_index, message.clone())); + continue; + } + _ => { + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + } + } + + match message { + // Python: role=system with a `tools` field renders a dynamic + // tool-declare (`## New Tools Available`). Map that to Developer + // messages that carry tools (OpenAI "developer" / system-tools). + ChatMessage::Developer { + content, + tools: Some(local_tools), + } if !local_tools.is_empty() => { + write_tool_declare(&mut out, local_tools, true)?; + if !content_is_empty(content) { + write_role_message(&mut out, "system", None, content)?; + } + } + ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { + write_role_message(&mut out, "system", None, content)?; + } + ChatMessage::User { content } => { + write_role_message(&mut out, "user", None, content)?; + } + ChatMessage::Assistant { content } => { + tool_call_id_index.clear(); + let mut call_position = 0usize; + for block in content { + if let AssistantContentBlock::ToolCall(call) = block { + call_position += 1; + if !call.id.is_empty() { + tool_call_id_index + .entry(call.id.clone()) + .or_insert((call_position, call.name.clone())); + } + } + } + + write_assistant_message(&mut out, content, thinking)?; + } + ChatMessage::ToolResponse { .. } => unreachable!("handled above"), + } + } + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + + match &request.tool_choice { + ChatToolChoice::Required => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.", + )?; + } + // Emit only when tools are present: Rust defaults tool_choice to None + // for tool-free requests, which must not inject a tool-choice message. + ChatToolChoice::None if !request.tools.is_empty() => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message.", + )?; + } + ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {} + } + + write_response_format(&mut out, request)?; + + if request.chat_options.add_generation_prompt() { + write_open_tag(&mut out, "message", &[("role", "assistant")])?; + write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[])?; + } + + Ok(out.finish()) +} + +fn request_tools(request: &ChatRequest) -> &[ChatTool] { + // Declare tools whenever the request carries them. K3 tool-declare is + // independent of tool_choice; tool_choice only injects control messages. + request.tools.as_slice() +} + +fn thinking_enabled(request: &ChatRequest) -> Result { + if let Some(thinking) = request.parse_template_bool("thinking")? { + return Ok(thinking); + } + if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { + return Ok(enable_thinking); + } + Ok(request + .chat_options + .reasoning_effort + .map(|effort| effort != crate::request::ReasoningEffort::None) + .unwrap_or(true)) +} + +fn thinking_effort(request: &ChatRequest) -> Result { + let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `thinking_effort` must be a string, got {value}" + )) + })? + } else if let Some(effort) = request.chat_options.reasoning_effort { + effort.as_str() + } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `reasoning_effort` must be a string, got {value}" + )) + })? + } else { + DEFAULT_THINKING_EFFORT + }; + + if !VALID_THINKING_EFFORTS.contains(&effort) { + return Err(Error::ChatTemplate(format!( + "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`" + ))); + } + Ok(effort.to_string()) +} + +fn content_is_empty(content: &ChatContent) -> bool { + match content { + ChatContent::Text(text) => text.is_empty(), + ChatContent::Parts(parts) => parts.iter().all(|part| match part { + ChatContentPart::Text { text } => text.is_empty(), + ChatContentPart::ImageUrl { .. } + | ChatContentPart::VideoUrl { .. } + | ChatContentPart::InputAudio { .. } + | ChatContentPart::AudioUrl { .. } => false, + }), + } +} + +fn write_tool_declare( + out: &mut K3TokenWriter<'_>, + tools: &[ChatTool], + dynamic: bool, +) -> Result<()> { + let mut specs = Vec::with_capacity(tools.len()); + for tool in tools { + let mut function = Map::new(); + function.insert( + "description".to_string(), + Value::String(tool.description.clone().unwrap_or_default()), + ); + function.insert("name".to_string(), Value::String(tool.name.clone())); + function.insert("parameters".to_string(), sort_json(&tool.parameters)); + specs.push(json!({ + "function": Value::Object(function), + "type": "function", + })); + } + let payload = compact_json(&sort_json(&Value::Array(specs)))?; + + let body = if dynamic { + format!( + "## New Tools Available\n\ + The system dynamically extends the toolset via lazy-loading.\n\ + You have access to all existing and extended tools.\n\ + Here are the specs for the extended tools.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + } else { + format!( + "# Tools\n\ + Here are the available tools, described in JSONSchema.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + }; + + write_internal_system(out, "tool-declare", &body) +} + +fn write_internal_system( + out: &mut K3TokenWriter<'_>, + message_type: &str, + body: &str, +) -> Result<()> { + write_open_tag( + out, + "message", + &[("role", "system"), ("type", message_type)], + )?; + out.ordinary(body.trim())?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_role_message( + out: &mut K3TokenWriter<'_>, + role: &str, + name: Option<&str>, + content: &ChatContent, +) -> Result<()> { + let mut attrs = vec![("role", role.to_string())]; + if let Some(name) = name.filter(|name| !name.is_empty()) { + attrs.push(("name", name.to_string())); + } + let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); + write_open_tag(out, "message", &attr_refs)?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_tool_message( + out: &mut K3TokenWriter<'_>, + tool_name: &str, + index: usize, + content: &ChatContent, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "message", + &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], + )?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_message( + out: &mut K3TokenWriter<'_>, + content: &[AssistantContentBlock], + thinking: bool, +) -> Result<()> { + write_open_tag(out, "message", &[("role", "assistant")])?; + + let mut reasoning = String::new(); + let mut response = String::new(); + let mut tool_calls = Vec::new(); + for block in content { + match block { + AssistantContentBlock::Reasoning { text } => reasoning.push_str(text), + AssistantContentBlock::Text { text } => response.push_str(text), + AssistantContentBlock::ToolCall(call) => tool_calls.push(call), + } + } + + // The think channel is structural: in thinking mode every assistant + // message carries open/close tags even when there is no reasoning content. + // In non-thinking mode the channel is dropped entirely. + if thinking { + write_open_tag(out, "think", &[])?; + if !reasoning.trim().is_empty() { + out.ordinary(&reasoning)?; + } + write_close_tag(out, "think")?; + } + + write_open_tag(out, "response", &[])?; + out.ordinary(&response)?; + write_close_tag(out, "response")?; + + if !tool_calls.is_empty() { + write_open_tag(out, "tools", &[])?; + for (index, tool_call) in tool_calls.into_iter().enumerate() { + write_assistant_tool_call(out, tool_call, index + 1)?; + } + write_close_tag(out, "tools")?; + } + + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_tool_call( + out: &mut K3TokenWriter<'_>, + tool_call: &AssistantToolCall, + index: usize, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "call", + &[ + ("tool", tool_call.name.as_str()), + ("index", index_str.as_str()), + ], + )?; + + let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?; + if let Some(raw) = json_block { + write_open_tag(out, "json", &[("type", "object")])?; + out.ordinary(&raw)?; + write_close_tag(out, "json")?; + } else { + for (key, value) in args { + let typ = xtml_type(&value); + write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)])?; + out.ordinary(&xtml_value(&value))?; + write_close_tag(out, "argument")?; + } + } + + write_close_tag(out, "call") +} + +fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { + match content { + ChatContent::Text(text) => write_text_with_images(out, text), + ChatContent::Parts(parts) => { + for part in parts { + match part { + ChatContentPart::Text { text } => write_text_with_images(out, text)?, + ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::VideoUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("video_url")); + } + ChatContentPart::InputAudio { .. } => { + return Err(Error::UnsupportedMultimodalContent("input_audio")); + } + ChatContentPart::AudioUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("audio_url")); + } + } + } + Ok(()) + } + } +} + +fn write_text_with_images(out: &mut K3TokenWriter<'_>, text: &str) -> Result<()> { + // Placeholder expansion is left as the literal K3 image token; multimodal + // preprocessing can replace it once image prompts are known. + out.ordinary(text) +} + +fn write_response_format(out: &mut K3TokenWriter<'_>, request: &ChatRequest) -> Result<()> { + let Some(rf) = request.chat_options.response_format.as_ref() else { + return Ok(()); + }; + + let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or(""); + + match rf_type { + "json_object" => { + write_internal_system( + out, + "response-format", + "The system is invoked with `response_format=json_object`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.", + )?; + } + "json_schema" => { + let schema = extract_response_schema(rf); + let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?; + write_internal_system( + out, + "response-format", + &format!( + "The system is invoked with `response_format=json_schema`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.\n\ + The JSON data must match the following schema:\n\ + ```json\n\ + {schema_json}\n\ + ```" + ), + )?; + } + _ => {} + } + Ok(()) +} + +fn extract_response_schema(response_format: &Value) -> Option { + let json_schema = response_format.get("json_schema")?; + if let Some(schema) = json_schema.get("schema") { + return Some(schema.clone()); + } + if let Some(schema) = json_schema.get("json_schema") { + return Some(schema.clone()); + } + Some(json_schema.clone()) +} + +fn normalize_tool_arguments(arguments: &str) -> Result<(Map, Option)> { + let trimmed = arguments.trim(); + if trimmed.is_empty() { + return Ok((Map::new(), None)); + } + match serde_json::from_str::(trimmed) { + Ok(Value::Object(map)) => Ok((map, None)), + Ok(_) => Err(Error::ChatTemplate( + "Kimi K3 tool call arguments must be a JSON object".to_string(), + )), + Err(_) => Ok((Map::new(), Some(arguments.to_string()))), + } +} + +fn xtml_type(value: &Value) -> &'static str { + match value { + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Object(_) => "object", + Value::Array(_) => "array", + } +} + +fn xtml_value(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + other => compact_json(other).unwrap_or_else(|_| other.to_string()), + } +} + +fn write_open_tag(out: &mut K3TokenWriter<'_>, tag: &str, attrs: &[(&str, &str)]) -> Result<()> { + out.control(OPEN)?; + out.ordinary(tag)?; + for (key, value) in attrs { + out.ordinary(&format!(" {key}"))?; + out.ordinary("=\"")?; + out.ordinary(&escape_attr_value(value))?; + out.ordinary("\"")?; + } + out.control(SEP) +} + +fn write_close_tag(out: &mut K3TokenWriter<'_>, tag: &str) -> Result<()> { + out.control(CLOSE)?; + out.ordinary(tag)?; + out.control(SEP) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +fn compact_json(value: &Value) -> Result { + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string())) +} + +fn sort_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()), + Value::Object(map) => { + let mut sorted = Map::new(); + let mut keys = map.keys().collect::>(); + keys.sort(); + for key in keys { + sorted.insert(key.clone(), sort_json(&map[key])); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json new file mode 100644 index 000000000000..569201f3b976 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json @@ -0,0 +1,15 @@ +{ + "messages": [ + { + "role": "user", + "content": "json pls" + } + ], + "add_generation_prompt": true, + "response_format": { + "type": "json_object" + }, + "template_kwargs": { + "thinking": false + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt new file mode 100644 index 000000000000..57ce4a687613 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="user"<|sep|>json pls<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="response-format"<|sep|>The system is invoked with `response_format=json_object`. +Your response must be raw JSON data without markdown code blocks (```json) or any additional formatting.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>response<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json new file mode 100644 index 000000000000..61ecebc268b8 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json @@ -0,0 +1,62 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "developer", + "tools": [ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ] + }, + { + "role": "user", + "content": "use calc" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt new file mode 100644 index 000000000000..edec2bebf2ab --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt @@ -0,0 +1,14 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>hi<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available +The system dynamically extends the toolset via lazy-loading. +You have access to all existing and extended tools. +Here are the specs for the extended tools. + +```json +[{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>use calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json new file mode 100644 index 000000000000..b7cf40c9b8b4 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json @@ -0,0 +1,31 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "see this" + }, + { + "type": "image_url", + "image_url": "https://example.com/a.png" + } + ] + }, + { + "role": "assistant", + "content": "old answer", + "reasoning_content": "old reasoning" + }, + { + "role": "user", + "content": "continue" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true, + "thinking_effort": "high" + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt new file mode 100644 index 000000000000..03582f1180a7 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>see this<|media_pad|><|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>old reasoning<|close|>think<|sep|><|open|>response<|sep|>old answer<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>continue<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json new file mode 100644 index 000000000000..3a64f8dd0bbc --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json @@ -0,0 +1,87 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Hangzhou weather and calc" + }, + { + "role": "assistant", + "content": "I'll check.", + "reasoning_content": "Need tools.", + "tool_calls": [ + { + "id": "get_weather:0", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\",\"days\":1}" + } + }, + { + "id": "calc:1", + "type": "function", + "function": { + "name": "calc", + "arguments": "{\"x\":2}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "calc:1", + "content": "2" + }, + { + "role": "tool", + "tool_call_id": "get_weather:0", + "content": "{\"city\":\"Hangzhou\",\"condition\":\"rain\"}" + } + ], + "add_generation_prompt": true, + "tool_choice": "required", + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt new file mode 100644 index 000000000000..1c24810072ad --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt @@ -0,0 +1,8 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"},{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Hangzhou weather and calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>Need tools.<|close|>think<|sep|><|open|>response<|sep|>I'll check.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool="get_weather" index="1"<|sep|><|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|><|open|>argument key="days" type="number"<|sep|>1<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool="calc" index="2"<|sep|><|open|>argument key="x" type="number"<|sep|>2<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="get_weather" index="1"<|sep|>{"city":"Hangzhou","condition":"rain"}<|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="calc" index="2"<|sep|>2<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`. +You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs new file mode 100644 index 000000000000..09bf356f5d64 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Native Kimi K3 XTML chat renderer. + +mod encoding; +#[cfg(test)] +mod tests; + +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; + +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::Result; +use crate::request::ChatRequest; + +/// Dedicated Kimi K3 XTML renderer. +#[derive(Clone)] +pub struct KimiK3ChatRenderer { + tokenizer: DynTokenizer, +} + +impl KimiK3ChatRenderer { + /// Create a Kimi K3 renderer. + pub fn new(tokenizer: DynTokenizer) -> Self { + Self { tokenizer } + } +} + +impl ChatRenderer for KimiK3ChatRenderer { + fn render(&self, request: &ChatRequest) -> Result { + request.validate()?; + + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs new file mode 100644 index 000000000000..d1fd6581f561 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Golden fixtures generated from HF remote-code `encoding_k3.py`. + +use std::path::PathBuf; +use std::sync::Arc; + +use expect_test::{expect, expect_file}; +use serde_json::json; +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; +use vllm_tokenizer::Tokenizer; +use vllm_tokenizer::test_utils::TestTokenizer; + +use super::KimiK3ChatRenderer; +use crate::ChatRenderer; +use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +const OPEN_ID: u32 = 256; +const CLOSE_ID: u32 = 257; +const SEP_ID: u32 = 258; +const END_OF_MSG_ID: u32 = 259; +const MEDIA_ID: u32 = 260; + +fn test_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token(CLOSE, CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID) +} + +fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec { + let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt; + let Prompt::TokenIds(token_ids) = prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + token_ids +} + +fn render_request(request: &crate::request::ChatRequest) -> String { + let tokenizer: DynTokenizer = Arc::new(test_tokenizer()); + let token_ids = render_token_ids(request, tokenizer.clone()); + tokenizer.decode(&token_ids, false).unwrap() +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/kimi_k3/fixtures") + .join(name) +} + +fn kimi_k3_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + // Fixture JSON owns thinking via `template_kwargs`. + enable_thinking: None, + no_generation_prompt_when_last_assistant: false, + } +} + +fn assert_golden(name: &str) { + let input_name = format!("{name}_input.json"); + let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + let rendered = render_request(&request); + expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered); +} + +#[test] +fn golden_history_preserve_and_image() { + assert_golden("history_preserve_and_image"); +} + +#[test] +fn golden_tools_history_and_required() { + assert_golden("tools_history_and_required"); +} + +#[test] +fn golden_controls_thinking_off() { + assert_golden("controls_thinking_off"); +} + +#[test] +fn golden_dynamic_system_tool_declare() { + assert_golden("dynamic_system_tool_declare"); +} + +#[test] +fn token_writer_protects_literal_control_and_media_markers() { + let tokenizer = Arc::new(test_tokenizer()); + let user_text = format!("literal {OPEN} and {}", super::encoding::IMAGE_PLACEHOLDER); + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ChatMessage::user(vec![ + ChatContentPart::text(user_text), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let token_ids = render_token_ids(&request, tokenizer.clone()); + + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 1 + ); + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 1 + ); + + let flattened = tokenizer.decode(&token_ids, false).unwrap(); + let flattened_ids = tokenizer.encode(&flattened, false).unwrap(); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 2 + ); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 2 + ); +} + +#[test] +fn thinking_history_renders_empty_think_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>think<|sep|><|close|>think<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn non_thinking_history_omits_reasoning_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { + text: "hidden reasoning".to_string(), + }, + AssistantContentBlock::Text { + text: "answer".to_string(), + }, + ]), + ChatMessage::user("follow-up"), + ]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(!rendered.contains("hidden reasoning")); + assert!(!rendered.contains("<|open|>think<|sep|>")); + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn partial_tool_results_keep_assistant_call_position() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "weather".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "search".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response("result-b", "call-b"), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.contains("<|open|>call tool=\"search\" index=\"2\"<|sep|>")); + assert!( + rendered + .contains("<|open|>message role=\"tool\" tool=\"search\" index=\"2\"<|sep|>result-b") + ); +} + +#[test] +fn defaults_thinking_effort_to_max() { + let rendered = render_request(&crate::request::ChatRequest::for_test()); + + expect![[r#"<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]] + .assert_eq(&rendered); +} + +#[test] +fn translates_standard_thinking_kwargs() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), json!(true)); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("high")); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=high")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); +} + +#[test] +fn native_k3_kwargs_take_precedence() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("thinking".to_string(), json!(true)), + ("enable_thinking".to_string(), json!(false)), + ("thinking_effort".to_string(), json!("low")), + ("reasoning_effort".to_string(), json!("high")), + ]); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=low")); + assert!(!rendered.contains("thinking_effort=high")); +} + +#[test] +fn standard_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("enable_thinking".to_string(), json!(false)), + ("reasoning_effort".to_string(), json!("none")), + ]); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn typed_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn rejects_removed_medium_thinking_effort() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("thinking_effort".to_string(), json!("medium")); + + let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())) + .render(&request) + .unwrap_err(); + + expect![[r#" + ChatTemplate( + "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`", + ) + "#]] + .assert_debug_eq(&error); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index 4d9c1581a01e..a83953f3318f 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -15,6 +15,7 @@ pub mod deepseek_v4; pub mod harmony; pub mod hf; mod inkling; +mod kimi_k3; mod selection; #[cfg(test)] mod test_utils; @@ -23,6 +24,7 @@ pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; pub use harmony::HarmonyChatRenderer; pub use inkling::InklingChatRenderer; +pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 711c0d87ddfb..5b6642c4c034 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -26,6 +26,8 @@ pub enum RendererSelection { Harmony, /// Force the Inkling native token renderer. Inkling, + /// Force the Kimi K3 XTML renderer. + KimiK3, } impl RendererSelection { @@ -37,6 +39,7 @@ impl RendererSelection { pub const HF_LITERAL: &str = "hf"; pub const INKLING_LITERAL: &str = "inkling"; pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model"; + pub const KIMI_K3_LITERAL: &str = "kimi_k3"; /// Resolve the renderer selection using the given model type string, if /// it's `Auto`. @@ -47,6 +50,7 @@ impl RendererSelection { Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, Self::GPT_OSS_MODEL_TYPE => Self::Harmony, Self::INKLING_MODEL_TYPE => Self::Inkling, + Self::KIMI_K3_LITERAL => Self::KimiK3, _ => Self::Hf, }, selection => selection, @@ -70,6 +74,8 @@ impl FromStr for RendererSelection { Ok(Self::Harmony) } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) { Ok(Self::Inkling) + } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) { + Ok(Self::KimiK3) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -88,6 +94,7 @@ impl fmt::Display for RendererSelection { Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), Self::Harmony => f.write_str(Self::HARMONY_LITERAL), Self::Inkling => f.write_str(Self::INKLING_LITERAL), + Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL), } } } @@ -114,7 +121,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 08eeca8176d4..0a4edc813798 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::fs; use std::path::Path; @@ -16,8 +17,9 @@ use crate::request::{ /// Options for constructing a [`ChatRequest`] from a fixture file. #[derive(Debug, Clone, Copy)] pub(crate) struct FixtureRequestOptions { - /// Whether to set the template kwarg `[enable_]thinking=true`. - pub enable_thinking: bool, + /// Optional thinking toggle applied only when the fixture does not already + /// set `thinking` / `enable_thinking` in [`FixtureRequest::template_kwargs`]. + pub enable_thinking: Option, /// Whether fixtures ending in an assistant message should omit the /// trailing generation prompt. pub no_generation_prompt_when_last_assistant: bool, @@ -46,6 +48,15 @@ pub(crate) struct FixtureRequest { messages: Vec, add_generation_prompt: Option, reasoning_effort: Option, + /// Standard response format passed to model-specific renderers. + #[serde(default)] + response_format: Option, + /// Extra chat-template kwargs (thinking, preserve_thinking, …). + #[serde(default)] + template_kwargs: HashMap, + /// When omitted, defaults to `auto` if tools are present, otherwise `none`. + #[serde(default)] + tool_choice: Option, } impl FixtureFile { @@ -57,6 +68,9 @@ impl FixtureFile { messages, add_generation_prompt: None, reasoning_effort: None, + response_format: None, + template_kwargs: HashMap::new(), + tool_choice: None, }, } } @@ -69,6 +83,7 @@ pub(crate) enum FixtureMessage { content: FixtureContent, }, Developer { + #[serde(default)] content: FixtureContent, #[serde(default)] tools: Vec, @@ -98,6 +113,12 @@ pub(crate) enum FixtureContent { Parts(Vec), } +impl Default for FixtureContent { + fn default() -> Self { + Self::Text(String::new()) + } +} + #[derive(Debug, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub(crate) enum FixtureContentPart { @@ -136,6 +157,13 @@ struct FixtureToolCallFunction { impl FixtureRequest { fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { + let tools = to_chat_tools(&self.tools); + let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + }); + let mut request = ChatRequest { request_id: "renderer-fixture".to_string(), messages: self @@ -144,12 +172,8 @@ impl FixtureRequest { .enumerate() .map(|(index, message)| fixture_message_to_chat_message(index, message)) .collect(), - tools: to_chat_tools(&self.tools), - tool_choice: if self.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, + tools, + tool_choice, ..ChatRequest::for_test() }; @@ -162,9 +186,15 @@ impl FixtureRequest { request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; } request.chat_options.reasoning_effort = self.reasoning_effort; - if options.enable_thinking { - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); + request.chat_options.response_format = self.response_format; + request.chat_options.template_kwargs.extend(self.template_kwargs); + + // Options supply a default thinking toggle only when the fixture did not. + if let Some(thinking) = options.enable_thinking { + let kwargs = &mut request.chat_options.template_kwargs; + if !kwargs.contains_key("thinking") && !kwargs.contains_key("enable_thinking") { + kwargs.insert("thinking".to_string(), Value::Bool(thinking)); + kwargs.insert("enable_thinking".to_string(), Value::Bool(thinking)); } } diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index eb72557c7d26..56a764f8987a 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -397,6 +397,10 @@ pub struct ChatOptions { /// Effort level exposed to chat templates for reasoning models. pub reasoning_effort: Option, + /// Standard response format available to model-specific renderers. + #[serde(default)] + pub response_format: Option, + /// Additional keyword arguments exposed to the chat template. pub template_kwargs: HashMap, } @@ -407,6 +411,7 @@ impl Default for ChatOptions { generation_prompt_mode: GenerationPromptMode::StartNewAssistant, chat_template: None, reasoning_effort: None, + response_format: None, template_kwargs: HashMap::new(), } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index ff15aa929ef1..9cd670d29d08 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -212,6 +212,23 @@ impl RoundtripCase { } } + /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser). + /// + /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` + + /// `tokenizer_config.json`). Weights are not required for this text-level + /// roundtrip. + fn kimi_k3() -> Self { + Self { + model_id: "moonshotai/Kimi-K3", + assistant_stop_suffix: "<|end_of_msg|>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// SeedOSS with `` / `` reasoning tags. fn seed_oss() -> Self { Self { @@ -255,7 +272,7 @@ impl RoundtripCase { fn gpt_oss() -> Self { Self { model_id: "openai/gpt-oss-20b", - assistant_stop_suffix: "", // not applicable for token-id cases + assistant_stop_suffix: "", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -268,7 +285,7 @@ impl RoundtripCase { fn inkling() -> Self { Self { model_id: "thinkingmachines/Inkling", - assistant_stop_suffix: "", + assistant_stop_suffix: "<|content_model_end_sampling|>", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -312,6 +329,8 @@ roundtrip_tests! { nemotron_v3 => [reasoning_and_content], gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history + // K3 drops plain-assistant reasoning in history; tool-call turns keep it. + kimi_k3 => [tool_call_mix], gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call inkling => [reasoning_and_content, tool_call_mix], } @@ -670,14 +689,23 @@ fn decoded_completion_stream( .collect() } Prompt::TokenIds(token_ids) => { - ensure!( - assistant_stop_suffix.is_empty(), - "token-id roundtrip cases do not support text stop suffixes" - ); + let body = if assistant_stop_suffix.is_empty() { + token_ids.as_slice() + } else { + let stop_token_ids = tokenizer + .encode(assistant_stop_suffix, false) + .context("failed to encode token-id completion stop suffix")?; + token_ids.strip_suffix(stop_token_ids.as_slice()).with_context(|| { + format!( + "token-id completion did not end with {:?}: {:?}", + assistant_stop_suffix, token_ids + ) + })? + }; incremental_decode_chunks( tokenizer, &prompt_token_ids, - token_ids, + body, TOKEN_COMPLETION_CHUNK_TOKENS, )? } diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index 030d4c6d1164..a326a0f9992a 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -23,16 +23,16 @@ serde.workspace = true serde_json.workspace = true serde_with.workspace = true thiserror-ext.workspace = true -time.workspace = true tokio = { workspace = true, features = ["signal"] } tokio-util.workspace = true tracing.workspace = true -tracing-subscriber.workspace = true uuid.workspace = true +vllm-bench.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true vllm-server.workspace = true +vllm-tracing.workspace = true [dev-dependencies] expect-test.workspace = true diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 314765bab26c..1572e20467ee 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -23,6 +23,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; use uuid::Uuid; use vllm_chat::ReasoningParserFactory; +use vllm_chat::multimodal::MmLimitPerPrompt; use vllm_engine_core_client::TransportMode; use vllm_managed_engine::ManagedEngineConfig; use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args}; @@ -79,13 +80,23 @@ impl Cli { } /// Supported top-level CLI commands. -#[derive(Debug, Subcommand, PartialEq, Eq)] +#[derive(Debug, Subcommand)] pub enum Command { /// Run the Rust OpenAI frontend as a Python-supervised worker. Frontend(FrontendArgs), /// Launch a managed Python headless engine, then run the Rust OpenAI /// frontend. Serve(ServeArgs), + /// Run vLLM benchmarks. + #[command(subcommand)] + Bench(BenchCommand), +} + +/// Supported benchmark commands. +#[derive(Debug, Subcommand)] +pub enum BenchCommand { + /// Benchmark online serving throughput. + Serve(vllm_bench::BenchServeArgs), } /// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for @@ -138,11 +149,6 @@ pub struct SharedRuntimeArgs { #[arg(long)] #[serde(default)] pub language_model_only: bool, - /// Override the maximum model context length. When set, the frontend uses - /// this value instead of the model's `max_position_embeddings` from - /// `config.json`. - #[arg(long)] - pub max_model_len: Option, /// Maximum number of log probabilities to return when `logprobs` is /// specified in sampling parameters. `-1` means no cap. #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] @@ -181,6 +187,17 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub default_chat_template_kwargs: Option>, + /// The maximum number of input items allowed per prompt for each + /// modality, as a JSON object (e.g. `{"image": 16, "video": 2}`). + /// + /// Also accepts the engine's configurable form + /// (e.g. `{"video": {"count": 1, "num_frames": 32}}`); the extra + /// profiling options are forwarded to the engine untouched. + /// Unspecified modalities are unlimited. + #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = "{}")] + #[serde(default)] + pub limit_mm_per_prompt: MmLimitPerPrompt, + /// The format to render message content within a chat template. /// /// * "auto" detects the format from the template @@ -343,6 +360,18 @@ impl SharedRuntimeArgs { .expect("profiler config serialization should not fail") } + /// Return the per-modality limits as JSON for managed Python engine + /// forwarding, or `None` when nothing is configured. + /// + /// Round-tripping the parsed map rather than the raw argument keeps the + /// engine's own profiling options (`num_frames`, `width`, ...) intact. + pub fn limit_mm_per_prompt_json(&self) -> Option { + (!self.limit_mm_per_prompt.is_empty()).then(|| { + serde_json::to_string(&self.limit_mm_per_prompt) + .expect("limit-mm-per-prompt serialization should not fail") + }) + } + /// Apply fallback logic for API key configuration from env variables. fn apply_env_api_key_fallback(&mut self) { if self.api_key.is_empty() @@ -394,6 +423,7 @@ impl SharedRuntimeArgs { language_model_only: self.language_model_only, chat_template: self.chat_template, default_chat_template_kwargs: self.default_chat_template_kwargs, + limit_mm_per_prompt: self.limit_mm_per_prompt, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, api_server_options, @@ -446,6 +476,7 @@ impl SharedRuntimeArgs { language_model_only: self.language_model_only, chat_template: self.chat_template, default_chat_template_kwargs: self.default_chat_template_kwargs, + limit_mm_per_prompt: self.limit_mm_per_prompt, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, api_server_options, @@ -654,7 +685,6 @@ impl ServeArgs { self.managed_engine.clone().into_config( self.runtime.model.clone(), - self.runtime.max_model_len, self.runtime.max_logprobs, profiler_config, reasoning_parser.as_deref(), @@ -662,6 +692,7 @@ impl ServeArgs { self.runtime.disable_log_stats, self.runtime.shutdown_timeout, handshake_port, + self.runtime.limit_mm_per_prompt_json(), ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index ee0d3ebf0486..e586477a897e 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -5,7 +5,27 @@ use expect_test::expect; use vllm_engine_core_client::TransportMode; use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection}; -use super::{Cli, Command}; +use super::{BenchCommand, Cli, Command}; + +#[test] +fn bench_serve_args_parse_without_managed_engine_repartition() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "bench", + "serve", + "--backend", + "openai-chat", + "--request-rate", + "inf", + ]) + .unwrap(); + + let Command::Bench(BenchCommand::Serve(args)) = cli.command else { + panic!("expected bench serve args"); + }; + assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat); + assert!(args.request_rate.is_infinite()); +} #[test] fn serve_args_forward_python_flags_with_separator() { @@ -38,15 +58,13 @@ fn serve_args_forward_python_flags_with_separator() { reasoning_parser: Auto, renderer: Auto, language_model_only: false, - max_model_len: Some( - 512, - ), max_logprobs: None, grpc_port: None, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -82,6 +100,9 @@ fn serve_args_forward_python_flags_with_separator() { handshake_port: None, data_parallel_size: 1, data_parallel_size_local: None, + max_model_len: Some( + "512", + ), python_args: [ "--dtype", "float16", @@ -646,7 +667,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3) For more information, try '--help'. "#]] @@ -736,13 +757,13 @@ fn frontend_args_accept_json() { reasoning_parser: None, renderer: Auto, language_model_only: false, - max_model_len: None, max_logprobs: None, grpc_port: None, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -803,11 +824,32 @@ fn frontend_args_json_applies_defaults() { assert_eq!(args.runtime.tool_call_parser, ParserSelection::None); assert_eq!(args.runtime.reasoning_parser, ParserSelection::None); assert_eq!(args.runtime.renderer, RendererSelection::Auto); - assert_eq!(args.runtime.max_model_len, None); assert_eq!(args.runtime.max_logprobs, None); assert_eq!(args.runtime.shutdown_timeout, 0); } +#[test] +fn frontend_args_json_ignores_engine_owned_max_model_len() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B","max_model_len":-1}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); +} + #[test] fn frontend_args_json_accepts_supported_non_default_fields() { let cli = Cli::try_parse_from([ @@ -820,7 +862,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "--output-address", "ipc:///tmp/output.sock", "--args-json", - r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"max_logprobs":-1,"shutdown_timeout":3}"#, + r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) .unwrap(); @@ -838,11 +880,38 @@ fn frontend_args_json_accepts_supported_non_default_fields() { ); assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32); assert!(args.runtime.language_model_only); - assert_eq!(args.runtime.max_model_len, Some(8192)); assert_eq!(args.runtime.max_logprobs, Some(-1)); assert_eq!(args.runtime.shutdown_timeout, 3); } +#[test] +fn serve_args_forward_auto_max_model_len_to_managed_engine() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--max-model-len", + "auto", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + assert_eq!(args.managed_engine.max_model_len.as_deref(), Some("auto")); + + let config = args.to_managed_engine_config(5555); + expect![[r#" + [ + "--max-model-len", + "auto", + "--reasoning-parser", + "qwen3", + ] + "#]] + .assert_debug_eq(&config.python_args); +} + #[test] fn serve_args_accept_none_reasoning_parser() { let cli = Cli::try_parse_from([ @@ -1050,6 +1119,24 @@ fn frontend_args_json_rejects_malformed_json() { "#]].assert_eq(&error.to_string()); } +#[test] +fn serve_args_reject_unsupported_modality_in_limit_mm_per_prompt() { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--limit-mm-per-prompt", + r#"{"unsupported_modality": 1}"#, + ]) + .unwrap_err(); + + expect![[r#" + error: invalid value '{"unsupported_modality": 1}' for '--limit-mm-per-prompt ': invalid JSON object: unknown variant `unsupported_modality`, expected one of `image`, `audio`, `video` at line 1 column 23 + + For more information, try '--help'. + "#]].assert_eq(&error.to_string()); +} + #[test] fn serve_args_reject_flags_before_model() { let error = Cli::try_parse_from(["vllm-rs", "serve", "--python", "python3", "Qwen/Qwen3-0.6B"]) @@ -1258,13 +1345,13 @@ fn serve_args_accept_handshake_aliases() { reasoning_parser: Auto, renderer: Auto, language_model_only: false, - max_model_len: None, max_logprobs: None, grpc_port: None, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1302,6 +1389,7 @@ fn serve_args_accept_handshake_aliases() { ), data_parallel_size: 4, data_parallel_size_local: None, + max_model_len: None, python_args: [], }, }, @@ -1407,6 +1495,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { @@ -1491,6 +1580,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { @@ -1593,6 +1683,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 66e8b2e9abb4..a48e61dbf7a6 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -299,11 +299,6 @@ pub struct EngineUnsupportedArgs { )] pub kv_sharing_fast_prefill: Option, - /// The maximum number of input items and options allowed per - /// prompt for each modality. - #[arg(long)] - pub limit_mm_per_prompt: Option, - /// Additional args passed to process media inputs, keyed by modalities. #[arg(long)] pub media_io_kwargs: Option, diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index 9637a3be5780..4e92ece70900 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -2,17 +2,18 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project mod cli; -mod logging; use std::env; -use std::process::ExitStatus; +use std::ffi::OsStr; +use std::process::{ExitCode, ExitStatus}; use anyhow::{Context, Result, anyhow, bail}; +use thiserror_ext::AsReport as _; use tokio_util::sync::CancellationToken; -use tracing::{info, warn}; +use tracing::{error, info, warn}; use vllm_managed_engine::ManagedEngineHandle; -use crate::cli::{Cli, Command}; +use crate::cli::{BenchCommand, Cli, Command}; #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; @@ -81,8 +82,15 @@ fn shutdown_signal() -> CancellationToken { token } -fn main() -> Result<()> { - logging::init_tracing(); +fn main() -> ExitCode { + let process_label = + match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() { + "bench" => "Bench", + "serve" | "frontend" => "RustFrontend", + _ => "Rust", + }; + vllm_tracing::init_tracing(process_label); + let cli = Cli::parse(); let mut runtime = tokio::runtime::Builder::new_multi_thread(); @@ -91,15 +99,27 @@ fn main() -> Result<()> { runtime.worker_threads(worker_threads); } - runtime + let result = runtime .build() - .context("failed to build Tokio runtime")? - .block_on(async_main(cli)) + .context("failed to build Tokio runtime") + .and_then(|runtime| runtime.block_on(async_main(cli))); + + match result { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + error!("process failed with error: {:#?}", error.as_report()); + ExitCode::FAILURE + } + } } async fn async_main(cli: Cli) -> Result<()> { match cli.command { Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await, + Command::Bench(BenchCommand::Serve(bench_args)) => { + vllm_bench::prepare_process(); + vllm_bench::run(bench_args).await + } Command::Serve(args) => { let handshake_port = args.managed_engine.resolve_handshake_port()?; diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 5b9cd500d399..a5709e52d67f 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -394,6 +394,17 @@ impl EngineCoreClient { self.engines.iter().map(|engine| &engine.ready_response).collect() } + /// Return the first engine's ready response. + /// + /// Per-engine fields such as `data_parallel_rank` should be read through + /// [`ready_responses`](Self::ready_responses). + pub fn ready_response(&self) -> &EngineCoreReadyResponse { + &self + .engines + .first() + .expect("engine core client requires at least one engine") + .ready_response + } /// Return the engine-reported effective model dtype. pub fn model_dtype(&self) -> ModelDtype { self.engines @@ -460,6 +471,12 @@ impl EngineCoreClient { self.inner.is_healthy() } + /// Subscribe to engine health changes. The current value is `true` while + /// the client is healthy and changes permanently to `false` on failure. + pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver { + self.inner.subscribe_health() + } + /// Return the first persistent health error observed by the client, if any. pub fn health_error(&self) -> Option> { self.inner.health_error() @@ -502,7 +519,7 @@ impl EngineCoreClient { "registered request to engine" ); - self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?; + self.inner.send_request_to_engine(&engine_id, req).await?; Ok(()) } .await; diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index a699f85c26ae..9c539bf180e9 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -6,10 +6,11 @@ use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use arc_swap::ArcSwapOption; +use bytes::Bytes; use parking_lot::Mutex; use thiserror_ext::AsReport as _; use tokio::runtime::Handle; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tracing::{debug, info, trace, warn}; use vllm_metrics::METRICS; use zeromq::RouterSendHalf; @@ -21,21 +22,35 @@ use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_outpu use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; -use crate::protocol::request::EngineCoreRequestType; +use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use crate::protocol::stats::SchedulerStats; use crate::protocol::utility::UtilityOutput; use crate::transport::{ConnectedEngine, EngineId}; use crate::{Error, Result, transport}; +const MSGPACK_ZERO_COPY_THRESHOLD_ENV: &str = "VLLM_MSGPACK_ZERO_COPY_THRESHOLD"; +const DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD: usize = 256; + +fn msgpack_zero_copy_threshold() -> usize { + std::env::var(MSGPACK_ZERO_COPY_THRESHOLD_ENV) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD) +} + pub(crate) struct ClientInner { input_send: RouterSendHalf, /// The runtime handle used for sending messages to the engine. handle: Handle, model_name: String, + /// Per-tensor byte threshold loaded from env variable + /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created. + msgpack_zero_copy_threshold: usize, scheduler_stats_recorder: SchedulerStatsRecorder, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, + health_tx: watch::Sender, } impl ClientInner { @@ -53,10 +68,12 @@ impl ClientInner { input_send, handle, model_name, + msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(), scheduler_stats_recorder, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), + health_tx: watch::Sender::new(true), } } @@ -169,6 +186,7 @@ impl ClientInner { /// persistent health error. pub fn close_registries(&self, error: Arc) { let persistent_error = self.record_health_error(error); + self.publish_unhealthy(); let request_senders = self.request_reg.lock().close(); let utility_senders = self.utility_reg.lock().close(); @@ -191,6 +209,12 @@ impl ClientInner { self.health_error.load().is_none() } + /// Subscribe to engine health changes. The current value is `true` while + /// the client is healthy and changes permanently to `false` on failure. + pub fn subscribe_health(&self) -> watch::Receiver { + self.health_tx.subscribe() + } + /// Resolve one utility output to the waiting caller. Returns `true` if a /// waiting caller existed. pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool { @@ -220,9 +244,29 @@ impl ClientInner { where T: serde::Serialize + std::fmt::Debug, { - // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux - // frames instead of always producing a single msgpack frame. - let payload = encode_msgpack(payload)?; + let payload = Bytes::from(encode_msgpack(payload)?); + self.send_encoded_to_engine(engine_id, request_type, payload, Vec::new()).await + } + + /// Send an add request, moving large tensor buffers into auxiliary frames. + pub async fn send_request_to_engine( + &self, + engine_id: &EngineId, + mut payload: EngineCoreRequest, + ) -> Result<()> { + let aux_frames = payload.extract_aux_frames(self.msgpack_zero_copy_threshold); + let payload = Bytes::from(encode_msgpack(&payload)?); + self.send_encoded_to_engine(engine_id, EngineCoreRequestType::Add, payload, aux_frames) + .await + } + + async fn send_encoded_to_engine( + &self, + engine_id: &EngineId, + request_type: EngineCoreRequestType, + payload: Bytes, + aux_frames: Vec, + ) -> Result<()> { let mut input_send = self.input_send.clone(); let engine_id = engine_id.clone(); @@ -233,6 +277,7 @@ impl ClientInner { &engine_id, request_type.to_frame(), payload, + aux_frames, ) .await }) @@ -280,6 +325,11 @@ impl ClientInner { .expect("health error must be recorded before registries close") } + /// Publish the sticky healthy-to-unhealthy transition. + fn publish_unhealthy(&self) { + self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false)); + } + /// Assert there is a recorded health error and return a `Shared` variant /// wrapping it for error returns when the client is already closed. fn closed_error(&self) -> Error { @@ -461,13 +511,18 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_registries_records_first_health_error_only() { let inner = test_inner().await; + let mut health = inner.subscribe_health(); + assert!(*health.borrow()); inner.close_registries(Arc::new(Error::EngineCoreDead)); + health.changed().await.expect("health sender remains open"); assert!(!inner.is_healthy()); + assert!(!*health.borrow()); assert!(matches!( inner.health_error().as_deref(), Some(Error::EngineCoreDead) )); + assert!(!*inner.subscribe_health().borrow()); inner.close_registries(Arc::new(client_closed!("shutdown"))); assert!(matches!( diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index fad117ac31d8..28e701ba7320 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -74,17 +74,13 @@ impl EngineRoutingState { /// /// Scheduler stats can raise the load estimate above the frontend-local /// view, but they should not lower it below requests this frontend has - /// already admitted. Waiting requests still get the same extra penalty - /// as the original `waiting * 4 + running` score. + /// already admitted. fn routing_score(&self) -> usize { - const WAITING_WEIGHT: usize = 4; - let Some(stats) = self.last_scheduler_stats else { return self.inflight; }; - let scheduler_total = stats.running + stats.waiting; - self.inflight.max(scheduler_total) + stats.waiting * (WAITING_WEIGHT - 1) + self.inflight.max(stats.running + stats.waiting) } /// Replace the local routing view with a fresh real scheduler snapshot. @@ -750,7 +746,7 @@ mod tests { } #[test] - fn routing_score_keeps_extra_waiting_penalty() { + fn routing_score_counts_waiting_without_extra_penalty() { let state = EngineRoutingState { inflight: 1, last_scheduler_stats: Some(EngineLoadSnapshot { @@ -759,7 +755,7 @@ mod tests { }), }; - assert_eq!(state.routing_score(), 14); + assert_eq!(state.routing_score(), 5); } #[test] diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index e8f277b128d0..ee066635ff0c 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -57,8 +57,16 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { vllm_version: "test-vllm-version".to_string(), world_size: 1, data_parallel_size: 1, + tensor_parallel_size: 1, + pipeline_parallel_size: 1, + decode_context_parallel_size: 1, + data_parallel_rank: 0, + max_num_seqs: 256, + max_num_batched_tokens: 8192, + instance_id: "test-instance".to_string(), kv_cache_size_tokens: None, kv_cache_max_concurrency: None, + kv_events_config: None, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index c8545017a96b..a4e1cebde67d 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -24,6 +24,19 @@ pub struct ReadyMessage { pub parallel_config_hash: Option, } +/// KV-event publisher configuration reported by EngineCore. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KvEventsConfig { + pub enable_kv_cache_events: bool, + pub publisher: String, + pub endpoint: String, + pub replay_endpoint: Option, + pub buffer_steps: u32, + pub hwm: u32, + pub max_queue_size: u32, + pub topic: String, +} + /// Post-initialization configuration sent from each engine on the input socket /// registration message, after the handshake completes. /// @@ -52,10 +65,28 @@ pub struct EngineCoreReadyResponse { pub world_size: u64, /// Data parallelism size from the parallel config. pub data_parallel_size: u64, + // Required discovery metadata; EngineCore and client versions must match. + /// Tensor-parallel size of this engine. + pub tensor_parallel_size: u32, + /// Pipeline-parallel size of this engine. + pub pipeline_parallel_size: u32, + /// Decode-context-parallel size of this engine. + pub decode_context_parallel_size: u32, + /// This engine's data-parallel rank. + pub data_parallel_rank: u32, + /// Scheduler cap on concurrently running sequences. + pub max_num_seqs: u64, + /// Scheduler cap on batched tokens per step. + pub max_num_batched_tokens: u64, + /// Unique identifier for this server instance. + pub instance_id: String, /// Total KV cache capacity in tokens, if reported. pub kv_cache_size_tokens: Option, /// Maximum achievable request concurrency given the KV cache, if reported. pub kv_cache_max_concurrency: Option, + /// KV-event publisher configuration, if configured. + #[serde(default)] + pub kv_events_config: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index a01186d0dcb9..28f6327d9b8a 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -209,17 +209,17 @@ impl WireLogprobs { logprob_token_ids: WireNdArray { dtype: "( field: &str, frames: &[Frame], expected_scalars: &[ScalarType], -) -> Result<(Vec, Vec, ScalarType, Endianness)> +) -> Result<(Vec, Bytes, ScalarType, Endianness)> where Frame: AsRef<[u8]>, { @@ -171,7 +172,7 @@ pub(super) fn resolve_array_bytes( value: WireArrayData, field: &str, frames: &[Frame], -) -> Result> +) -> Result where Frame: AsRef<[u8]>, { @@ -187,7 +188,7 @@ where ), ) })?; - Ok(frame.as_ref().to_vec()) + Ok(Bytes::copy_from_slice(frame.as_ref())) } } } diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index b105ea8d4a08..42f7ec78b599 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() { "messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]" ); } + +#[test] +fn decodes_zero_row_logprobs_as_empty() { + for shape in [[0usize, 0], [0, 3]] { + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields( + None, + Some(Value::Array(vec![ + ndarray_value("), } +impl MmFeatureSpec { + /// Extract large tensor buffers from this feature in serialized field order. + pub(crate) fn extract_aux_frames(&mut self, aux_frames: &mut Vec, threshold: usize) { + if let Some(data) = &mut self.data { + for elem in data.values_mut() { + elem.extract_aux_frames(aux_frames, threshold); + } + } + if let Some(is_embed) = &mut self.mm_position.is_embed { + is_embed.extract_aux_frame(aux_frames, threshold); + } + } +} + +impl MmFieldElem { + /// Extract large tensor buffers from this field element. + fn extract_aux_frames(&mut self, aux_frames: &mut Vec, threshold: usize) { + if let Some(data) = &mut self.data { + data.extract_aux_frames(aux_frames, threshold); + } + } +} + +impl MmKwargValue { + /// Recursively extract large tensor buffers from this nested value. + fn extract_aux_frames(&mut self, aux_frames: &mut Vec, threshold: usize) { + match self { + Self::Tensor(tensor) => tensor.extract_aux_frame(aux_frames, threshold), + Self::List(values) => { + for value in values { + value.extract_aux_frames(aux_frames, threshold); + } + } + Self::Int(_) | Self::Float(_) => {} + } + } +} + /// Defines how to interpret tensor data belonging to a keyword argument for /// `MultiModalKwargsItems`, and vice versa. /// diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs index ee1a6ac0b2ec..99e1b588e175 100644 --- a/rust/src/engine-core-client/src/protocol/request.rs +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -137,6 +137,17 @@ impl EngineCoreRequest { } Ok(()) } + + /// Extract large request tensors into ordered auxiliary frames. + pub(crate) fn extract_aux_frames(&mut self, threshold: usize) -> Vec { + let mut aux_frames = Vec::new(); + if let Some(features) = &mut self.mm_features { + for feature in features { + feature.extract_aux_frames(&mut aux_frames, threshold); + } + } + aux_frames + } } #[cfg(test)] @@ -144,9 +155,15 @@ mod tests { use rmpv::Value; use super::*; + use crate::protocol::multimodal::{ + MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange, + }; use crate::protocol::sampling::EngineCoreSamplingParams; + use crate::protocol::tensor::{WireArrayData, WireTensor}; use crate::protocol::{decode_value, encode_msgpack}; + const AUX_FRAME_THRESHOLD: usize = 256; + #[test] fn engine_core_request_serializes_as_full_array() { let request = EngineCoreRequest { @@ -175,4 +192,84 @@ mod tests { assert_eq!(array[10], Value::Nil); assert_eq!(array[11], Value::from(7)); } + + #[test] + fn engine_core_request_extracts_large_nested_tensors_in_wire_order() { + let inline = vec![1_u8; AUX_FRAME_THRESHOLD - 1]; + let first_aux = vec![2_u8; AUX_FRAME_THRESHOLD]; + let second_aux = vec![3_u8; AUX_FRAME_THRESHOLD + 1]; + let first_aux_ptr = first_aux.as_ptr(); + let second_aux_ptr = second_aux.as_ptr(); + let mut request = EngineCoreRequest { + mm_features: Some(vec![MmFeatureSpec { + data: Some(BTreeMap::from([ + ( + "inline".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(WireTensor::from_raw( + "uint8", + vec![inline.len()], + inline, + ))), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + ), + ( + "nested".to_string(), + MmFieldElem { + data: Some(MmKwargValue::List(vec![ + MmKwargValue::Int(7), + MmKwargValue::Tensor(WireTensor::from_raw( + "uint8", + vec![first_aux.len()], + first_aux, + )), + ])), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + ), + ])), + modality: "image".to_string(), + identifier: "id".to_string(), + mm_position: PlaceholderRange { + offset: 0, + length: second_aux.len(), + is_embed: Some(WireTensor::from_raw( + "bool", + vec![second_aux.len()], + second_aux, + )), + }, + mm_hash: None, + }]), + ..EngineCoreRequest::default() + }; + + let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD); + + assert_eq!(aux_frames.len(), 2); + assert_eq!(aux_frames[0].as_ptr(), first_aux_ptr); + assert_eq!(aux_frames[1].as_ptr(), second_aux_ptr); + let feature = &request.mm_features.as_ref().unwrap()[0]; + let MmKwargValue::Tensor(inline) = + feature.data.as_ref().unwrap()["inline"].data.as_ref().unwrap() + else { + panic!("expected inline tensor"); + }; + assert!(matches!(inline.data, WireArrayData::RawView(_))); + let MmKwargValue::List(nested) = + feature.data.as_ref().unwrap()["nested"].data.as_ref().unwrap() + else { + panic!("expected nested tensor list"); + }; + let MmKwargValue::Tensor(nested_tensor) = &nested[1] else { + panic!("expected nested tensor"); + }; + assert_eq!(nested_tensor.data, WireArrayData::AuxIndex(1)); + assert_eq!( + feature.mm_position.is_embed.as_ref().unwrap().data, + WireArrayData::AuxIndex(2) + ); + assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty()); + } } diff --git a/rust/src/engine-core-client/src/protocol/tensor.rs b/rust/src/engine-core-client/src/protocol/tensor.rs index 5eb1668e7528..01540dab34b2 100644 --- a/rust/src/engine-core-client/src/protocol/tensor.rs +++ b/rust/src/engine-core-client/src/protocol/tensor.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use bytemuck::allocation::pod_collect_to_vec; +use bytemuck::{Pod, cast_slice}; +use bytes::Bytes; use enum_as_inner::EnumAsInner; use half::{bf16, f16}; use rmpv::Value; @@ -20,6 +21,21 @@ struct MsgpackExtRef<'a>((i8, ByteSlice<'a>)); struct ByteSlice<'a>(&'a [u8]); +struct PodVec(Vec); + +impl AsRef<[u8]> for PodVec { + fn as_ref(&self) -> &[u8] { + cast_slice(&self.0) + } +} + +fn bytes_from_pod_vec(data: Vec) -> Bytes +where + T: Pod + Send + 'static, +{ + Bytes::from_owner(PodVec(data)) +} + impl Serialize for ByteSlice<'_> { fn serialize(&self, serializer: S) -> std::result::Result where @@ -55,53 +71,63 @@ pub struct WireNdArray { impl WireNdArray { /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. pub fn from_f32(shape: Vec, data: Vec) -> Result { validate_element_count(&shape, data.len())?; - Ok(Self { - dtype: "float32".to_string(), + Ok(Self::from_raw_bytes( + "float32", shape, - data: WireArrayData::RawView(pod_collect_to_vec::(&data)), - }) + bytes_from_pod_vec(data), + )) } /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. pub fn from_f16(shape: Vec, data: Vec) -> Result { validate_element_count(&shape, data.len())?; - Ok(Self { - dtype: "float16".to_string(), + Ok(Self::from_raw_bytes( + "float16", shape, - data: WireArrayData::RawView(pod_collect_to_vec::(&data)), - }) + bytes_from_pod_vec(data), + )) } /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. pub fn from_bf16(shape: Vec, data: Vec) -> Result { validate_element_count(&shape, data.len())?; - Ok(Self { - dtype: "bfloat16".to_string(), + Ok(Self::from_raw_bytes( + "bfloat16", shape, - data: WireArrayData::RawView(pod_collect_to_vec::(&data)), - }) + bytes_from_pod_vec(data), + )) } /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. pub fn from_i64(shape: Vec, data: Vec) -> Result { validate_element_count(&shape, data.len())?; - Ok(Self { - dtype: "int64".to_string(), + Ok(Self::from_raw_bytes( + "int64", shape, - data: WireArrayData::RawView(pod_collect_to_vec::(&data)), - }) + bytes_from_pod_vec(data), + )) } /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. pub fn from_u32(shape: Vec, data: Vec) -> Result { validate_element_count(&shape, data.len())?; - Ok(Self { - dtype: "uint32".to_string(), + Ok(Self::from_raw_bytes( + "uint32", shape, - data: WireArrayData::RawView(pod_collect_to_vec::(&data)), - }) + bytes_from_pod_vec(data), + )) } /// Build a bool tensor/ndarray backed by raw-view bytes. @@ -113,7 +139,9 @@ impl WireNdArray { Ok(Self { dtype: "bool".to_string(), shape, - data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()), + data: WireArrayData::RawView(Bytes::from( + data.into_iter().map(u8::from).collect::>(), + )), }) } @@ -122,12 +150,22 @@ impl WireNdArray { /// Use this as an escape hatch when the caller already owns bytes that /// match the requested `dtype` and `shape`. pub fn from_raw(dtype: impl Into, shape: Vec, data: Vec) -> Self { + Self::from_raw_bytes(dtype, shape, Bytes::from(data)) + } + + /// Build a tensor/ndarray from an owned immutable raw-view buffer. + pub fn from_raw_bytes(dtype: impl Into, shape: Vec, data: Bytes) -> Self { Self { dtype: dtype.into(), shape, data: WireArrayData::RawView(data), } } + + /// Move a sufficiently large inline buffer into the ordered auxiliary-frame list. + pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec, threshold: usize) { + self.data.extract_aux_frame(aux_frames, threshold); + } } /// Validate that the number of elements implied by the shape matches the length @@ -165,7 +203,25 @@ pub enum WireArrayData { /// stored. AuxIndex(usize), /// The raw bytes of this array/tensor. - RawView(Vec), + RawView(Bytes), +} + +impl WireArrayData { + /// Replace a sufficiently large raw view with its one-based auxiliary-frame index. + fn extract_aux_frame(&mut self, aux_frames: &mut Vec, threshold: usize) { + let Self::RawView(bytes) = self else { + return; + }; + if bytes.len() < threshold { + return; + } + + let index = aux_frames.len() + 1; + let bytes = std::mem::replace(self, Self::AuxIndex(index)) + .into_raw_view() + .expect("raw view was matched above"); + aux_frames.push(bytes); + } } impl<'de> Deserialize<'de> for WireArrayData { @@ -175,7 +231,9 @@ impl<'de> Deserialize<'de> for WireArrayData { { let value = Value::deserialize(deserializer)?; match value { - Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)), + Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => { + Ok(Self::RawView(Bytes::from(bytes))) + } Value::Ext(tag, _) => Err(serde::de::Error::custom(format!( "unsupported extension type code {tag}" ))), @@ -196,9 +254,6 @@ impl Serialize for WireArrayData { where S: Serializer, { - // TODO: outbound request serialization currently only supports inline - // raw-view bytes. Emitting aux frames needs transport-level plumbing; - // serializing `AuxIndex` here only preserves an already-built reference. match self { Self::AuxIndex(index) => serializer.serialize_u64(*index as u64), Self::RawView(bytes) => { @@ -216,7 +271,7 @@ mod tests { 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"); + rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone().into())).expect("encode"); let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone())) .expect("encode expected"); @@ -229,11 +284,15 @@ mod tests { #[test] fn constructors_build_raw_view_tensors() { - let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap(); + let f32_data = vec![1.0, 2.5]; + let f32_data_ptr = f32_data.as_ptr().cast::(); + let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap(); assert_eq!(f32_tensor.dtype, "float32"); assert_eq!(f32_tensor.shape, vec![2]); + let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view"); + assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr); assert_eq!( - f32_tensor.data.into_raw_view().expect("raw view"), + f32_raw_view, [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::>() ); @@ -253,15 +312,15 @@ mod tests { let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap(); assert_eq!(i64_tensor.dtype, "int64"); assert_eq!( - i64_tensor.data.into_raw_view().expect("raw view"), - (-7_i64).to_ne_bytes() + i64_tensor.data.into_raw_view().expect("raw view").as_ref(), + (-7_i64).to_ne_bytes().as_ref() ); let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap(); assert_eq!(u32_tensor.dtype, "uint32"); assert_eq!( - u32_tensor.data.into_raw_view().expect("raw view"), - 42_u32.to_ne_bytes() + u32_tensor.data.into_raw_view().expect("raw view").as_ref(), + 42_u32.to_ne_bytes().as_ref() ); let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap(); diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 64b80185c32c..11d08db646ba 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -19,7 +19,7 @@ use zeromq::prelude::{Socket, SocketRecv, SocketSend}; use zeromq::util::PeerIdentity; use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage}; -use crate::protocol::handshake::{HandshakeInitMessage, ReadyMessage}; +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; use crate::protocol::logprobs::MaybeWireLogprobs; use crate::protocol::multimodal::{ MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange, @@ -32,7 +32,7 @@ use crate::protocol::output::{ use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::stats::SchedulerStats; -use crate::protocol::tensor::WireTensor; +use crate::protocol::tensor::{WireArrayData, WireTensor}; use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; use crate::test_utils::{ IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, @@ -1728,6 +1728,86 @@ async fn client_decodes_multipart_logprob_outputs() { client.shutdown().await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn client_sends_large_multimodal_tensor_as_aux_frame() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-multimodal-aux".to_vec(); + let tensor_data = (0..64).map(|value| value as f32).collect::>(); + let expected_bytes = + tensor_data.iter().flat_map(|value| value.to_ne_bytes()).collect::>(); + let mut request = sample_multimodal_request(); + request.mm_features.as_mut().unwrap()[0] + .data + .as_mut() + .unwrap() + .get_mut("pixel_values") + .unwrap() + .data = Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![64], tensor_data).unwrap(), + )); + + let (shutdown_tx, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + move |dealer, push| { + Box::pin(async move { + let add = recv_engine_message(dealer).await; + assert_eq!(add.len(), 3); + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + let MmKwargValue::Tensor(tensor) = + request.mm_features.as_ref().unwrap()[0].data.as_ref().unwrap()["pixel_values"] + .data + .as_ref() + .unwrap() + else { + panic!("expected tensor"); + }; + assert_eq!(tensor.data, WireArrayData::AuxIndex(1)); + assert_eq!(add[2].as_ref(), expected_bytes); + + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output( + "req-mm", + vec![], + Some(EngineCoreFinishReason::Length), + )], + finished_requests: Some(BTreeSet::from(["req-mm".to_string()])), + ..Default::default() + } + .into(), + ) + .await; + }) + }, + ); + + let client = connect_client_with_ipc( + handshake_test_config( + handshake_address, + 1, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await; + + let outputs = client.call(request).await.unwrap().collect::>().await; + assert_eq!(outputs.len(), 1); + assert!(outputs[0].is_ok()); + + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + client.shutdown().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { init_tracing(); @@ -2597,6 +2677,21 @@ fn python_msgpack_fixtures_match_rust_encoding() { rust_ready_keys, python_ready_keys, "EngineCoreReadyResponse drifted from the Python dataclass", ); + + let ready_response: EngineCoreReadyResponse = + rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); + assert!(kv_events_config.enable_kv_cache_events); + assert_eq!(kv_events_config.publisher, "zmq"); + assert_eq!(kv_events_config.endpoint, "tcp://127.0.0.1:5557"); + assert_eq!( + kv_events_config.replay_endpoint.as_deref(), + Some("tcp://127.0.0.1:5558") + ); + assert_eq!(kv_events_config.topic, "kv"); + assert_eq!(kv_events_config.buffer_steps, 10_000); + assert_eq!(kv_events_config.hwm, 100_000); + assert_eq!(kv_events_config.max_queue_size, 100_000); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 0005ff20f920..c68e2b6b471e 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -353,6 +353,18 @@ def engine_outputs_wire(output): ) +@dataclass +class KVEventsConfig: + enable_kv_cache_events: bool + publisher: str + endpoint: str + replay_endpoint: str | None + buffer_steps: int + hwm: int + max_queue_size: int + topic: str + + @dataclass class EngineCoreReadyResponse: max_model_len: int @@ -363,8 +375,16 @@ class EngineCoreReadyResponse: vllm_version: str world_size: int data_parallel_size: int + tensor_parallel_size: int + pipeline_parallel_size: int + decode_context_parallel_size: int + data_parallel_rank: int + max_num_seqs: int + max_num_batched_tokens: int + instance_id: str kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None + kv_events_config: KVEventsConfig | None = None ready_response = EngineCoreReadyResponse( @@ -376,6 +396,23 @@ class EngineCoreReadyResponse: vllm_version="0.0.0", data_parallel_size=1, world_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", + kv_events_config=KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint="tcp://127.0.0.1:5557", + replay_endpoint="tcp://127.0.0.1:5558", + buffer_steps=10_000, + hwm=100_000, + max_queue_size=100_000, + topic="kv", + ), ) print(msgspec.msgpack.encode(request).hex()) diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index d2eb57aa148e..d3e3ecae2bf2 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -512,14 +512,14 @@ pub async fn send_message( input_send: &mut RouterSendHalf, engine_id: &EngineId, request_type: Bytes, - payload: Vec, + payload: Bytes, + aux_frames: Vec, ) -> Result<()> { - let message = ZmqMessage::try_from(vec![ - engine_id.to_frame(), - request_type, - Bytes::from(payload), - ]) - .expect("router messages must contain identity and payload"); + let mut frames = Vec::with_capacity(3 + aux_frames.len()); + frames.extend([engine_id.to_frame(), request_type, payload]); + frames.extend(aux_frames); + let message = + ZmqMessage::try_from(frames).expect("router messages must contain identity and payload"); trace!( ?engine_id, diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index 95d61f652c1e..eadefced236e 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -39,6 +39,12 @@ pub struct ManagedEngineArgs { /// Number of data parallel replicas to run on this node. #[arg(long)] pub data_parallel_size_local: Option, + /// Maximum model context length forwarded to the managed Python engine. + /// + /// Rust leaves validation to Python so values such as `auto` and + /// human-readable integers retain their engine-owned semantics. + #[arg(long)] + pub max_model_len: Option, /// Additional arguments forwarded to `python -m vllm.entrypoints.cli.main /// serve ...`. @@ -78,7 +84,6 @@ impl ManagedEngineArgs { pub fn into_config( self, model: String, - max_model_len: Option, max_logprobs: Option, profiler_config: Option, reasoning_parser: Option<&str>, @@ -86,12 +91,13 @@ impl ManagedEngineArgs { disable_log_stats: bool, shutdown_timeout: u64, handshake_port: u16, + limit_mm_per_prompt: Option, ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. - if let Some(max_model_len) = max_model_len { + if let Some(max_model_len) = self.max_model_len { python_args.push("--max-model-len".to_string()); - python_args.push(max_model_len.to_string()); + python_args.push(max_model_len); } if let Some(max_logprobs) = max_logprobs { python_args.push("--max-logprobs".to_string()); @@ -121,6 +127,10 @@ impl ManagedEngineArgs { python_args.push("--data-parallel-size-local".to_string()); python_args.push(data_parallel_size_local.to_string()); } + if let Some(limit_mm_per_prompt) = limit_mm_per_prompt { + python_args.push("--limit-mm-per-prompt".to_string()); + python_args.push(limit_mm_per_prompt); + } ManagedEngineConfig { python: self.python, diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs index ab017b5136c0..9cba325116c5 100644 --- a/rust/src/parser/benches/utils/adapter.rs +++ b/rust/src/parser/benches/utils/adapter.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use vllm_parser::tool::{ - Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput, + Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput, }; use vllm_parser::unified::{ UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, @@ -19,6 +19,10 @@ impl Tokenizer for BenchTokenizer { Ok(text.chars().map(|_| u32::MAX).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], @@ -85,8 +89,8 @@ impl ToolParser for UnifiedToolParserAdapter { self.inner.preserve_special_tokens() } - fn structural_tag_model(&self) -> Option { - self.inner.structural_tag_model() + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + self.inner.structural_tag_builder() } fn tool_call_id(&self, tool_index: usize) -> Option<&str> { diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs index aebcde8a036e..09bb0f189bfb 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{DeepSeekDsmlToolParser, DsmlTokens}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3.2 models. /// @@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser { true } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::DeepSeekV32) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::DeepSeekV32.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs index 31470abcdfc8..f4540ea47cab 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{DeepSeekDsmlToolParser, DsmlTokens}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V4 models. /// @@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser { true } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::DeepSeekV4) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::DeepSeekV4.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { @@ -73,7 +73,7 @@ mod tests { use super::DeepSeekV4ToolParser; use crate::tool::test_utils::{collect_stream, test_tools}; - use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _}; + use crate::tool::{ToolParser, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -91,13 +91,10 @@ mod tests { } #[test] - fn deepseek_v4_exposes_structural_tag_model() { + fn deepseek_v4_exposes_structural_tag_builder() { let parser = DeepSeekV4ToolParser::new(&test_tools()); - assert_eq!( - parser.structural_tag_model(), - Some(StructuralTagModel::DeepSeekV4) - ); + assert!(parser.structural_tag_builder().is_some()); } #[test] diff --git a/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs b/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs index 3e3c16624d08..f932cb60196d 100644 --- a/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs +++ b/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3 JSON-fenced tool calls. /// @@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::DeepSeekR1) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::DeepSeekR1.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs b/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs index 1a1ec3634af5..34d5eed78489 100644 --- a/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs +++ b/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3.1 raw JSON tool calls. /// @@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::DeepSeekV31) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::DeepSeekV31.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/glm_xml/glm47_moe.rs b/rust/src/parser/src/tool/glm_xml/glm47_moe.rs index f13cd7983e50..06477b404913 100644 --- a/rust/src/parser/src/tool/glm_xml/glm47_moe.rs +++ b/rust/src/parser/src/tool/glm_xml/glm47_moe.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{GlmXmlToolParser, Separator}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Tool parser for GLM-4.7 MoE XML-style tool calls. /// @@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Glm47) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Glm47.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/hy_v3.rs b/rust/src/parser/src/tool/hy_v3.rs index 227f4aa28742..a1426d6a3bd2 100644 --- a/rust/src/parser/src/tool/hy_v3.rs +++ b/rust/src/parser/src/tool/hy_v3.rs @@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until}; use super::parameters::ToolSchemas; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagBuilder, Tool}; const TOOL_CALLS_START: &str = ""; const TOOL_CALLS_END: &str = ""; @@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::HyV3) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::HyV3.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/json/hermes.rs b/rust/src/parser/src/tool/json/hermes.rs index 64c594ee0ff3..aed27a8cd4dd 100644 --- a/rust/src/parser/src/tool/json/hermes.rs +++ b/rust/src/parser/src/tool/json/hermes.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Hermes", @@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Hermes) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Hermes.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/json/llama.rs b/rust/src/parser/src/tool/json/llama.rs index 23e49724f07d..262340a9e806 100644 --- a/rust/src/parser/src/tool/json/llama.rs +++ b/rust/src/parser/src/tool/json/llama.rs @@ -12,7 +12,9 @@ use super::{ argument_delta_event, tool_call_header_event, }; use crate::tool::utils::{JsonObjectScanState, parse_buffered_event}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; +use crate::tool::{ + Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput, +}; #[derive(Debug, Clone, PartialEq, Eq)] enum LlamaJsonMode { @@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Llama) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Llama.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs index cf7c1fbe237b..7d5b71078aaa 100644 --- a/rust/src/parser/src/tool/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Qwen XML", @@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Qwen3) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Qwen3.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/kimi_k2.rs b/rust/src/parser/src/tool/kimi_k2.rs index aa4dc0e6ad65..540b5a3d37bd 100644 --- a/rust/src/parser/src/tool/kimi_k2.rs +++ b/rust/src/parser/src/tool/kimi_k2.rs @@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while}; use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagBuilder, Tool}; const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>"; const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>"; @@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser { true } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Kimi) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Kimi.builder()) } fn tool_call_id(&self, tool_index: usize) -> Option<&str> { diff --git a/rust/src/parser/src/tool/minimax_m2.rs b/rust/src/parser/src/tool/minimax_m2.rs index 903b14eb6a5e..d90e6676f0b6 100644 --- a/rust/src/parser/src/tool/minimax_m2.rs +++ b/rust/src/parser/src/tool/minimax_m2.rs @@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until}; use super::parameters::ToolSchemas; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagBuilder, Tool}; const TOOL_CALL_START: &str = ""; const TOOL_CALL_END: &str = ""; @@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Minimax) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Minimax.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index 216d3dff4cf0..e25eae882f9b 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser; pub use seed_oss::SeedOssToolParser; use serde::{Deserialize, Serialize}; use serde_json::Value; -pub use xgrammar_structural_tag::Model as StructuralTagModel; +pub use xgrammar_structural_tag::builders::StructuralTagBuilder; use crate::utils; @@ -187,8 +187,8 @@ pub trait ToolParser: Send { false } - /// Return the xgrammar structural-tag model used for strict tool calling. - fn structural_tag_model(&self) -> Option { + /// Return the xgrammar structural-tag builder used for strict tool calling. + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { None } diff --git a/rust/src/parser/src/tool/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs index 361de524e8ba..0a9385a85ce4 100644 --- a/rust/src/parser/src/tool/qwen_coder.rs +++ b/rust/src/parser/src/tool/qwen_coder.rs @@ -9,8 +9,8 @@ use winnow::token::{literal, take_until}; use super::parameters::ToolSchemas; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; -use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::Tool; +use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; +use crate::tool::{StructuralTagBuilder, Tool}; const TOOL_CALL_START: &str = ""; const TOOL_CALL_END: &str = ""; @@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser { Ok(Box::new(Self::new(tools))) } - fn structural_tag_model(&self) -> Option { - Some(StructuralTagModel::Qwen3Coder) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(xgrammar_structural_tag::Model::Qwen3Coder.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { @@ -294,7 +294,7 @@ mod tests { use serde_json::{Value, json}; use thiserror_ext::AsReport; - use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser}; + use super::{Qwen3CoderToolParser, ToolParser}; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; @@ -308,13 +308,10 @@ mod tests { } #[test] - fn qwen_coder_exposes_structural_tag_model() { + fn qwen_coder_exposes_structural_tag_builder() { let parser = Qwen3CoderToolParser::new(&test_tools()); - assert_eq!( - parser.structural_tag_model(), - Some(StructuralTagModel::Qwen3Coder) - ); + assert!(parser.structural_tag_builder().is_some()); } #[test] diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index 13183192c6c9..afecc31665bb 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer; use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use crate::reasoning::ReasoningParser; -use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; /// Unified parser that composes existing reasoning and tool parsers. pub struct CombinedParser { @@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser { || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens()) } - fn structural_tag_model(&self) -> Option { - self.tool.as_ref().and_then(|parser| parser.structural_tag_model()) + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + self.tool.as_ref().and_then(|parser| parser.structural_tag_builder()) } fn tool_call_id(&self, tool_index: usize) -> Option<&str> { @@ -269,10 +269,7 @@ mod tests { fn combined_parser_emits_tool_calls_from_visible_content() { let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap(); let mut parser = CombinedParser::new(None, Some(tool)); - assert!(matches!( - parser.structural_tag_model(), - Some(crate::tool::StructuralTagModel::Qwen3) - )); + assert!(parser.structural_tag_builder().is_some()); let output = collect( &mut parser, diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs index eb78321d1269..d35fb48d9cc0 100644 --- a/rust/src/parser/src/unified/inkling.rs +++ b/rust/src/parser/src/unified/inkling.rs @@ -61,7 +61,7 @@ enum InklingEvent { TextStart, ReasoningStart, MessageStart, - Header, + Header { len: usize }, ToolJsonStart, ToolJsonHeader { name: String }, ToolJsonArgs { len: usize, complete: bool }, @@ -72,7 +72,9 @@ enum InklingEvent { enum InklingMode { #[default] Idle, - MessageHeader, + MessageHeader { + pending: String, + }, Text, Reasoning, ToolJsonHeader, @@ -117,7 +119,9 @@ impl InklingUnifiedParser { self.mode = InklingMode::Idle; for token_id in prompt_token_ids.iter().rev().copied() { if token_id == self.message_model_token_id { - self.mode = InklingMode::MessageHeader; + self.mode = InklingMode::MessageHeader { + pending: String::new(), + }; return; } if token_id == self.content_thinking_token_id { @@ -140,8 +144,19 @@ impl InklingUnifiedParser { InklingEvent::Reasoning { len } => { output.push_reasoning(self.buffer[..len].to_string()); } - InklingEvent::MessageStart => self.mode = InklingMode::MessageHeader, - InklingEvent::Header => {} + InklingEvent::MessageStart => { + self.mode = InklingMode::MessageHeader { + pending: String::new(), + }; + } + InklingEvent::Header { len } => { + let InklingMode::MessageHeader { pending } = &mut self.mode else { + return Err(parsing_failed!( + "Inkling header text outside a message header" + )); + }; + pending.push_str(&self.buffer[..len]); + } InklingEvent::TextStart => self.mode = InklingMode::Text, InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning, InklingEvent::ToolJsonStart => self.mode = InklingMode::ToolJsonHeader, @@ -174,7 +189,10 @@ impl InklingUnifiedParser { } } InklingEvent::BlockEnd => { - self.mode = InklingMode::Idle; + let mode = std::mem::take(&mut self.mode); + if let InklingMode::MessageHeader { pending } = mode { + output.push_text(pending); + } self.active_tool_index = None; } } @@ -182,10 +200,14 @@ impl InklingUnifiedParser { } fn reset(&mut self) -> String { - self.mode = InklingMode::Idle; + let mut uncommitted = match std::mem::take(&mut self.mode) { + InklingMode::MessageHeader { pending } => pending, + _ => String::new(), + }; self.active_tool_index = None; self.emitted_tool_count = 0; - std::mem::take(&mut self.buffer) + uncommitted.push_str(&std::mem::take(&mut self.buffer)); + uncommitted } } @@ -225,11 +247,15 @@ impl UnifiedParser for InklingUnifiedParser { fn finish(&mut self) -> Result { let mut output = UnifiedParserOutput::default(); - match &self.mode { + match &mut self.mode { InklingMode::Idle | InklingMode::Text => { output.push_text(std::mem::take(&mut self.buffer)) } - InklingMode::MessageHeader => self.buffer.clear(), + InklingMode::MessageHeader { pending } => { + pending.push_str(&self.buffer); + self.buffer.clear(); + output.push_text(std::mem::take(pending)); + } InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), InklingMode::ToolJsonHeader | InklingMode::ToolJsonArgs { .. } @@ -254,7 +280,7 @@ fn parse_next_inkling_event( ) -> ModalResult { match mode { InklingMode::Idle => parse_idle_event(input), - InklingMode::MessageHeader => parse_message_header_event(input), + InklingMode::MessageHeader { .. } => parse_message_header_event(input), InklingMode::Text => parse_text_event(input), InklingMode::Reasoning => parse_reasoning_event(input), InklingMode::ToolJsonHeader => parse_tool_json_header_event(input), @@ -346,7 +372,7 @@ fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult) -> ModalResult { - safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header) + safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Header { len }) } /// Parse safe text before the end of a Inkling text block. @@ -414,6 +440,10 @@ mod tests { Ok(text.chars().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], @@ -668,15 +698,58 @@ mod tests { let mut parser = test_parser(); parser.initialize(&[200001]).unwrap(); + let mut output = parser.parse_chunk("get_").unwrap(); + output.append( + parser + .parse_complete(concat!( + "weather<|content_invoke_tool_json|>", + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>" + )) + .unwrap(), + ); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + } + + #[test] + fn inkling_initialize_model_opener_flushes_bare_text_at_finish() { + let mut parser = test_parser(); + parser.initialize(&[200001]).unwrap(); + + let mut output = parser.parse_chunk("plain ").unwrap(); + output.append(parser.parse_chunk("answer").unwrap()); + assert!(output.normal_text().is_empty()); + + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "plain answer"); + } + + #[test] + fn inkling_initialize_model_opener_flushes_bare_text_before_end_marker() { + let mut parser = test_parser(); + parser.initialize(&[200001]).unwrap(); + let output = parser .parse_complete(concat!( - "get_weather<|content_invoke_tool_json|>", - "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>" + "plain answer", + "<|end_message|>", + "<|content_model_end_sampling|>" )) .unwrap(); - assert!(output.normal_text().is_empty()); - assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.normal_text(), "plain answer"); + } + + #[test] + fn inkling_reset_returns_pending_message_header() { + let mut parser = test_parser(); + parser.initialize(&[200001]).unwrap(); + parser.parse_chunk("plain ").unwrap(); + parser.parse_chunk("answer").unwrap(); + + assert_eq!(parser.reset(), "plain answer"); } #[test] @@ -733,6 +806,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs new file mode 100644 index 000000000000..a721b29f3f5d --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -0,0 +1,1126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for the Kimi K3 (XTML) chat format. +//! +//! Original Python implementations: +//! - `vllm/reasoning/kimi_k3_reasoning_parser.py` +//! - `vllm/tool_parsers/kimi_k3_tool_parser.py` +//! +//! K3 wraps one assistant message into XTML channels built from the dedicated +//! special tokens `<|open|>`, `<|close|>`, and `<|sep|>`: +//! +//! ```text +//! <|open|>think<|sep|>reasoning<|close|>think<|sep|> +//! <|open|>response<|sep|>visible answer<|close|>response<|sep|> +//! <|open|>tools<|sep|> +//! <|open|>call tool="get_weather" index="1"<|sep|> +//! <|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|> +//! <|close|>call<|sep|> +//! <|close|>tools<|sep|> +//! <|close|>message<|sep|> +//! ``` +//! +//! In chat serving the generation prompt ends with `<|open|>think<|sep|>` +//! (thinking) or `<|open|>response<|sep|>` (instruct), so the model output +//! starts *inside* that channel without re-emitting the open tag; +//! [`UnifiedParser::initialize`] detects this from the prompt token IDs. +//! +//! Argument decoding mirrors the renderer's type tagging (inverse encoding): +//! `type="string"` values pass the raw text through, other types are +//! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute +//! values reverse the renderer escaping (`"` before `&`). +//! +//! Known limitation (shared with the Python parser): string argument and +//! response bodies are emitted raw, so a value that literally contains +//! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is +//! indistinguishable from a real closing marker. + +mod structural_tag; + +pub use structural_tag::KimiK3StructuralTagBuilder; + +use serde_json::{Map, Value}; +use vllm_tokenizer::DynTokenizer; +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; +use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated}; +use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::{literal, rest, take_till, take_until, take_while}; + +use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::{StructuralTagBuilder, Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +const OPEN: &str = "<|open|>"; +const SEP: &str = "<|sep|>"; +const END_OF_MSG: &str = "<|end_of_msg|>"; + +const THINK_OPEN: &str = "<|open|>think<|sep|>"; +const THINK_CLOSE: &str = "<|close|>think<|sep|>"; +const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; +const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; +const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; +const CALL_OPEN: &str = "<|open|>call"; +const CALL_CLOSE: &str = "<|close|>call<|sep|>"; +const ARG_OPEN: &str = "<|open|>argument"; +const ARG_CLOSE: &str = "<|close|>argument<|sep|>"; +const JSON_OPEN: &str = "<|open|>json"; +const JSON_CLOSE: &str = "<|close|>json<|sep|>"; + +const IDLE_MARKERS: &[&str] = &[ + THINK_OPEN, + RESPONSE_OPEN, + TOOLS_OPEN, + MESSAGE_CLOSE, + END_OF_MSG, +]; +const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG]; +const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG]; + +/// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in +/// the prompt tail (attribute-bearing message opens, message bodies) never name +/// a generation channel. +const MAX_PREFILL_TAG_TOKENS: usize = 8; + +type KimiK3Input<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +enum KimiK3Event { + Text { + len: usize, + }, + Reasoning { + len: usize, + }, + /// Structural noise consumed without emitting anything. + Skip, + ThinkOpen, + ThinkClose, + ResponseOpen, + ResponseClose, + ToolsOpen, + ToolsClose, + /// The assistant message closed; everything after it is ignored. + MessageEnd, + CallOpen { + name: String, + index: Option, + }, + CallComplete { + arguments: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +enum KimiK3Mode { + /// Before any channel opens: channel opens are expected, but raw text + /// falls through as visible text so marker-free output still streams. + #[default] + Idle, + /// Inside the `think` channel. + Reasoning, + /// Inside the `response` channel. + Response, + /// After the `response` (or `tools`) channel closed: only a `tools` + /// channel or the message close may follow, and noise is never content. + Epilogue, + /// Inside the `tools` channel, between `call` blocks. + Tools, + /// Inside one `call` block, buffering its body until the close marker. + Call { + name: String, + index: Option, + scan: MarkerScanState, + }, + /// After the message closed: ignore the rest (EOS leakage guard). + Done, +} + +/// Unified parser for Kimi K3 XTML think / response / tools channels. +pub struct KimiK3UnifiedParser { + buffer: String, + mode: KimiK3Mode, + /// Number of calls emitted in the current response. + emitted_call_count: usize, + tokenizer: DynTokenizer, + open_token_id: u32, + sep_token_id: u32, +} + +impl KimiK3UnifiedParser { + /// Create a Kimi K3 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let open_token_id = token_id(tokenizer.as_ref(), OPEN)?; + let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; + + Ok(Self { + buffer: String::new(), + mode: KimiK3Mode::default(), + emitted_call_count: 0, + tokenizer, + open_token_id, + sep_token_id, + }) + } + + /// Detect the prefilled generation channel from the prompt tail. + /// + /// `add_generation_prompt` ends the prompt with `<|open|>think<|sep|>` or + /// `<|open|>response<|sep|>`, so generation starts inside that channel and + /// never re-emits the open tag. Locate the last `<|open|>…<|sep|>` pair and + /// decode the short tag between the two structural tokens. + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = KimiK3Mode::Idle; + + let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else { + return; + }; + let Some(open_pos) = + prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id) + else { + return; + }; + let tag_ids = &prompt_token_ids[open_pos + 1..sep_pos]; + if tag_ids.is_empty() || tag_ids.len() > MAX_PREFILL_TAG_TOKENS { + return; + } + let Ok(tag) = self.tokenizer.decode(tag_ids, /* skip_special_tokens */ false) else { + return; + }; + self.mode = match tag.trim() { + "think" => KimiK3Mode::Reasoning, + "response" => KimiK3Mode::Response, + _ => KimiK3Mode::Idle, + }; + } + + fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> { + match event { + KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()), + KimiK3Event::Reasoning { len } => { + output.push_reasoning(self.buffer[..len].to_string()); + } + KimiK3Event::Skip => {} + KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning, + KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle, + KimiK3Event::ResponseOpen => self.mode = KimiK3Mode::Response, + KimiK3Event::ResponseClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::ToolsOpen => self.mode = KimiK3Mode::Tools, + KimiK3Event::ToolsClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::MessageEnd => self.mode = KimiK3Mode::Done, + KimiK3Event::CallOpen { name, index } => { + self.mode = KimiK3Mode::Call { + name, + index, + scan: MarkerScanState::default(), + }; + } + KimiK3Event::CallComplete { arguments } => { + let mode = std::mem::replace(&mut self.mode, KimiK3Mode::Tools); + let KimiK3Mode::Call { name, .. } = mode else { + return Err(parsing_failed!( + "Kimi K3 call completion without an active tool call" + )); + }; + // An empty/garbage call block without a tool name is dropped, + // matching the Python parser. + if name.is_empty() { + return Ok(()); + } + + let tool_index = self.emitted_call_count; + self.emitted_call_count += 1; + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments, + }); + } + } + Ok(()) + } + + fn reset_state(&mut self) -> String { + self.mode = KimiK3Mode::Idle; + self.emitted_call_count = 0; + std::mem::take(&mut self.buffer) + } +} + +impl UnifiedParser for KimiK3UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.emitted_call_count = 0; + self.initialize_mode(prompt_token_ids); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER) + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + parse_next_kimi_k3_event(input, &mut self.mode) + })? { + self.apply_event(event, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + KimiK3Mode::Idle | KimiK3Mode::Response => { + output.push_text(std::mem::take(&mut self.buffer)); + } + KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(), + // A tools channel truncated between complete calls loses only its + // closing markers; keep the calls already emitted. + KimiK3Mode::Tools if self.buffer.is_empty() => {} + KimiK3Mode::Tools | KimiK3Mode::Call { .. } => { + return Err(parsing_failed!("incomplete Kimi K3 tool call")); + } + } + + self.mode = KimiK3Mode::Idle; + Ok(output) + } + + fn reset(&mut self) -> String { + self.reset_state() + } +} + +/// Parse one Kimi K3 event from buffered streaming input. +fn parse_next_kimi_k3_event( + input: &mut KimiK3Input<'_>, + mode: &mut KimiK3Mode, +) -> ModalResult { + match mode { + KimiK3Mode::Idle => parse_idle_event(input), + KimiK3Mode::Reasoning => parse_reasoning_event(input), + KimiK3Mode::Response => parse_response_event(input), + KimiK3Mode::Epilogue => parse_epilogue_event(input), + KimiK3Mode::Tools => parse_tools_event(input), + KimiK3Mode::Call { scan, .. } => call_body_event(input, scan), + KimiK3Mode::Done => parse_done_event(input), + } +} + +/// Parse an event while waiting for the next channel open. +fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_OPEN).value(KimiK3Event::ThinkOpen), + literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen), + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_idle_text_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `think` channel. +fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_CLOSE).value(KimiK3Event::ThinkClose), + // `<|end_of_msg|>` can reach the parser under `ignore_eos` or + // `include_stop_str_in_output`; never leak it into reasoning. + literal(END_OF_MSG).value(KimiK3Event::MessageEnd), + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `response` channel. +fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose), + // The response body also implicitly ends at a `tools` channel. + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_response_text_event, + )) + .parse_next(input) +} + +/// Parse an event after the response channel closed. +fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + skip_epilogue_noise_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `tools` channel, between `call` blocks. +fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + call_open_event, + literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose), + // Defensive: an unterminated tools channel still ends with the message. + message_end_event, + skip_tools_noise_event, + )) + .parse_next(input) +} + +/// Parse a message close or end-of-message marker. +fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG))) + .value(KimiK3Event::MessageEnd) + .parse_next(input) +} + +/// Ignore everything after the assistant message closed. +fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult { + rest.value(KimiK3Event::Skip).parse_next(input) +} + +/// Parse safe text while waiting for the next channel marker. +fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Parse safe reasoning before the think close marker. +fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len }) +} + +/// Parse safe response text before the next channel marker. +fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Skip non-content noise after the response channel closed. +fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Skip non-content noise between `call` blocks. +fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Parse a `call` open tag into its tool name and one-based index. +fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let (attrs,) = seq!( + _: literal(CALL_OPEN), + take_until(0.., SEP), + _: literal(SEP), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + Ok(KimiK3Event::CallOpen { + name: attr_value(&attrs, "tool").unwrap_or_default().to_string(), + index: attr_value(&attrs, "index") + .filter(|index| !index.is_empty()) + .map(str::to_string), + }) +} + +/// Parse the buffered call body through `<|close|>call<|sep|>` into a +/// completed call. +fn call_body_event( + input: &mut KimiK3Input<'_>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (body,) = seq!( + take_until_marker(CALL_CLOSE, scan), + _: literal(CALL_CLOSE), + ) + .parse_next(input)?; + let arguments = parse_call_arguments(body)?; + + Ok(KimiK3Event::CallComplete { arguments }) +} + +/// Parse a complete `call` body into the OpenAI-style arguments JSON string. +/// +/// The body is either one raw `json` block (passed through unmodified) or a +/// sequence of typed `argument` blocks (converted per their `type` tags). +fn parse_call_arguments(body: &str) -> ModalResult { + let mut input = body; + terminated( + delimited(ws0, alt((json_block_arguments, typed_arguments)), ws0), + eof, + ) + .parse_next(&mut input) + .map_err(|_| xtml_error("Kimi K3 call body")) +} + +/// Parse one raw `json` argument block, passing its body through unmodified. +fn json_block_arguments(input: &mut &str) -> ModalResult { + let (raw,) = seq!( + _: literal(JSON_OPEN), + _: take_until(0.., SEP), // attrs (`type="object"`), unused on decode + _: literal(SEP), + take_until(0.., JSON_CLOSE), + _: literal(JSON_CLOSE), + ) + .parse_next(input)?; + Ok(raw.to_string()) +} + +/// Parse typed `argument` blocks into a serialized JSON object, preserving +/// argument order. +fn typed_arguments(input: &mut &str) -> ModalResult { + let pairs: Vec<(String, Value)> = + repeat(0.., terminated(argument_block, ws0)).parse_next(input)?; + let arguments = pairs.into_iter().collect::>(); + serde_json::to_string(&arguments).map_err(|_| xtml_error("Kimi K3 arguments")) +} + +/// Parse one typed `argument` block into its key/value pair. +fn argument_block(input: &mut &str) -> ModalResult<(String, Value)> { + let (attrs, raw) = seq!( + _: literal(ARG_OPEN), + take_until(0.., SEP), + _: literal(SEP), + take_until(0.., ARG_CLOSE), + _: literal(ARG_CLOSE), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + let key = attr_value(&attrs, "key").unwrap_or_default().to_string(); + let arg_type = attr_value(&attrs, "type").unwrap_or("string"); + Ok((key, decode_argument_value(arg_type, raw))) +} + +/// Decode one typed argument value per its XTML `type` tag. +/// +/// `string` values pass the raw text through (the renderer emits them +/// unescaped); other types are JSON-decoded, falling back to the raw text on +/// malformed payloads so one quirky value does not fail the whole call. +fn decode_argument_value(arg_type: &str, raw: &str) -> Value { + if arg_type == "string" { + return Value::String(raw.to_string()); + } + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// Parse a complete XTML attribute string like ` tool="get_weather" index="1"`. +fn parse_tag_attrs(attrs: &str) -> ModalResult> { + let mut input = attrs; + terminated(repeat(0.., preceded(ws1, tag_attr)), (ws0, eof)) + .parse_next(&mut input) + .map_err(|_| xtml_error("XTML tag attributes")) +} + +/// Parse one XTML `key="value"` attribute pair. +fn tag_attr(input: &mut &str) -> ModalResult<(String, String)> { + seq!( + take_while(1.., |char: char| char.is_alphanumeric() || char == '_').map(str::to_string), + _: literal("=\""), + take_till(0.., '"').map(unescape_attr_value), + _: literal("\""), + ) + .parse_next(input) +} + +/// Reverse XTML attribute escaping: `"` first, then `&` (the inverse +/// of the encode order). +fn unescape_attr_value(value: &str) -> String { + value.replace(""", "\"").replace("&", "&") +} + +/// Look up one parsed attribute value by key. +fn attr_value<'a>(attrs: &'a [(String, String)], key: &str) -> Option<&'a str> { + attrs.iter().find(|(name, _)| name == key).map(|(_, value)| value.as_str()) +} + +/// Build a cut error for determinably malformed XTML structure. +fn xtml_error(label: &'static str) -> ErrMode { + let mut error = ContextError::new(); + error.push(StrContext::Label(label)); + ErrMode::Cut(error) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer as _; + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::{ + END_OF_MSG, KimiK3UnifiedParser, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, + }; + use crate::tool::ToolCallDelta; + use crate::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + }; + + const OPEN_ID: u32 = 256; + const CLOSE_ID: u32 = 257; + const SEP_ID: u32 = 258; + const END_OF_MSG_ID: u32 = 259; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token("<|close|>", CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for KimiK3UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_parser() -> KimiK3UnifiedParser { + KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() + } + + fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + /// Split `text` into small chunks to stress marker-split handling. + fn char_chunks(text: &str, size: usize) -> Vec { + let chars: Vec = text.chars().collect(); + chars.chunks(size).map(|chunk| chunk.iter().collect()).collect() + } + + fn arg(key: &str, arg_type: &str, value: &str) -> String { + format!( + "{OPEN}argument key=\"{key}\" type=\"{arg_type}\"{SEP}{value}<|close|>argument{SEP}" + ) + } + + fn call(attrs: &str, body: &str) -> String { + format!("{OPEN}call {attrs}{SEP}{body}<|close|>call{SEP}") + } + + fn thinking_output(reasoning: &str, response: &str, tools_body: &str) -> String { + let mut output = format!("{THINK_OPEN}{reasoning}{THINK_CLOSE}"); + output.push_str(&format!("{RESPONSE_OPEN}{response}{RESPONSE_CLOSE}")); + if !tools_body.is_empty() { + output.push_str(&format!("{TOOLS_OPEN}{tools_body}{TOOLS_CLOSE}")); + } + output.push_str("<|close|>message<|sep|>"); + output + } + + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { + output.calls().first().expect("expected one tool call").clone() + } + + #[test] + fn kimi_k3_create_requires_structural_tokens() { + let error = match KimiK3UnifiedParser::new(&[], Arc::new(TestTokenizer::new())) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == OPEN + )); + } + + #[test] + fn kimi_k3_parses_reasoning_response_and_typed_tool_call() { + let body = [ + arg("city", "string", "Hangzhou"), + arg("days", "number", "1.5"), + arg("detailed", "boolean", "true"), + arg("filters", "object", r#"{"kind":"rain"}"#), + arg("hours", "array", "[8,20]"), + ] + .concat(); + let text = thinking_output( + "Need the weather tool.", + "I'll check.", + &call("tool=\"get_weather\" index=\"1\"", &body), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "Need the weather tool."); + assert_eq!(output.normal_text(), "I'll check."); + let call = first_call(&output); + assert_eq!(call.tool_index, 0); + assert_eq!(call.name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&call.arguments).unwrap(), + json!({ + "city": "Hangzhou", + "days": 1.5, + "detailed": true, + "filters": { "kind": "rain" }, + "hours": [8, 20], + }) + ); + } + + #[test] + fn kimi_k3_arguments_preserve_order_and_number_formatting() { + let body = [ + arg("y", "number", "1.0"), + arg("x", "number", "2"), + arg("items", "array", r#"["left","right"]"#), + ] + .concat(); + let text = thinking_output("t", "", &call("tool=\"add\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + first_call(&output).arguments, + r#"{"y":1.0,"x":2,"items":["left","right"]}"# + ); + } + + #[test] + fn kimi_k3_streaming_splits_markers_across_chunks() { + let text = thinking_output( + "step by step", + "the answer", + &call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")), + ); + + for size in [1, 3, 7] { + let chunks = char_chunks(&text, size); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let output = collect_stream(&mut test_parser(), &chunk_refs); + + assert_eq!(output.reasoning_text(), "step by step", "chunk size {size}"); + assert_eq!(output.normal_text(), "the answer", "chunk size {size}"); + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":42}"#); + } + } + + #[test] + fn kimi_k3_streaming_emits_text_incrementally() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let first = parser.parse_chunk("Hel").unwrap(); + assert_eq!(first.normal_text(), "Hel"); + + let second = parser.parse_chunk("lo<|close|>resp").unwrap(); + assert_eq!(second.normal_text(), "lo"); + + let mut output = parser.parse_chunk("onse<|sep|>").unwrap(); + output.append(parser.finish().unwrap()); + assert_eq!(output.normal_text(), ""); + } + + #[test] + fn kimi_k3_initialize_think_prefill_starts_in_reasoning() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!( + "reasoning{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}" + )) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reasoning"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_initialize_response_prefill_starts_in_response() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>response<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("answer{RESPONSE_CLOSE}<|close|>message{SEP}")) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn kimi_k3_initialize_message_open_prefill_starts_idle() { + let mut parser = test_parser(); + let prompt = + tokenizer().encode("<|open|>message role=\"assistant\"<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("{THINK_OPEN}reason{THINK_CLOSE}{RESPONSE_OPEN}hi")) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "hi"); + } + + #[test] + fn kimi_k3_plain_text_falls_through_as_text() { + let output = collect_stream(&mut test_parser(), &["plain ", "answer"]); + + assert_eq!(output.normal_text(), "plain answer"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + + #[test] + fn kimi_k3_tool_call_waits_for_close_marker() { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + + let argument = arg("x", "number", "1"); + for chunk in [ + TOOLS_OPEN, + "<|open|>call tool=\"calc\" index=\"1\"<|sep|>", + argument.as_str(), + ] { + output.append(parser.parse_chunk(chunk).unwrap()); + assert!(output.calls().is_empty()); + } + + output.append(parser.parse_chunk("<|close|>call<|sep|>").unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#); + } + + #[test] + fn kimi_k3_parses_multiple_tool_calls() { + let tools_body = format!( + "{}{}", + call( + "tool=\"get_weather\" index=\"1\"", + &arg("city", "string", "SF") + ), + call("tool=\"get_time\" index=\"2\"", ""), + ); + let text = thinking_output("t", "r", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(calls[0].arguments, r#"{"city":"SF"}"#); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("get_time")); + assert_eq!(calls[1].arguments, "{}"); + } + + #[test] + fn kimi_k3_json_block_arguments_pass_through_raw() { + // Spacing and key order must survive unmodified: raw `json` blocks are + // not validated or normalized. + let raw = r#"{"b": 1, "a": [2 , 3]}"#; + let body = format!("{OPEN}json type=\"object\"{SEP}{raw}<|close|>json{SEP}"); + let text = thinking_output("t", "", &call("tool=\"run\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, raw); + } + + #[test] + fn kimi_k3_string_argument_passes_raw_text_through() { + let value = "line one\nline two {\"not\": \"json\"} & "; + let text = thinking_output( + "t", + "", + &call( + "tool=\"write\" index=\"1\"", + &arg("content", "string", value), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": value }) + ); + } + + #[test] + fn kimi_k3_malformed_typed_argument_falls_back_to_raw_text() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"calc\" index=\"1\"", + &arg("x", "number", "not a number"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, r#"{"x":"not a number"}"#); + } + + #[test] + fn kimi_k3_attribute_values_are_unescaped() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"a"b&c\" index=\"1\"", + &arg("key", "string", "value"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).name.as_deref(), Some("a\"b&c")); + } + + #[test] + fn kimi_k3_call_without_tool_name_is_dropped() { + let tools_body = format!( + "{}{}", + call("index=\"1\"", &arg("x", "number", "1")), + call("tool=\"real\" index=\"2\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("real")); + } + + #[test] + fn kimi_k3_tool_indices_ignore_xtml_index_attribute() { + let tools_body = format!( + "{}{}{}", + call("tool=\"first\" index=\"3\"", ""), + call("tool=\"second\"", ""), + call("tool=\"third\" index=\"x\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + output.calls().iter().map(|call| call.tool_index).collect::>(), + [0, 1, 2] + ); + } + + #[test] + fn kimi_k3_ignores_output_after_message_close() { + let mut parser = test_parser(); + let output = parser + .parse_complete(&format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}junk{END_OF_MSG}" + )) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_epilogue_noise_is_not_content() { + let text = format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}\n{TOOLS_OPEN}{}{TOOLS_CLOSE}\n<|close|>message{SEP}", + call("tool=\"calc\" index=\"1\"", ""), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_finish_flushes_unclosed_reasoning() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!("{THINK_OPEN}still thinking")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "still thinking"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn kimi_k3_finish_flushes_partial_marker_as_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk("answer<|clo").unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "answer<|clo"); + } + + #[test] + fn kimi_k3_finish_fails_mid_tool_call() { + let mut parser = test_parser(); + parser + .parse_chunk(&format!( + "{TOOLS_OPEN}<|open|>call tool=\"calc\" index=\"1\"<|sep|>{}", + arg("x", "number", "1") + )) + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Kimi K3 tool call")); + } + + #[test] + fn kimi_k3_finish_after_truncated_tools_keeps_complete_calls() { + let mut parser = test_parser(); + let mut output = parser + .parse_chunk(&format!( + "{TOOLS_OPEN}{}", + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "1")) + )) + .unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + } + + #[test] + fn kimi_k3_malformed_call_attributes_fail_fast() { + let mut parser = test_parser(); + let error = parser + .parse_chunk(&format!("{TOOLS_OPEN}<|open|>call garbage attrs<|sep|>")) + .unwrap_err(); + + assert!(error.to_report_string().contains("XTML tag attributes")); + } + + #[test] + fn kimi_k3_empty_response_channel_emits_nothing() { + let text = thinking_output("t", "", &call("tool=\"calc\" index=\"1\"", "")); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.reasoning_text(), "t"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_reset_returns_buffered_text() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + parser.parse_chunk("answer<|close|>resp").unwrap(); + + let raw = parser.reset(); + + assert_eq!(raw, "<|close|>resp"); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs new file mode 100644 index 000000000000..d3515ad2dc33 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Structural-tag grammar for Kimi K3 XTML tool calls. + +use serde_json::{Map, Value}; +use xgrammar_structural_tag::Result; +use xgrammar_structural_tag::builders::{ + StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, +}; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; +use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; + +use super::{ + ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, + RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, +}; + +pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = + KimiK3StructuralTagBuilder; + +const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"]; + +/// Kimi K3 XTML structural-tag builder. +#[derive(Debug, Clone, Copy, Default)] +pub struct KimiK3StructuralTagBuilder; + +impl StructuralTagBuilder for KimiK3StructuralTagBuilder { + fn build(&self, ctx: StructuralTagContext<'_>) -> Result { + let mut elements = response_prefix(ctx.options.reasoning); + + let tools = match ctx.tool_choice { + // Serving lowering filters empty tools before calling the builder, + // while direct `build_structural_tag(..., auto, ...)` calls do not. + BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, + BuilderToolChoice::Auto => Some(Format::optional(tools_channel( + ctx.function_tools, + ctx.options, + ))), + BuilderToolChoice::Forced | BuilderToolChoice::Required => { + Some(tools_channel(ctx.function_tools, ctx.options)) + } + }; + if let Some(tools) = tools { + elements.push(tools); + } + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(StructuralTag::new(Format::sequence(elements))) + } +} + +fn response_prefix(reasoning: bool) -> Vec { + let mut elements = Vec::new(); + if reasoning { + elements.push(Format::tag( + "", + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), + THINK_CLOSE, + )); + elements.push(Format::const_string(RESPONSE_OPEN)); + } else { + elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); + } + elements.push(Format::tag( + "", + Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]), + RESPONSE_CLOSE, + )); + elements +} + +fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { + let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); + Format::tag( + TOOLS_OPEN, + Format::tags_with_separator(calls, "", true, false), + TOOLS_CLOSE, + ) +} + +fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { + let parameters = function_parameters(&tool.function); + let call_body = Format::or(vec![ + typed_arguments(¶meters, options), + raw_json_arguments(¶meters, options), + ]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + let Some(schema) = parameters.as_object() else { + return if parameters == &Value::Bool(false) { + Format::const_string("") + } else { + Format::star(permissive_argument()) + }; + }; + let Some(properties) = schema.get("properties").and_then(Value::as_object) else { + return Format::star(permissive_argument()); + }; + if properties.is_empty() { + return Format::star(permissive_argument()); + } + + let root_defs = root_definitions(schema); + let arguments = properties + .iter() + .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options)) + .map(Format::Tag) + .collect::>(); + let arguments = match arguments.as_slice() { + [argument] => argument.clone(), + _ => Format::or(arguments), + }; + // Keep typed arguments order-agnostic and non-unique, but do not allow an + // empty call when the root schema declares required properties. + if schema + .get("required") + .and_then(Value::as_array) + .is_some_and(|required| !required.is_empty()) + { + Format::plus(arguments) + } else { + Format::star(arguments) + } +} + +fn argument_tags( + key: &str, + schema: &Value, + root_defs: &Map, + options: StructuralTagOptions, +) -> Vec { + let types = schema_types(schema); + types + .into_iter() + .map(|xtml_type| { + let content = if xtml_type == "string" { + string_argument_content(schema) + } else { + json_schema( + attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs), + options, + ) + }; + TagFormat::new( + format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr_value(key) + ), + content, + ARG_CLOSE, + ) + }) + .collect() +} + +fn schema_types(schema: &Value) -> Vec<&'static str> { + let Some(schema) = schema.as_object() else { + return XTML_TYPES.to_vec(); + }; + let mut types = Vec::new(); + match schema.get("type") { + Some(Value::String(value)) => push_schema_type(&mut types, value), + Some(Value::Array(values)) => { + for value in values.iter().filter_map(Value::as_str) { + push_schema_type(&mut types, value); + } + } + _ => {} + } + if types.is_empty() + && let Some(value) = schema.get("const") + { + push_value_type(&mut types, value); + } + if types.is_empty() + && let Some(values) = schema.get("enum").and_then(Value::as_array) + { + for value in values { + push_value_type(&mut types, value); + } + } + if types.is_empty() { + XTML_TYPES.to_vec() + } else { + types + } +} + +fn push_value_type(types: &mut Vec<&'static str>, value: &Value) { + let xtml_type = match value { + Value::String(_) => "string", + Value::Number(_) => "number", + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Object(_) => "object", + Value::Array(_) => "array", + }; + if !types.contains(&xtml_type) { + types.push(xtml_type); + } +} + +fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) { + "integer" + } else { + xtml_type + }; + schema.insert("type".to_string(), Value::String(json_type.to_string())); + Value::Object(schema) +} + +fn explicitly_integer_only(schema: &Map) -> bool { + match schema.get("type") { + Some(Value::String(value)) => value == "integer", + Some(Value::Array(values)) => { + let values = values.iter().filter_map(Value::as_str).collect::>(); + values.contains(&"integer") && !values.contains(&"number") + } + _ => false, + } +} + +fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) { + let xtml_type = match json_type { + "string" => Some("string"), + "integer" | "number" => Some("number"), + "boolean" => Some("boolean"), + "null" => Some("null"), + "object" => Some("object"), + "array" => Some("array"), + _ => None, + }; + if let Some(xtml_type) = xtml_type + && !types.contains(&xtml_type) + { + types.push(xtml_type); + } +} + +fn string_argument_content(schema: &Value) -> Format { + let Some(schema) = schema.as_object() else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + let values = schema + .get("enum") + .and_then(Value::as_array) + .cloned() + .or_else(|| schema.get("const").cloned().map(|value| vec![value])); + let Some(values) = values else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + if values.is_empty() + || values.len() > 256 + || values + .iter() + .any(|value| value.as_str().is_none_or(|value| value.contains("<|"))) + { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + } + let values = values.iter().filter_map(Value::as_str).collect::>(); + match values.as_slice() { + [value] => Format::const_string(*value), + _ => Format::or(values.into_iter().map(Format::const_string).collect()), + } +} + +fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + json_schema(parameters.clone(), options), + JSON_CLOSE, + ) +} + +fn permissive_argument() -> Format { + let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#); + let alternatives = XTML_TYPES + .iter() + .map(|xtml_type| { + Format::sequence(vec![ + key.clone(), + Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")), + if *xtml_type == "string" { + Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]) + } else { + Format::json_schema(Value::Bool(true)) + }, + ]) + }) + .collect(); + Format::tag( + format!("{OPEN}argument key=\""), + Format::or(alternatives), + ARG_CLOSE, + ) +} + +fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { + Format::JsonSchema( + JsonSchemaFormat::new(schema) + .with_any_order(options.any_order) + .with_max_whitespace_cnt(options.max_whitespace_cnt), + ) +} + +fn root_definitions(schema: &Map) -> Map { + ["$defs", "definitions"] + .into_iter() + .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone()))) + .collect() +} + +fn attach_root_definitions(schema: &Value, root_defs: &Map) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + for (key, value) in root_defs { + schema.entry(key.clone()).or_insert_with(|| value.clone()); + } + Value::Object(schema) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::json; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::KimiK3StructuralTagBuilder; + + fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { + ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new(name).with_parameters(parameters), + )) + } + + #[test] + fn required_structural_tag_matches_xtml_channels() { + let tools = vec![tool( + "get_weather", + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap(); + + expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + } + + #[test] + fn typed_arguments_require_one_tag_only_for_nonempty_required() { + let required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query"] + }), + StructuralTagOptions::default(), + ); + let optional = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } } + }), + StructuralTagOptions::default(), + ); + let empty_required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": [] + }), + StructuralTagOptions::default(), + ); + + assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus"); + assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star"); + assert_eq!( + serde_json::to_value(empty_required).unwrap()["type"], + "star" + ); + } + + #[test] + fn reasoning_grammar_starts_inside_prefilled_think_channel() { + let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::auto(), + StructuralTagOptions::default().with_reasoning(true), + ) + .unwrap(); + let value = serde_json::to_value(tag).unwrap(); + + assert_eq!( + value["format"]["elements"][0]["end"], + "<|close|>think<|sep|>" + ); + assert_eq!( + value["format"]["elements"][1]["value"], + "<|open|>response<|sep|>" + ); + assert_eq!(value["format"]["elements"][3]["type"], "optional"); + } + + #[test] + fn forced_choice_keeps_only_the_named_tool() { + let tools = vec![ + tool("search", json!({ "type": "object" })), + tool("lookup", json!({ "type": "object" })), + ]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::function("lookup"), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains("lookup")); + assert!(!tag.contains("search")); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + let tools = vec![tool( + "set_count", + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default(), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains(r#"type=\"number\""#)); + assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); + assert!(tag.contains(r#"type=\"null\""#)); + assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); + } + + #[test] + fn unsafe_string_enum_falls_back_as_a_whole() { + let format = super::string_argument_content(&json!({ + "type": "string", + "enum": ["safe", "<|unsafe"] + })); + + assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text"); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index d5ecb6492904..ae5524759df3 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -6,17 +6,19 @@ mod combined; mod gemma4; mod inkling; +mod kimi_k3; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; pub use inkling::InklingUnifiedParser; +pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; use crate::reasoning::ReasoningError; use crate::tool::{ - StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput, + StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput, }; /// Result alias for unified parser operations. @@ -171,8 +173,8 @@ pub trait UnifiedParser: Send { false } - /// Return the xgrammar structural-tag model used for strict tool calling. - fn structural_tag_model(&self) -> Option { + /// Return the xgrammar structural-tag builder used for strict tool calling. + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { None } diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index f3e03863d495..1ed397d589cc 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -35,6 +35,7 @@ tokio-openssl.workspace = true tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true +tonic-health.workspace = true tonic-prost.workspace = true tower.workspace = true tower-http.workspace = true diff --git a/rust/src/server/build.rs b/rust/src/server/build.rs index c20ff1c86b8f..585c3c70b990 100644 --- a/rust/src/server/build.rs +++ b/rust/src/server/build.rs @@ -9,7 +9,13 @@ fn main() -> Result<(), Box> { .build_server(true) .build_client(true) .protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers - .compile_protos(&[format!("{proto_dir}/vllm_grpc.proto")], &[proto_dir])?; + .compile_protos( + &[ + format!("{proto_dir}/control.proto"), + format!("{proto_dir}/inference.proto"), + ], + &[proto_dir], + )?; Ok(()) } diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 2c2779569ac7..387776263059 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -70,6 +71,7 @@ async fn main() -> Result<()> { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: HashMap::new(), chat_template_content_format: ChatTemplateContentFormatOption::Auto, max_logprobs: None, api_server_options: ApiServerOptions::default(), diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 91373c82f2d8..1fb87f2478eb 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -10,6 +10,7 @@ use axum::http::{HeaderName, HeaderValue, Method}; use educe::Educe; use serde::Serialize; use serde_json::Value; +use vllm_chat::multimodal::MmLimitPerPrompt; use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; @@ -184,6 +185,9 @@ pub struct Config { pub chat_template: Option, /// Server-default keyword arguments merged into every chat-template render. pub default_chat_template_kwargs: Option>, + /// Maximum number of input items allowed per prompt for each modality. + /// Unspecified modalities are unlimited. + pub limit_mm_per_prompt: MmLimitPerPrompt, /// How to serialize `message.content` for chat-template rendering. pub chat_template_content_format: ChatTemplateContentFormatOption, /// Optional maximum number of top log probabilities accepted by the diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index fbd13b77bafb..3a8c0ae5c92a 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -141,6 +141,22 @@ mod tests { assert!(response.error.message.contains("max_tokens=4")); } + #[test] + fn sampling_params_validation_maps_to_invalid_request() { + let api_error = text_submit_error( + "failed to submit completion request", + vllm_text::Error::SamplingParams(vllm_text::SamplingParamsError::OutOfRange { + parameter: "top_p", + value: 0.0, + expected: "(0, 1]", + }), + ); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + let response = api_error.to_error_response(); + assert_eq!(response.error.error_type, "invalid_request_error"); + assert!(response.error.message.contains("top_p")); + } + #[test] fn chat_wrapped_prompt_too_long_maps_to_invalid_request() { let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong { diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs new file mode 100644 index 000000000000..897209822252 --- /dev/null +++ b/rust/src/server/src/grpc/control.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use thiserror_ext::AsReport as _; +use tonic::{Request, Response, Status}; +use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; + +use super::{ControlServer, pb}; +use crate::state::AppState; + +pub(crate) type ControlGrpcService = ControlServer; + +/// gRPC control service backed by the shared application state. +pub struct ControlServiceImpl { + state: Arc, +} + +impl ControlServiceImpl { + pub fn new(state: Arc) -> Self { + Self { state } + } + + fn ready(&self) -> &EngineCoreReadyResponse { + self.state.engine_core_client().ready_response() + } + + fn parallelism_info(&self) -> pb::ParallelismInfo { + let ready = self.ready(); + pb::ParallelismInfo { + tensor_parallel_size: ready.tensor_parallel_size, + pipeline_parallel_size: ready.pipeline_parallel_size, + data_parallel_size: ready.data_parallel_size.min(u64::from(u32::MAX)) as u32, + data_parallel_rank: ready.data_parallel_rank, + decode_context_parallel_size: ready.decode_context_parallel_size, + } + } +} + +const GRPC_API_VERSION: &str = "vllm"; + +#[tonic::async_trait] +impl pb::control_server::Control for ControlServiceImpl { + async fn get_server_info( + &self, + _request: Request, + ) -> Result, Status> { + let ready = self.ready(); + Ok(Response::new(pb::ServerInfo { + engine_version: ready.vllm_version.clone(), + api_version: GRPC_API_VERSION.to_string(), + instance_id: ready.instance_id.clone(), + parallelism: Some(self.parallelism_info()), + max_model_len: self.state.engine_core_client().max_model_len(), + kv_block_size: ready.block_size.min(u64::from(u32::MAX)) as u32, + total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), + max_running_requests: ready.max_num_seqs, + max_batched_tokens: ready.max_num_batched_tokens, + })) + } + + async fn get_model_info( + &self, + _request: Request, + ) -> Result, Status> { + let served = self.state.served_model_names(); + Ok(Response::new(pb::ModelInfo { + model_id: self.state.chat.text().model_id().to_string(), + served_model_name: self.state.primary_model_name().to_string(), + served_model_aliases: served.iter().skip(1).cloned().collect(), + // GenerateRequest accepts both prompt representations. + supports_text_input: true, + supports_token_ids_input: true, + supports_multimodal: self.state.chat.supports_multimodal(), + reasoning_parser: self + .state + .chat + .reasoning_parser_name() + .unwrap_or_default() + .to_string(), + tool_call_parser: self + .state + .chat + .tool_call_parser_name() + .unwrap_or_default() + .to_string(), + })) + } + + async fn abort( + &self, + request: Request, + ) -> Result, Status> { + let request_ids = request.into_inner().request_ids; + if request_ids.is_empty() { + return Ok(Response::new(pb::AbortResponse {})); + } + self.state + .chat + .abort(&request_ids) + .await + .map_err(|error| Status::internal(error.to_report_string()))?; + Ok(Response::new(pb::AbortResponse {})) + } + + async fn get_kv_event_sources( + &self, + _request: Request, + ) -> Result, Status> { + let client = self.state.engine_core_client(); + let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect(); + Ok(Response::new(pb::GetKvEventSourcesResponse { sources })) + } +} + +pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { + let config = response.kv_events_config.as_ref()?; + if !config.enable_kv_cache_events || config.publisher != "zmq" { + return None; + } + + Some(pb::KvEventSource { + transport: "zmq".to_string(), + endpoint: config.endpoint.clone(), + topic: config.topic.clone(), + replay_endpoint: config.replay_endpoint.clone().unwrap_or_default(), + data_parallel_rank: Some(response.data_parallel_rank), + encoding: "msgpack".to_string(), + schema_version: 1, + buffer_steps: config.buffer_steps, + hwm: config.hwm, + max_queue_size: config.max_queue_size, + }) +} diff --git a/rust/src/server/src/grpc/health.rs b/rust/src/server/src/grpc/health.rs new file mode 100644 index 000000000000..cd35aec972e8 --- /dev/null +++ b/rust/src/server/src/grpc/health.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use tonic::server::NamedService; +use tonic_health::ServingStatus; +use tonic_health::server::HealthReporter; +use tracing::{info, warn}; + +use super::{ControlGrpcService, InferenceGrpcService}; + +pub(crate) async fn monitor_health( + mut health_reporter: HealthReporter, + mut engine_health: watch::Receiver, + shutdown: CancellationToken, +) { + let inference_service = InferenceGrpcService::NAME; + let control_service = ControlGrpcService::NAME; + let status = ServingStatus::NotServing; + let health_event_first = tokio::select! { + result = engine_health.wait_for(|healthy| !*healthy) => { + match result { + Ok(_) => warn!( + status = ?status, + reason = "engine_unhealthy", + "marking gRPC health services as not serving" + ), + Err(error) => warn!( + %error, + status = ?status, + reason = "health_channel_closed", + "engine health channel closed; marking gRPC health services as not serving" + ), + } + true + } + _ = shutdown.cancelled() => { + info!( + status = ?status, + reason = "server_shutdown", + "server shutting down; marking gRPC health services as not serving" + ); + false + } + }; + + health_reporter.set_not_serving::().await; + health_reporter.set_not_serving::().await; + // Both gRPC services use the same engine client, so overall server health + // mirrors their shared engine health. + health_reporter.set_service_status("", status).await; + + if health_event_first { + shutdown.cancelled().await; + info!( + reason = "server_shutdown", + "server shutting down; closing gRPC health watches" + ); + } + + health_reporter.clear_service_status(inference_service).await; + health_reporter.clear_service_status(control_service).await; + health_reporter.clear_service_status("").await; +} diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs new file mode 100644 index 000000000000..56fa40924cf1 --- /dev/null +++ b/rust/src/server/src/grpc/inference.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::pin::Pin; +use std::sync::Arc; + +use futures::{Stream, StreamExt as _}; +use thiserror_ext::AsReport as _; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{Request, Response, Status}; +use tracing::info; +use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _}; + +use super::convert::{self, ResponseOpts}; +use super::{InferenceServer, pb}; +use crate::state::AppState; + +pub(crate) type InferenceGrpcService = InferenceServer; + +/// gRPC inference service backed by the shared application state. +pub struct InferenceServiceImpl { + state: Arc, +} + +impl InferenceServiceImpl { + pub fn new(state: Arc) -> Self { + Self { state } + } +} + +#[tonic::async_trait] +impl pb::inference_server::Inference for InferenceServiceImpl { + type GenerateStreamStream = + Pin> + Send>>; + + /// Unary generate: collect all output and return a single response. + async fn generate( + &self, + request: Request, + ) -> Result, Status> { + let proto_req = request.into_inner(); + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); + let text_request = + convert::to_text_request(proto_req, false, self.state.served_model_names())?; + + let request_id = text_request.request_id.clone(); + info!(%request_id, "grpc generate (unary)"); + + let stream = self.state.chat.text().generate(text_request).await; + let stream = stream.map_err(text_error_to_status)?; + + let collected = stream.collect_output().await.map_err(text_error_to_status)?; + + // Build the single aggregated response. + let prompt_info = convert::to_prompt_info( + &collected.prompt_token_ids, + collected.prompt_logprobs.as_ref(), + &response_opts, + ); + + let finish_info = vllm_text::Finished { + usage: collected.usage, + finish_reason: collected.finish_reason, + kv_transfer_params: collected.kv_transfer_params, + ec_transfer_params: collected.ec_transfer_params, + }; + + let outputs = convert::to_sequence_output( + &collected.text, + &collected.token_ids, + collected.logprobs.as_ref(), + Some(&finish_info), + &response_opts, + ); + + Ok(Response::new(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: Some(outputs), + })) + } + + /// Streaming generate: yield incremental responses as tokens are produced. + async fn generate_stream( + &self, + request: Request, + ) -> Result, Status> { + let proto_req = request.into_inner(); + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); + let text_request = + convert::to_text_request(proto_req, true, self.state.served_model_names())?; + + let request_id = text_request.request_id.clone(); + info!(%request_id, "grpc generate (stream)"); + + let stream = self.state.chat.text().generate(text_request).await; + let stream = stream.map_err(text_error_to_status)?; + + let (tx, rx) = mpsc::channel(32); + + tokio::spawn(async move { + futures::pin_mut!(stream); + while let Some(event) = stream.next().await { + let response = match event { + Err(e) => Err(text_error_to_status(e)), + Ok(DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) => { + let prompt_info = convert::to_prompt_info( + &prompt_token_ids, + prompt_logprobs.as_ref(), + &response_opts, + ); + Ok(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: None, + }) + } + Ok(DecodedTextEvent::TextDelta { + delta, + token_ids, + logprobs, + finished, + }) => Ok(pb::GenerateResponse { + prompt_info: None, + outputs: Some(convert::to_sequence_output( + &delta, + &token_ids, + logprobs.as_ref(), + finished.as_ref(), + &response_opts, + )), + }), + }; + + if tx.send(response).await.is_err() { + // Client disconnected. + break; + } + } + }); + + let response_stream = ReceiverStream::new(rx); + Ok(Response::new(Box::pin(response_stream))) + } +} + +fn text_error_to_status(error: vllm_text::Error) -> Status { + let message = error.to_report_string(); + if error.is_request_validation_error() { + Status::invalid_argument(message) + } else { + Status::internal(message) + } +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 023aae38ac07..06c7c9259eb8 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -1,167 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade. +//! gRPC services backed by the shared application state. +mod control; mod convert; - -use std::pin::Pin; -use std::sync::Arc; - -use futures::{Stream, StreamExt as _}; -use thiserror_ext::AsReport as _; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; -use tonic::{Request, Response, Status}; -use tracing::info; -use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _}; - -use self::convert::ResponseOpts; -use crate::state::AppState; +mod health; +mod inference; /// Generated protobuf/gRPC types for the `vllm` package. pub mod pb { tonic::include_proto!("vllm"); } -pub use pb::generate_server::GenerateServer; +pub(crate) use control::ControlGrpcService; +pub use control::ControlServiceImpl; +pub(crate) use health::monitor_health; +pub(crate) use inference::InferenceGrpcService; +pub use inference::InferenceServiceImpl; +pub use pb::control_server::ControlServer; +pub use pb::inference_server::InferenceServer; #[cfg(test)] mod tests; - -/// gRPC Generate service implementation backed by the shared application state. -pub struct GenerateServiceImpl { - state: Arc, -} - -impl GenerateServiceImpl { - pub fn new(state: Arc) -> Self { - Self { state } - } -} - -#[tonic::async_trait] -impl pb::generate_server::Generate for GenerateServiceImpl { - type GenerateStreamStream = - Pin> + Send>>; - - /// Unary generate: collect all output and return a single response. - async fn generate( - &self, - request: Request, - ) -> Result, Status> { - let proto_req = request.into_inner(); - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, false, self.state.served_model_names())?; - - let request_id = text_request.request_id.clone(); - info!(%request_id, "grpc generate (unary)"); - - let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(text_error_to_status)?; - - let collected = stream.collect_output().await.map_err(text_error_to_status)?; - - // Build the single aggregated response. - let prompt_info = convert::to_prompt_info( - &collected.prompt_token_ids, - collected.prompt_logprobs.as_ref(), - &response_opts, - ); - - let finish_info = vllm_text::Finished { - usage: collected.usage, - finish_reason: collected.finish_reason, - kv_transfer_params: collected.kv_transfer_params, - ec_transfer_params: collected.ec_transfer_params, - }; - - let outputs = convert::to_sequence_output( - &collected.text, - &collected.token_ids, - collected.logprobs.as_ref(), - Some(&finish_info), - &response_opts, - ); - - Ok(Response::new(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: Some(outputs), - })) - } - - /// Streaming generate: yield incremental responses as tokens are produced. - async fn generate_stream( - &self, - request: Request, - ) -> Result, Status> { - let proto_req = request.into_inner(); - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, true, self.state.served_model_names())?; - - let request_id = text_request.request_id.clone(); - info!(%request_id, "grpc generate (stream)"); - - let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(text_error_to_status)?; - - let (tx, rx) = mpsc::channel(32); - - tokio::spawn(async move { - futures::pin_mut!(stream); - while let Some(event) = stream.next().await { - let response = match event { - Err(e) => Err(text_error_to_status(e)), - Ok(DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) => { - let prompt_info = convert::to_prompt_info( - &prompt_token_ids, - prompt_logprobs.as_ref(), - &response_opts, - ); - Ok(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: None, - }) - } - Ok(DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - }) => Ok(pb::GenerateResponse { - prompt_info: None, - outputs: Some(convert::to_sequence_output( - &delta, - &token_ids, - logprobs.as_ref(), - finished.as_ref(), - &response_opts, - )), - }), - }; - - if tx.send(response).await.is_err() { - // Client disconnected. - break; - } - } - }); - - let response_stream = ReceiverStream::new(rx); - Ok(Response::new(Box::pin(response_stream))) - } -} - -fn text_error_to_status(error: vllm_text::Error) -> Status { - let message = error.to_report_string(); - if error.is_request_validation_error() { - Status::invalid_argument(message) - } else { - Status::internal(message) - } -} diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 56d0c1c2f5a9..8265fefa050e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -16,17 +16,28 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::TcpStream; use tokio_openssl::SslStream; use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri}; +use tonic_health::pb::HealthCheckRequest; +use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus; +use tonic_health::pb::health_client::HealthClient; +use tonic_health::server::health_reporter; use tower::service_fn; use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt, }; +use vllm_engine_core_client::mock_engine::{ + DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, + default_ready_response, +}; +use vllm_engine_core_client::protocol::handshake::KvEventsConfig; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; -use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; +use vllm_engine_core_client::test_utils::{ + IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready, +}; +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode}; use vllm_llm::Llm; use vllm_text::tokenizer::DynTokenizer; use vllm_text::{Prompt, TextBackend}; @@ -34,8 +45,10 @@ use vllm_tokenizer::test_utils::TestTokenizer; use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; -use super::pb::generate_client::GenerateClient; -use super::{GenerateServer, GenerateServiceImpl, pb}; +use super::control::kv_event_source; +use super::pb::control_client::ControlClient; +use super::pb::inference_client::InferenceClient; +use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb}; use crate::listener::{Listener, MaybeTlsListener}; use crate::state::AppState; use crate::tls; @@ -149,10 +162,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.expect("recv engine message").into_vec() } -fn test_llm(client: EngineCoreClient) -> Llm { - Llm::new(client).with_request_id_randomization(false) -} - #[derive(Clone, Debug)] struct FakeTextBackend; @@ -200,7 +209,12 @@ impl ChatRenderer for FakeTextBackend { async fn setup_grpc_service( engine_id: impl Into, output_specs: Vec<(Vec, Option)>, -) -> (GenerateServer, MockEngineTask) { +) -> ( + InferenceServer, + ControlServer, + tokio::sync::watch::Receiver, + MockEngineTask, +) { let ipc = IpcNamespace::new().expect("create ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let engine_id = engine_id.into(); @@ -232,14 +246,17 @@ async fn setup_grpc_service( ) .await .expect("connect client"); + let engine_health = client.subscribe_health(); let chat = ChatLlm::from_shared_backend( - test_llm(client), + Llm::new(client), Arc::new(FakeTextBackend) as Arc, ); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); ( - GenerateServer::new(GenerateServiceImpl::new(state)), + InferenceServer::new(InferenceServiceImpl::new(state.clone())), + ControlServer::new(ControlServiceImpl::new(state)), + engine_health, engine_task, ) } @@ -250,29 +267,60 @@ async fn grpc_test_server( engine_id: impl Into, output_specs: Vec<(Vec, Option)>, ) -> ( - GenerateClient, + InferenceClient, tokio::task::JoinHandle<()>, MockEngineTask, ) { - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service(engine_id, output_specs).await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + (InferenceClient::new(channel), server_task, engine_task) +} + +async fn start_grpc_test_server( + inference_service: InferenceServer, + control_service: ControlServer, + engine_health: tokio::sync::watch::Receiver, + shutdown: tokio_util::sync::CancellationToken, +) -> (Channel, tokio::task::JoinHandle<()>) { + let (health_reporter, health_service) = health_reporter(); + health_reporter.set_serving::>().await; + health_reporter.set_serving::>().await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let addr = listener.local_addr().expect("local addr"); let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); - TonicServer::builder() - .add_service(svc) - .serve_with_incoming(incoming) - .await - .expect("grpc server"); + let server = TonicServer::builder() + .add_service(health_service) + .add_service(control_service) + .add_service(inference_service) + .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); + let health_monitor = + super::monitor_health(health_reporter, engine_health, shutdown.clone()); + let server = async move { + let result = server.await; + shutdown.cancel(); + result + }; + let (server_result, ()) = tokio::join!(server, health_monitor); + server_result.expect("grpc server"); }); - let grpc_client = GenerateClient::connect(format!("http://{addr}")) + let channel = Endpoint::from_shared(format!("http://{addr}")) + .expect("grpc endpoint") + .connect() .await - .expect("connect grpc client"); + .expect("connect grpc channel"); - (grpc_client, server_task, engine_task) + (channel, server_task) } /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode). @@ -283,7 +331,8 @@ async fn grpc_tls_test_server( certs: &TestCerts, cert_reqs: i32, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let (inference_service, control_service, _engine_health, engine_task) = + setup_grpc_service(engine_id, output_specs).await; let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs)) .expect("build grpc tls config"); @@ -293,7 +342,8 @@ async fn grpc_tls_test_server( let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context); TonicServer::builder() - .add_service(svc) + .add_service(control_service) + .add_service(inference_service) .serve_with_incoming(incoming) .await .expect("grpc tls server"); @@ -309,7 +359,7 @@ async fn grpc_tls_client( certs: &TestCerts, addr: &str, identity: Option<&str>, -) -> Result, tonic::transport::Error> { +) -> Result, tonic::transport::Error> { let ca = certs.path("ca.pem"); let identity = identity.map(|name| { ( @@ -346,7 +396,7 @@ async fn grpc_tls_client( .expect("grpc endpoint") .connect_with_connector(connector) .await?; - Ok(GenerateClient::new(channel)) + Ok(InferenceClient::new(channel)) } /// Complete a raw TLS handshake against the gRPC port (offering ALPN `h2`) for @@ -373,7 +423,8 @@ async fn grpc_server_with_keepalive( engine_id: impl Into, keepalive: Option, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await; + let (inference_service, control_service, _engine_health, engine_task) = + setup_grpc_service(engine_id, default_stream_output_specs()).await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let addr = listener.local_addr().expect("local addr").to_string(); @@ -388,7 +439,8 @@ async fn grpc_server_with_keepalive( let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); builder - .add_service(svc) + .add_service(control_service) + .add_service(inference_service) .serve_with_incoming(incoming) .await .expect("grpc server"); @@ -598,6 +650,35 @@ async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = grpc_test_server( + b"engine-grpc-invalid-sampling", + default_stream_output_specs(), + ) + .await; + + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-sampling".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + sampling: Some(pb::RandomSampling { + top_p: 2.0, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("should fail when top_p is out of range"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("top_p")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { @@ -1035,3 +1116,373 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_abort_resolves_external_id_and_empty_is_noop() { + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut inference_client = InferenceClient::new(channel.clone()); + let mut control_client = ControlClient::new(channel); + let request_id = "test-abort-active"; + + let mut stream = inference_client + .generate_stream(pb::GenerateRequest { + request_id: request_id.to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect("start generation") + .into_inner(); + + loop { + let response = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for active generation output") + .expect("read active generation output") + .expect("generation ended before producing output"); + if let Some(output) = response.outputs { + assert!( + output.finish_info.is_none(), + "generation finished before abort behavior was exercised" + ); + break; + } + } + + control_client + .abort(pb::AbortRequest::default()) + .await + .expect("empty abort should be a no-op"); + assert!( + tokio::time::timeout(Duration::from_millis(100), stream.message()) + .await + .is_err(), + "empty abort unexpectedly ended the active generation" + ); + + control_client + .abort(pb::AbortRequest { + request_ids: vec![ + request_id.to_string(), + request_id.to_string(), + "unknown".to_string(), + ], + }) + .await + .expect("abort active generation"); + + let finish_reason = loop { + let response = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for aborted generation") + .expect("read aborted generation") + .expect("generation ended without an aborted response"); + if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) { + break finish_info.finish_reason; + } + }; + assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32); + + control_client + .abort(pb::AbortRequest { + request_ids: vec![request_id.to_string()], + }) + .await + .expect("repeated abort should be idempotent"); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_reports_server_and_model_info() { + let (generate_service, control_service, engine_health, _engine_task) = + setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await; + let (channel, server_task) = start_grpc_test_server( + generate_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = ControlClient::new(channel); + + let server = client + .get_server_info(pb::GetServerInfoRequest {}) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.engine_version, "test-vllm-version"); + assert_eq!(server.api_version, "vllm"); + assert_eq!(server.instance_id, "test-instance"); + assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32); + assert_eq!(server.kv_block_size, DEFAULT_MOCK_BLOCK_SIZE as u32); + assert_eq!(server.total_kv_blocks, DEFAULT_MOCK_NUM_GPU_BLOCKS); + assert_eq!(server.max_running_requests, 256); + assert_eq!(server.max_batched_tokens, 8_192); + let parallelism = server.parallelism.expect("parallelism metadata"); + assert_eq!(parallelism.tensor_parallel_size, 1); + assert_eq!(parallelism.pipeline_parallel_size, 1); + assert_eq!(parallelism.data_parallel_size, 1); + assert_eq!(parallelism.data_parallel_rank, 0); + assert_eq!(parallelism.decode_context_parallel_size, 1); + + let model = client + .get_model_info(pb::GetModelInfoRequest {}) + .await + .expect("get model info") + .into_inner(); + assert_eq!(model.model_id, "test-model"); + assert_eq!(model.served_model_name, "test-model"); + assert!(model.served_model_aliases.is_empty()); + assert!(model.supports_text_input); + assert!(model.supports_token_ids_input); + assert!(!model.supports_multimodal); + assert!(model.reasoning_parser.is_empty()); + assert!(model.tool_call_parser.is_empty()); + + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn control_aggregates_multi_engine_capacity() { + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + + let mut ready_0 = default_ready_response(); + ready_0.max_model_len = 8_192; + ready_0.num_gpu_blocks = 10; + ready_0.data_parallel_size = 2; + + let mut ready_1 = default_ready_response(); + ready_1.max_model_len = 4_096; + ready_1.num_gpu_blocks = 20; + ready_1.data_parallel_size = 2; + ready_1.data_parallel_rank = 1; + + let engine_tasks = [ready_0, ready_1].map(|ready| { + let engine_id = EngineId::from_engine_index(ready.data_parallel_rank); + MockEngineTask::new(spawn_mock_engine_task_with_ready( + handshake_address.clone(), + engine_id, + ready, + |_, _| boxed_test_future(async {}), + )) + }); + + let client = EngineCoreClient::connect(EngineCoreClientConfig { + transport_mode: TransportMode::HandshakeOwner { + handshake_address, + advertised_host: "127.0.0.1".to_string(), + engine_count: 2, + ready_timeout: Duration::from_secs(2), + local_input_address: Some(ipc.input_endpoint()), + local_output_address: Some(ipc.output_endpoint()), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + }) + .await + .expect("connect multi-engine client"); + let chat = ChatLlm::from_shared_backend( + Llm::new(client), + Arc::new(FakeTextBackend) as Arc, + ); + let service = ControlServiceImpl::new(Arc::new(AppState::new( + vec!["test-model".to_string()], + chat, + ))); + + let server = pb::control_server::Control::get_server_info( + &service, + tonic::Request::new(pb::GetServerInfoRequest {}), + ) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.max_model_len, 4_096); + assert_eq!(server.total_kv_blocks, 30); + + drop(engine_tasks); +} + +#[test] +fn kv_event_source_filters_and_exposes_zmq_publisher() { + let mut ready = default_ready_response(); + ready.data_parallel_rank = 2; + ready.kv_events_config = Some(KvEventsConfig { + enable_kv_cache_events: false, + publisher: "null".to_string(), + endpoint: "tcp://*:5559".to_string(), + replay_endpoint: Some("tcp://*:5560".to_string()), + buffer_steps: 10_000, + hwm: 100_000, + max_queue_size: 100_000, + topic: "kv".to_string(), + }); + + assert!(kv_event_source(&ready).is_none()); + + let config = ready.kv_events_config.as_mut().unwrap(); + config.enable_kv_cache_events = true; + config.publisher = "zmq".to_string(); + let source = kv_event_source(&ready).expect("configured ZMQ event source"); + assert_eq!(source.transport, "zmq"); + assert_eq!(source.endpoint, "tcp://*:5559"); + assert_eq!(source.topic, "kv"); + assert_eq!(source.replay_endpoint, "tcp://*:5560"); + assert_eq!(source.data_parallel_rank, Some(2)); + assert_eq!(source.encoding, "msgpack"); + assert_eq!(source.schema_version, 1); + assert_eq!(source.buffer_steps, 10_000); + assert_eq!(source.hwm, 100_000); + assert_eq!(source.max_queue_size, 100_000); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() { + let (inference_service, control_service, _connected_engine_health, _engine_task) = + setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await; + let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true); + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut health_client = HealthClient::new(channel); + + let mut health_streams = Vec::new(); + for service in ["vllm.Inference", "vllm.Control", ""] { + let service_label = if service.is_empty() { + "overall" + } else { + service + }; + let mut stream = health_client + .watch(HealthCheckRequest { + service: service.to_string(), + }) + .await + .unwrap_or_else(|error| { + panic!("failed to start health watch for {service_label}: {error}") + }) + .into_inner(); + let initial = stream + .message() + .await + .unwrap_or_else(|error| { + panic!("failed to read initial health status for {service_label}: {error}") + }) + .unwrap_or_else(|| { + panic!("health watch for {service_label} ended before its initial status") + }); + assert_eq!( + initial.status, + HealthServingStatus::Serving as i32, + "unexpected initial health status for {service_label}" + ); + health_streams.push((service_label, stream)); + } + + engine_health_tx.send(false).expect("publish unhealthy engine state"); + + for (service_label, mut stream) in health_streams { + let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}")) + .unwrap_or_else(|error| { + panic!("failed to read health update for {service_label}: {error}") + }) + .unwrap_or_else(|| panic!("health watch for {service_label} ended before its update")); + assert_eq!( + update.status, + HealthServingStatus::NotServing as i32, + "unexpected health status for {service_label}" + ); + } + + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn grpc_health_watch_closes_on_graceful_shutdown() { + let (inference_service, control_service, engine_health, _engine_task) = setup_grpc_service( + b"engine-grpc-health-shutdown", + default_stream_output_specs(), + ) + .await; + let shutdown = tokio_util::sync::CancellationToken::new(); + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + shutdown.clone(), + ) + .await; + let mut health_client = HealthClient::new(channel); + let mut stream = health_client + .watch(HealthCheckRequest { + service: "vllm.Inference".to_string(), + }) + .await + .expect("start health watch for vllm.Inference") + .into_inner(); + + let initial = stream + .message() + .await + .expect("read initial health status for vllm.Inference") + .expect("health watch ended before its initial status"); + assert_eq!( + initial.status, + HealthServingStatus::Serving as i32, + "unexpected initial health status for vllm.Inference" + ); + + shutdown.cancel(); + + let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for shutdown health update for vllm.Inference") + .expect("failed to read shutdown health update for vllm.Inference") + .expect("health watch ended before its shutdown update"); + assert_eq!( + update.status, + HealthServingStatus::NotServing as i32, + "unexpected shutdown health status for vllm.Inference" + ); + + let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for vllm.Inference health watch to close") + .expect("failed while closing vllm.Inference health watch"); + assert!( + stream_end.is_none(), + "vllm.Inference health watch remained open" + ); + + tokio::time::timeout(Duration::from_secs(2), server_task) + .await + .expect("timed out waiting for gRPC server shutdown") + .expect("gRPC server task failed"); +} diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 7f1ec91c6f2d..224ec672aec7 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -39,6 +39,7 @@ use tokio::net::TcpListener; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tonic::transport::Server as TonicServer; +use tonic_health::server::health_reporter; use tower::ServiceExt as _; use tracing::{info, trace, warn}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; @@ -100,6 +101,7 @@ async fn build_state(config: &Config) -> Result> { .default_chat_template_kwargs .clone() .unwrap_or_default(), + limit_mm_per_prompt: config.limit_mm_per_prompt.clone(), }, ) .await @@ -179,14 +181,17 @@ where result = build_state(&config) => result?, _ = shutdown.cancelled() => return Ok(()), }; + let model = state.primary_model_name().to_owned(); + let app = extend_router(build_router(state.clone())); + + info!(model, "starting vLLM server"); + let listener = Listener::bind(&config.listener_mode) .await .context("failed to bind listener for OpenAI server")?; let bind_address = listener.local_addr_display()?; - let model = state.primary_model_name().to_owned(); - let app = extend_router(build_router(state.clone())); - // Optionally bind the gRPC Generate server on a separate port. Bind + // Optionally bind the gRPC Inference server on a separate port. Bind // synchronously here so bind errors (port in use, permission denied, ...) // surface before serving rather than being deferred until shutdown. let grpc_setup = if let Some(grpc_port) = config.grpc_port { @@ -203,14 +208,29 @@ where .map(tls::build_grpc_server_config) .transpose() .context("invalid gRPC TLS configuration")?; - let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone())); + let (health_reporter, health_service) = health_reporter(); + let engine_health = state.engine_core_client().subscribe_health(); + health_reporter.set_serving::().await; + health_reporter.set_serving::().await; + let control_service = + grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone())); + let inference_service = + grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone())); let svc = TonicServer::builder() .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) .layer(middleware::request_runtime_layer(state.clone())) - .add_service(svc); - info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); - Some((grpc_listener, svc, grpc_tls)) + .add_service(health_service) + .add_service(control_service) + .add_service(inference_service); + Some(( + addr, + grpc_listener, + svc, + grpc_tls, + health_reporter, + engine_health, + )) } else { None }; @@ -220,7 +240,7 @@ where } else { "http" }; - info!(%bind_address, %scheme, %model, "starting OpenAI server"); + let model = model.as_str(); // Run HTTP and gRPC concurrently under a child token of the caller's shutdown // token. Caller cancellation propagates into both protocols; if either @@ -274,6 +294,11 @@ where }; let server = serve_connections(listener, app, shutdown.cancelled_owned(), timeouts); + info!( + bind_address, + scheme, model, "OpenAI server is ready to accept requests" + ); + let result = tokio::select! { result = server => { result.context("HTTP server failed") @@ -294,29 +319,41 @@ where let server_shutdown = server_shutdown.clone(); let force_shutdown = force_shutdown.clone(); async move { - let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else { + let Some((addr, grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = + grpc_setup + else { // No gRPC configured: just wait for shutdown so we do not race the // join! by resolving early and tripping the cancellation token. shutdown.cancelled().await; return Ok(()); }; + let tls = grpc_tls.is_some(); let incoming = match grpc_tls { Some(context) => MaybeTlsListener::tls(grpc_listener, context), None => MaybeTlsListener::plain(grpc_listener), }; - let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()); - - let result = tokio::select! { - result = server => { - result.context("gRPC server failed") - } - _ = force_shutdown.cancelled() => { - warn!("gRPC graceful shutdown deadline elapsed; aborting server"); - Ok(()) - } + let server = + svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); + let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown); + + info!(%addr, tls, model, "gRPC server is ready to accept requests"); + + let server = async move { + let result = tokio::select! { + result = server => { + result.context("gRPC server failed") + } + _ = force_shutdown.cancelled() => { + warn!("gRPC graceful shutdown deadline elapsed; aborting server"); + Ok(()) + } + }; + + server_shutdown.cancel(); + result }; - server_shutdown.cancel(); + let (result, ()) = tokio::join!(server, health_monitor); result } }; diff --git a/rust/src/server/src/middleware/offload.rs b/rust/src/server/src/middleware/offload.rs index edc6560bf20a..c310ab442fc7 100644 --- a/rust/src/server/src/middleware/offload.rs +++ b/rust/src/server/src/middleware/offload.rs @@ -30,8 +30,8 @@ const OFFLOADED_PATHS: &[&str] = &[ "/detokenize", "/inference/v1/generate", // gRPC routes: - "/vllm.Generate/Generate", - "/vllm.Generate/GenerateStream", + "/vllm.Inference/Generate", + "/vllm.Inference/GenerateStream", ]; /// Return a Tower layer that runs selected data-plane requests on the request runtime, @@ -124,8 +124,8 @@ mod tests { assert!(should_offload("/tokenize")); assert!(should_offload("/detokenize")); assert!(should_offload("/inference/v1/generate")); - assert!(should_offload("/vllm.Generate/Generate")); - assert!(should_offload("/vllm.Generate/GenerateStream")); + assert!(should_offload("/vllm.Inference/Generate")); + assert!(should_offload("/vllm.Inference/GenerateStream")); } #[test] diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 6fcdb427ccf4..e176938db400 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -238,13 +238,18 @@ fn collect_generate( None }; let prompt_logprobs = if include_prompt_logprobs { - let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| { - ApiError::server_error( - "raw generate response requested prompt_logprobs but generation returned none" - .to_string(), - ) - })?; - Some(raw_prompt_logprobs_to_maps(prompt_logprobs)) + match collected.prompt_logprobs.as_ref() { + Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)), + // A single-token prompt has no scored positions; same mapping + // as /v1/completions. + None if collected.prompt_token_ids.len() == 1 => Some(vec![None]), + None => { + return Err(ApiError::server_error( + "raw generate response requested prompt_logprobs but generation returned none" + .to_string(), + )); + } + } } else { None }; @@ -472,4 +477,48 @@ mod tests { Some(2) ); } + + #[test] + fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() { + let output_without_payload = |prompt_token_ids: Vec| CollectedGenerateOutput { + request_id: "raw-1".to_string(), + prompt_logprobs: None, + token_ids: vec![3], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: vllm_llm::TokenUsage { + prompt_token_count: prompt_token_ids.len(), + output_token_count: 1, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids, + }; + + let response = collect_generate( + output_without_payload(vec![9707]), + "raw-1".to_string(), + ApiServerOptions::default(), + ResponseOptions { + include_prompt_logprobs: true, + ..Default::default() + }, + ) + .expect("single-token prompt without payload maps to [None]"); + let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present"); + assert_eq!(prompt_logprobs.len(), 1); + assert!(prompt_logprobs[0].is_none()); + + collect_generate( + output_without_payload(vec![9707, 11]), + "raw-2".to_string(), + ApiServerOptions::default(), + ResponseOptions { + include_prompt_logprobs: true, + ..Default::default() + }, + ) + .expect_err("multi-token prompt without payload is an engine failure"); + } } diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index b08c533a51bd..8a1e9c4383a1 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{ ChatMessageDelta, }; use crate::routes::openai::utils::logprobs::{ - decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps, + decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps, }; use crate::routes::openai::utils::types::{ ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage, @@ -131,6 +131,7 @@ async fn collect_chat_completion( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, ) -> Result { let collected = stream.collect_message().await.map_err(|error| { @@ -157,7 +158,9 @@ async fn collect_chat_completion( // When reasoning is hidden, omit them rather than leaking hidden reasoning // tokens through per-token metadata. let include_output_metadata = include_reasoning || reasoning.is_none(); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?.to_string(); + let finish_reason = + chat_finish_reason_to_openai(&finish_reason, saw_tool_calls && !is_named_tool_choice)? + .to_string(); let tool_calls = message .tool_calls() .map(|call| ToolCall { @@ -181,14 +184,11 @@ async fn collect_chat_completion( None }; let prompt_logprobs = if include_prompt_logprobs { - Some(decoded_prompt_logprobs_to_maps( - prompt_logprobs.as_ref().ok_or_else(|| { - server_error!( - "chat response requested prompt_logprobs but generation returned none" - ) - })?, + Some(prompt_logprobs_to_maps( + prompt_logprobs.as_ref(), + &prompt_token_ids, return_tokens_as_token_ids, - )) + )?) } else { None }; @@ -257,6 +257,7 @@ async fn chat_completion_chunk_stream( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, mut y: TryYielder, ) -> Result<(), ApiError> { @@ -457,7 +458,7 @@ async fn chat_completion_chunk_stream( &response_model, created, finish_reason, - saw_tool_calls, + saw_tool_calls && !is_named_tool_choice, ) { Ok(chunk) => yield_chunk!(chunk), Err(error) => { @@ -790,10 +791,10 @@ fn final_chunk( response_model: &str, created: u64, finish_reason: FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result { let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?; + let finish_reason = chat_finish_reason_to_openai(&finish_reason, use_tool_calls_finish_reason)?; debug!( finish_reason = %finish_reason, @@ -812,10 +813,10 @@ fn final_chunk( fn chat_finish_reason_to_openai( finish_reason: &FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result<&'static str, ApiError> { match finish_reason { - FinishReason::Stop(_) if saw_tool_calls => Ok("tool_calls"), + FinishReason::Stop(_) if use_tool_calls_finish_reason => Ok("tool_calls"), FinishReason::Stop(_) => Ok("stop"), FinishReason::Length => Ok("length"), FinishReason::Abort => Ok("abort"), diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 5b2fa3c19ed7..2c4050f8f1e3 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -54,6 +54,8 @@ pub(super) struct ResponseOptions { pub return_token_ids: bool, /// Whether to format logprob tokens as `token_id:{id}`. pub return_tokens_as_token_ids: bool, + /// Whether the request forces one named function tool. + pub is_named_tool_choice: bool, } /// Validate and lower one OpenAI chat completion request into the internal chat @@ -87,6 +89,15 @@ pub(super) fn prepare_chat_request( )?; let template_kwargs = request.chat_template_kwargs.unwrap_or_default(); + let response_format = + request.response_format.as_ref().map(serde_json::to_value).transpose().map_err( + |error| { + ApiError::invalid_request( + format!("failed to serialize response_format: {error}"), + Some("response_format"), + ) + }, + )?; let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) @@ -98,6 +109,7 @@ pub(super) fn prepare_chat_request( .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); let requested_logprobs = request.logprobs; + let is_named_tool_choice = matches!(&request.tool_choice, Some(ToolChoice::Function { .. })); // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's // behavior. @@ -147,6 +159,7 @@ pub(super) fn prepare_chat_request( generation_prompt_mode, chat_template: request.chat_template, reasoning_effort: request.reasoning_effort, + response_format, template_kwargs, }, tools: convert_tools(request.tools)?, @@ -180,6 +193,7 @@ pub(super) fn prepare_chat_request( echo, return_token_ids: request.return_token_ids.unwrap_or(false), return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false), + is_named_tool_choice, }, chat_request, }) @@ -396,6 +410,7 @@ fn convert_tool_choice(tool_choice: Option<&ToolChoice>) -> Result, - prompt_token_ids: &[u32], - return_tokens_as_token_ids: bool, -) -> Result>>, ApiError> { - if let Some(prompt_logprobs) = prompt_logprobs { - return Ok(decoded_prompt_logprobs_to_maps( - prompt_logprobs, - return_tokens_as_token_ids, - )); - } - - if let [_token_id] = prompt_token_ids { - return Ok(vec![None]); - } - - Err(server_error!( - "completion response requested prompt_logprobs but generation returned none" - )) -} - fn usage_chunk( request_id: &str, response_model: &str, diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs index 078c42b507fd..e300e124b1ab 100644 --- a/rust/src/server/src/routes/openai/utils/logprobs.rs +++ b/rust/src/server/src/routes/openai/utils/logprobs.rs @@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai( }) } -/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response -/// shape. -pub fn decoded_prompt_logprobs_to_maps( - prompt_logprobs: &DecodedPromptLogprobs, +/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a +/// missing single-token payload as `[None]`. +pub fn prompt_logprobs_to_maps( + prompt_logprobs: Option<&DecodedPromptLogprobs>, + prompt_token_ids: &[u32], return_tokens_as_token_ids: bool, -) -> Vec>> { - std::iter::once(None) - .chain(prompt_logprobs.scored_positions.iter().map(|position| { - Some(position_top_logprobs_map( - position, - return_tokens_as_token_ids, - )) - })) - .collect() +) -> Result>>, ApiError> { + if let Some(prompt_logprobs) = prompt_logprobs { + return Ok(std::iter::once(None) + .chain(prompt_logprobs.scored_positions.iter().map(|position| { + Some(position_top_logprobs_map( + position, + return_tokens_as_token_ids, + )) + })) + .collect()); + } + + if let [_token_id] = prompt_token_ids { + return Ok(vec![None]); + } + + Err(server_error!( + "prompt_logprobs were requested but generation returned none" + )) } /// Convert decoded token-position logprobs into the OpenAI chat `logprobs` @@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 { mod tests { use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; - use super::decoded_logprobs_to_openai_chat; + use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps}; + + #[test] + fn prompt_logprobs_maps_reject_missing_multi_token_payload() { + prompt_logprobs_to_maps(None, &[9707, 11], false) + .expect_err("multi-token prompt without payload is an engine failure"); + } fn sample_logprobs() -> DecodedLogprobs { DecodedLogprobs { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 64e904b562b6..e4e7fed9960d 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -153,6 +153,20 @@ fn default_stream_output_specs() -> Vec<(Vec, Option Vec<(Vec, Option)> { + vec![ + (bytes_to_token_ids(b"Need tool."), None), + ( + bytes_to_token_ids(b"\n{\"name\":\"get_weather\", "), + None, + ), + ( + bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n"), + Some(EngineCoreFinishReason::Stop), + ), + ] +} + fn assert_adapter_a_lora_request(request: &EngineCoreRequest) { let lora = request.lora_request.as_ref().expect("lora request"); assert_eq!(lora.lora_name, "adapter-a"); @@ -563,6 +577,12 @@ fn render_fake_content(content: &ChatContent, placeholder: &str) -> vllm_chat::R } fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo { + qwen_multimodal_model_info_with_limits(std::collections::HashMap::new()) +} + +fn qwen_multimodal_model_info_with_limits( + limit_mm_per_prompt: vllm_chat::multimodal::MmLimitPerPrompt, +) -> vllm_chat::multimodal::MultimodalModelInfo { let config_path = std::env::temp_dir().join(format!( "vllm-server-qwen-config-{}.json", uuid::Uuid::new_v4() @@ -580,6 +600,7 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo { ..Default::default() }, Arc::new(fake_chat_tokenizer()), + limit_mm_per_prompt, ) .expect("load multimodal info") .expect("qwen multimodal info is registered"); @@ -2293,6 +2314,74 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() { assert_eq!(json["choices"][0]["message"]["content"], "hi"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt() { + // The request is rejected by `--limit-mm-per-prompt` validation before + // ever reaching the engine, so the mock engine task is never awaited. + let (chat, _engine_task) = test_models_with_engine_outputs_and_backend( + b"engine-openai-mm-limit", + default_stream_output_specs(), + Arc::new(FakeChatBackend::with_multimodal_model_info( + qwen_multimodal_model_info_with_limits(std::collections::HashMap::from([( + vllm_chat::multimodal::MmLimitModality::Image, + vllm_chat::multimodal::MmLimitSpec::Count(1), + )])), + )), + ) + .await; + let app = build_router(Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + ))); + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": false, + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "describe "}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" + } + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" + } + } + ] + }] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["error"]["type"], "invalid_request_error"); + assert_eq!( + json["error"]["message"], + "At most 1 image(s) may be provided in one prompt." + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() { @@ -4554,17 +4643,7 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() { async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { let (app, engine_task) = test_app_with_backend_and_stream_output_specs( Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), - vec![ - (bytes_to_token_ids(b"Need tool."), None), - ( - bytes_to_token_ids(b"\n{\"name\":\"get_weather\", "), - None, - ), - ( - bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n"), - Some(EngineCoreFinishReason::Stop), - ), - ], + weather_tool_call_output_specs(), ) .await; @@ -4613,6 +4692,63 @@ async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { assert!(text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn named_tool_choice_uses_stop_finish_reason() { + for stream in [false, true] { + let (app, engine_task) = test_app_with_backend_and_stream_output_specs( + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), + weather_tool_call_output_specs(), + ) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": stream, + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}} + } + } + }], + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"} + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + + assert!(text.contains("\"tool_calls\":"), "{text}"); + assert!(text.contains("\"name\":\"get_weather\""), "{text}"); + assert!(text.contains("\"finish_reason\":\"stop\""), "{text}"); + assert!(!text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn tool_call_sse_chunks_can_carry_logprobs() { diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 242a7af47335..6684aa55d26a 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -82,6 +82,7 @@ impl TokenizeChatRequest { generation_prompt_mode, chat_template: self.chat_template, reasoning_effort: None, + response_format: None, template_kwargs: self.chat_template_kwargs.unwrap_or_default(), }, tools: convert_tools(self.tools)?, diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 49ae5dbd6b95..4fc7a18753a4 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -177,6 +177,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index e72e68196c31..6b385fad68d4 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -6,6 +6,7 @@ use vllm_engine_core_client::Error as EngineCoreError; use vllm_llm::Error as LlmError; pub use crate::lower::logprobs::LogprobsError; +pub use crate::lower::sampling::SamplingParamsError; pub use crate::lower::token_ids::TokenIdsError; #[derive(Debug, Error)] @@ -23,6 +24,8 @@ pub enum Error { Logprobs(#[from] LogprobsError), #[error(transparent)] TokenIds(#[from] TokenIdsError), + #[error(transparent)] + SamplingParams(#[from] SamplingParamsError), #[error( "`min_tokens` must be less than or equal to `max_tokens`, \ got min_tokens={min_tokens}, max_tokens={max_tokens}" @@ -50,6 +53,7 @@ impl Error { | Self::EmptyPromptTokenIds { .. } | Self::Logprobs(_) | Self::TokenIds(_) + | Self::SamplingParams(_) | Self::MinTokensExceedsMaxTokens { .. } | Self::InvalidThinkingTokenBudget | Self::InvalidRepetitionDetection { .. } diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 1ccd53127acd..646e4bac9f3e 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -10,7 +10,7 @@ use std::mem::take; pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend}; -pub use error::{Error, LogprobsError, Result, TokenIdsError}; +pub use error::{Error, LogprobsError, Result, SamplingParamsError, TokenIdsError}; use futures::Stream; pub use lower::{ PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens, @@ -38,22 +38,90 @@ trait_set! { pub trait TextOutputStream = Stream> + Send + 'static; } -/// Raw text facade above [`Llm`]. -/// -/// This layer stays below chat semantics: prompt text or prompt token IDs flow -/// in, decoded text deltas and terminal metadata flow out. -pub struct TextLlm { - /// Generate-only client owned by this text facade. - llm: Llm, +/// Text request preparation shared by inference and render-only frontends. +pub struct TextRequestProcessor { /// Tokenizer/model metadata backend responsible for prompt encode/decode /// and sampling hints. backend: DynTextBackend, /// Runtime context window size reported by the engine startup handshake. + /// Render-only frontends supply the downstream engine's effective value. max_model_len: u32, /// Maximum number of top log probabilities accepted by this text facade. max_logprobs: i32, } +impl TextRequestProcessor { + /// Create a processor with the effective model context length. + pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self { + Self { + backend, + max_model_len, + max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, + } + } + + /// Override the maximum accepted logprobs count. + pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { + if let Some(max_logprobs) = max_logprobs { + self.max_logprobs = max_logprobs; + } + self + } + + /// Return the tokenizer used by this processor. + pub fn tokenizer(&self) -> DynTokenizer { + self.backend.tokenizer() + } + + /// Return the effective model context length. + pub fn max_model_len(&self) -> u32 { + self.max_model_len + } + + /// Tokenize and lower one request without submitting it to an engine. + pub fn prepare(&self, mut request: TextRequest) -> Result { + 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)?, + // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark + // and infra workloads bypass chat rendering and tokenizer overhead entirely. + Prompt::TokenIds(token_ids) => token_ids, + }; + let sampling_hints = self.backend.sampling_hints()?; + let sampling_limits = SamplingLimits { + max_model_len: self.max_model_len, + max_logprobs: self.max_logprobs, + model_vocab_size: self.backend.model_vocab_size(), + tokenizer_vocab_size: self.backend.tokenizer_vocab_size(), + }; + + lower_text_request( + request, + prompt_token_ids, + sampling_hints, + sampling_limits, + tokenizer.as_ref(), + ) + } +} + +/// Raw text facade above [`Llm`]. +/// +/// This layer stays below chat semantics: prompt text or prompt token IDs flow +/// in, decoded text deltas and terminal metadata flow out. +pub struct TextLlm { + /// Generate-only client owned by this text facade. + llm: Llm, + /// Shared engine-free request preparation. + processor: TextRequestProcessor, +} + impl TextLlm { /// Create a new text-generation facade from a shared LLM client plus a text /// backend. @@ -64,23 +132,19 @@ impl TextLlm { Self { llm, - backend, - max_model_len, - max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, + processor: TextRequestProcessor::new(backend, max_model_len), } } /// Override the maximum accepted logprobs count. pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { - if let Some(max_logprobs) = max_logprobs { - self.max_logprobs = max_logprobs; - } + self.processor = self.processor.with_max_logprobs(max_logprobs); self } /// Return the backend model ID. pub fn model_id(&self) -> &str { - self.backend.model_id() + self.processor.backend.model_id() } /// Expose the underlying engine-core client for low-level utility/admin @@ -91,19 +155,19 @@ impl TextLlm { /// Return the tokenizer used by this text backend. pub fn tokenizer(&self) -> DynTokenizer { - self.backend.tokenizer() + self.processor.tokenizer() } /// Tokenizer vocabulary size (the number of tokens the tokenizer knows), /// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`. pub fn tokenizer_vocab_size(&self) -> usize { - self.backend.tokenizer_vocab_size() + self.processor.backend.tokenizer_vocab_size() } /// Model vocabulary size from the model config, used to bound generated /// token IDs and logits-domain sampling controls. pub fn model_vocab_size(&self) -> usize { - self.backend.model_vocab_size() + self.processor.backend.model_vocab_size() } /// Tokenize if needed, lower to a generate request, and return the raw @@ -117,7 +181,7 @@ impl TextLlm { /// incrementally decoded text. pub async fn generate(&self, request: TextRequest) -> Result { let (text_request, raw_stream) = self.generate_inner(request).await?; - let tokenizer = self.backend.tokenizer(); + let tokenizer = self.processor.tokenizer(); let decoded_stream = output::decoded_text_event_stream( text_request.request_id, tokenizer, @@ -131,40 +195,12 @@ impl TextLlm { async fn generate_inner( &self, - mut request: TextRequest, + request: TextRequest, ) -> 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)?, - // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark - // and infra workloads bypass chat rendering and tokenizer overhead entirely. - Prompt::TokenIds(token_ids) => token_ids, - }; - - let sampling_hints = self.backend.sampling_hints()?; - let sampling_limits = SamplingLimits { - max_model_len: self.max_model_len, - max_logprobs: self.max_logprobs, - model_vocab_size: self.backend.model_vocab_size(), - tokenizer_vocab_size: self.backend.tokenizer_vocab_size(), - }; - let PreparedTextRequest { text_request, generate_request, - } = lower_text_request( - request, - prompt_token_ids, - sampling_hints, - sampling_limits, - &*tokenizer, - )?; + } = self.processor.prepare(request)?; let raw_stream = self.llm.generate(generate_request).await?; Ok((text_request, raw_stream)) diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index bd43a1d141ca..aa54595acdaa 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -4,9 +4,11 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; +pub(crate) mod sampling; pub(crate) mod token_ids; use logprobs::validate_logprobs; +use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ EngineCoreSamplingParams, RepetitionDetectionParams, @@ -186,6 +188,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, }; + validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; Ok(params) } @@ -319,7 +322,7 @@ mod tests { use super::*; use crate::backend::hf::HfTextBackend; use crate::backend::{SamplingHints, TextBackend as _}; - use crate::error::{LogprobsError, TokenIdsError}; + use crate::error::{LogprobsError, SamplingParamsError, TokenIdsError}; use crate::request::{Prompt, TextRequest}; fn stub_tokenizer() -> TestTokenizer { @@ -482,6 +485,120 @@ mod tests { assert!(message.contains("min_count=1")); } + #[test] + fn lower_sampling_params_rejects_invalid_sampling_ranges() { + let cases = [ + ( + "temperature", + SamplingParams { + temperature: Some(5.0), + ..SamplingParams::default() + }, + ), + ( + "top_p", + SamplingParams { + top_p: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "min_p", + SamplingParams { + min_p: Some(2.0), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "frequency_penalty", + SamplingParams { + frequency_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ( + "presence_penalty", + SamplingParams { + presence_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ]; + + for (expected_parameter, sampling_params) in cases { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::OutOfRange { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should be rejected" + ); + } + } + + #[test] + fn lower_sampling_params_rejects_non_finite_sampling_values() { + for (expected_parameter, sampling_params) in [ + ( + "temperature", + SamplingParams { + temperature: Some(f32::INFINITY), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(f32::NAN), + ..SamplingParams::default() + }, + ), + ] { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::NotFinite { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should reject non-finite values" + ); + } + } + + #[test] + fn lower_sampling_params_accepts_python_compatible_repetition_penalty_above_two() { + let params = lower_sampling_params_with_limits( + SamplingParams { + repetition_penalty: Some(2.5), + ..SamplingParams::default() + }, + sample_sampling_limits(), + ) + .unwrap(); + + assert_eq!(params.repetition_penalty, 2.5); + } + #[test] fn lower_text_request_applies_python_style_eos_hints() { let prepared = lower_text_request( diff --git a/rust/src/text/src/lower/sampling.rs b/rust/src/text/src/lower/sampling.rs new file mode 100644 index 000000000000..edcdc4b0492f --- /dev/null +++ b/rust/src/text/src/lower/sampling.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use thiserror::Error; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; + +#[derive(Debug, Error, PartialEq)] +pub enum SamplingParamsError { + #[error("{parameter} must be a finite number, got {value}")] + NotFinite { parameter: &'static str, value: f32 }, + #[error("{parameter} must be in {expected}, got {value}")] + OutOfRange { + parameter: &'static str, + value: f32, + expected: &'static str, + }, +} + +fn validate_frequency_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("frequency_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_presence_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("presence_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_temperature(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("temperature", value)?; + validate_closed_range("temperature", value, 0.0, 2.0, "[0, 2]") +} + +fn validate_top_p(value: f32) -> Result<(), SamplingParamsError> { + if value > 0.0 && value <= 1.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "top_p", + value, + expected: "(0, 1]", + }) +} + +fn validate_min_p(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("min_p", value, 0.0, 1.0, "[0, 1]") +} + +fn validate_repetition_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("repetition_penalty", value)?; + if value > 0.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "repetition_penalty", + value, + expected: "(0, inf)", + }) +} + +pub(crate) fn validate_resolved_sampling_params( + params: &EngineCoreSamplingParams, +) -> Result<(), SamplingParamsError> { + validate_temperature(params.temperature)?; + validate_top_p(params.top_p)?; + validate_min_p(params.min_p)?; + validate_frequency_penalty(params.frequency_penalty)?; + validate_presence_penalty(params.presence_penalty)?; + validate_repetition_penalty(params.repetition_penalty) +} + +fn validate_finite(parameter: &'static str, value: f32) -> Result<(), SamplingParamsError> { + if value.is_finite() { + return Ok(()); + } + Err(SamplingParamsError::NotFinite { parameter, value }) +} + +fn validate_closed_range( + parameter: &'static str, + value: f32, + min: f32, + max: f32, + expected: &'static str, +) -> Result<(), SamplingParamsError> { + if value >= min && value <= max { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter, + value, + expected, + }) +} diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 9d19741b7b01..cb0fe9c18e7f 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -306,6 +306,12 @@ pub async fn decoded_text_event_stream( /// If stop string matches, returns tuple /// (index into stop string vec, byte index of first byte of stop string in /// output) +/// +/// When several stop strings match within the newly generated text (for example +/// when a step appends multiple tokens, or one token decodes to several +/// characters), the stop string that completes earliest in the text is selected, +/// so the result matches appending one character at a time. Ties are broken by +/// stop-list order. fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Option<(usize, usize)> { // We compare byte subslices to avoid utf8 boundary problem let output = output.as_bytes(); @@ -314,12 +320,15 @@ fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Opti .iter() .map(|ss| (ss.as_bytes(), ss.len(), next_off.saturating_sub(ss.len()))) .enumerate() - .find_map(|(ss_idx, (ss, len, start_off))| { + .filter_map(|(ss_idx, (ss, len, start_off))| { output[start_off..] .windows(len) .position(|w| w == ss) - .map(|pos| (ss_idx, start_off + pos)) + .map(|pos| (ss_idx, start_off + pos, start_off + pos + len)) }) + // `min_by_key` keeps the first minimum, so ties fall back to stop-list order. + .min_by_key(|&(_, _, end)| end) + .map(|(ss_idx, start, _)| (ss_idx, start)) } #[cfg(test)] @@ -524,11 +533,49 @@ mod tests { #[test] fn stop_string_matches_first_of_multiple() { let stops = vec!["wor".to_string(), "say".to_string()]; - // "say" appears earlier but "wor" is checked first (index 0) + // Only "wor" can complete in the 1-new-byte window; the "say" at index 0 + // is behind the window start and is not a candidate. let result = matches_stop_string(&stops, "say wor", 1); assert_eq!(result, Some((0, 4))); } + /// Several stop strings can land in the same window when a step appends + /// multiple tokens (speculative decoding) or when one token decodes to + /// several characters. The earliest-completing one must win over stop-list + /// order, so that a batched step agrees with character-at-a-time appending. + #[test] + fn stop_string_earliest_completing_wins_regardless_of_list_order() { + // " The user is a": " is a" (5 bytes) was appended in one step. Both + // "is" (index 10, completes at 12) and " a" (index 12, completes at 14) + // land in the window. "is" completes earlier, so it wins either order. + for stops in [vec!["a", "is"], vec!["is", "a"]] { + let owned: Vec = stops.iter().map(|s| s.to_string()).collect(); + let is_idx = stops.iter().position(|s| *s == "is").unwrap(); + let result = matches_stop_string(&owned, " The user is a", " is a".len()); + assert_eq!(result, Some((is_idx, 10)), "stop list order {stops:?}"); + } + } + + /// Both stops complete at the same offset, so stop-list order decides. This + /// pins `min_by_key`'s first-minimum behavior, which is what implements the + /// tie-break. + #[test] + fn stop_string_ties_broken_by_list_order() { + for (stops, expected) in [(vec!["ab", "b"], (0, 0)), (vec!["b", "ab"], (0, 1))] { + let owned: Vec = stops.iter().map(|s| s.to_string()).collect(); + let result = matches_stop_string(&owned, "ab", "ab".len()); + assert_eq!(result, Some(expected), "stop list order {stops:?}"); + } + } + + #[test] + fn stop_string_completion_position_not_start_position() { + // "b" starts later than "abc" but completes earlier, so it must win. + let stops = vec!["abc".to_string(), "b".to_string()]; + let result = matches_stop_string(&stops, "abc", 3); + assert_eq!(result, Some((1, 1))); + } + #[test] fn stop_string_matches_second_of_multiple() { let stops = vec!["xyz".to_string(), "wor".to_string()]; diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index 08ec5a22d6be..eb527294a5b0 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -1,13 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::borrow::Cow; use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use fastokens::Tokenizer as FastokensTokenizer; use fastokens::decoders::Decoder as FastokensDecoder; +use fastokens::pre_tokenized::{ + PreTokenizedString as FastokensPreTokenizedString, Split as FastokensSplit, +}; +use fastokens::{PreTokenizer as FastokensPreTokenizer, Split as FastokensSplitPreTokenizer}; use thiserror_ext::AsReport as _; -use tokenizers::Tokenizer as HfTokenizer; +use tokenizers::{ + AddedVocabulary, Model as _, OffsetType, PreTokenizer as _, Tokenizer as HfTokenizer, +}; use tracing::{info, warn}; use crate::byte_level_decode::decode_byte_level; @@ -16,6 +23,8 @@ use crate::{Result, Tokenizer}; mod added_tokens; +static EMPTY_HF_ADDED_VOCABULARY: LazyLock = LazyLock::new(AddedVocabulary::new); + enum Backend { Hf(Box), Fastokens(Box), @@ -53,6 +62,85 @@ fn decode_fastokens_byte_level( Ok(decode_byte_level(tokens)) } +fn encode_hf_ordinary(tokenizer: &HfTokenizer, text: &str) -> tokenizers::Result> { + let mut pretokenized = + EMPTY_HF_ADDED_VOCABULARY.extract_and_normalize(tokenizer.get_normalizer(), text); + + if let Some(pre_tokenizer) = tokenizer.get_pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized.tokenize(|normalized| tokenizer.get_model().tokenize(normalized.get()))?; + let encoding = pretokenized.into_encoding(None, 0, OffsetType::Byte)?; + let encoding = tokenizer.post_process(encoding, None, false)?; + Ok(encoding.get_ids().to_vec()) +} + +fn fastokens_fused_split(tokenizer: &FastokensTokenizer) -> Option<&FastokensSplitPreTokenizer> { + // Keep this predicate aligned with fastokens::Tokenizer::detect_fused_byte_level. + let FastokensPreTokenizer::Sequence(steps) = tokenizer.pre_tokenizer()? else { + return None; + }; + let [ + FastokensPreTokenizer::Split(split), + FastokensPreTokenizer::ByteLevel(byte_level), + ] = steps.as_slice() + else { + return None; + }; + byte_level.is_bulk_only().then_some(split) +} + +fn fastokens_pre_tokenized_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> FastokensPreTokenizedString { + // This is fastokens::Tokenizer::build_pre_tokenized with added_tokens = None. + let normalized = tokenizer + .normalizer() + .map_or(Cow::Borrowed(text), |normalizer| normalizer.normalize(text)); + match normalized { + Cow::Borrowed(_) => FastokensPreTokenizedString::from_text(text), + Cow::Owned(text) => { + let len = text.len(); + FastokensPreTokenizedString::new( + text, + vec![FastokensSplit { + range: 0..len, + token_id: None, + }], + ) + } + } +} + +fn encode_fastokens_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> std::result::Result, fastokens::Error> { + if text.is_empty() { + return Ok(Vec::new()); + } + + let mut pretokenized = fastokens_pre_tokenized_ordinary(tokenizer, text); + let ids = if let Some(split) = fastokens_fused_split(tokenizer) { + split.pre_tokenize(&mut pretokenized)?; + pretokenized + .tokenize_batched(|buffer, splits, output| { + tokenizer.model().tokenize_batch_fused(buffer, splits, output) + }) + .map_err(fastokens::Error::Model)? + } else { + if let Some(pre_tokenizer) = tokenizer.pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized + .tokenize(|text, output| tokenizer.model().tokenize_into(text, output)) + .map_err(fastokens::Error::Model)? + }; + + Ok(tokenizer.post_process(ids, false)) +} + /// Tokenizer from `tokenizer.json` in HuggingFace format. /// /// This tries to load with `fastokens` first for better performance, then falls @@ -156,6 +244,17 @@ impl Tokenizer for HuggingFaceTokenizer { } } + fn encode_ordinary(&self, text: &str) -> Result> { + match &self.backend { + Backend::Hf(tokenizer) => encode_hf_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())), + Backend::Fastokens(tokenizer) | Backend::FastokensByteLevel(tokenizer) => { + encode_fastokens_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())) + } + } + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { match &self.backend { Backend::Hf(t) => t @@ -200,12 +299,19 @@ impl Tokenizer for HuggingFaceTokenizer { #[cfg(test)] mod tests { + use std::path::{Path, PathBuf}; + + use serde_json::{Value, json}; use tempfile::tempdir; use tokenizers::models::bpe::BPE; + use tokenizers::pre_tokenizers::byte_level::ByteLevel; use tokenizers::{AddedToken, Tokenizer as HfTokenizer}; use super::{HuggingFaceTokenizer, Tokenizer}; + const REGULAR_TOKEN: &str = "<|regular|>"; + const SPECIAL_TOKEN: &str = "<|special|>"; + fn tiny_bpe_tokenizer() -> HfTokenizer { let vocab = [ ("".to_string(), 0), @@ -232,6 +338,186 @@ mod tests { HfTokenizer::new(model) } + fn ordinary_test_tokenizer_json(fused: bool, with_added_tokens: bool) -> Value { + let mut alphabet: Vec = ByteLevel::alphabet().into_iter().collect(); + alphabet.sort_unstable(); + let vocab = alphabet + .into_iter() + .enumerate() + .map(|(id, token)| (token.to_string(), json!(id))) + .collect::>(); + + let pre_tokenizer = if fused { + json!({ + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": "\\S+|\\s+"}, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + }) + } else { + json!({ + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }) + }; + let added_tokens = with_added_tokens.then(|| { + json!([ + { + "id": 256, + "content": REGULAR_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": true, + "special": false + }, + { + "id": 257, + "content": SPECIAL_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ]) + }); + + json!({ + "version": "1.0", + "truncation": { + "direction": "Right", + "max_length": 24, + "strategy": "LongestFirst", + "stride": 0 + }, + "padding": null, + "added_tokens": added_tokens.unwrap_or_else(|| json!([])), + "normalizer": {"type": "NFC"}, + "pre_tokenizer": pre_tokenizer, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": null, + "end_of_word_suffix": null, + "fuse_unk": false, + "byte_fallback": false, + "ignore_merges": false, + "vocab": vocab, + "merges": [] + } + }) + } + + fn write_tokenizer_json(dir: &Path, name: &str, value: &Value) -> PathBuf { + let path = dir.join(name); + std::fs::write( + &path, + serde_json::to_vec(value).expect("serialize tokenizer"), + ) + .expect("write tokenizer"); + path + } + + fn assert_ordinary_matches_added_empty( + constructor: fn(&Path) -> crate::Result, + fused: bool, + ) { + let dir = tempdir().expect("create temp dir"); + let added_path = write_tokenizer_json( + dir.path(), + "with-added.json", + &ordinary_test_tokenizer_json(fused, true), + ); + let empty_path = write_tokenizer_json( + dir.path(), + "added-empty.json", + &ordinary_test_tokenizer_json(fused, false), + ); + let tokenizer = constructor(&added_path).expect("load tokenizer with added tokens"); + let added_empty = constructor(&empty_path).expect("load tokenizer with empty added tokens"); + + if let super::Backend::Fastokens(inner) | super::Backend::FastokensByteLevel(inner) = + &tokenizer.backend + { + assert_eq!(super::fastokens_fused_split(inner).is_some(), fused); + } + + assert_eq!( + tokenizer.encode(REGULAR_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(REGULAR_TOKEN).unwrap()] + ); + assert_eq!( + tokenizer.encode(SPECIAL_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(SPECIAL_TOKEN).unwrap()] + ); + + for text in [ + "", + "hello", + "Cafe\u{301}", + REGULAR_TOKEN, + SPECIAL_TOKEN, + "hello <|regular|> Cafe\u{301} <|special|> tail", + ] { + assert_eq!( + tokenizer.encode_ordinary(text).unwrap(), + added_empty.encode(text, false).unwrap(), + "fused={fused}, text={text:?}", + ); + } + if matches!(&tokenizer.backend, super::Backend::Hf(_)) { + assert_eq!( + tokenizer + .encode_ordinary("hello <|regular|> Cafe\u{301} <|special|> tail") + .unwrap() + .len(), + 24, + "HF post-processing must retain configured truncation", + ); + } + } + + #[test] + fn hf_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_hf, fused); + } + } + + #[test] + fn fastokens_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_fastokens, fused); + } + } + #[test] fn hf_constructor_resolves_added_token_ids() { let mut tokenizer = tiny_bpe_tokenizer(); diff --git a/rust/src/tokenizer/src/incremental.rs b/rust/src/tokenizer/src/incremental.rs index 5e470ae4b005..f608fa874b82 100644 --- a/rust/src/tokenizer/src/incremental.rs +++ b/rust/src/tokenizer/src/incremental.rs @@ -199,6 +199,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result { let bytes = token_ids.iter().map(|id| *id as u8).collect::>(); Ok(String::from_utf8_lossy(&bytes).into_owned()) @@ -273,6 +277,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let mut text = String::new(); for &token_id in token_ids { @@ -410,6 +418,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result { match token_ids { [1] => Ok("abc".into()), diff --git a/rust/src/tokenizer/src/lib.rs b/rust/src/tokenizer/src/lib.rs index 6c9fcd3fdea2..0f8c7dc16e59 100644 --- a/rust/src/tokenizer/src/lib.rs +++ b/rust/src/tokenizer/src/lib.rs @@ -25,6 +25,10 @@ pub trait Tokenizer: Send + Sync { /// Encode one prompt string into token IDs. fn encode(&self, text: &str, add_special_tokens: bool) -> Result>; + /// Equivalent to `encode(text, false)`, except that every added, + /// special, and control-token matcher is bypassed. + fn encode_ordinary(&self, text: &str) -> Result>; + /// Decode one token sequence into text. fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result; diff --git a/rust/src/tokenizer/src/tekken.rs b/rust/src/tokenizer/src/tekken.rs index 20b6c26ffc8c..5f9342e4a5e6 100644 --- a/rust/src/tokenizer/src/tekken.rs +++ b/rust/src/tokenizer/src/tekken.rs @@ -35,6 +35,12 @@ impl Tokenizer for TekkenTokenizer { .map_err(|error| tokenizer_error!("encoding failed: {error}")) } + fn encode_ordinary(&self, text: &str) -> Result> { + self.inner + .encode(text, false, false) + .map_err(|error| tokenizer_error!("encoding failed: {error}")) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let policy = if skip_special_tokens { tekken::SpecialTokenPolicy::Ignore @@ -67,3 +73,51 @@ impl Tokenizer for TekkenTokenizer { self.inner.is_special_token(token_id) } } + +#[cfg(test)] +mod tests { + use base64::Engine as _; + use tekken::config::TokenizerVersion; + use tekken::{SpecialTokenInfo, TokenInfo}; + + use super::*; + + fn test_tokenizer() -> TekkenTokenizer { + let vocab = (0_u8..=255) + .map(|byte| TokenInfo { + rank: byte as usize, + token_bytes: base64::engine::general_purpose::STANDARD.encode([byte]), + token_str: None, + }) + .collect(); + let special_tokens = vec![SpecialTokenInfo { + rank: 0, + token_str: "".to_string(), + is_control: true, + }]; + let inner = Tekkenizer::new( + vocab, + &special_tokens, + r"(?s).", + 257, + 1, + TokenizerVersion::V3, + None, + ) + .expect("build Tekken tokenizer"); + TekkenTokenizer { inner } + } + + #[test] + fn ordinary_matches_tekkens_empty_special_encoding() { + let tokenizer = test_tokenizer(); + let text = "user text"; + let control_id = tokenizer.token_to_id("").unwrap(); + let ordinary_ids = tokenizer.encode_ordinary(text).unwrap(); + + assert_eq!(control_id, 0); + assert_eq!(ordinary_ids, tokenizer.encode(text, false).unwrap()); + assert!(!ordinary_ids.contains(&control_id)); + assert_eq!(tokenizer.decode(&ordinary_ids, false).unwrap(), text); + } +} diff --git a/rust/src/tokenizer/src/test_utils.rs b/rust/src/tokenizer/src/test_utils.rs index 36d1f3d18aa1..6fdb6e027211 100644 --- a/rust/src/tokenizer/src/test_utils.rs +++ b/rust/src/tokenizer/src/test_utils.rs @@ -208,6 +208,10 @@ impl Tokenizer for TestTokenizer { Ok(ids) } + fn encode_ordinary(&self, text: &str) -> Result> { + Ok(text.as_bytes().iter().copied().map(u32::from).collect()) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let mut output = String::new(); let mut pending_bytes = Vec::new(); @@ -374,6 +378,32 @@ mod tests { assert!(!tokenizer.is_special_id(0xF002)); } + #[test] + fn ordinary_encoding_bypasses_all_configured_tokens() { + let tokenizer = TestTokenizer::new() + .with_bos_token("", 256) + .with_special_token("", 257) + .with_regular_token("", 258); + let ordinary_text = "user and "; + + assert_eq!(tokenizer.encode("", false).unwrap(), vec![257]); + assert_eq!(tokenizer.encode("", false).unwrap(), vec![258]); + assert_eq!( + tokenizer.encode_ordinary(ordinary_text).unwrap(), + ordinary_text.as_bytes().iter().copied().map(u32::from).collect::>() + ); + + let mut segmented = tokenizer.encode("", false).unwrap(); + segmented.extend(tokenizer.encode_ordinary(ordinary_text).unwrap()); + segmented.extend(tokenizer.encode("", false).unwrap()); + assert_eq!(segmented.first(), Some(&257)); + assert_eq!(segmented.last(), Some(&258)); + assert_eq!( + tokenizer.decode(&segmented, false).unwrap(), + format!("{ordinary_text}") + ); + } + #[test] #[should_panic(expected = "configured test token id 255 overlaps byte fallback range 0..=255")] fn configured_token_id_must_stay_outside_byte_range() { diff --git a/rust/src/tokenizer/src/tiktoken.rs b/rust/src/tokenizer/src/tiktoken.rs index d5ca119b8080..fb569d69c35b 100644 --- a/rust/src/tokenizer/src/tiktoken.rs +++ b/rust/src/tokenizer/src/tiktoken.rs @@ -462,6 +462,13 @@ impl Tokenizer for TiktokenTokenizer { }) } + fn encode_ordinary(&self, text: &str) -> Result> { + Ok(match &self.backend { + Backend::Riptoken(backend) => backend.inner.encode_ordinary(text), + Backend::TiktokenRs(backend) => backend.inner.encode_ordinary(text), + }) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { // Filter passes: // @@ -525,7 +532,7 @@ fn detect_bpe_pattern(config: &TiktokenModelConfig) -> &'static str { let model_type = config.effective_model_type(); match model_type { - Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN, + Some("kimi" | "kimi_k2" | "kimi_k25" | "kimi_k3" | "deepseek_v3") => KIMI_PATTERN, _ => CL100K_BASE_PATTERN, } } @@ -752,6 +759,39 @@ mod tests { } } + #[test] + fn tiktoken_ordinary_bypasses_every_registered_added_token() { + let dir = tempfile::tempdir().expect("create temp dir"); + let bpe_path = write_synthetic_bpe_file(dir.path()); + fs::write( + dir.path().join("tokenizer_config.json"), + r#"{ + "added_tokens_decoder": { + "257": { "content": "<|im_end|>", "special": true }, + "258": { "content": "<|tool_call_begin|>", "special": false } + } + }"#, + ) + .expect("write tokenizer_config.json"); + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 260}"#) + .expect("write config.json"); + + let input = "<|im_end|><|tool_call_begin|><|reserved_token_259|>"; + let expected: Vec = input.as_bytes().iter().copied().map(u32::from).collect(); + for backend in explicit_backends(&bpe_path) { + assert_eq!(backend.encode("<|im_end|>", false).unwrap(), vec![257]); + assert_eq!( + backend.encode("<|tool_call_begin|>", false).unwrap(), + vec![258] + ); + assert_eq!( + backend.encode("<|reserved_token_259|>", false).unwrap(), + vec![259] + ); + assert_eq!(backend.encode_ordinary(input).unwrap(), expected); + } + } + /// `vocab_size` may live under `text_config` for composite (e.g. /// multimodal) configs. #[test] @@ -777,6 +817,7 @@ mod tests { #[test] fn tiktoken_detects_kimi_pattern_from_model_type() { let kimi = config_json!({ "model_type": "kimi_k25" }); + let kimi_k3 = config_json!({ "model_type": "kimi_k3" }); let baseten_kimi = config_json!({ "model_type": "deepseek_v3" }); let nested_kimi = config_json!({ "model_type": "composite_wrapper", @@ -790,6 +831,7 @@ mod tests { let missing = config_json!({ "text_config": {} }); assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN); + assert_eq!(detect_bpe_pattern(&kimi_k3), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN); assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN); diff --git a/rust/src/tracing/Cargo.toml b/rust/src/tracing/Cargo.toml new file mode 100644 index 000000000000..e003720814e1 --- /dev/null +++ b/rust/src/tracing/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "vllm-tracing" +version.workspace = true +edition.workspace = true +description = "Shared tracing subscriber and log formatting for vLLM Rust binaries" +license.workspace = true + +[dependencies] +time.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/rust/src/cmd/src/logging.rs b/rust/src/tracing/src/lib.rs similarity index 93% rename from rust/src/cmd/src/logging.rs rename to rust/src/tracing/src/lib.rs index 4d31b92ee238..eb6da9c0b247 100644 --- a/rust/src/cmd/src/logging.rs +++ b/rust/src/tracing/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +//! Shared tracing subscriber and log formatting for vLLM Rust binaries. + use std::{env, fmt, process}; use time::UtcOffset; @@ -26,18 +28,21 @@ const RESET: &str = "\x1b[0m"; const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] = format_description!("[month]-[day] [hour]:[minute]:[second]"); -const PROCESS_LABEL: &str = "RustFrontend"; - -/// Install the process-wide vLLM-style tracing subscriber for the CLI binary. -pub(crate) fn init_tracing() { +/// Install the process-wide vLLM-style tracing subscriber. +pub fn init_tracing(process_label: &str) { let filter = build_targets_filter( env::var("VLLM_LOGGING_LEVEL").ok().as_deref(), env::var("RUST_LOG").ok().as_deref(), ); - let formatter = VllmEventFormatter::new(); + let formatter = VllmEventFormatter::new(process_label); let _ = tracing_subscriber::registry() - .with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter)) + .with( + tracing_subscriber::fmt::layer() + .event_format(formatter) + .with_writer(std::io::stderr) + .with_filter(filter), + ) .try_init(); } @@ -94,9 +99,9 @@ struct VllmEventFormatter { } impl VllmEventFormatter { - fn new() -> Self { + fn new(process_label: &str) -> Self { Self { - prefix: format!("({} pid={})", PROCESS_LABEL, process::id()), + prefix: format!("({process_label} pid={})", process::id()), timer: VllmLocalTimer::default(), } } @@ -291,6 +296,13 @@ fn map_python_log_level(level: &str) -> LevelFilter { mod tests { use super::*; + #[test] + fn formatter_prefix_uses_process_label() { + let formatter = VllmEventFormatter::new("Bench"); + + assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id())); + } + #[test] fn rust_log_target_overrides_are_merged_with_vllm_default_level() { let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error")); diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py index a2640aa52ffa..fea30d6f9cfa 100644 --- a/scripts/benchmark_helion_kernels.py +++ b/scripts/benchmark_helion_kernels.py @@ -5,8 +5,10 @@ Benchmark a registered Helion kernel against a baseline. For each input case produced by the kernel's registered input generator, this -measures the latency of the Helion kernel and a chosen baseline, then reports -the speedup. +checks the Helion kernel's numerics once against its eager reference, measures +its latency against a chosen performance baseline, then reports the speedup. +Use ``--numerics-with-perf-baseline`` to check numerics against the performance +baseline instead. Two baselines are supported (``--baseline``): @@ -30,9 +32,17 @@ python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ --baseline cuda + # Check numerics against the performance baseline instead of eager + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --baseline cuda --numerics-with-perf-baseline + # Disable CUDA graph capture and save results python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ --no-cudagraph --output results.json + + # Only verify numerics, skipping the timing runs + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --numerics-only """ import argparse @@ -43,15 +53,19 @@ import sys from collections.abc import Callable from dataclasses import asdict, dataclass +from typing import Any import torch +from torch.utils._pytree import tree_flatten from vllm.triton_utils import triton try: + from helion.autotuner.accuracy import assert_close as helion_assert_close + from helion.autotuner.accuracy import is_fp8_dtype + from vllm.benchmarks.lib.utils import default_vllm_config from vllm.kernels.helion import get_kernel_by_name, get_registered_kernels - from vllm.kernels.helion.ops import import_all_kernels from vllm.logger import init_logger from vllm.utils.import_utils import has_helion except ImportError as e: @@ -59,6 +73,27 @@ print("Please ensure vLLM is installed and in your Python path") sys.exit(1) + +def import_all_kernels() -> None: + """Trigger Helion op registration, tolerating cross-version name drift. + + Current vLLM registers every Helion kernel as a side effect of importing + ``vllm.kernels.helion.ops``; some builds instead expose an explicit importer + whose name has drifted (``import_all_kernels`` / ``import_all_ops``). Call + whichever exists; if none does, importing the module already registered + them. + """ + try: + import vllm.kernels.helion.ops as ops + except ImportError: + return + for fn_name in ("import_all_kernels", "import_all_ops"): + fn = getattr(ops, fn_name, None) + if callable(fn): + fn() + return + + logger = init_logger("vllm.scripts.benchmark_helion_kernels") @@ -121,6 +156,21 @@ def fmt(row: list[str]) -> str: print(fmt(row)) +def log_versions() -> None: + """Log torch/helion/triton versions at the head of the output.""" + from importlib.metadata import PackageNotFoundError, version + + def pkg_version(name: str) -> str: + try: + return version(name) + except PackageNotFoundError: + return "not installed" + + logger.info("torch: %s", torch.__version__) + logger.info("helion: %s", pkg_version("helion")) + logger.info("triton: %s", getattr(triton, "__version__", pkg_version("triton"))) + + def list_kernels() -> None: kernels = get_registered_kernels() @@ -176,12 +226,8 @@ def make_cuda_baseline(kernel_name: str) -> Callable: return cuda_op -def make_autotune_baseline(kernel_name: str) -> Callable: - """Return the kernel's autotune baseline wrapped in ``torch.compile``. - - The baseline is the native-torch reference the kernel is tuned against, - registered via ``helion_settings.autotune_baseline_fn``. - """ +def make_eager_baseline(kernel_name: str) -> Callable: + """Return the kernel's registered native-torch reference.""" wrapper = get_kernel_by_name(kernel_name) settings = wrapper.helion_settings baseline_fn = getattr(settings, "autotune_baseline_fn", None) @@ -195,6 +241,13 @@ def make_autotune_baseline(kernel_name: str) -> Callable: ) sys.exit(1) + return baseline_fn + + +def make_autotune_baseline(kernel_name: str) -> Callable: + """Return the kernel's autotune baseline wrapped in ``torch.compile``.""" + baseline_fn = make_eager_baseline(kernel_name) + return torch.compile( baseline_fn, fullgraph=True, @@ -204,6 +257,25 @@ def make_autotune_baseline(kernel_name: str) -> Callable: ) +def make_correctness_baseline( + kernel_name: str, + timed_baseline_fn: Callable, + numerics_with_perf_baseline: bool, +) -> Callable: + """Choose the numerical reference independently from the timed baseline.""" + if not numerics_with_perf_baseline: + logger.info( + "Using the eager reference for '%s' correctness", + kernel_name, + ) + return make_eager_baseline(kernel_name) + logger.info( + "Using the selected performance baseline for '%s' correctness", + kernel_name, + ) + return timed_baseline_fn + + def cleanup_gpu_resources() -> None: try: torch.accelerator.empty_cache() @@ -227,6 +299,118 @@ def _reduce(times: list[float], return_mode: str) -> float: return _REDUCERS[return_mode](times) +def _assert_close(actual: object, expected: object, atol: float, rtol: float) -> None: + """Compare pytrees, allowing the one-ULP FP8 variance used by kernel tests.""" + actual_flat, actual_spec = tree_flatten(actual) + expected_flat, expected_spec = tree_flatten(expected) + if actual_spec != expected_spec: + raise AssertionError( + f"Output structure mismatch: {actual_spec} != {expected_spec}" + ) + + for actual_leaf, expected_leaf in zip(actual_flat, expected_flat, strict=True): + is_fp8 = isinstance(actual_leaf, torch.Tensor) and is_fp8_dtype( + actual_leaf.dtype + ) + helion_assert_close( + actual_leaf, + expected_leaf, + atol=1 if is_fp8 else atol, + rtol=0 if is_fp8 else rtol, + ) + + +def check_correctness( + kernel: Any, + baseline_fn: Callable, + inputs: tuple[Any, ...], + case: str, +) -> None: + """Run one numerical comparison on copies separate from benchmark inputs.""" + kernel_inputs = copy.deepcopy(inputs) + baseline_inputs = copy.deepcopy(inputs) + + kernel_output = kernel(*kernel_inputs) + baseline_output = baseline_fn(*baseline_inputs) + + settings = kernel.helion_settings + try: + custom_check = getattr(settings, "autotune_baseline_accuracy_check_fn", None) + if custom_check is not None: + custom_check(kernel_output, baseline_output) + custom_check(kernel_inputs, baseline_inputs) + return + + configured_atol = getattr(settings, "autotune_baseline_atol", None) + configured_rtol = getattr(settings, "autotune_baseline_rtol", None) + atol = 1e-2 if configured_atol is None else configured_atol + rtol = 1e-2 if configured_rtol is None else configured_rtol + _assert_close( + kernel_output, + baseline_output, + atol=atol, + rtol=rtol, + ) + _assert_close( + kernel_inputs, + baseline_inputs, + atol=atol, + rtol=rtol, + ) + except AssertionError as e: + raise AssertionError(f"Numerics check failed for case {case}:\n{e}") from e + + +@dataclass +class CorrectnessResult: + """Outcome of the numerics check for a single shape case.""" + + case: str + passed: bool + error: str | None = None + + +def check_kernel_correctness( + kernel: Any, + baseline_fn: Callable, + inputs_dict: dict[Any, tuple[Any, ...]] | None = None, +) -> list[CorrectnessResult]: + """Run the per-shape numerics check for a kernel, continuing past failures. + + Runs the same comparison as ``check_correctness`` for every shape case + produced by the kernel's input generator, but records the outcome per case + instead of raising on the first mismatch. This lets callers (e.g. a CI gate) + report every failing shape in one pass rather than aborting early. + + Args: + kernel: The Helion kernel wrapper to check. + baseline_fn: Reference callable sharing the kernel's argument interface. + inputs_dict: Optional mapping of case key to input tuple. Defaults to + ``kernel.get_inputs()``. + + Returns: + One ``CorrectnessResult`` per shape case, in iteration order. A case that + raises (compile/run error or numerics mismatch) is marked + ``passed=False`` with the exception text in ``error``; iteration + continues regardless. An empty input mapping returns an empty list, which + the CLI reports as a skipped check. + """ + if inputs_dict is None: + inputs_dict = kernel.get_inputs() + + results: list[CorrectnessResult] = [] + for key, inputs in inputs_dict.items(): + case = str(key) + try: + check_correctness(kernel, baseline_fn, inputs, case) + except Exception as e: # noqa: BLE001 - any failure is recorded, not fatal + results.append(CorrectnessResult(case=case, passed=False, error=str(e))) + else: + results.append(CorrectnessResult(case=case, passed=True)) + cleanup_gpu_resources() + return results + + def do_bench_cudagraph_l2_clear( fn: Callable, rep: int = 100, return_mode: str = "mean" ) -> float: @@ -298,6 +482,7 @@ def do_bench_cudagraph_l2_clear( def benchmark( kernel_name: str, baseline_fn: Callable, + correctness_fn: Callable, repeat: int, cudagraph: bool, return_mode: str, @@ -313,6 +498,9 @@ def benchmark( for key, inputs in inputs_dict.items(): logger.info("Benchmarking case %s", key) + check_correctness(kernel, correctness_fn, inputs, str(key)) + logger.info("Numerics check passed for case %s", key) + # Kernels may mutate their inputs in place; give each side its own copy. kernel_inputs = copy.deepcopy(inputs) baseline_inputs = copy.deepcopy(inputs) @@ -374,7 +562,7 @@ def main() -> None: choices=["cuda", "autotune"], default="autotune", help=( - "Baseline to compare against: 'autotune' uses the kernel's " + "Performance baseline: 'autotune' uses the kernel's " "autotune_baseline_fn under torch.compile; 'cuda' uses the mapped " "torch.ops._C op (default: autotune)" ), @@ -390,9 +578,24 @@ def main() -> None: type=str, help="Path to save benchmark results as JSON (default: log only)", ) + parser.add_argument( + "--numerics-only", + action="store_true", + help="Only run the per-case numerics check; skip timing and reporting", + ) + parser.add_argument( + "--numerics-with-perf-baseline", + action="store_true", + help=( + "Compare numerics against the selected performance baseline instead " + "of the eager reference" + ), + ) args = parser.parse_args() + log_versions() + import_all_kernels() if args.list: @@ -425,10 +628,43 @@ def main() -> None: baseline_fn = make_cuda_baseline(args.kernel) else: baseline_fn = make_autotune_baseline(args.kernel) + correctness_fn = make_correctness_baseline( + args.kernel, + baseline_fn, + args.numerics_with_perf_baseline, + ) + + if args.numerics_only: + results = check_kernel_correctness(wrapper, correctness_fn) + if not results: + logger.warning( + "No input cases generated for '%s'; skipping numerics check", + args.kernel, + ) + return + for r in results: + if r.passed: + logger.info("Numerics check passed for case %s", r.case) + else: + logger.error( + "Numerics check FAILED for case %s: %s", r.case, r.error + ) + failed = [r for r in results if not r.passed] + if failed: + logger.error( + "%d/%d case(s) failed numerics for '%s'", + len(failed), + len(results), + args.kernel, + ) + sys.exit(1) + logger.info("Numerics check passed for all cases of '%s'", args.kernel) + return rows = benchmark( args.kernel, baseline_fn, + correctness_fn, args.repeat, args.cudagraph, args.return_mode, @@ -442,6 +678,9 @@ def main() -> None: { "kernel": args.kernel, "baseline": args.baseline, + "numerics_baseline": ( + args.baseline if args.numerics_with_perf_baseline else "eager" + ), "cudagraph": args.cudagraph, "repeat": args.repeat, "return_mode": args.return_mode, diff --git a/setup.py b/setup.py index 40d4ca103bea..a8817a960893 100644 --- a/setup.py +++ b/setup.py @@ -783,6 +783,7 @@ def extract_precompiled_and_patch_package( "vllm/_qutlass_C.abi3.so", "vllm/_flashmla_C.abi3.so", "vllm/_flashmla_extension_C.abi3.so", + "vllm/_flashkda_C.abi3.so", "vllm/_sparse_flashmla_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", @@ -1121,7 +1122,7 @@ def _read_requirements(filename: str) -> list[str]: # copying the relevant .py files from the source repository. ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True)) -if sys.version_info >= (3, 11): +if not _is_xpu() and sys.version_info >= (3, 11): ext_modules.append(CMakeExtension(name="vllm.spinloop")) ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) @@ -1150,6 +1151,10 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append( CMakeExtension(name="vllm._flashmla_extension_C", optional=True) ) + if USE_PRECOMPILED_EXTENSIONS or ( + CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0") + ): + ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True)) if envs.VLLM_USE_PRECOMPILED or ( CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") ): @@ -1286,7 +1291,7 @@ def add_vllm_package_data(filename: str) -> None: # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml # - .buildkite/test-amd.yaml - "helion": ["helion==1.1.0"], + "helion": ["helion==1.4.0"], # Optional deps for gRPC server (vllm serve --grpc) "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"], # Optional deps for OpenTelemetry tracing diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 810a3a0aeed3..4f99120c0c39 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/tests/basic_correctness/test_basic_correctness.py @@ -78,6 +78,15 @@ def _resolve_target_test_suite() -> str: TARGET_TEST_SUITE = _resolve_target_test_suite() +# ROCm can occasionally retain the object until fixture teardown. Retry only +# that assertion after cleanup; collecting cyclic garbage here would mask the +# reference cycles this test is intended to catch. +@pytest.mark.flaky( + reruns=2, + reruns_delay=5, + only_rerun="AssertionError", + condition=current_platform.is_rocm(), +) def test_vllm_gc_ed(): """Verify vllm instance is GC'ed when it is deleted""" llm = LLM("hmellor/tiny-random-LlamaForCausalLM") diff --git a/tests/benchmarks/test_skip_tokenizer_init.py b/tests/benchmarks/test_skip_tokenizer_init.py new file mode 100644 index 000000000000..741db2e0d626 --- /dev/null +++ b/tests/benchmarks/test_skip_tokenizer_init.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for --skip-tokenizer-init with --dataset-name custom. + +Before the fix (introduced by #39896), running: + + vllm bench serve \ + --backend vllm-pooling \ + --dataset-name custom \ + --dataset-path \ + --model ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11 \ + --endpoint /pooling \ + --skip-tokenizer-init \ + ... + +raised immediately with: + + AssertionError: Tokenizer must be initialized before loading dataset + +even though CustomDataset.sample() already handles tokenizer=None. +This test exercises main_async() directly so it catches any regression +re-introduced at the serve.py level, not just inside get_samples(). +""" + +import argparse +import asyncio +import json +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +import vllm.benchmarks.serve as serve_module + +# Exact prompt payload from the failing benchmark run against a +# Prithvi-EO-2.0 pooling endpoint (URL-in / base64-out format). +_PRITHVI_PROMPT = { + "data": { + "data": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif", + "data_format": "url", + "out_data_format": "b64_json", + "indices": [1, 2, 3, 8, 11, 12], + }, + "priority": 0, + "softmax": False, +} + + +def _write_dataset(path: Path) -> None: + path.write_text(json.dumps({"prompt": _PRITHVI_PROMPT}) + "\n") + + +def _args(dataset_path: str) -> argparse.Namespace: + """Reproduce the argparse.Namespace that serve.py builds from the + failing command, including skip_tokenizer_init=True.""" + return argparse.Namespace( + # dataset + dataset_name="custom", + dataset_path=dataset_path, + disable_shuffle=False, + num_prompts=1, + custom_output_len=256, + skip_chat_template=True, + chat_template_kwargs=None, + no_oversample=False, + seed=0, + request_id_prefix="bench-", + # model / tokenizer + model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", + served_model_name=None, + tokenizer=None, + tokenizer_mode="auto", + trust_remote_code=False, + skip_tokenizer_init=True, # <-- the flag under test + # backend / endpoint + backend="vllm-pooling", + base_url="http://127.0.0.1:8000", + host="127.0.0.1", + port=8000, + endpoint="/pooling", + header=None, + insecure=False, + # traffic + request_rate=16.0, + burstiness=1.0, + max_concurrency=None, + probe_request_rate=0.0, + # misc serve args that main_async reads before reaching get_samples + plot_timeline=False, + plot_dataset_stats=False, + self_timed=None, + metadata=None, + label=None, + logprobs=None, + use_beam_search=False, + ignore_eos=False, + goodput=None, + percentile_metrics="ttft,tpot,itl,e2el", + metric_percentiles="25,50,75,99", + save_result=False, + append_result=False, + result_dir=".", + result_filename=None, + num_warmups=0, + profile=False, + disable_tqdm=True, + lora_modules=None, + lora_assignment="random", + ramp_up_strategy=None, + ramp_up_start_rps=None, + ramp_up_end_rps=None, + ready_check_timeout_sec=0, + extra_body=None, + top_p=None, + top_k=None, + min_p=None, + temperature=None, + frequency_penalty=None, + presence_penalty=None, + repetition_penalty=None, + save_detailed=False, + input_len=None, + output_len=None, + ) + + +@pytest.mark.benchmark +def test_main_async_skip_tokenizer_init_does_not_raise(tmp_path: Path) -> None: + """main_async must not raise AssertionError when skip_tokenizer_init=True + and dataset_name='custom'. + + On main (before the fix) this test fails with: + AssertionError: Tokenizer must be initialized before loading dataset + """ + dataset_path = tmp_path / "dataset_url_input_india.jsonl" + _write_dataset(dataset_path) + + args = _args(str(dataset_path)) + + # Patch benchmark() so we never make real HTTP requests — the regression + # triggers before benchmark() is ever called, so this just keeps the test + # fast and hermetic. + mock_result = { + "completed": 1, + "failed": 0, + "total_input_tokens": 1, + "total_output_tokens": 1, + "request_throughput": 1.0, + "output_throughput": 1.0, + "total_token_throughput": 1.0, + "input_lens": [], + "output_lens": [], + "ttfts": [], + "itls": [], + "generated_texts": [], + "errors": [], + "duration": 1.0, + } + with patch.object( + serve_module, "benchmark", new=AsyncMock(return_value=mock_result) + ): + # Must NOT raise AssertionError + asyncio.run(serve_module.main_async(args)) diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py index e320f5a11208..9c9c59619824 100644 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ b/tests/compile/correctness_e2e/test_sequence_parallel.py @@ -348,10 +348,14 @@ def test_tp_sp_generation( if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): pytest.skip("inductor graph partition is only available in PyTorch 2.9+") - # Skip FP8 SP-only test on sm89 (compute capability 8.9) + # Skip FP8 SP-only test on sm89 (compute capability 8.9). + # An unknown capability (None) is left to fail in the test body rather than + # being silently skipped here. + capability = current_platform.get_device_capability() if ( "fp8" in model_id.lower() - and current_platform.get_device_capability() < (9, 0) + and capability is not None + and capability < (9, 0) and (not fuse_gemm_comms) ): pytest.skip("FP8 reduction support begins with sm90 capable devices.") diff --git a/tests/compile/fullgraph/test_basic_correctness.py b/tests/compile/fullgraph/test_basic_correctness.py index 35989dcde1dc..7471564340dd 100644 --- a/tests/compile/fullgraph/test_basic_correctness.py +++ b/tests/compile/fullgraph/test_basic_correctness.py @@ -64,6 +64,8 @@ class TestSetting: "bfloat16", "--max-model-len", "2048", + "--gpu-memory-utilization", + "0.98", ], pp_size=1, tp_size=1, diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py index 95306e2062f8..c2af7be048af 100644 --- a/tests/compile/fullgraph/test_full_cudagraph.py +++ b/tests/compile/fullgraph/test_full_cudagraph.py @@ -9,7 +9,7 @@ from tests.utils import wait_for_gpu_memory_to_clear from tests.v1.attention.utils import full_cg_backend_configs as backend_configs from vllm import LLM, SamplingParams -from vllm.config import CompilationConfig +from vllm.config import CompilationConfig, CUDAGraphMode from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -97,7 +97,9 @@ def llm_pair(request): trust_remote_code=True, max_model_len=1024, max_num_seqs=128, - compilation_config=CompilationConfig(cudagraph_mode="PIECEWISE"), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.PIECEWISE + ), generation_config="vllm", seed=42, ) diff --git a/tests/compile/fullgraph/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py index cc138454802b..0f3368c0dbc7 100644 --- a/tests/compile/fullgraph/test_full_graph.py +++ b/tests/compile/fullgraph/test_full_graph.py @@ -75,7 +75,7 @@ def test_full_graph( monkeypatch: pytest.MonkeyPatch, model: str, model_kwargs: dict[str, Any], - compilation_mode: int, + compilation_mode: CompilationMode, ): if ( "w8a8" in model @@ -197,7 +197,7 @@ def test_custom_compile_config( ], ) def test_fp8_kv_scale_compile( - compilation_mode: int, + compilation_mode: CompilationMode, model: str, backend: AttentionBackendEnum | None, ): @@ -213,7 +213,9 @@ def test_fp8_kv_scale_compile( run_model(compilation_mode, model, **model_kwargs) -def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs): +def run_model( + compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs +): compilation_config = ( compile_config if isinstance(compile_config, CompilationConfig) diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index a4ed63ffe7b5..929b27c11c12 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -12,7 +12,9 @@ from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches -def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs): +def run_model( + compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs +): """Run a model with the given compilation config for E2E fusion tests.""" compilation_config = ( compile_config @@ -54,7 +56,7 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg ) # Fetch match table from each worker via RPC and sum across workers. - worker_tables = llm.llm_engine.engine_core.collective_rpc( + worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc( "get_compilation_match_table" ) combined: defaultdict[str, int] = defaultdict(int) diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index b18c41658fde..b6ad4e2e6e85 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -181,8 +181,13 @@ def test_tp2_ar_rms_fp4_fusions( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b, gpt_oss_20b], + "model_name, matches_fn, model_kwargs, hf_overrides, model_impl", + [ + (*llama3_8b, "auto"), + (*llama3_8b, "transformers"), + (*qwen3_a3b, "auto"), + (*gpt_oss_20b, "auto"), + ], ) @pytest.mark.parametrize( "attn_backend", @@ -202,17 +207,26 @@ def test_tp2_ar_rms_fusions( matches_fn: Callable[[int], Matches], model_kwargs: dict, hf_overrides: Callable[[int], dict], + model_impl: str, attn_backend: AttentionBackendCase, n_layers: int, custom_ops: str, inductor_graph_partition: bool, run_e2e_fusion_test, ): + if model_impl == "transformers" and not current_platform.is_rocm(): + pytest.skip("Transformers 3D AR+RMS regression is ROCm-only") + matches = matches_fn(n_layers) + if model_impl == "transformers": + # Transformers add+RMSNorm canonicalization exposes every generic + # AR+RMS fusion site, including the final norm. + matches = matches._replace(aiter_ar_rms_fusion=matches.ar_rms_fusion) # Reduce size of model and skip weight loading time model_kwargs["hf_overrides"] = hf_overrides(n_layers) model_kwargs["load_format"] = "dummy" + model_kwargs["model_impl"] = model_impl model_kwargs["max_model_len"] = 1024 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} model_kwargs["disable_custom_all_reduce"] = False diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1aac4b2bec49..e9c8d0deaa75 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -272,12 +272,10 @@ def __init__( token_num=16, eps=1e-6, dtype: torch.dtype = torch.bfloat16, - use_triton_quant: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps - self.use_triton_quant = use_triton_quant assert hidden_size % self.quant_group_size == 0, ( f"hidden_size ({hidden_size}) must be a multiple of " f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant" @@ -289,10 +287,6 @@ def __init__( ] def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - if self.use_triton_quant: - return torch.ops.vllm.triton_per_token_group_quant_fp8( - rms, self.quant_group_size - ) return torch.ops.vllm.rocm_aiter_group_fp8_quant.default( rms, self.quant_group_size ) @@ -339,11 +333,7 @@ def forward(self, hidden_states): def ops_in_model_before(self): return [ torch.ops.vllm.all_reduce.default, - ( - torch.ops.vllm.triton_per_token_group_quant_fp8.default - if self.use_triton_quant - else torch.ops.vllm.rocm_aiter_group_fp8_quant.default - ), + torch.ops.vllm.rocm_aiter_group_fp8_quant.default, ] def ops_in_model_after(self): @@ -646,7 +636,6 @@ def all_reduce_fusion_pass_on_test_model( @multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize("use_triton_quant", [True, False]) @pytest.mark.parametrize("batch_size", [8]) @pytest.mark.parametrize("seq_len", [8]) @pytest.mark.parametrize("hidden_size", [128]) @@ -663,7 +652,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -676,9 +664,9 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( * ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual, single ``rms`` consumer) * ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with- - residual, DSv3.2 indexer fan-out; parametrized over both - ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant`` - producers). + residual, DSv3.2 indexer fan-out; parametrized over + ``rocm_aiter_group_fp8_quant`` + producer). """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -703,7 +691,6 @@ def run_torch_spawn(fn, nprocs): hidden_size, dtype, enable_rms_norm_custom_op, - use_triton_quant, monkeypatch, ), nprocs=nprocs, @@ -721,7 +708,6 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): set_random_seed(0) @@ -749,10 +735,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( custom_ops = [] if enable_rms_norm_custom_op: custom_ops.append("+rms_norm") - # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip`` - # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at - # the call site). The patterns in this PR are robust to both Triton and - # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example + # We always enable +quant_fp8 so the matcher's example # trace finds the same form the test model uses. custom_ops.append("+quant_fp8") @@ -783,9 +766,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - model = test_model_cls( - hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant - ) + model = test_model_cls(hidden_size, token_num, dtype=dtype) hidden_states = torch.randn((token_num, hidden_size), requires_grad=False) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 92d1902b2c2f..591b014d9e25 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -195,8 +195,6 @@ def ops_in_model_before(self): # Blockwise path if self.use_aiter_fusion and self.use_aiter_quant_op: return [rocm_aiter_ops.get_group_quant_op()] - if self.use_aiter_fusion: - return [torch.ops.vllm.triton_per_token_group_quant_fp8.default] else: if self.use_aiter_quant_op: return [rocm_aiter_ops.get_per_token_quant_op()] diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 76536f3387c5..da335b91557f 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -290,12 +290,9 @@ def test_attention_quant_pattern( model_class: type[AttentionQuantPatternModel], backend: AttentionBackendEnum, dist_init, - monkeypatch, - use_fresh_inductor_cache, + disable_vllm_compile_cache, ): """Test AttentionStaticQuantPattern fusion pass""" - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - if backend == AttentionBackendEnum.FLASHINFER and ( not current_platform.is_device_capability((10, 0)) or not has_flashinfer() ): diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 0a38ffca483a..5b2ef5bfd95c 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -419,8 +419,7 @@ def test_mla_attention_quant_pattern( model_class: type[MLAAttentionQuantPatternModel], backend: AttentionBackendEnum, dist_init, - monkeypatch, - use_fresh_inductor_cache, + disable_vllm_compile_cache, ): """Test MLA AttentionQuantPattern fusion pass""" if ( @@ -429,8 +428,6 @@ def test_mla_attention_quant_pattern( ): pytest.skip("NVFP4 is not supported on this GPU (requires SM 100+).") - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - custom_ops_list = custom_ops.split(",") if custom_ops else [] device = torch.device(f"{DEVICE_TYPE}:0") diff --git a/tests/compile/passes/test_rmsnorm_reshape_fusion.py b/tests/compile/passes/test_rmsnorm_reshape_fusion.py new file mode 100644 index 000000000000..410b2e97eb83 --- /dev/null +++ b/tests/compile/passes/test_rmsnorm_reshape_fusion.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.ir.ops +from tests.compile.backend import TestBackend +from vllm.compilation.passes.fusion.add_rms_fusion import ( + AddRMSNormFusionPass, + RMSNormReshapeFusionPass, +) +from vllm.compilation.passes.fx_utils import find_op_nodes, is_func +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass +from vllm.config import CompilationConfig, CompilationMode, VllmConfig +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm" +) + + +class RMSNormModel(torch.nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(hidden_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rms = vllm.ir.ops.rms_norm(x, self.weight, 1e-6) + return rms.reshape(-1, rms.shape[-1]) + + +class AddRMSNormModel(torch.nn.Module): + def __init__(self, hidden_size: int, residual_first: bool) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(hidden_size)) + self.residual_first = residual_first + + def forward( + self, x: torch.Tensor, residual: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = residual + x if self.residual_first else x + residual + rms = vllm.ir.ops.rms_norm(residual, self.weight, 1e-6) + return rms.reshape(-1, rms.shape[-1]), residual + + def ops_in_model_before(self): + return [torch.ops.aten.add, torch.ops.vllm_ir.rms_norm] + + def ops_in_model_after(self): + return [torch.ops.vllm_ir.fused_add_rms_norm] + + +def _run_fusion_test(model, config, passes, *inputs): + backend = TestBackend(NoOpEliminationPass(config), *passes, PostCleanupPass(config)) + outputs_unfused = model(*inputs) + outputs_fused = torch.compile(model, backend=backend)(*inputs) + torch.testing.assert_close(outputs_fused, outputs_unfused) + return backend + + +def _is_reshape(node): + return isinstance(node, torch.fx.Node) and is_func( + node, torch.ops.aten.reshape.default + ) + + +@pytest.fixture +def vllm_config(): + config = VllmConfig( + compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE) + ) + with vllm.config.set_current_vllm_config(config): + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(0) + yield config + + +def test_rmsnorm_reshape_fusion(vllm_config): + fusion_pass = RMSNormReshapeFusionPass(vllm_config) + model = RMSNormModel(hidden_size=32) + x = torch.randn(2, 7, 32) + backend = _run_fusion_test(model, vllm_config, [fusion_pass], x) + + assert fusion_pass.matched_count == 1 + (rms_node,) = find_op_nodes(torch.ops.vllm_ir.rms_norm, backend.graph_post_pass) + assert _is_reshape(rms_node.args[0]) + + +@pytest.mark.parametrize("residual_first", [True, False]) +def test_add_rmsnorm_reshape_fusion(vllm_config, residual_first): + add_fusion = AddRMSNormFusionPass(vllm_config) + reshape_fusion = RMSNormReshapeFusionPass(vllm_config) + model = AddRMSNormModel(hidden_size=32, residual_first=residual_first) + x = torch.randn(2, 7, 32) + residual = torch.randn_like(x) + backend = _run_fusion_test( + model, vllm_config, [add_fusion, reshape_fusion], x, residual + ) + + assert add_fusion.matched_count == 1 + assert reshape_fusion.matched_count == 1 + backend.check_before_ops(model.ops_in_model_before()) + backend.check_after_ops(model.ops_in_model_after()) diff --git a/tests/compile/passes/test_rope_kvcache_fusion.py b/tests/compile/passes/test_rope_kvcache_fusion.py index 709490f1972b..2f68061749d7 100644 --- a/tests/compile/passes/test_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rope_kvcache_fusion.py @@ -29,6 +29,7 @@ from vllm.config.utils import Range from vllm.forward_context import get_forward_context, set_forward_context from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention import attention as attention_module from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding from vllm.platforms import current_platform from vllm.utils.torch_utils import _encode_layer_name @@ -43,6 +44,37 @@ FP8_DTYPE = current_platform.fp8_dtype() +def test_kv_cache_update_ops_fake_tensor_metadata(monkeypatch: pytest.MonkeyPatch): + query = torch.empty(1, dtype=torch.bfloat16) + kv_cache = torch.empty(1, dtype=torch.uint8) + context = (None, None, kv_cache, None) + monkeypatch.setattr(attention_module, "get_attention_context", lambda _: context) + monkeypatch.setattr(rope_kvcache_fusion, "get_attention_context", lambda _: context) + + runtime_output = attention_module.unified_kv_cache_update(query, query, "layer") + fake_output = attention_module.unified_kv_cache_update_fake(query, query, "layer") + assert runtime_output.shape == fake_output.shape + assert runtime_output.dtype == fake_output.dtype + assert runtime_output.device == fake_output.device + + args = ( + query, + query, + query, + torch.empty(1, dtype=torch.int64), + torch.empty(1), + True, + "layer", + ) + runtime_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_impl( + *args + ) + fake_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_fake(*args) + assert runtime_output.shape == fake_output.shape + assert runtime_output.dtype == fake_output.dtype + assert runtime_output.device == fake_output.device + + def test_rope_kvcache_fusion_default_keeps_large_ranges_unfused(): vllm_config = VllmConfig( compilation_config=CompilationConfig( diff --git a/tests/compile/passes/test_silu_mul_quant_fusion.py b/tests/compile/passes/test_silu_mul_quant_fusion.py index bc134ed427a1..7d291cc50448 100644 --- a/tests/compile/passes/test_silu_mul_quant_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_fusion.py @@ -158,13 +158,6 @@ def __init__(self, hidden_size: int, dtype: torch.dtype, **kwargs): input_dtype=dtype, ) - if not current_platform.is_fp8_fnuz(): - kernel = self.w8a8_block_fp8_linear.kernel - orig_quant = kernel.quant_fp8 - kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant( - *a, use_triton=True, **kw - ) - self.enable_silu_mul_custom_op = self.silu_and_mul.enabled() def forward(self, x): @@ -175,9 +168,7 @@ def forward(self, x): def ops_in_model_before(self): return [ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul, - rocm_aiter_ops.get_group_quant_op() - if current_platform.is_fp8_fnuz() - else torch.ops.vllm.triton_per_token_group_quant_fp8.default, + rocm_aiter_ops.get_group_quant_op(), ] def ops_in_model_after(self): diff --git a/tests/compile/test_compile_ranges.py b/tests/compile/test_compile_ranges.py index 9fd8e9577ba0..4dfea42a6b4a 100644 --- a/tests/compile/test_compile_ranges.py +++ b/tests/compile/test_compile_ranges.py @@ -66,7 +66,7 @@ def uuid(self) -> str: return InductorPass.hash_dict(state) -def test_compile_ranges(use_fresh_inductor_cache): +def test_compile_ranges(disable_vllm_compile_cache): post_grad_range_checker = PostGradRangeChecker( [ Range(start=1, end=8), @@ -168,7 +168,7 @@ def uuid(self) -> str: return InductorPass.hash_dict(state) -def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache): +def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache): """Verify that compile_sizes entries are compiled with fully concrete shapes (no SymInts), while compile_ranges entries retain dynamic shapes.""" checker = PostGradStaticShapeChecker() @@ -209,10 +209,9 @@ def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache): ) -def test_inductor_cache_compile_ranges(monkeypatch, use_fresh_inductor_cache): - # To force multiple compilations, we disable the compile cache - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - +def test_inductor_cache_compile_ranges(disable_vllm_compile_cache): + # disable_vllm_compile_cache sets VLLM_DISABLE_COMPILE_CACHE=1 to force + # multiple compilations by disabling vLLM's on-disk compile cache. post_grad_range_checker = PostGradRangeChecker( ranges=[ Range(start=1, end=8), diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index d822b68c5036..c45907ee6e12 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -83,9 +83,33 @@ def test_copy_pass(): def test_custom_op(): # proper syntax _ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"]) + _ = CompilationConfig(custom_ops=["none", "+rms_norm"]) + _ = CompilationConfig(custom_ops=["+rms_norm", "+rms_norm"]) - with pytest.raises(ValueError, match="Invalid syntax '"): - _ = CompilationConfig(custom_ops=["quant_fp8"]) + for custom_ops in (["quant_fp8"], ["+"], ["-"]): + with pytest.raises(ValueError, match="Invalid syntax '"): + CompilationConfig(custom_ops=custom_ops) + + +@pytest.mark.parametrize( + ("custom_ops", "config_kwargs", "match"), + [ + (["all", "none"], {}, "can contain only one base mode"), + ( + ["none", "+rms_norm", "-rms_norm"], + {}, + "cannot both enable and disable.*rms_norm", + ), + ( + ["-rotary_embedding"], + {"pass_config": PassConfig(enable_qk_norm_rope_fusion=True)}, + "cannot both enable and disable.*rotary_embedding", + ), + ], +) +def test_reject_contradictory_custom_ops(custom_ops, config_kwargs, match): + with pytest.raises(ValueError, match=match): + CompilationConfig(custom_ops=custom_ops, **config_kwargs) # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 96c3f49aba3e..3260d5aecc5a 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import gc import tempfile from contextlib import contextmanager @@ -9,7 +8,7 @@ import torch from tests.models.utils import check_logprobs_close -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.compilation.decorators import support_torch_compile from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config from vllm.config.compilation import ( @@ -48,6 +47,7 @@ def get_test_models(): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_dynamic_shapes_compilation( monkeypatch, + vllm_runner, model_name, shapes_type, use_aot_compile, @@ -78,9 +78,13 @@ def test_dynamic_shapes_compilation( print(f"Testing {shapes_type.name} dynamic shapes...") - # Initialize the model with specific dynamic shapes configuration - model = LLM( - model=model_name, + sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) + test_prompts = [prompt, "The capital of France is"] + + # VllmRunner shuts down the engine core on exit, so the eager model + # below never races a lingering compiled engine for GPU memory. + with vllm_runner( + model_name, compilation_config={ "mode": CompilationMode.VLLM_COMPILE, "dynamic_shapes_config": { @@ -89,32 +93,25 @@ def test_dynamic_shapes_compilation( }, }, max_model_len=1024, - ) - - sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) - test_prompts = [prompt, "The capital of France is"] - - compiled_outputs = [] - for p in test_prompts: - output = model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Compiled model produced empty output" - compiled_outputs.append((output.token_ids, output.text, output.logprobs)) - - del model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - - eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024) - eager_outputs = [] - for p in test_prompts: - output = eager_model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Eager model produced empty output" - eager_outputs.append((output.token_ids, output.text, output.logprobs)) - del eager_model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + enable_chunked_prefill=None, + ) as vllm_model: + compiled_outputs = [] + for p in test_prompts: + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Compiled model produced empty output" + compiled_outputs.append((output.token_ids, output.text, output.logprobs)) + + with vllm_runner( + model_name, + enforce_eager=True, + max_model_len=1024, + enable_chunked_prefill=None, + ) as vllm_model: + eager_outputs = [] + for p in test_prompts: + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Eager model produced empty output" + eager_outputs.append((output.token_ids, output.text, output.logprobs)) check_logprobs_close( outputs_0_lst=eager_outputs, @@ -239,44 +236,39 @@ def test(model_class, input1, input2, is_01_specialization=False): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") -def test_piecewise_backend_empty_sym_shape_indices(): +def test_piecewise_backend_empty_sym_shape_indices(vllm_runner): """Test that PiecewiseBackend handles empty sym_shape_indices correctly. When all inputs have static shapes (no torch.SymInt), sym_shape_indices will be empty. The fix in PiecewiseBackend.__call__ handles this case by using the first compiled range_entry. """ - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - # Use small max_model_len and max_num_batched_tokens to encourage # static shape compilation with empty sym_shape_indices - llm = LLM( - model="Qwen/Qwen3-0.6B", + with vllm_runner( + "Qwen/Qwen3-0.6B", max_model_len=512, max_num_batched_tokens=1, + enable_chunked_prefill=None, compilation_config={ "mode": CompilationMode.VLLM_COMPILE, "dynamic_shapes_config": { "type": DynamicShapesType.BACKED.value, }, }, - ) - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) + ) as vllm_model: + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) - # Generate with static shape inputs - output = llm.generate("Hello, my name is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output" - - # Generate again to verify compilation works with empty sym_shape_indices - output = llm.generate("The capital of France is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output on second run" + # Generate with static shape inputs + output = vllm_model.llm.generate( + "Hello, my name is", sampling_params=sampling_params + ) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output" - del llm - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + # Generate again to verify compilation works with empty sym_shape_indices + output = vllm_model.llm.generate( + "The capital of France is", sampling_params=sampling_params + ) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output on second run" diff --git a/tests/config/test_bailing_mtp_config.py b/tests/config/test_bailing_mtp_config.py index 8fae29959f23..bb7cf3ad5631 100644 --- a/tests/config/test_bailing_mtp_config.py +++ b/tests/config/test_bailing_mtp_config.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import get_args + from transformers import PretrainedConfig from vllm.config.speculative import MTPModelTypes, SpeculativeConfig @@ -34,7 +36,7 @@ def test_bailing_hybrid_mtp_hf_config_override(): assert overridden.model_type == "bailing_hybrid_mtp" assert overridden.architectures == ["BailingMoeV25MTPModel"] assert overridden.n_predict == 1 - assert "bailing_hybrid_mtp" in MTPModelTypes.__args__ + assert "bailing_hybrid_mtp" in get_args(MTPModelTypes) def test_bailing_hybrid_mtp_model_arch_config(): diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 3cc26e6e4761..35bc1e167b52 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -214,3 +214,65 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): base_hash = CacheConfig().compute_hash() assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash + + +def test_envs_compile_factors_relocation_invariant(tmp_path): + """Relocating HOME or the XDG roots must not change the compile-cache + env hash. + + Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME, + VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about + compiled artifacts, only about where directories live. When they leak + into compile_factors(), a cache produced under one HOME/XDG layout + silently misses under another - which defeats copying or pre-baking a + compile cache into a container image. + """ + import os + import subprocess + import sys + + code = """ +import sys +import logging +logging.disable(logging.CRITICAL) +from vllm import envs +from vllm.config.utils import hash_factors +print(hash_factors(envs.compile_factors())) +""" + + def hash_with(extra_env): + env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"} + # Drop explicit overrides so the derived defaults are what is + # exercised, then apply the relocation under test. + for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"): + env.pop(key, None) + env.update(extra_env) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + check=True, + env=env, + ) + return result.stdout.strip() + + xdg_cache = tmp_path / "relocated-xdg-cache" + xdg_config = tmp_path / "relocated-xdg-config" + new_home = tmp_path / "relocated-home" + for d in (xdg_cache, xdg_config, new_home): + d.mkdir() + + base = hash_with({}) + relocated_xdg = hash_with( + {"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)} + ) + relocated_home = hash_with({"HOME": str(new_home)}) + + assert relocated_xdg == base, ( + "XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache " + "env hash - a location-only derived var is leaking into the key" + ) + assert relocated_home == base, ( + "HOME relocation changed the compile-cache env hash - a " + "location-only derived var is leaking into the key" + ) diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 9720d84672fd..5d0d755e5106 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -1,21 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + import pytest +from transformers import PretrainedConfig from vllm.config.model import ModelConfig from vllm.config.multimodal import MultiModalConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + ModelArchConfigConvertorBase, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum def test_mm_encoder_attn_backend_str_conversion(): - config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN") + config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN") # type: ignore[arg-type] assert config.mm_encoder_attn_backend == AttentionBackendEnum.FLASH_ATTN def test_mm_encoder_attn_backend_invalid(): with pytest.raises(ValueError): - MultiModalConfig(mm_encoder_attn_backend="not_a_backend") + MultiModalConfig(mm_encoder_attn_backend="not_a_backend") # type: ignore[arg-type] + + +def test_mm_hasher_algorithm_invalid(): + with pytest.raises(ValueError, match="mm_hasher_algorithm"): + MultiModalConfig(mm_hasher_algorithm="md5") # type: ignore[arg-type] def test_mm_encoder_attn_backend_hash_updates(): @@ -43,6 +54,15 @@ def test_language_model_only_affects_model_hash(): assert base_hash != lm_only_hash +@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"]) +def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str): + config = MultiModalConfig( + media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}} + ) + + assert config.use_gpu_video_backend() + + def test_mm_encoder_fp8_scale_path_requires_fp8(): with pytest.raises(ValueError, match="mm_encoder_attn_dtype"): MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json") @@ -59,3 +79,97 @@ def test_mm_encoder_attn_dtype_hash_updates(tmp_path): ).compute_hash() assert base_hash != fp8_hash assert fp8_hash != fp8_static_hash + + +def _make_mm_prefix_model_config( + *, + language_model_only: bool = False, +) -> ModelConfig: + model_config = MagicMock(spec=ModelConfig) + model_config.multimodal_config = MultiModalConfig( + language_model_only=language_model_only + ) + # Bind real helper methods onto the mock. + model_config._supports_multimodal_for_mm_prefix = ( + ModelConfig._supports_multimodal_for_mm_prefix.__get__( + model_config, ModelConfig + ) + ) + return model_config + + +@pytest.mark.parametrize("supports_mm", [True, False]) +def test_supports_multimodal_for_mm_prefix_uses_registry(supports_mm: bool): + model_config = _make_mm_prefix_model_config() + + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=supports_mm, + ) as mocked: + assert model_config._supports_multimodal_for_mm_prefix() is supports_mm + mocked.assert_called_once_with(model_config) + + # Sticky cache — registry must not be consulted again. + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + side_effect=AssertionError("should use cache"), + ): + assert model_config._supports_multimodal_for_mm_prefix() is supports_mm + + +def test_supports_multimodal_for_mm_prefix_before_multimodal_config(): + model_config = _make_mm_prefix_model_config() + model_config.multimodal_config = None + + assert model_config._supports_multimodal_for_mm_prefix() is True + assert not hasattr(model_config, "_supports_multimodal_inputs_cached") + + +def test_language_model_only_disables_via_supports_multimodal_inputs(): + """language_model_only zeros all limits, so registry reports text-only.""" + model_config = _make_mm_prefix_model_config(language_model_only=True) + + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=False, + ): + assert model_config._supports_multimodal_for_mm_prefix() is False + + +def test_convertor_clears_mm_prefix_when_multimodal_disabled(): + hf_config = PretrainedConfig( + model_type="gemma3", + architectures=["Gemma3ForConditionalGeneration"], + ) + hf_config.is_mm_prefix_lm = True + convertor = ModelArchConfigConvertorBase(hf_config, hf_config) + + assert convertor.is_mm_prefix_lm(supports_multimodal=True) is True + assert convertor.is_mm_prefix_lm(supports_multimodal=False) is False + + enabled = convertor.convert(supports_multimodal=True) + disabled = convertor.convert(supports_multimodal=False) + assert enabled.is_mm_prefix_lm is True + assert disabled.is_mm_prefix_lm is False + + +def test_sticky_cache_survives_text_subconfig_regeneration(): + """with_hf_config deepcopies the cached decision onto text submodules.""" + model_config = _make_mm_prefix_model_config() + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=False, + ): + assert model_config._supports_multimodal_for_mm_prefix() is False + + # Simulate deepcopy onto a Gemma4ForCausalLM-like config that would + # otherwise fail registry lookup / return False incorrectly. + text_config = _make_mm_prefix_model_config() + text_config._supports_multimodal_inputs_cached = ( + model_config._supports_multimodal_inputs_cached + ) + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + side_effect=AssertionError("must not re-query registry"), + ): + assert text_config._supports_multimodal_for_mm_prefix() is False diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py new file mode 100644 index 000000000000..117ae5a66d53 --- /dev/null +++ b/tests/config/test_speculative_draft_max_position_embeddings.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the EAGLE draft ``max_position_embeddings`` override (#48894). + +EAGLE drafts share the target's positional space, but some draft +checkpoints (e.g. ``yuhuili/EAGLE3-LLaMA3.1-Instruct-8B``) ship a +``max_position_embeddings`` (2048) far smaller than the target's context. +That value sizes the draft's rotary ``cos_sin_cache`` while the proposer +feeds positions up to the target's ``max_model_len``, so the cache gather +goes out of bounds — a device-side assert under torch.compile and silent +garbage reads in eager mode. ``SpeculativeConfig`` must raise the draft's +value to the target's ``max_model_len``, with a log, for the eagle/eagle3 +methods only. +""" + +import logging +from typing import Literal + +import pytest +from transformers import PretrainedConfig + +from vllm.config.model import ModelConfig +from vllm.config.parallel import ParallelConfig +from vllm.config.speculative import SpeculativeConfig + +# All repos are public; only config/tokenizer-config files are fetched. +EAGLE3_DRAFT = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" # max_position_embeddings=2048 +LLAMA3_TARGET = "unsloth/Meta-Llama-3.1-8B-Instruct" # max_position_embeddings=131072 +AR_MODEL = "JackFram/llama-68m" # max_position_embeddings=2048 + +_LOGGER = "vllm.config.speculative" +_OVERRIDE_MSG = "Overriding draft model max_position_embeddings" + + +@pytest.fixture +def vllm_caplog(caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch): + """Make caplog see vLLM logger records (vLLM sets propagate=False).""" + monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True) + with caplog.at_level(logging.INFO, logger=_LOGGER): + yield caplog + + +def _override_logged(caplog: pytest.LogCaptureFixture) -> bool: + return any(_OVERRIDE_MSG in record.getMessage() for record in caplog.records) + + +@pytest.mark.cpu_test +def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture): + hf_config = PretrainedConfig(max_position_embeddings=2048) + SpeculativeConfig._maybe_override_draft_max_position_embeddings( + hf_config, target_max_model_len=8192 + ) + assert hf_config.max_position_embeddings == 8192 + assert _override_logged(vllm_caplog) + + +@pytest.mark.cpu_test +def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture): + hf_config = PretrainedConfig(max_position_embeddings=8192) + SpeculativeConfig._maybe_override_draft_max_position_embeddings( + hf_config, target_max_model_len=8192 + ) + assert hf_config.max_position_embeddings == 8192 + assert not _override_logged(vllm_caplog) + + +@pytest.mark.cpu_test +def test_override_ignores_missing_attribute(vllm_caplog: pytest.LogCaptureFixture): + hf_config = PretrainedConfig() + hf_config.__dict__.pop("max_position_embeddings", None) + SpeculativeConfig._maybe_override_draft_max_position_embeddings( + hf_config, target_max_model_len=8192 + ) + assert not hasattr(hf_config, "max_position_embeddings") + assert not _override_logged(vllm_caplog) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("method", ["eagle", "eagle3"]) +def test_eagle_draft_inherits_target_max_model_len( + method: Literal["eagle", "eagle3"], vllm_caplog: pytest.LogCaptureFixture +): + target_model_config = ModelConfig(LLAMA3_TARGET) + assert target_model_config.max_model_len > 2048 + speculative_config = SpeculativeConfig( + target_model_config=target_model_config, + target_parallel_config=ParallelConfig(), + model=EAGLE3_DRAFT, + method=method, + num_speculative_tokens=3, + ) + draft_hf_config = speculative_config.draft_model_config.hf_config + assert draft_hf_config.max_position_embeddings == target_model_config.max_model_len + assert _override_logged(vllm_caplog) + + +@pytest.mark.cpu_test +def test_independent_draft_model_keeps_its_own_limit( + vllm_caplog: pytest.LogCaptureFixture, +): + """An independent AR draft may genuinely have a smaller context than the + target; its max_position_embeddings must not be resized.""" + target_model_config = ModelConfig( + AR_MODEL, hf_overrides={"max_position_embeddings": 8192} + ) + assert target_model_config.max_model_len == 8192 + speculative_config = SpeculativeConfig( + target_model_config=target_model_config, + target_parallel_config=ParallelConfig(), + model=AR_MODEL, + method="draft_model", + num_speculative_tokens=3, + ) + draft_hf_config = speculative_config.draft_model_config.hf_config + assert draft_hf_config.max_position_embeddings == 2048 + assert not _override_logged(vllm_caplog) diff --git a/tests/conftest.py b/tests/conftest.py index 05a81c75eac2..f9003e55e738 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,7 +31,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from huggingface_hub import snapshot_download +from vllm.transformers_utils.repo_utils import hf_api from PIL import Image from transformers import ( AutoConfig, @@ -66,6 +66,7 @@ from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import BeamSearchParams +from vllm.transformers_utils.repo_utils import with_retry from vllm.transformers_utils.utils import maybe_model_redirect from vllm.utils.collection_utils import is_list_of from vllm.utils.torch_utils import set_default_torch_num_threads @@ -349,6 +350,20 @@ def audio_assets() -> AudioTestAssets: _R = TypeVar("_R") +def _fix_v4_tied_weights_keys(model_cls: type) -> None: + """Convert a v4 list-format _tied_weights_keys to the transformers v5 dict form.""" + tied = getattr(model_cls, "_tied_weights_keys", None) + if not isinstance(tied, list) or not tied: + return + result = { + k: "model.embed_tokens.weight" + for k in tied + if "lm_head" in k and k.endswith(".weight") + } + if result: + setattr(model_cls, "_tied_weights_keys", result) + + class HfRunner: def get_default_device(self): from vllm.platforms import current_platform @@ -474,6 +489,22 @@ def _init( trust_remote_code=trust_remote_code, ) else: + if trust_remote_code and hasattr(self.config, "auto_map"): + cls_ref = self.config.auto_map.get(auto_cls.__name__) + if cls_ref is not None: + from vllm.transformers_utils.dynamic_module import ( + try_get_class_from_dynamic_module, + ) + + model_cls = try_get_class_from_dynamic_module( + cls_ref, + model_name, + trust_remote_code=trust_remote_code, + warn_on_fail=False, + ) + if model_cls is not None: + _fix_v4_tied_weights_keys(model_cls) + model = cast( nn.Module, auto_cls.from_pretrained( @@ -513,9 +544,15 @@ def _init( # it will call torch.accelerator.device_count() from transformers import AutoProcessor - self.processor = AutoProcessor.from_pretrained( - model_name, - trust_remote_code=trust_remote_code, + # A concurrent refresh of the shared HF cache can briefly hide + # processor configuration files. Retry just as model config loading + # does in vllm.transformers_utils.config. + self.processor = with_retry( + lambda: AutoProcessor.from_pretrained( + model_name, + trust_remote_code=trust_remote_code, + ), + f"Error loading processor for {model_name}", ) if skip_tokenizer_init: if self.processor is None: @@ -1466,7 +1503,7 @@ def num_gpus_available(): def dummy_opt_path(): json_path = os.path.join(_dummy_opt_path, "config.json") if not os.path.exists(_dummy_opt_path): - snapshot_download( + hf_api().snapshot_download( repo_id="facebook/opt-125m", local_dir=_dummy_opt_path, ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"], @@ -1484,7 +1521,7 @@ def dummy_opt_path(): def dummy_llava_path(): json_path = os.path.join(_dummy_llava_path, "config.json") if not os.path.exists(_dummy_llava_path): - snapshot_download( + hf_api().snapshot_download( repo_id="llava-hf/llava-1.5-7b-hf", local_dir=_dummy_llava_path, ignore_patterns=[ @@ -1509,7 +1546,7 @@ def dummy_llava_path(): def dummy_gemma2_embedding_path(): json_path = os.path.join(_dummy_gemma2_embedding_path, "config.json") if not os.path.exists(_dummy_gemma2_embedding_path): - snapshot_download( + hf_api().snapshot_download( repo_id="BAAI/bge-multilingual-gemma2", local_dir=_dummy_gemma2_embedding_path, ignore_patterns=[ @@ -1698,33 +1735,46 @@ def disable_deepgemm_ue8m0(monkeypatch): is_deep_gemm_e8m0_used.cache_clear() +def _should_clean_gpu_memory_between_tests() -> bool: + # This must stay opt-in: a function-scoped fixture cannot distinguish + # stale VRAM from allocations owned by longer-lived module/session fixtures. + return os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") == "1" + + @pytest.fixture(autouse=True) def clean_gpu_memory_between_tests(): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1": + if not _should_clean_gpu_memory_between_tests(): yield return - # Wait for GPU memory to be cleared before starting the test import gc - from tests.utils import wait_for_gpu_memory_to_clear + from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle num_gpus = torch.accelerator.device_count() - if num_gpus > 0: + + def _wait_for_settled_gpu_memory() -> None: + if num_gpus <= 0: + return try: - wait_for_gpu_memory_to_clear( - devices=list(range(num_gpus)), - threshold_ratio=0.1, - ) + if current_platform.is_rocm(): + wait_for_rocm_memory_to_settle() + else: + wait_for_gpu_memory_to_clear( + devices=list(range(num_gpus)), + threshold_ratio=0.1, + ) except ValueError as e: logger.info("Failed to clean GPU memory: %s", e) + _wait_for_settled_gpu_memory() + yield - # Clean up GPU memory after the test if torch.cuda.is_available(): torch.accelerator.empty_cache() gc.collect() + _wait_for_settled_gpu_memory() @pytest.fixture @@ -1738,6 +1788,22 @@ def use_fresh_inductor_cache(): yield +@pytest.fixture +def disable_vllm_compile_cache(monkeypatch, use_fresh_inductor_cache): + """ + Use a fresh inductor cache AND disable vLLM's on-disk torch.compile cache. + + This forces compilation (and any custom compile passes) to actually run + instead of being served from a warm cache left behind by previous runs + (e.g. on persistent CI agents). Use this for tests that inspect what + happens during compilation; use ``use_fresh_inductor_cache`` (or + ``fresh_vllm_cache``) instead when the vLLM compile cache must stay + enabled (e.g. cache save/load tests). + """ + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + yield + + @pytest.fixture def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache): """Temporary VLLM_CACHE_ROOT combined with a fresh inductor cache.""" diff --git a/tests/detokenizer/test_check_stop_strings.py b/tests/detokenizer/test_check_stop_strings.py new file mode 100644 index 000000000000..2fae373f60b5 --- /dev/null +++ b/tests/detokenizer/test_check_stop_strings.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for check_stop_strings. + +These are pure-function tests (no model / GPU). They pin down which stop +string is selected when several stop strings match within the text that was +appended in a single step -- which happens under speculative decoding, where +multiple tokens (and therefore multiple stop strings) can be appended at once. +""" + +import pytest + +from vllm.v1.engine.detokenizer import check_stop_strings + + +@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]]) +def test_earliest_completing_stop_wins_regardless_of_list_order(stop): + # " The user is a": " is a" (5 chars) was appended in one step. Both "is" + # (index 10) and " a" (index 13) land in the same window. "is" completes + # earlier in the text, so it must win over list order. + text = " The user is a" + new_char_count = len(" is a") + + assert check_stop_strings(text, new_char_count, stop, include_in_output=False) == ( + "is", + 10, + ) + + +@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]]) +def test_earliest_completing_stop_include_in_output(stop): + text = " The user is a" + new_char_count = len(" is a") + + # Truncate to the end of "is" (index 12) -> " The user is". + assert check_stop_strings(text, new_char_count, stop, include_in_output=True) == ( + "is", + 12, + ) + + +def test_completion_position_not_start_position(): + # "b" starts later than "abc" but completes earlier, so it must win. + text = "abc" + assert check_stop_strings( + text, len(text), ["abc", "b"], include_in_output=False + ) == ("b", 1) + + +@pytest.mark.parametrize( + "stop,expected", + [ + (["ab", "b"], ("ab", 0)), + (["b", "ab"], ("b", 1)), + ], +) +def test_ties_broken_by_list_order(stop, expected): + # "ab" and "b" both complete at index 2; list order decides the winner. + text = "ab" + assert ( + check_stop_strings(text, len(text), stop, include_in_output=False) == expected + ) + + +def test_single_stop_in_window_unchanged(): + # The common case (one stop in the window) is unaffected by the change. + text = "hello world." + assert check_stop_strings(text, 1, ["."], include_in_output=False) == (".", 11) + # Stop completes at the very end -> no truncation needed (-1). + assert check_stop_strings(text, 1, ["."], include_in_output=True) == (".", -1) + + +def test_no_match_and_empty_inputs_return_none(): + assert check_stop_strings("hello", 5, ["zzz"], include_in_output=False) is None + assert check_stop_strings("hello", 0, ["h"], include_in_output=False) is None + assert check_stop_strings("hello", 5, [], include_in_output=False) is None diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 48b007da664c..19a095c7c6f7 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -7,6 +7,7 @@ from collections.abc import Callable from typing import Any +from unittest.mock import Mock import pytest import ray @@ -19,6 +20,8 @@ tensor_model_parallel_all_reduce, tensor_model_parallel_reduce_scatter, ) +from vllm.distributed.device_communicators import flashinfer_all_reduce +from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors @@ -278,6 +281,43 @@ def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32)) +@pytest.mark.parametrize("aliased", [False, True]) +def test_cuda_communicator_checkpoints_flashinfer_workspaces( + monkeypatch: pytest.MonkeyPatch, + aliased: bool, +) -> None: + group = object() + normal_workspace = Mock() + quant_workspace = normal_workspace if aliased else Mock() + unique_workspaces = ( + [normal_workspace] if aliased else [normal_workspace, quant_workspace] + ) + + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal_workspace) + monkeypatch.setattr( + flashinfer_all_reduce, "_fi_ar_quant_workspace", quant_workspace + ) + monkeypatch.setattr( + flashinfer_all_reduce, + "_fi_ar_workspace_groups", + {id(workspace): group for workspace in unique_workspaces}, + ) + monkeypatch.setattr( + flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False + ) + + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.cpu_group = group + communicator.fi_ar_comm = None + communicator.all2all_manager = None + communicator.checkpoint_prepare() + communicator.checkpoint_restore() + + for workspace in unique_workspaces: + workspace.checkpoint_prepare.assert_called_once_with() + workspace.checkpoint_restore.assert_called_once_with(group) + + def test_async_intermediate_tensors_lazy_wait() -> None: work = _DummyWork() post_calls = {"n": 0} diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 4ce7497598ad..01c254c2d039 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -3,7 +3,9 @@ import os import subprocess +import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest import requests @@ -40,6 +42,106 @@ def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: return False +def _traffic_loop( + server: RemoteOpenAIServer, + dp_rank: int | None, + ready: threading.Barrier, + stop: threading.Event, + finished: threading.Event, + is_probe: bool = False, +) -> list[tuple[float, float, int | None]]: + url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") + payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} + headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} + request_payload = None if is_probe else payload + responses = [] + is_ready = False + while not stop.is_set(): + request_start = time.perf_counter() + try: + response = requests.post( + url, json=request_payload, headers=headers, timeout=120 + ) + status_code = response.status_code + except requests.exceptions.RequestException: + status_code = None + responses.append((request_start, time.perf_counter(), status_code)) + if status_code == 200: + if not is_ready: + ready.wait(timeout=120) + is_ready = True + if finished.is_set(): + return responses + time.sleep(0.05) + return responses + + +def _downtime(responses: list[tuple[float, float, int | None]]) -> float: + rejected = [end for _, end, status in responses if status == 503] + if not rejected: + return 0 + recovered = next( + end for _, end, status in responses if status == 200 and end > rejected[-1] + ) + return recovered - rejected[0] + + +def _scale_with_traffic( + server: RemoteOpenAIServer, + source_dp_size: int, + new_dp_size: int, + traffic_mode: str, +) -> None: + traffic_clients: list[int | None] = [] + if traffic_mode == "light": + traffic_clients = [0] + elif traffic_mode == "heavy": + traffic_clients = [None] * source_dp_size + clients = [(None, True)] + [(rank, False) for rank in traffic_clients] + ready = threading.Barrier(len(clients) + 1) + stop = threading.Event() + finished = threading.Event() + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + futures = [ + executor.submit( + _traffic_loop, server, rank, ready, stop, finished, is_probe + ) + for rank, is_probe in clients + ] + try: + ready.wait(timeout=120) + start_time = time.perf_counter() + assert _send_scale_command(server, new_dp_size) + scale_seconds = time.perf_counter() - start_time + finished.set() + probe_result, *results = [future.result(timeout=120) for future in futures] + finally: + stop.set() + + bad_statuses = { + status + for responses in [probe_result, *results] + for _, _, status in responses + if status not in (200, 503) + } + assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" + probe_503 = [start for start, _, status in probe_result if status == 503] + assert probe_503, "Scaling probe did not observe commit" + assert not results or any( + status == 200 and start_time <= request_start and request_end < probe_503[0] + for responses in results + for request_start, request_end, status in responses + ), "No request completed successfully during preparation" + + print( + f"[Elastic EP timing][{source_dp_size}->{new_dp_size}]" + f"[traffic={traffic_mode}] " + f"scale_seconds={scale_seconds:.3f} " + f"downtime_seconds={_downtime(probe_result):.3f}" + ) + + def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: assert server.port is not None result = evaluate_gsm8k( @@ -59,7 +161,7 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: return accuracy -def _base_serve_args(use_async_eplb: bool = False) -> list[str]: +def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]: args = [ "--trust-remote-code", "--tensor-parallel-size", @@ -78,57 +180,65 @@ def _base_serve_args(use_async_eplb: bool = False) -> list[str]: "--eplb-config.num_redundant_experts", "0", "--eplb-config.use_async", - "true" if use_async_eplb else "false", + "true", "--eplb-config.step_interval", - "10", + "300", "--eplb-config.window_size", "5", "--data-parallel-backend", "ray", "--data-parallel-size", - "2", + str(dp_size), "--api-server-count", "1", + "--disable-access-log-for-endpoints", + "/is_scaling_elastic_ep", ] leader_address = os.environ.get("LEADER_ADDRESS") if leader_address: args.extend(["--data-parallel-address", leader_address]) + if enforce_eager: + args.append("--enforce-eager") return args @pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] + ("enforce_eager", "traffic_mode"), + [ + pytest.param(True, "none", id="enforce_eager_none"), + pytest.param(True, "light", id="enforce_eager_light"), + pytest.param(True, "heavy", id="enforce_eager_heavy"), + pytest.param(False, "heavy", id="cuda_graphs_heavy"), + ], ) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling(use_async_eplb: bool): - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl +def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) + target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) + assert target_dp_size > initial_dp_size + vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager) with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 ) as server: - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)") - - assert _send_scale_command(server, 4) - time.sleep(10) - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)") + initial_accuracy = _run_gsm8k_eval(server, "Initial") + _scale_with_traffic(server, initial_dp_size, target_dp_size, traffic_mode) + scale_up_accuracy = _run_gsm8k_eval(server, "After scale up") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" ) - assert _send_scale_command(server, 2) - time.sleep(5) - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") - + _scale_with_traffic(server, target_dp_size, initial_dp_size, traffic_mode) + scale_down_accuracy = _run_gsm8k_eval(server, "After scale down") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" @@ -147,24 +257,20 @@ def test_elastic_ep_scaling(use_async_eplb: bool): print(f" Tolerance: {ACCURACY_TOL:.3f}") -@pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] -) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling_uneven(use_async_eplb: bool): +def test_elastic_ep_scaling_uneven(): """Test scale up with uneven worker distribution. This tests the case where num_new_workers % old_dp_size != 0, specifically 2 -> 3 where remainder = 1 % 2 = 1. This exercises the remainder handling in sender-receiver pairing. """ - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + vllm_serve_args = _base_serve_args() with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 @@ -174,7 +280,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale 2 -> 3: This has remainder = 1 % 2 = 1 # Tests uneven sender-receiver pairing assert _send_scale_command(server, 3) - time.sleep(10) scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( @@ -184,7 +289,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale back down to 2 assert _send_scale_command(server, 2) - time.sleep(5) scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py index 7d5e58b26ef8..d8d870bc98c8 100644 --- a/tests/distributed/test_eplb_fused_moe_layer.py +++ b/tests/distributed/test_eplb_fused_moe_layer.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Test that the interaction between EPLB and FusedMoE Layer is okay +# Test that the interaction between EPLB and FusedMoEFactory Layer is okay from dataclasses import dataclass @@ -16,7 +16,7 @@ get_eplb_group, get_tp_group, ) -from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner from .eplb_utils import distributed_run, set_env_vars_and_device @@ -69,8 +69,8 @@ def make_fused_moe_layer( rank: int, layer_idx: int, test_config: TestConfig, -) -> FusedMoE: - fml = FusedMoE( +) -> MoERunner: + fml = FusedMoEFactory( num_experts=test_config.num_experts, top_k=test_config.num_topk, hidden_size=test_config.hidden_size, diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index e2d54821ce9c..fe11e13cd60b 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Test that the interaction between EPLB and FusedMoE Layer is okay for DP w/ NVFP4 +# Test that the interaction between EPLB and MoERunner Layer is okay for DP w/ NVFP4 from dataclasses import dataclass @@ -19,7 +19,7 @@ get_eplb_group, ) from vllm.forward_context import set_forward_context -from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory, MoERunner from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, ModelOptNvFp4FusedMoE, @@ -44,7 +44,7 @@ def make_fused_moe_layer( rank: int, layer_idx: int, test_config: TestConfig, -) -> FusedMoE: +) -> MoERunner: quant_config = None device = torch.device(f"cuda:{rank}") @@ -55,7 +55,7 @@ def make_fused_moe_layer( exclude_modules=[], ) - fml = FusedMoE( + fml = FusedMoEFactory( num_experts=test_config.num_experts, top_k=test_config.num_topk, hidden_size=test_config.hidden_size, diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index 9b5601ad1d93..894f3f7f8b83 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -284,7 +284,7 @@ def test_data_parallel_rank_tagging(publisher_config): sub_1.close() -def test_event_publisher_factory(): +def test_event_publisher_factory(random_port): """Test event publisher factory creation behavior under different configurations""" from vllm.config.kv_events import KVEventsConfig from vllm.distributed.kv_events import ZmqEventPublisher @@ -308,10 +308,15 @@ def test_event_publisher_factory(): config = KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", - endpoint="inproc://test-factory-true", + endpoint=f"tcp://*:{random_port}", + replay_endpoint=f"tcp://*:{random_port + 100}", ) - publisher = EventPublisherFactory.create(config, DP_RANK) + publisher = EventPublisherFactory.create(config, DP_RANK + 1) assert isinstance(publisher, ZmqEventPublisher) + resolved_config = publisher.get_publisher_config() + assert resolved_config.endpoint == f"tcp://*:{random_port + 1}" + assert resolved_config.replay_endpoint == f"tcp://*:{random_port + 101}" + assert config.endpoint == f"tcp://*:{random_port}" publisher.shutdown() # test unknown publisher diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index fb2d9bb832e7..f9b76ea1b030 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -73,7 +73,10 @@ def _spawn_workers(worker_fn, world_size, *, dp_size=None): err_queue.close() err_queue.join_thread() if errors: - pytest.fail("Worker(s) failed:\n" + "\n---\n".join(errors)) + combined = "\n---\n".join(errors) + if "NCCL GIN" in combined: + pytest.skip("NCCL GIN not available on this system") + pytest.fail("Worker(s) failed:\n" + combined) def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue): @@ -916,8 +919,11 @@ def _deepep_v2_lifecycle_worker(rank, world_size): DeepEPV2All2AllManager, ) - cpu_group = get_ep_group().cpu_group - manager = DeepEPV2All2AllManager(cpu_group) + ep_group = get_ep_group() + manager = DeepEPV2All2AllManager( + ep_group.cpu_group, + device_group=ep_group.device_group, + ) assert manager.rank == rank assert manager.world_size == world_size diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index cbdc02527064..042349e4fbb1 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -254,8 +254,8 @@ def test_multiproc_executor_shutdown_cleanup(): for worker in executor.workers: assert not worker.proc.is_alive(), "Worker processes should be terminated" - # Verify shutdown event is set - assert executor.shutdown_event.is_set(), "Shutdown event should be set" + # Verify shutdown flag is set + assert executor.shutting_down, "Shutdown flag should be set" # Multiple shutdowns should be safe (idempotent) executor.shutdown() @@ -292,7 +292,6 @@ def test_multiproc_executor_pipeline_parallel(): "Max concurrent batches should follow the configured PP/async " "scheduling policy" ) - finally: # Clean up executor.shutdown() @@ -338,6 +337,10 @@ def test_multiproc_executor_multi_node(): - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 Total world_size = 4, nnodes = 2 """ + # Python 3.14+ changed default multiprocessing start method to 'forkserver' + # which cannot pickle nested functions. Use 'fork' for this test. + mp_ctx = multiprocessing.get_context("fork") + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) port = s.getsockname()[1] @@ -405,12 +408,12 @@ def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int): executor.shutdown() # Create a queue to collect results from both processes - result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue() + result_queue: multiprocessing.Queue[dict[str, int | bool]] = mp_ctx.Queue() # Start both node processes processes = [] for node_rank in range(2): - p = multiprocessing.Process( + p = mp_ctx.Process( target=run_node, args=(node_rank, result_queue, port), name=f"Node{node_rank}", diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index fd4a1d0f570b..22e0321da60e 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -121,7 +121,6 @@ def iter_params(self, model_id: str): "ibm/PowerMoE-3b": PPTestSettings.fast(), "internlm/internlm2-chat-7b": PPTestSettings.fast(), "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(), - "pfnet/plamo-2-1b": PPTestSettings.fast(), "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(), "meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(), # Tests TransformersForCausalLM diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index bfa28cc5c444..fb2b04abf0d5 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -38,6 +38,14 @@ def on_gfx942() -> bool: return False +def _quickreduce_tolerance(quant_level: QuickReduceRegime) -> dict: + # With INT3 inputs of absmax 23, per-rank error peaks at 2.75, + # doubled by the TP2 sum = 5.5. + if quant_level == QuickReduceRegime.INT3: + return {"atol": 5.5, "rtol": 0.1} + return {"atol": 2.5, "rtol": 0.1} + + set_random_seed(42) _test_size_rng = random.Random(44) # Size over 8MB is sufficient for custom quick allreduce. @@ -280,8 +288,9 @@ def graph_quickreduce( out2 = tensor_model_parallel_all_reduce(inp2) dist.all_reduce(inp2, group=group) graph.replay() - torch.testing.assert_close(out1, inp1, atol=2.5, rtol=0.1) - torch.testing.assert_close(out2, inp2, atol=2.5, rtol=0.1) + tol = _quickreduce_tolerance(fa.qr_quant_level) + torch.testing.assert_close(out1, inp1, **tol) + torch.testing.assert_close(out2, inp2, **tol) @ray.remote(num_gpus=1, max_calls=1) @@ -309,14 +318,15 @@ def eager_quickreduce( ) _assert_quickreduce(fa, inp) out = fa.quick_all_reduce(inp) - torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1) + tol = _quickreduce_tolerance(fa.qr_quant_level) + torch.testing.assert_close(out, inp * tp_size, **tol) inp = torch.tensor( [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device ) _assert_quickreduce(fa, inp) out = fa.quick_all_reduce(inp) - torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1) + torch.testing.assert_close(out, inp * tp_size, **tol) @ray.remote(num_gpus=1, max_calls=1) diff --git a/tests/distributed/test_rocm_quick_reduce.py b/tests/distributed/test_rocm_quick_reduce.py index 081698483a43..e422b032e257 100644 --- a/tests/distributed/test_rocm_quick_reduce.py +++ b/tests/distributed/test_rocm_quick_reduce.py @@ -15,14 +15,16 @@ import pytest import torch import torch.distributed as dist -from huggingface_hub import snapshot_download import vllm.envs as envs from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.network_utils import get_open_port +from ..utils import multi_gpu_test + pytestmark = pytest.mark.skipif( not current_platform.is_rocm(), reason="ROCm-only quick-reduce tests", @@ -60,6 +62,8 @@ def _make_quick_allreduce( qar.use_fp16_kernels = use_fp16_kernels qar.qr_quant_level = QuickReduceRegime[quant_level] qar.qr_max_size = qr_max_size + qar.qr_min_size = None + qar.qr_quantization_min_size = None return qar @@ -149,6 +153,116 @@ def _run_two_gpu_quick_allreduce_test( ) +CUDAGRAPH_WORLD_SIZE = 2 +CUDAGRAPH_ROUNDS = 10 +CUDAGRAPH_NUM_ELEMENTS = 1 << 21 # 2M fp16 = 4 MB, above the quick-reduce thresholds + + +def _quick_allreduce_cudagraph_worker( + rank: int, + world_size: int, + port: int, + quant_level: str, +): + # FP keeps the all-reduce bit-exact for small fp16 integers, so every + # replayed round can be checked exactly. + os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level + os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0" + _log(f"cudagraph worker start: rank={rank} quant={quant_level}") + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + ) + + qar = None + try: + from vllm.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + ) + + qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) + assert not qar.disabled + + N = CUDAGRAPH_NUM_ELEMENTS + inp = torch.empty(N, dtype=torch.float16, device=device) + out = torch.empty(N, dtype=torch.float16, device=device) + assert qar.should_quick_allreduce(inp) + + # Every rank contributes the same value v in a round, so the true + # cross-rank all-reduce sum is simply world_size * v. + def expected(v): + return float(world_size * v) + + if rank == 0: + print( + f"[repro] world_size={world_size} elems={N} regime={quant_level} fp16", + flush=True, + ) + + # Warmup, then capture a graph with EXACTLY ONE quick-reduce (isolated + # qr, so it is the sole writer of its flag slot -- the condition that + # triggers the stale-flag bug). + inp.fill_(1.0) + qar.quick_all_reduce(inp, out=out) + torch.accelerator.synchronize() + dist.barrier() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + qar.quick_all_reduce(inp, out=out) + torch.accelerator.synchronize() + dist.barrier() + + for v in range(CUDAGRAPH_ROUNDS): + inp.fill_(float(v)) # in-place: same value on every rank + dist.barrier() + g.replay() + torch.accelerator.synchronize() + dist.barrier() + got = out.float() + expect = expected(v) + if rank == 0: + print(f"round {v}: got={got[:10]}, expected={expect}", flush=True) + mismatch = ~torch.isclose(got, torch.full_like(got, expect)) + num_mismatch = int(mismatch.sum().item()) + assert num_mismatch == 0, ( + f"rank={rank} round={v} expected={expect} " + f"mismatched {num_mismatch}/{got.numel()} elements; " + f"unique wrong values={torch.unique(got[mismatch])[:8].tolist()}" + ) + _log(f"cudagraph worker complete: rank={rank} rounds={CUDAGRAPH_ROUNDS}") + finally: + if qar is not None: + qar.close() + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_cudagraph_replay_test(*, world_size: int, quant_level: str): + _log(f"launch {world_size}-GPU cudagraph replay case: quant={quant_level}") + ctx = mp.get_context("spawn") + port = get_open_port() + procs = [] + + for rank in range(world_size): + proc = ctx.Process( + target=_quick_allreduce_cudagraph_worker, + args=(rank, world_size, port, quant_level), + ) + proc.start() + procs.append(proc) + + for proc in procs: + proc.join(timeout=120) + assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" + _log(f"finished {world_size}-GPU cudagraph replay case: quant={quant_level}") + + MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" E2E_PREFILL_TOKENS = 1024 E2E_MAX_MODEL_LEN = 1536 @@ -208,11 +322,11 @@ def _log_prompt_summaries() -> None: @lru_cache(maxsize=1) def _get_model_path() -> str: try: - path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True) + path = hf_api().snapshot_download(repo_id=MODEL_NAME, local_files_only=True) _log(f"using cached model snapshot: {path}") return path except Exception: - path = snapshot_download(repo_id=MODEL_NAME) + path = hf_api().snapshot_download(repo_id=MODEL_NAME) _log(f"downloaded model snapshot: {path}") return path @@ -511,13 +625,21 @@ def test_quick_reduce_regime_values(): assert QuickReduceRegime.INT8.value == 1 assert QuickReduceRegime.INT6.value == 2 assert QuickReduceRegime.INT4.value == 3 - assert QuickReduceRegime.NONE.value == 4 + assert QuickReduceRegime.INT3.value == 4 + assert QuickReduceRegime.NONE.value == 5 def test_quick_reduce_regime_names(): from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime - assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"} + assert set(QuickReduceRegime.__members__) == { + "FP", + "INT8", + "INT6", + "INT4", + "INT3", + "NONE", + } @pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"]) @@ -693,7 +815,7 @@ def test_quick_allreduce_min_size_table(): for dtype in [torch.float16, torch.bfloat16]: for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES: min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)] - assert len(min_sizes) == 4 + assert len(min_sizes) == 5 assert all(size > 0 for size in min_sizes) @@ -719,6 +841,18 @@ def test_quick_allreduce_two_gpu_correctness(quant_level): ) +@multi_gpu_test(num_gpus=CUDAGRAPH_WORLD_SIZE) +def test_quick_allreduce_cudagraph_replay(): + # Regression test for the stale flag_color bug: a quick-reduce captured in a + # CUDA graph must return correct results on every replay, not the data from + # a previous round. + _log("cudagraph replay case") + _run_cudagraph_replay_test( + world_size=CUDAGRAPH_WORLD_SIZE, + quant_level="FP", + ) + + @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 7cf3b01e75c7..0b4139650321 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,17 +1,28 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io +import pickle import random import threading import time +from types import SimpleNamespace from unittest import mock import multiprocess as mp import numpy as np import pytest +import torch import torch.distributed as dist -from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.distributed.device_communicators import shm_broadcast +from vllm.distributed.device_communicators.shm_broadcast import ( + MessageQueue, + ShmRingBuffer, + _rebuild_tensor, + _reduce_tensor, + check_shm_free_space, +) from vllm.distributed.utils import StatelessProcessGroup from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -348,6 +359,280 @@ def test_message_queue_busy_to_idle(): distributed_run(worker_fn_test_busy_to_idle, 4) +@worker_fn_wrapper +def worker_fn_tensor_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + message_queue = MessageQueue.create_from_process_group( + dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank + ) + + # Both ranks construct the identical reference payload. + torch.manual_seed(42) + payload = { + # 2MiB: rides the shm ring as an out-of-band buffer (the receiving + # side must copy out of the reusable ring chunk). + "mid": torch.randn(1024, 512), + # 16MiB > max_chunk_bytes: overflows to the zmq socket (the + # receiving side aliases the zmq.Frame zero-copy). + "big": torch.randn(4096, 2048, dtype=torch.bfloat16), + "nested": ["plain", 123, {"inner": torch.arange(5)}], + } + + if rank == writer_rank: + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + assert wrapped_reduce.call_count == 3 + # Cycle the ring (max_chunks=4) several times over so that aliased + # ring chunks would be overwritten. + for i in range(16): + message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) + else: + received = message_queue.dequeue(timeout=30) + for key in ("mid", "big"): + assert torch.equal(received[key], payload[key]), key + assert received[key].dtype == payload[key].dtype, key + assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) + + snapshot = received["mid"].clone() + for i in range(16): + junk = message_queue.dequeue(timeout=30) + assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) + # Tensors received via the shm ring must not alias chunk memory + # that the writer has reused for subsequent messages. + assert torch.equal(received["mid"], snapshot) + # Rebuilt tensors must be writable, like regular tensors. + received["mid"] += 1.0 + received["big"][0, 0] = 1.0 + + dist.barrier() + print(f"tensor broadcast passed the test! Rank {rank}") + + +def test_tensor_broadcast(): + distributed_run(worker_fn_tensor_broadcast, 2) + + +def _dumps_oob(obj) -> tuple[bytes, list]: + """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor + dispatch table + out-of-band buffers >= 1MiB.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + bio = io.BytesIO() + pickler = pickle.Pickler( + bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback + ) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + pickler.dump(obj) + return bio.getvalue(), buffers + + +@pytest.mark.parametrize( + "case", + [ + "small", + "mid", + "bf16", + "fp8", + "empty", + "scalar", + "noncontig", + "requires_grad", + "conj", + "param", + ], +) +def test_tensor_pickle_roundtrip(case: str): + tensor = { + # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers. + "small": lambda: torch.randn(100, 10), + "mid": lambda: torch.randn(1024, 512), + # Dtypes numpy doesn't recognize. + "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16), + "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn), + # Shape edge cases. + "empty": lambda: torch.empty(0, 8), + "scalar": lambda: torch.tensor(3.14), + "noncontig": lambda: torch.randn(64, 64).t(), + # These fall back to torch's default reducer. + "requires_grad": lambda: torch.randn(8, 8, requires_grad=True), + "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(), + "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False), + }[case]() + + data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))}) + received = pickle.loads(data, buffers=buffers)["tensor"] + + assert received.shape == tensor.shape + assert received.dtype == tensor.dtype + if tensor.dtype == torch.float8_e4m3fn: + assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8)) + else: + assert torch.equal(received, tensor) + assert received.requires_grad == tensor.requires_grad + assert isinstance(received, type(tensor)) + if tensor.numel() and not tensor.requires_grad: + # Rebuilt tensors must be writable, like regular tensors. + received.view(-1)[0] = 1.0 + + +@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) +def test_reduce_tensor_fallback(case: str): + """Tensors the zero-copy reducer can't safely alias must fall back to + torch's default reduction.""" + if case == "cuda": + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + tensor = torch.randn(4, device="cuda") + elif case == "requires_grad": + tensor = torch.randn(8, requires_grad=True) + else: + tensor = torch.randn(4, dtype=torch.complex64).conj() + + reduced = _reduce_tensor(tensor) + assert reduced[0] is not _rebuild_tensor + + +@pytest.mark.parametrize("should_warn", [False, True]) +def test_reader_timeout_caps_indefinite_waits(should_warn): + with ( + mock.patch( + "vllm.distributed.device_communicators.shm_broadcast." + "SHM_READER_RECHECK_INTERVAL_MS", + new=7, + ), + mock.patch( + "vllm.distributed.device_communicators.shm_broadcast." + "VLLM_RINGBUFFER_WARNING_INTERVAL", + new=60, + ), + ): + timeout = MessageQueue.ReadTimeoutWithWarnings( + timeout=None, should_warn=should_warn + ) + assert timeout.timeout_ms() == 7 + + +def test_reader_rechecks_shm_after_idle_wait_timeout_without_notify(): + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=1024 * 1024, + max_chunks=1, + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + payload = 123 + poll_started = threading.Event() + allow_timeout = threading.Event() + result = {} + + def acquire_read_in_thread(): + try: + with reader.acquire_read(indefinite=True) as buf: + result["value"] = buf[0] + except Exception as exc: + result["exc"] = exc + + def poll_timeout(*, timeout: int | None = None): + poll_started.set() + assert allow_timeout.wait(timeout=5) + return [] + + try: + writer.wait_until_ready() + reader.wait_until_ready() + reader._spin_condition.last_read = 0 + reader._spin_condition.busy_loop_s = 0 + + with ( + mock.patch( + "vllm.distributed.device_communicators.shm_broadcast." + "SHM_READER_RECHECK_INTERVAL_MS", + new=50, + ), + mock.patch( + "vllm.distributed.device_communicators.shm_broadcast." + "VLLM_RINGBUFFER_WARNING_INTERVAL", + new=60, + ), + mock.patch.object( + reader._spin_condition.poller, + "poll", + side_effect=poll_timeout, + ) as poll, + ): + read_thread = threading.Thread(target=acquire_read_in_thread, daemon=True) + read_thread.start() + assert poll_started.wait(timeout=5) + with writer.acquire_write(timeout=0.1) as buf: + buf[0] = payload + allow_timeout.set() + read_thread.join(timeout=5) + + assert not read_thread.is_alive() + poll.assert_called_once_with(timeout=50) + + if "exc" in result: + raise result["exc"] + assert result["value"] == payload + with writer.buffer.get_metadata(0) as metadata_buffer: + assert metadata_buffer[0] == 1 + assert metadata_buffer[1] == 1 + finally: + writer.shutdown() + reader.shutdown() + for socket in ( + writer.local_socket, + writer._spin_condition.local_notify_socket, + reader.local_socket, + reader._spin_condition.local_notify_socket, + reader._spin_condition.read_cancel_socket, + reader._spin_condition.write_cancel_socket, + ): + socket.close(linger=0) + + +def test_acquire_read_releases_slot_when_reader_raises(): + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=1024 * 1024, + max_chunks=1, + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + try: + writer.wait_until_ready() + reader.wait_until_ready() + + writer.enqueue({"payload": "first"}) + + with ( + pytest.raises(RuntimeError, match="reader failed"), + reader.acquire_read(timeout=0.1), + ): + raise RuntimeError("reader failed") + + with writer.buffer.get_metadata(0) as metadata_buffer: + assert metadata_buffer[0] == 1 + assert metadata_buffer[1] == 1 + + with writer.acquire_write(timeout=0.1) as buf: + buf[0] = 0 + finally: + writer.shutdown() + reader.shutdown() + + def test_warning_logs(caplog_vllm): """ Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals @@ -392,3 +677,39 @@ def test_warning_logs(caplog_vllm): # Clean up when done writer.shutdown() reader.shutdown() + + +def _fake_disk_usage(free_bytes: int): + return SimpleNamespace(total=free_bytes, used=0, free=free_bytes) + + +def test_check_shm_free_space_raises_when_insufficient(tmp_path): + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(32 << 20) + ), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + check_shm_free_space(240 << 20, shm_path=str(tmp_path)) + + +def test_check_shm_free_space_passes_when_sufficient(tmp_path): + with mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(512 << 20) + ): + check_shm_free_space(240 << 20, shm_path=str(tmp_path)) + + +def test_check_shm_free_space_skipped_when_path_missing(tmp_path): + check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) + + +def test_shm_ring_buffer_creation_checks_free_space(): + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) + ), + mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index b79aa1974d13..ddc6dec04e81 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -27,10 +27,13 @@ WeightTransferTrainerFactory, ) from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, WeightTransferInitRequest, WeightTransferUpdateRequest, ) from vllm.distributed.weight_transfer.ipc_engine import ( + IPCTrainerInitInfo, + IPCTrainerWeightTransferEngine, IPCWeightTransferEngine, IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo, @@ -1071,7 +1074,7 @@ def inference_receive_ipc_tensor( import torch - _set_ray_assigned_device() + device = _set_ray_assigned_device() from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig @@ -1088,6 +1091,8 @@ def load_weights(self, weights): for name, tensor in weights: self.received.append((name, tensor.clone())) + # Trainer sends unpacked IPC handles; the worker learns packed=False from + # the init handshake below (IPCWeightTransferInitInfo defaults to False). config = WeightTransferConfig(backend="ipc") vllm_config = MagicMock() parallel_config = MagicMock(spec=ParallelConfig) @@ -1099,9 +1104,7 @@ def load_weights(self, weights): vllm_config.model_config = MagicMock() recorder = Recorder() - engine = IPCWeightTransferEngine( - config, vllm_config, _get_ray_assigned_device(), recorder - ) + engine = IPCWeightTransferEngine(config, vllm_config, 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 @@ -1211,6 +1214,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises(): torch.device("cuda:0"), MagicMock(spec=torch.nn.Module), ) + # No init handshake here, so the engine keeps its default packed=False. dummy_tensor = torch.ones(10, 10, device="cuda:0") _, ipc_handle = reduce_tensor(dummy_tensor) @@ -1247,7 +1251,7 @@ def update_weights(self, update_info: dict) -> None: self.order.append("update") self.last_update_info = update_info - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: self.order.append("finish") @@ -1264,8 +1268,8 @@ class _DummyTrainerEngine(TrainerWeightTransferEngine): """Minimal concrete trainer engine to exercise base-class + factory.""" @classmethod - def trainer_init(cls, config, init_info, *, client, source): - return cls(config, client=client, source=source) + def trainer_init(cls, init_info, *, client, source): + return cls(client=client, source=source) def send_weights(self): pass @@ -1303,6 +1307,10 @@ def test_ray_client_sends_typed_requests(self, monkeypatch): assert isinstance(update_req, WeightTransferUpdateRequest) assert update_req.update_info == {"names": ["w"]} + client.finish_weight_update("step-42") + handle.finish_weight_update.remote.assert_called_once_with() + handle.update_weight_version.remote.assert_called_once_with("step-42") + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" captured = {} @@ -1334,6 +1342,9 @@ def fake_post(self, path, json=None): client.update_weights(update_info) assert captured["json"]["update_info"] == update_info + client.finish_weight_update("step-42") + assert captured["json"] == {"weight_version": "step-42"} + class TestModuleSource: """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" @@ -1362,18 +1373,17 @@ def test_source_is_reiterable(self): class TestTrainerFactory: """WeightTransferTrainerFactory registry mechanics.""" - def test_builtin_registry_has_no_trainer_backends_yet(self): - # Concrete backends register in the per-backend migration PRs. - assert WeightTransferTrainerFactory._registry == {} + def test_registry_has_ipc(self): + # IPC is the first backend migrated to the trainer engine; NCCL / + # sparse NCCL register in later PRs. + assert "ipc" in WeightTransferTrainerFactory._registry def test_register_and_dispatch(self): saved = dict(WeightTransferTrainerFactory._registry) try: WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine) engine = WeightTransferTrainerFactory.trainer_init( - "dummy", - WeightTransferConfig(backend="dummy"), - MagicMock(), + MagicMock(backend="dummy"), # backend read from the init info client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) @@ -1388,20 +1398,26 @@ def test_register_and_dispatch(self): def test_unknown_backend_raises(self): with pytest.raises(ValueError, match="Invalid weight transfer backend"): WeightTransferTrainerFactory.trainer_init( - "nope", - WeightTransferConfig(backend="nope"), - MagicMock(), + MagicMock(backend="nope"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) + def test_ipc_init_info_declares_backend(self): + assert IPCTrainerInitInfo.backend == "ipc" + + def test_trainer_init_info_subclass_must_set_backend(self): + with pytest.raises(TypeError, match="class-level `backend`"): + + class _NoBackend(TrainerInitInfo): + pass + class TestTrainerEngineBase: """Base-class construction (no GPU).""" def test_source_stored_and_sender_by_default(self): engine = _DummyTrainerEngine( - WeightTransferConfig(backend="nccl"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) @@ -1410,10 +1426,52 @@ def test_source_stored_and_sender_by_default(self): def test_shutdown_default_is_noop(self): engine = _DummyTrainerEngine( - WeightTransferConfig(backend="nccl"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), is_sender=False, ) assert engine.is_sender is False engine.shutdown() # must not raise + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 1, + reason="Need at least 1 GPU (CUDA IPC handles).", +) +def test_ipc_trainer_send_weights_drives_client_in_order(): + """send_weights issues start -> update -> finish and ships per-round metadata; + the packed wire param rides the init info, not the per-round update_info.""" + client = RecordingClient() + engine = IPCTrainerWeightTransferEngine( + client=client, + source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))), + packed=False, + ) + + engine.send_weights() + + assert client.order == ["start", "update", "finish"] + assert client.last_update_info is not None + assert client.last_update_info["names"] == ["w"] + assert client.last_update_info["shapes"] == [[4]] + assert "packed" not in client.last_update_info + + +def test_ipc_trainer_init_ships_packed_to_worker(): + """trainer_init drives the inference-side init handshake and propagates the + must-agree `packed` flag to the worker.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU (CUDA IPC handles).") + + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=True), # backend from init info + client=client, + source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))), + ) + + assert isinstance(engine, IPCTrainerWeightTransferEngine) + assert engine.is_sender is True + assert engine.packed is True + assert client.order == ["init"] + assert client.last_init_info == {"packed": True} diff --git a/tests/engine/test_short_mm_context.py b/tests/engine/test_short_mm_context.py index 23489c213332..940709c8e53f 100644 --- a/tests/engine/test_short_mm_context.py +++ b/tests/engine/test_short_mm_context.py @@ -3,6 +3,8 @@ import pytest +from vllm.exceptions import VLLMValidationError + from ..conftest import IMAGE_ASSETS HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( @@ -19,7 +21,9 @@ def test_context_length_too_short(vllm_runner, image_assets, model): images = [asset.pil_image for asset in image_assets] - with pytest.raises(ValueError, match="longer than the maximum model length"): + with pytest.raises( + VLLMValidationError, match="longer than the maximum model length" + ): vllm_model = vllm_runner( model, # LLaVA has a feature size of 576 diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 94f5bb6068ef..db792f8c7327 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -13,10 +13,16 @@ """ import json +from argparse import Namespace +from http import HTTPStatus from unittest.mock import MagicMock import pytest +from fastapi import FastAPI +from fastapi.exceptions import RequestValidationError +from fastapi.testclient import TestClient +from vllm.entrypoints.anthropic.api_router import attach_router from vllm.entrypoints.anthropic.protocol import ( AnthropicMessagesRequest, ) @@ -38,6 +44,7 @@ PromptTokenUsageInfo, UsageInfo, ) +from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler _convert = AnthropicServingMessages._convert_anthropic_to_openai_request _img_url = AnthropicServingMessages._convert_image_source_to_url @@ -1385,3 +1392,67 @@ def test_empty_completion_emits_one_text_block(self): assert len(result.content) == 1 assert result.content[0].type == "text" assert result.content[0].text == "" + + +# ====================================================================== +# cache_salt pass-through (Issue #46688) +# ====================================================================== + + +class TestCacheSalt: + def test_cache_salt_passed_through(self): + """cache_salt on the Anthropic request reaches the converted + ChatCompletionRequest so prefix-cache isolation works via /v1/messages.""" + request = _make_request( + [{"role": "user", "content": "Hello"}], + cache_salt="tenant-abc-secret-salt", + ) + result = _convert(request) + assert result.cache_salt == "tenant-abc-secret-salt" + + def test_cache_salt_defaults_to_none(self): + """Omitting cache_salt leaves it unset (unchanged default behavior).""" + request = _make_request([{"role": "user", "content": "Hello"}]) + result = _convert(request) + assert result.cache_salt is None + + @staticmethod + def _make_api_app(): + app = FastAPI() + attach_router(app) + app.state.args = Namespace(log_error_stack=False) + app.exception_handler(RequestValidationError)(validation_exception_handler) + + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = AssertionError( + "invalid requests must not reach the serving handler" + ) + app.state.anthropic_serving_messages = handler + return app, handler + + def test_cache_salt_openapi_requires_non_empty_string(self): + app, _ = self._make_api_app() + field_schema = app.openapi()["components"]["schemas"][ + "AnthropicMessagesRequest" + ]["properties"]["cache_salt"] + string_schema = next( + option for option in field_schema["anyOf"] if option.get("type") == "string" + ) + + assert string_schema["minLength"] == 1 + + def test_empty_cache_salt_returns_bad_request(self): + app, handler = self._make_api_app() + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + "/v1/messages", + json={ + "model": "test-model", + "max_tokens": 1, + "messages": [{"role": "user", "content": "Hello"}], + "cache_salt": "", + }, + ) + + assert response.status_code == HTTPStatus.BAD_REQUEST + handler.create_messages.assert_not_awaited() diff --git a/tests/entrypoints/cohere/__init__.py b/tests/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py new file mode 100644 index 000000000000..8380faa29ebf --- /dev/null +++ b/tests/entrypoints/cohere/test_api_router.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/api_router.py``. + +Covers: + +* The optional-import guard: ``attach_router`` is a no-op when the + ``cohere`` SDK isn't installed. +* The env-var opt-in gate: ``attach_router`` is a no-op unless + ``VLLM_ENABLE_COHERE_API=1`` is set. +* The router wiring: response shapes (JSON + SSE), error translation, + and the ``cohere_serving_chat_v2 is None`` fallback (501 Not + Implemented). +""" + +import json +from argparse import Namespace +from collections.abc import AsyncGenerator +from http import HTTPStatus + +import pytest +from fastapi import FastAPI, HTTPException +from fastapi.exceptions import RequestValidationError +from fastapi.testclient import TestClient + +from vllm.entrypoints.cohere import api_router as api_router_mod +from vllm.entrypoints.cohere.api_router import attach_router +from vllm.entrypoints.cohere.protocol import ( + AssistantMessageResponse, + CohereChatV2Response, +) +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.entrypoints.serve.utils.server_utils import ( + http_exception_handler, + validation_exception_handler, +) + + +@pytest.fixture(autouse=True) +def _enable_cohere_api(monkeypatch): + """Auto-enable the Cohere API gate for every test in this module. + + The endpoint is opt-in in production (``VLLM_ENABLE_COHERE_API=1``); + every test in this file exercises the enabled path *except* the + dedicated gate test in :class:`TestEnvVarGate`, which unsets the + flag inside the test body. + """ + monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1") + + +# ---------------------------------------------------------------------- +# Fakes +# ---------------------------------------------------------------------- + + +class _Handler: + """Minimal stand-in for :class:`CohereServingChatV2` used by the + router. Each test sets ``self.result`` to either: + + * a :class:`CohereChatV2Response` (non-streaming JSON path); + * an async generator yielding SSE frames (streaming path); + * an :class:`ErrorResponse` (error envelope path); or + * an exception (router-level 500 path). + """ + + def __init__(self, result): + self.result = result + + async def create_chat_v2(self, request, raw_request): + if isinstance(self.result, Exception): + raise self.result + return self.result + + +def _build_app(handler: _Handler | None) -> FastAPI: + app = FastAPI() + attach_router(app) + app.state.cohere_serving_chat_v2 = handler + return app + + +def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI: + """Build a FastAPI app that mirrors the real vLLM setup by installing + ``validation_exception_handler`` and ``http_exception_handler``. The + :class:`CohereErrorEnvelopeMiddleware` registered by ``attach_router`` + is expected to translate any resulting vLLM ``ErrorResponse`` body + into the ``CohereError`` wire shape. + """ + app = FastAPI() + attach_router(app) + app.state.cohere_serving_chat_v2 = handler + # ``validation_exception_handler`` reads ``req.app.state.args``; the + # real cli builds this via argparse. + app.state.args = Namespace(log_error_stack=False) + app.exception_handler(RequestValidationError)(validation_exception_handler) + app.exception_handler(HTTPException)(http_exception_handler) + return app + + +def _minimal_request_body() -> dict: + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + } + + +# ---------------------------------------------------------------------- +# Optional-import guard +# ---------------------------------------------------------------------- + + +class TestOptionalCohereImport: + """``attach_router`` probes for the SDK once at module load (because + the route handler uses types imported from ``cohere``) and stashes + the result in ``_SDK_AVAILABLE``. Tests simulate the "SDK missing" + state by flipping that flag for the duration of the test. + + ``attach_router`` checks the env-var gate before the SDK probe, so + the SDK-missing branch is only reachable when the operator opts in + via ``VLLM_ENABLE_COHERE_API=1``. The flag-off-and-SDK-missing case + below exists to pin down that ordering — the flag-off short-circuits + """ + + def test_flag_off_and_sdk_missing_stays_silent_about_sdk(self, monkeypatch, caplog): + """Flag off doesn't do SDK-missing check. + + When the operator hasn't opted in, ``attach_router`` must not + warn about the ``cohere`` SDK being missing: they never asked + for the endpoint, so surfacing the SDK gap is misleading noise. + Only the flag-off DEBUG message should fire. + """ + monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False) + monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False) + + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"): + app = FastAPI() + attach_router(app) + + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" not in paths + # The flag-off short-circuit ran; the SDK check never did. + assert not any( + "SDK is not installed" in rec.message for rec in caplog.records + ), "SDK-missing log leaked despite the flag being off" + + def test_flag_on_but_sdk_missing_logs_warning(self, monkeypatch, caplog): + """Misconfiguration path: the operator explicitly opted into the + endpoint via ``VLLM_ENABLE_COHERE_API=1`` (already set by the + autouse fixture) but forgot to install ``cohere``. + """ + monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False) + + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"): + app = FastAPI() + attach_router(app) + + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" not in paths + warn_sdk_records = [ + rec + for rec in caplog.records + if "VLLM_ENABLE_COHERE_API=1" in rec.message + and "SDK is not installed" in rec.message + ] + assert warn_sdk_records, ( + "expected a WARNING that pairs the opt-in flag with the " + "missing SDK so operators notice the misconfiguration" + ) + assert all(rec.levelname == "WARNING" for rec in warn_sdk_records) + + def test_attach_router_registers_route_when_cohere_present(self): + app = _build_app(handler=None) + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" in paths + + +# ---------------------------------------------------------------------- +# VLLM_ENABLE_COHERE_API gate +# ---------------------------------------------------------------------- + + +class TestEnvVarGate: + """The Cohere v2 endpoint is opt-in via ``VLLM_ENABLE_COHERE_API``. + + Even with the SDK installed, :func:`attach_router` must skip route + registration and middleware installation unless the env flag is + set. The autouse fixture on this module enables the flag by + default, so each test here explicitly disables it. + """ + + def test_attach_router_noop_when_flag_unset(self, monkeypatch, caplog): + monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False) + + # The flag-off skip logs at DEBUG on purpose: this is the default + # state for every non-Cohere vLLM deployment, so an INFO log on + # every server startup would be pointless noise. The test raises + # caplog's level accordingly. + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"): + app = FastAPI() + attach_router(app) + + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" not in paths + debug_flag_records = [ + rec + for rec in caplog.records + if "VLLM_ENABLE_COHERE_API is not set" in rec.message + ] + assert debug_flag_records, ( + "expected a DEBUG message that the cohere flag is off" + ) + assert all(rec.levelname == "DEBUG" for rec in debug_flag_records) + + def test_attach_router_noop_when_flag_zero(self, monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "0") + + app = FastAPI() + attach_router(app) + + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" not in paths + + +# ---------------------------------------------------------------------- +# Endpoint behavior +# ---------------------------------------------------------------------- + + +class TestEndpoint: + def test_501_when_handler_missing(self): + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.NOT_IMPLEMENTED + body = r.json() + assert "does not support" in body["message"] + assert "id" not in body # excluded by ``exclude_none=True`` + + def test_non_streaming_response_is_json(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hello"}]) + result = CohereChatV2Response(id="r1", finish_reason="COMPLETE", message=msg) + app = _build_app(handler=_Handler(result)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("application/json") + body = r.json() + assert body["id"] == "r1" + assert body["finish_reason"] == "COMPLETE" + assert body["message"]["content"][0]["text"] == "hello" + + def test_streaming_response_is_sse(self): + async def _gen() -> AsyncGenerator[str, None]: + yield 'data: {"type":"message-start"}\n\n' + yield "data: [DONE]\n\n" + + app = _build_app(handler=_Handler(_gen())) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={**_minimal_request_body(), "stream": True}, + ) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("text/event-stream") + body = r.text + assert "message-start" in body + assert body.rstrip().endswith("[DONE]") + + def test_error_response_translated_to_cohere_envelope(self): + err = ErrorResponse( + error=ErrorInfo( + message="bad request", + type="bad_request", + code=400, + ) + ) + app = _build_app(handler=_Handler(err)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + assert body == {"message": "bad request"} + + def test_handler_exception_returns_500_envelope(self): + app = _build_app(handler=_Handler(RuntimeError("kaboom"))) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + body = r.json() + assert body == {"message": "kaboom"} + + def test_non_json_content_type_rejected(self): + """The ``validate_json_request`` dependency raises + ``RequestValidationError`` (HTTP 422) for non-JSON content + types, matching the behavior of the other vLLM API routers. + """ + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + content=json.dumps(_minimal_request_body()), + headers={"content-type": "text/plain"}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + + def test_invalid_body_returns_422(self): + # ``model`` is required; omit it to trip Pydantic validation. + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + + +# ---------------------------------------------------------------------- +# CohereErrorEnvelopeMiddleware +# ---------------------------------------------------------------------- + + +class TestCohereErrorEnvelope: + """When the app installs vLLM's global exception handlers, validation + and HTTP errors escape as ``ErrorResponse`` bodies before the route + handler runs. The middleware installed by ``attach_router`` must + normalise those bodies to the ``CohereError`` shape declared on the + endpoint's OpenAPI ``responses`` map so schema-conformance tests + (``test_openai_schema.py``) don't see a mismatch on ``/cohere/*`` + responses. + """ + + def test_validation_error_body_is_cohere_shaped(self): + # ``model=""`` and ``messages=[]`` trip our custom field + # validators, which raise pydantic ValueErrors and are routed + # through ``validation_exception_handler`` in the real vLLM + # server (producing the ``{"error": {...}}`` shape). + app = _build_app_with_vllm_handlers(handler=None) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json={"messages": [], "model": ""}) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + # ``CohereError`` has ``message`` at the top level, not nested + # under an ``error`` envelope. + assert "error" not in body + assert "message" in body + assert isinstance(body["message"], str) and body["message"] + + def test_http_error_body_is_cohere_shaped(self): + # A raised ``HTTPException`` from anywhere in the request cycle + # is routed through ``http_exception_handler`` (producing the + # ``ErrorResponse`` shape) and must be translated. + app = _build_app_with_vllm_handlers(handler=None) + + @app.get("/cohere/v2/boom") + async def _boom(): + raise HTTPException(status_code=418, detail="teapot") + + with TestClient(app) as client: + r = client.get("/cohere/v2/boom") + assert r.status_code == 418 + body = r.json() + assert body == {"message": "teapot"} + + def test_non_cohere_path_is_not_translated(self): + app = _build_app_with_vllm_handlers(handler=None) + + @app.get("/v1/other") + async def _other(): + raise HTTPException(status_code=400, detail="nope") + + with TestClient(app) as client: + r = client.get("/v1/other") + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + # Non-cohere paths keep the vLLM ``ErrorResponse`` shape. + assert "error" in body + assert body["error"]["message"] == "nope" + + def test_streaming_response_passes_through(self): + # SSE responses have content-type text/event-stream; the + # middleware must never buffer these (which would break + # streaming) even though they're on ``/cohere/*``. + async def _gen() -> AsyncGenerator[str, None]: + yield 'data: {"type":"message-start"}\n\n' + yield "data: [DONE]\n\n" + + app = _build_app_with_vllm_handlers(handler=_Handler(_gen())) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={**_minimal_request_body(), "stream": True}, + ) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("text/event-stream") + assert "message-start" in r.text + assert r.text.rstrip().endswith("[DONE]") + + def test_already_cohere_shaped_body_passes_through(self): + # When the handler returns an ``ErrorResponse`` the route + # itself translates it to ``CohereError``; the middleware sees + # the ``CohereError`` shape and must leave it alone. + err = ErrorResponse( + error=ErrorInfo(message="already cohere", type="Bad Request", code=400) + ) + app = _build_app_with_vllm_handlers(handler=_Handler(err)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + # No ``error`` wrapper: the route already emitted the wire shape. + assert body == {"message": "already cohere"} + + def test_request_id_preserved_in_translated_body(self): + # Client-provided ``X-Request-Id`` should be echoed as + # ``CohereError.id`` so callers can correlate failures. + app = _build_app_with_vllm_handlers(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={"messages": [], "model": ""}, + headers={"X-Request-Id": "req-abc"}, + ) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + assert body.get("id") == "req-abc" diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py new file mode 100644 index 000000000000..9e0ec90a746f --- /dev/null +++ b/tests/entrypoints/cohere/test_chat_v2.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end integration tests for the ``POST /cohere/v2/chat`` endpoint. + +These tests spin up a real ``vllm serve`` process via +:class:`tests.utils.RemoteOpenAIServer` and exercise the Cohere Chat v2 +contract over HTTP. They mirror the pattern used by +:mod:`tests.entrypoints.anthropic.test_messages`. + +Two layers of integration are covered: + +1. Raw HTTP via :mod:`httpx` — always runs, verifies the wire contract. +2. Cohere SDK (``pip install cohere``) — auto-skipped when the optional + dependency isn't installed, verifies SDK-level interop. + +We use a tiny generic chat model (``HuggingFaceTB/SmolLM2-135M-Instruct``, +≈135M params) so the fixture stays runnable on CPU-only laptops. The +Cohere v2 endpoint is model-agnostic and just performs the v2 ↔ OpenAI +chat-completion translation, so the choice of model only matters for +response *shape* — the translation logic itself is unit-tested +elsewhere. +""" + +import json + +import httpx +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer + +# Tiny chat-tuned model so the fixture is cheap to boot on CPU-only +# hosts (Mac, CI runners without GPUs). The Cohere v2 translation is +# completely independent of which underlying chat model is loaded. +MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct" +SERVED_MODEL_NAME = "command-r-plus-08-2024" + + +@pytest.fixture(scope="module") +def server(): + args = [ + # 1k tokens is plenty for the prompts these tests send; trimming + # it from the default keeps the KV-cache footprint tiny. + "--max-model-len", + "1024", + "--max-num-seqs", + "4", + "--dtype", + "bfloat16", + "--enforce-eager", + # Advertise a Cohere model name so Cohere SDK ``model=`` calls + # round-trip without the ``model not found`` check. + "--served-model-name", + SERVED_MODEL_NAME, + # Disable the reasoning-model path so the test doesn't require + # the conversation to surface a thinking block (SmolLM2 doesn't + # emit Cohere-style reasoning tokens out of the box). + "--no-cohere-is-reasoning-model", + ] + + # Cap the CPU KV-cache pool the vLLM CPU backend reserves at + # startup. The default (4 GB) trips Mac's RAM check on smaller + # machines; 1 GB is more than enough for max_model_len=1024 and + # max_num_seqs=4. + # ``VLLM_ENABLE_COHERE_API=1`` flips the opt-in gate that + # ``vllm.entrypoints.cohere.api_router.attach_router`` reads to + # decide whether to register ``POST /cohere/v2/chat`` at server + # startup. Without it the fixture would boot a server that returns + # 404 for every request the tests below make. + env_dict = { + "VLLM_CPU_KVCACHE_SPACE": "1", + "VLLM_ENABLE_COHERE_API": "1", + } + + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server: + yield remote_server + + +# ---------------------------------------------------------------------- +# Layer 1: raw HTTP contract (no optional deps) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def httpx_client(server): + async with httpx.AsyncClient( + base_url=server.url_root, timeout=httpx.Timeout(120.0) + ) as client: + yield client + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_non_streaming(httpx_client: httpx.AsyncClient): + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Say hi."}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + payload = resp.json() + # The response envelope follows Cohere v2's schema. + assert "message" in payload + assert payload["message"]["role"] == "assistant" + # ``content`` is a list of content blocks; the synthesized + # ``CohereServingChatV2`` should always emit at least one ``text`` + # block (it falls back to an empty block when the model returned + # nothing). + content = payload["message"]["content"] + assert isinstance(content, list) and len(content) >= 1 + assert content[0]["type"] == "text" + # ``usage`` is always populated by the translator. + assert "usage" in payload + assert "billed_units" in payload["usage"] + assert "tokens" in payload["usage"] + # ``finish_reason`` is one of Cohere's enum values. + assert payload["finish_reason"] in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_streaming(httpx_client: httpx.AsyncClient): + """Streaming returns SSE frames in the v2 message-lifecycle shape.""" + events: list[dict] = [] + async with httpx_client.stream( + "POST", + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 8, + "stream": True, + }, + ) as resp: + assert resp.status_code == 200, await resp.aread() + assert resp.headers["content-type"].startswith("text/event-stream") + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[len("data: ") :] + if data == "[DONE]": + events.append({"type": "_DONE_"}) + continue + events.append(json.loads(data)) + + types = [ev["type"] for ev in events] + # The lifecycle always starts with message-start and ends with [DONE] + # preceded by message-end. + assert types[0] == "message-start" + assert types[-1] == "_DONE_" + assert types[-2] == "message-end" + # message-start must carry the chunk/message id. + assert events[0].get("id") + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_validation_error_returns_400( + httpx_client: httpx.AsyncClient, +): + # Missing required ``model`` field → FastAPI/Pydantic returns 400. + resp = await httpx_client.post( + "/cohere/v2/chat", + json={"messages": []}, + ) + assert resp.status_code == 400 + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_documents_field_accepted( + httpx_client: httpx.AsyncClient, +): + """The v2 endpoint forwards ``documents`` into chat_template_kwargs. + + We only assert the request is accepted and produces a 200 response — + the renderer-level effect is covered by ``tests/renderers/test_cohere.py``. + """ + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Summarize."}], + "documents": [{"id": "d1", "data": {"title": "T", "snippet": "S"}}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + + +# ---------------------------------------------------------------------- +# Layer 2: Cohere SDK round-trip (auto-skipped if SDK absent) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def cohere_async_client(server): + cohere = pytest.importorskip("cohere") + # The vLLM endpoint is mounted at ``/cohere/v2/chat`` while the + # cohere SDK targets ``${base_url}/v2/chat``; point base_url at the + # ``/cohere`` prefix so paths line up. + client = cohere.AsyncClientV2( + api_key="dummy", + base_url=server.url_for("cohere"), + ) + try: + yield client + finally: + # ``AsyncClientV2`` exposes a sync close; if a future version + # adds aclose we still close cleanly. + close = getattr(client, "aclose", None) or getattr(client, "close", None) + if close is not None: + result = close() + if hasattr(result, "__await__"): + await result + + +@pytest.mark.asyncio +async def test_cohere_sdk_non_streaming(cohere_async_client): + resp = await cohere_async_client.chat( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Say hi."}], + max_tokens=16, + ) + # SDK parses our JSON into typed objects. + assert resp.message.role == "assistant" + assert resp.message.content is not None + assert len(resp.message.content) >= 1 + assert resp.message.content[0].type == "text" + assert resp.finish_reason in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_sdk_streaming(cohere_async_client): + events: list[str] = [] + stream = cohere_async_client.chat_stream( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Hi"}], + max_tokens=8, + ) + async for ev in stream: + events.append(ev.type) + + assert events, "SDK stream yielded no events" + assert events[0] == "message-start" + assert events[-1] == "message-end" diff --git a/tests/entrypoints/cohere/test_cohere_chat_message.py b/tests/entrypoints/cohere/test_cohere_chat_message.py new file mode 100644 index 000000000000..8c297bf19ae9 --- /dev/null +++ b/tests/entrypoints/cohere/test_cohere_chat_message.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for :mod:`vllm.entrypoints.cohere.cohere_chat_message`. + +The module hosts the citation-carrying subclasses of the OpenAI chat +completion protocol (``CohereChatMessage`` / ``CohereDeltaMessage``) plus +the shared ``Citation`` / ``CitationSource`` shapes. These tests pin: + +1. Serialization behavior we own (empty-``citations`` cleanup on the + subclass ``_serialize``, ``exclude_none`` shape, round-trip). +2. The constraint that ``type`` discriminators reject unknown values. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation, + CitationSource, + CohereChatMessage, + CohereDeltaMessage, +) + +# ====================================================================== +# CitationSource +# ====================================================================== + + +class TestCitationSource: + def test_invalid_type_rejected(self): + with pytest.raises(ValidationError): + CitationSource(type="other") # type: ignore[arg-type] + + def test_none_fields_excluded_from_dump(self): + s = CitationSource(type="document", id="d1") + assert s.model_dump(exclude_none=True) == { + "type": "document", + "id": "d1", + } + + def test_resolved_source_survives_json_round_trip(self): + # The parser produces fully-resolved sources (see + # ``_melody_sources_to_vllm`` in + # ``vllm/reasoning/cohere_command_reasoning_parser.py``) which + # then flow through the internal streaming pipeline (parser -> + # OpenAI stream chunk -> ``_chat_completion_stream_to_v2``). + # Every field must round-trip through ``model_dump_json`` at + # that layer or the wire event ends up missing pieces. + src = CitationSource( + type="tool", + id="res_a0", + tool_output={"id": "res_a0", "text": "r"}, + ) + src2 = CitationSource.model_validate_json(src.model_dump_json()) + assert src2.type == "tool" + assert src2.id == "res_a0" + assert src2.tool_output == {"id": "res_a0", "text": "r"} + + +# ====================================================================== +# Citation +# ====================================================================== + + +class TestCitation: + def test_invalid_type_rejected(self): + with pytest.raises(ValidationError): + Citation(type="OTHER") # type: ignore[arg-type] + + def test_dump_excludes_none_fields(self): + c = Citation(start=0, end=5, text="hello") + dumped = c.model_dump(exclude_none=True) + assert dumped == { + "start": 0, + "end": 5, + "text": "hello", + "sources": [], + } + + +# ====================================================================== +# CohereDeltaMessage +# ====================================================================== + + +class TestCohereDeltaMessage: + def test_citations_omitted_from_dump_when_none(self): + # Non-grounded deltas should never carry ``citations: null`` on + # the wire even when routed through the Cohere subclass. + d = CohereDeltaMessage(content="hello") + dumped = d.model_dump(exclude_none=True) + assert "citations" not in dumped + + def test_citations_dump_round_trip(self): + d = CohereDeltaMessage(citations=[Citation(start=0, end=5, text="hi")]) + dumped = d.model_dump(exclude_none=True) + assert dumped["citations"][0]["text"] == "hi" + # Round-trip through the subclass to confirm the field validates. + d2 = CohereDeltaMessage.model_validate(dumped) + assert d2.citations[0].text == "hi" + + +# ====================================================================== +# CohereChatMessage +# ====================================================================== + + +class TestCohereChatMessage: + def test_citations_omitted_from_dump_when_none(self): + m = CohereChatMessage(role="assistant", content="hi") + dumped = m.model_dump(exclude_none=True) + assert "citations" not in dumped + + def test_citations_dump_round_trip(self): + m = CohereChatMessage( + role="assistant", + content="hello", + citations=[ + Citation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="document", id="d1")], + ) + ], + ) + dumped = m.model_dump(exclude_none=True) + assert dumped["role"] == "assistant" + assert dumped["content"] == "hello" + assert dumped["citations"][0]["text"] == "hello" + assert dumped["citations"][0]["sources"][0]["id"] == "d1" + # Round-trip through the subclass to confirm the field validates. + m2 = CohereChatMessage.model_validate(dumped) + assert m2.citations[0].text == "hello" + assert m2.citations[0].sources[0].id == "d1" diff --git a/tests/entrypoints/cohere/test_protocol.py b/tests/entrypoints/cohere/test_protocol.py new file mode 100644 index 000000000000..7d9d811047a5 --- /dev/null +++ b/tests/entrypoints/cohere/test_protocol.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/protocol.py``. + +The module is mostly a thin wrapper around the official ``cohere`` SDK +types plus a few local additions: + +* :class:`CohereError` envelope. +* :class:`CohereChatV2Request` (model required, ``max_tokens`` non-negative). +* :class:`CohereChatV2Response` plus the usage / logprob helpers. +* The streaming event subclasses that bake a wire-format ``type`` + discriminator into ``model_dump()`` so SSE consumers can demux on it. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.cohere.protocol import ( + AssistantMessageResponse, + CitationEndEvent, + CitationStartEvent, + CohereChatV2Request, + CohereChatV2Response, + CohereError, + CohereUsageTokens, + ContentDeltaEvent, + ContentEndEvent, + ContentStartEvent, + MessageEndEvent, + MessageStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolPlanDeltaEvent, +) + +# ====================================================================== +# CohereError +# ====================================================================== + + +class TestCohereError: + def test_model_dump_excludes_none(self): + err = CohereError(message="boom") + assert err.model_dump(exclude_none=True) == {"message": "boom"} + + +# ====================================================================== +# CohereChatV2Request +# ====================================================================== + + +class TestCohereChatV2Request: + def test_empty_model_rejected(self): + with pytest.raises(ValidationError, match="model is required"): + CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}]) + + def test_negative_max_tokens_rejected(self): + with pytest.raises(ValidationError, match="non-negative"): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=-1, + ) + + def test_zero_max_tokens_allowed(self): + # Zero is allowed (the docs allow 0 -> return prompt only). + req = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=0, + ) + assert req.max_tokens == 0 + + def test_invalid_tool_choice_rejected(self): + with pytest.raises(ValidationError): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + tool_choice="ANY", # not REQUIRED/NONE + ) + + +# ====================================================================== +# Usage / Logprobs +# ====================================================================== + + +class TestUsage: + def test_tokens_serialization(self): + u = CohereUsageTokens(input_tokens=10, output_tokens=5) + assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0} + + +# ====================================================================== +# CohereChatV2Response +# ====================================================================== + + +class TestCohereChatV2Response: + def test_invalid_finish_reason_rejected(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}]) + with pytest.raises(ValidationError): + CohereChatV2Response( + id="r1", + finish_reason="NOT_A_REASON", # type: ignore[arg-type] + message=msg, + ) + + +# ====================================================================== +# Streaming event ``type`` discriminator baked into model_dump() +# ====================================================================== + + +class TestStreamingEventTypeField: + """Each event subclass adds a ``type: Literal[...]`` field with a + default so ``model_dump()`` always emits the wire-format discriminator + (the parent SDK classes don't declare ``type`` as a Pydantic field). + """ + + @pytest.mark.parametrize( + "cls, expected_type, kwargs", + [ + ( + MessageStartEvent, + "message-start", + {"id": "a", "delta": {"message": {"role": "assistant"}}}, + ), + ( + ContentStartEvent, + "content-start", + { + "index": 0, + "delta": {"message": {"content": {"type": "text", "text": ""}}}, + }, + ), + ( + ContentDeltaEvent, + "content-delta", + { + "index": 0, + "delta": {"message": {"content": {"text": "hi"}}}, + }, + ), + (ContentEndEvent, "content-end", {"index": 0}), + ( + ToolPlanDeltaEvent, + "tool-plan-delta", + {"delta": {"message": {"tool_plan": "thinking"}}}, + ), + ( + ToolCallStartEvent, + "tool-call-start", + { + "index": 0, + "delta": { + "message": { + "tool_calls": { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": ""}, + } + } + }, + }, + ), + ( + ToolCallDeltaEvent, + "tool-call-delta", + { + "index": 0, + "delta": { + "message": {"tool_calls": {"function": {"arguments": "{}"}}} + }, + }, + ), + (ToolCallEndEvent, "tool-call-end", {"index": 0}), + ( + CitationStartEvent, + "citation-start", + { + "index": 0, + "delta": { + "message": { + "citations": {"start": 0, "end": 5, "text": "hello"} + } + }, + }, + ), + (CitationEndEvent, "citation-end", {"index": 0}), + ( + MessageEndEvent, + "message-end", + {"id": "a", "delta": {"finish_reason": "COMPLETE"}}, + ), + ], + ) + def test_type_field_default(self, cls, expected_type, kwargs): + ev = cls(**kwargs) + # type field is auto-populated from the Literal default. + assert ev.type == expected_type + # The discriminator must be present in the serialized payload so + # clients reading the stream can demux on it. + dumped = ev.model_dump(exclude_none=True) + assert dumped["type"] == expected_type + # Same in JSON form (what ``_emit`` serializes). + assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True) + + def test_type_field_cannot_be_overridden_to_wrong_value(self): + # Literal types reject any value other than the bake-in default. + with pytest.raises(ValidationError): + MessageStartEvent( + id="a", + delta={"message": {"role": "assistant"}}, + type="other", # type: ignore[arg-type] + ) diff --git a/tests/entrypoints/cohere/test_registry_and_args.py b/tests/entrypoints/cohere/test_registry_and_args.py new file mode 100644 index 000000000000..5e907b0d9340 --- /dev/null +++ b/tests/entrypoints/cohere/test_registry_and_args.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sanity tests for the small registry / CLI-arg / config additions made +by the Cohere v2 chat API: + +* ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry. +* ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry + (aliased to the cached HF tokenizer). +* ``vllm/entrypoints/openai/cli_args.py``: new + ``cohere_is_reasoning_model`` field. +* ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode`` + Literal. +""" + +import dataclasses +import typing + +import pytest + +from vllm.config.model import TokenizerMode +from vllm.entrypoints.openai.cli_args import BaseFrontendArgs, make_arg_parser +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import TokenizerRegistry +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class TestRendererRegistry: + def test_cohere_renderer_registered(self): + # The registry resolves to the importable ``CohereRenderer`` class. + cls = RENDERER_REGISTRY.load_renderer_cls("cohere") + assert cls.__name__ == "CohereRenderer" + # Sanity: class lives in the cohere renderer module. + assert cls.__module__ == "vllm.renderers.cohere" + + +class TestTokenizerRegistry: + def test_cohere_aliased_to_cached_hf_tokenizer(self): + # ``cohere`` mode uses the standard HF tokenizer; only the + # renderer stage is replaced. This test guards against accidental + # divergence. + cls = TokenizerRegistry.load_tokenizer_cls("cohere") + assert cls.__name__ == "CachedHfTokenizer" + + +class TestTokenizerModeLiteral: + def test_cohere_is_a_valid_tokenizer_mode(self): + # The Literal must enumerate ``"cohere"`` so engine arg parsing + # accepts ``--tokenizer-mode cohere``. + modes = typing.get_args(TokenizerMode) + assert "cohere" in modes + + +# ---------------------------------------------------------------------- +# ``--cohere-is-reasoning-model`` CLI flag +# ---------------------------------------------------------------------- + + +class TestCohereCliArg: + """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end + through ``make_arg_parser`` — mirrors the pattern used in + :mod:`tests.entrypoints.openai.test_cli_args`. + """ + + def test_default_value_on_dataclass_is_true(self): + fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)} + assert "cohere_is_reasoning_model" in fields + field = fields["cohere_is_reasoning_model"] + assert field.default is True + assert field.type is bool + + @pytest.fixture + def serve_parser(self) -> FlexibleArgumentParser: + parser = FlexibleArgumentParser() + return make_arg_parser(parser) + + def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser): + # No flag supplied → dataclass default (True) wins. + args = serve_parser.parse_args(["--model", "m"]) + assert args.cohere_is_reasoning_model is True + + def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser): + # Boolean dataclass fields are wired up as ``--flag value`` / + # ``--no-flag`` pairs by FlexibleArgumentParser. + args = serve_parser.parse_args( + ["--model", "m", "--no-cohere-is-reasoning-model"] + ) + assert args.cohere_is_reasoning_model is False + + def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser): + args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"]) + assert args.cohere_is_reasoning_model is True diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py new file mode 100644 index 000000000000..6ee6ffef289e --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -0,0 +1,2084 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 -> OpenAI request / response conversion +implemented in ``vllm/entrypoints/cohere/serving.py``. + +These cover the pure-Python classmethods so we don't need an engine. +For the instance methods that read ``self._is_reasoning_model`` we +build a lightweight :class:`_FakeServing` subclass that skips the +heavy ``OpenAIServingChat.__init__`` chain (which would otherwise need +a real engine client, model registry, etc.) — the same pattern used in +``test_serving_streaming.py``. +""" + +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.cohere.cohere_chat_message import ( + CitationSource, +) +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, +) +from vllm.entrypoints.cohere.serving import ( + _FINISH_REASON_MAP, + CohereServingChatV2, + _map_finish_reason, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +def _make_request(**kwargs) -> CohereChatV2Request: + kwargs.setdefault("model", "m") + kwargs.setdefault("messages", [{"role": "user", "content": "hi"}]) + return CohereChatV2Request(**kwargs) + + +def _convert(request: CohereChatV2Request) -> ChatCompletionRequest: + return CohereServingChatV2._convert_v2_to_chat_completion(request) + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain. + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_to_v2`` and friends); the rest of the parent + state is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _build_chat_completion_response( + *, + response_id: str = "resp_1", + content: str | None = "hello", + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = "stop", + citations: list[Any] | None = None, + usage: dict[str, Any] | None = None, + kv_transfer_params: dict[str, Any] | None = None, +) -> ChatCompletionResponse: + message: dict[str, Any] = {"role": "assistant"} + if content is not None: + message["content"] = content + if reasoning is not None: + message["reasoning"] = reasoning + if tool_calls is not None: + message["tool_calls"] = tool_calls + if citations is not None: + message["citations"] = citations + kwargs: dict[str, Any] = dict( + id=response_id, + object="chat.completion", + created=0, + model="m", + choices=[{"index": 0, "message": message, "finish_reason": finish_reason}], + # ``usage`` is a required field on ChatCompletionResponse, but the + # production code defensively handles ``None`` -> no usage block; + # we round-trip that behavior by post-setting the attribute below. + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + if kv_transfer_params is not None: + kwargs["kv_transfer_params"] = kv_transfer_params + resp = ChatCompletionResponse(**kwargs) + if usage is None: + resp.usage = None + else: + # Replace the placeholder with the caller-provided usage. + resp = ChatCompletionResponse.model_validate( + {**resp.model_dump(), "usage": usage} + ) + if kv_transfer_params is not None: + resp.kv_transfer_params = kv_transfer_params + return resp + + +# ====================================================================== +# _map_finish_reason +# ====================================================================== + + +class TestMapFinishReason: + @pytest.mark.parametrize( + "openai, cohere", + [ + ("stop", "COMPLETE"), + ("length", "MAX_TOKENS"), + ("tool_calls", "TOOL_CALL"), + ("stop_sequence", "STOP_SEQUENCE"), + ("error", "ERROR"), + (None, "COMPLETE"), + ], + ) + def test_known_reasons(self, openai, cohere): + assert _map_finish_reason(openai) == cohere + + def test_unknown_reason_defaults_to_complete(self): + assert _map_finish_reason("not_a_real_reason") == "COMPLETE" + + def test_finish_reason_map_is_complete(self): + # Sanity check that the lookup table covers all documented states. + assert set(_FINISH_REASON_MAP) == { + "stop", + "length", + "tool_calls", + "stop_sequence", + "error", + None, + } + + +# ====================================================================== +# _coerce_text_content (system / tool string fallback) +# ====================================================================== + + +class TestCoerceTextContent: + def test_string_passthrough(self): + assert CohereServingChatV2._coerce_text_content("hi") == "hi" + + def test_concatenates_text_blocks(self): + from cohere.types import SystemChatMessageV2 + + sys_msg = SystemChatMessageV2( + content=[ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"}, + ] + ) + assert CohereServingChatV2._coerce_text_content(sys_msg.content) == "ab" + + +# ====================================================================== +# User message conversion +# ====================================================================== + + +class TestConvertUserMessage: + def test_string_content(self): + req = _make_request(messages=[{"role": "user", "content": "hi"}]) + result = _convert(req) + assert result.messages == [{"role": "user", "content": "hi"}] + + def test_text_only_list_flattened_to_string(self): + # Single-text-block list is flattened back to a string for maximum + # downstream-template compatibility. + req = _make_request( + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi"}], + } + ] + ) + result = _convert(req) + assert result.messages[0] == {"role": "user", "content": "hi"} + + def test_image_url_content_with_detail(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ], + } + ] + ) + result = _convert(req) + msg = result.messages[0] + assert msg["role"] == "user" + assert msg["content"] == [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ] + + def test_image_url_without_detail_omits_field(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + } + ], + } + ] + ) + result = _convert(req) + assert result.messages[0]["content"][0]["image_url"] == { + "url": "https://x/i.png" + } + + def test_text_plus_image_keeps_list(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + }, + ], + } + ] + ) + result = _convert(req) + content = result.messages[0]["content"] + assert isinstance(content, list) + assert len(content) == 2 + assert content[0] == {"type": "text", "text": "describe"} + assert content[1]["type"] == "image_url" + + +# ====================================================================== +# Assistant message conversion +# ====================================================================== + + +class TestConvertAssistantMessage: + def test_string_content(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst == {"role": "assistant", "content": "hello"} + + def test_text_and_thinking_blocks(self): + # ``thinking`` blocks collapse back into the ``reasoning`` field on + # the OpenAI message; ``text`` blocks become ``content``. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "let me think"}, + {"type": "text", "text": "Hi!"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["role"] == "assistant" + assert asst["content"] == "Hi!" + assert asst["reasoning"] == "let me think" + + def test_thinking_only(self): + # Thinking-only assistant messages have no ``content`` set. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "ponder"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst.get("reasoning") == "ponder" + assert "content" not in asst + + def test_multiple_thinking_blocks_concatenated(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "first."}, + {"type": "thinking", "thinking": "second."}, + {"type": "text", "text": "done."}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["reasoning"] == "first.second." + assert asst["content"] == "done." + + def test_tool_plan_collapses_into_reasoning(self): + # Cohere's ``tool_plan`` is the older chain-of-thought field; it + # should be appended to ``reasoning`` so the rendered template + # preserves the planning context. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "I'll call a tool."}, + ], + "tool_plan": "plan: use calculator", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["content"] == "I'll call a tool." + assert asst["reasoning"] == "plan: use calculator" + assert asst["tool_calls"][0]["function"] == { + "name": "calc", + "arguments": "{}", + } + + def test_tool_calls_with_missing_function_pieces_get_defaults(self): + # The conversion defends against missing function name/arguments + # by emitting empty string / "{}" defaults so downstream + # validation never sees a None. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + ], + }, + ] + ) + result = _convert(req) + tc = result.messages[1]["tool_calls"][0] + assert tc["id"] == "c1" + assert tc["type"] == "function" + assert tc["function"] == {"name": "", "arguments": "{}"} + + +# ====================================================================== +# Tool message conversion +# ====================================================================== + + +class TestConvertToolMessage: + def _request_with_tool_message(self, content: Any) -> CohereChatV2Request: + return _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": content}, + ] + ) + + def test_string_content(self): + req = self._request_with_tool_message("result text") + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg == { + "role": "tool", + "tool_call_id": "c1", + "content": "result text", + } + + def test_text_only_list_flattened_to_newline_string(self): + # Text-only tool results are flattened to a single newline-joined + # string for compatibility with standard chat templates. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "line 1"}, + {"type": "text", "text": "line 2"}, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg["content"] == "line 1\nline 2" + + def test_with_document_preserves_structured_content(self): + # When documents appear in the tool result, we keep the list shape + # so the Cohere renderer can lift them into grounding sources. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "see attachment"}, + { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + }, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert isinstance(tool_msg["content"], list) + assert tool_msg["content"][0] == {"type": "text", "text": "see attachment"} + assert tool_msg["content"][1] == { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + } + + +# ====================================================================== +# System message +# ====================================================================== + + +class TestSystemMessage: + def test_system_string(self): + req = _make_request( + messages=[ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0] == { + "role": "system", + "content": "be helpful", + } + + def test_system_text_blocks_concatenated(self): + req = _make_request( + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "part1 "}, + {"type": "text", "text": "part2"}, + ], + }, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0]["content"] == "part1 part2" + + +# ====================================================================== +# Base ChatCompletionRequest field mapping +# ====================================================================== + + +class TestBuildBaseChatCompletion: + def test_sampling_and_limits_mapped(self): + req = _make_request( + max_tokens=128, + stop_sequences=["", "STOP"], + temperature=0.5, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + k=50, + p=0.95, + logprobs=True, + priority=2, + kv_transfer_params={"x": 1}, + chat_template_kwargs={"y": 2}, + ) + result = _convert(req) + assert result.model == "m" + # ``max_tokens`` is deprecated in favor of ``max_completion_tokens`` + # but the serving code intentionally sets both for compatibility. + assert result.max_completion_tokens == 128 + assert result.stop == ["", "STOP"] + assert result.temperature == 0.5 + assert result.seed == 42 + assert result.frequency_penalty == 0.1 + assert result.presence_penalty == 0.2 + assert result.top_k == 50 + assert result.top_p == 0.95 + assert result.logprobs is True + assert result.priority == 2 + assert result.kv_transfer_params == {"x": 1} + # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the + # base build at least preserves what the caller passed. + assert (result.chat_template_kwargs or {}).get("y") == 2 + + def test_priority_defaults_to_zero(self): + # ChatCompletionRequest.priority defaults to 0; ``None`` Cohere + # priority must be coerced rather than passed through. + req = _make_request() + result = _convert(req) + assert result.priority == 0 + + +# ====================================================================== +# Streaming options +# ====================================================================== + + +class TestStreamingOptions: + def test_no_stream_leaves_defaults(self): + result = _convert(_make_request(stream=False)) + assert not result.stream + assert result.stream_options is None + + def test_stream_enables_usage_options(self): + # The v2 translator forces ``include_usage=True`` so the + # ``message-end`` event can surface ``billed_units`` / ``tokens``; + # ``continuous_usage_stats`` is intentionally left at its + # ``StreamOptions`` default (False) — Cohere v2 only reports + # usage on the terminal event. + result = _convert(_make_request(stream=True)) + assert result.stream is True + assert result.stream_options is not None + assert result.stream_options.include_usage is True + + +# ====================================================================== +# Response format +# ====================================================================== + + +class TestResponseFormat: + def test_text_is_passthrough(self): + result = _convert(_make_request(response_format={"type": "text"})) + assert result.response_format is None + + def test_json_object(self): + result = _convert(_make_request(response_format={"type": "json_object"})) + assert result.response_format is not None + assert result.response_format.type == "json_object" + assert result.response_format.json_schema is None + + def test_json_schema(self): + schema = {"type": "object", "properties": {"a": {"type": "string"}}} + result = _convert( + _make_request( + response_format={"type": "json_object", "json_schema": schema} + ) + ) + assert result.response_format is not None + assert result.response_format.type == "json_schema" + assert result.response_format.json_schema is not None + assert result.response_format.json_schema.name == "cohere_v2_json_schema" + # ``JsonSchemaResponseFormat.json_schema`` has alias=``schema`` on + # the Pydantic field, so we observe the value via the serialized + # payload (which is what downstream consumers actually read). + dumped = result.response_format.json_schema.model_dump(exclude_none=True) + assert dumped["json_schema"] == schema + + +# ====================================================================== +# Tools / tool_choice +# ====================================================================== + + +class TestApplyTools: + def test_no_tools(self): + result = _convert(_make_request()) + assert result.tools is None + + def test_basic_tool(self): + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": {"type": "object"}, + }, + } + ] + ) + ) + assert result.tools is not None + assert len(result.tools) == 1 + tool = result.tools[0] + assert tool.type == "function" + assert tool.function.name == "calc" + assert tool.function.description == "calculator" + # ``strict`` is an extra attribute on FunctionDefinition (the + # field is only stamped onto the OpenAI tool when strict_tools is + # set on the request). The default path must not set it. + assert getattr(tool.function, "strict", None) is None + + def test_strict_tools_propagates_to_function(self): + result = _convert( + _make_request( + strict_tools=True, + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tools[0].function.strict is True + + +class TestApplyToolChoice: + def test_required(self): + result = _convert( + _make_request( + tool_choice="REQUIRED", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "required" + + def test_none(self): + result = _convert( + _make_request( + tool_choice="NONE", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "none" + + def test_default_to_auto_when_tools_present(self): + # No explicit ``tool_choice`` + tools present → auto, mirroring + # Cohere's documented "free choice" default. + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ] + ) + ) + assert result.tool_choice == "auto" + + def test_no_tools_no_choice_left_unset(self): + # When there are no tools the underlying ChatCompletionRequest + # default applies; we must not stamp ``auto``. + result = _convert(_make_request()) + assert result.tool_choice != "auto" + + +# ====================================================================== +# Cohere-specific template kwargs forwarding +# ====================================================================== + + +class TestApplyCohereTemplateKwargs: + def test_string_documents_wrapped(self): + result = _convert(_make_request(documents=["doc 1", "doc 2"])) + docs = (result.chat_template_kwargs or {}).get("documents") + assert docs == [ + {"id": "doc_0", "data": {"text": "doc 1"}}, + {"id": "doc_1", "data": {"text": "doc 2"}}, + ] + + def test_document_with_explicit_id_preserved(self): + result = _convert( + _make_request( + documents=[ + {"id": "custom", "data": {"text": "t"}}, + {"data": {"text": "t2"}}, # no id -> synthesized + ] + ) + ) + docs = result.chat_template_kwargs["documents"] + assert docs[0] == {"id": "custom", "data": {"text": "t"}} + assert docs[1]["id"] == "doc_1" + + def test_safety_mode_normalized_to_lowercase(self): + result = _convert(_make_request(safety_mode="CONTEXTUAL")) + assert result.chat_template_kwargs["safety_mode"] == "contextual" + + def test_citation_options_forwarded_as_dict(self): + result = _convert(_make_request(citation_options={"mode": "accurate"})) + assert result.chat_template_kwargs["citation_options"] == {"mode": "accurate"} + + def test_thinking_forwarded_as_dict(self): + result = _convert( + _make_request(thinking={"type": "enabled", "token_budget": 16}) + ) + assert result.chat_template_kwargs["thinking"] == { + "type": "enabled", + "token_budget": 16, + } + + def test_existing_chat_template_kwargs_preserved(self): + # User-supplied kwargs should not be clobbered by the v2 fields + # (setdefault semantics). + result = _convert( + _make_request( + chat_template_kwargs={ + "safety_mode": "user-explicit", + "extra": "x", + }, + safety_mode="CONTEXTUAL", + ) + ) + assert result.chat_template_kwargs["safety_mode"] == "user-explicit" + assert result.chat_template_kwargs["extra"] == "x" + + def test_no_template_kwargs_when_no_cohere_fields(self): + # Without any of the Cohere-specific fields and no caller-supplied + # kwargs, we must leave ``chat_template_kwargs`` as None so other + # renderers see a clean request. + result = _convert(_make_request()) + assert result.chat_template_kwargs is None + + +# ====================================================================== +# Message-citation forwarding (chat_template_kwargs["_messages_citations"]) +# ====================================================================== + + +class TestMessageCitations: + """Covers the ``_messages_citations`` chat_template_kwargs entry + that carries ``AssistantChatMessageV2.citations`` from the request + through to the Cohere renderer. Response-shape ``Citation`` -> + melody ``FilterCitation`` conversion is best-effort: some + melody-side fields have no equivalent on the Cohere wire and get + defaulted. + """ + + def test_absent_when_no_assistant_citations(self): + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + ) + ) + assert result.chat_template_kwargs is None or ( + "_messages_citations" not in (result.chat_template_kwargs or {}) + ) + + def test_document_citation_forwarded_by_index(self): + # Assistant at request-index 1 cites the *second* top-level + # document (``doc_x``, position 1). Melody addresses documents + # as ``tool_call_index=0`` (the reserved documents bucket) with + # ``tool_result_indices`` selecting positions inside the + # top-level ``documents`` list. + result = _convert( + _make_request( + documents=[ + {"id": "doc_a", "data": {"text": "irrelevant"}}, + {"id": "doc_x", "data": {"text": "cited doc"}}, + ], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "document", "id": "doc_x"}], + "type": "TEXT_CONTENT", + } + ], + }, + ], + ) + ) + citations_by_index = result.chat_template_kwargs["_messages_citations"] + assert list(citations_by_index.keys()) == [1] + cite = citations_by_index[1][0] + assert cite["start_index"] == 0 + assert cite["end_index"] == 1 + assert cite["text"] == "a" + assert cite["is_thinking"] is False + # ``document_ids`` is a parser-*output* field in melody (see + # ``PromptRenderIds`` in src/templating/util.rs); on the input + # side melody expects the id resolved into + # ``tool_result_indices``. Setting a ``document_ids`` list here + # would have no effect at all -- the renderer would still + # render ```` (no anchor). + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [1]}] + + def test_document_citation_with_unresolvable_id_drops_citation(self): + # A citation source pointing at an id that doesn't appear in + # the request's ``documents`` array leaves the citation with + # nothing to anchor to. Emitting it with ``sources=[]`` would + # render a malformed ``text`` marker, and silently + # misattributing to ``documents[0]`` would be worse -- so we + # drop the whole citation, matching the cohere api's behavior. + # Since this message's only citation was dropped, the message + # index shouldn't appear in the forwarded map at all. + result = _convert( + _make_request( + documents=[{"id": "doc_present", "data": {"text": "x"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "document", "id": "doc_ghost"}], + "type": "TEXT_CONTENT", + } + ], + }, + ], + ) + ) + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations") + assert not forwarded + + def test_document_citation_resolves_auto_assigned_fallback_id(self): + # Requests may include documents without an explicit ``id``; + # ``_apply_cohere_template_kwargs`` synthesizes ``doc_{idx}`` + # ids for them. A citation that targets that fallback id must + # still resolve to the right position. + result = _convert( + _make_request( + documents=[ + {"data": {"text": "unnamed 0"}}, + {"data": {"text": "unnamed 1"}}, + ], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "document", "id": "doc_1"}], + "type": "TEXT_CONTENT", + } + ], + }, + ], + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][1] + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [1]}] + + def test_tool_citation_resolves_to_bucket_and_result_index(self): + # ``ToolSource.id`` on the wire is the id of a specific document + # inside a tool result (NOT the tool_call_id -- ``type`` is a + # payload-shape hint, not an id-space discriminator). Here the + # citation points at the second doc in the second tool call's + # results, which melody addresses as bucket ``1`` (first tool + # call, no top-level documents present so it takes bucket 0 + # ... wait: buckets are per unique tool_call_id, so with no + # top-level docs, ``call_a`` = 0 and ``call_b`` = 1) and + # ``tool_result_indices=[1]`` (2nd doc inside that bucket). + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res_a0", "data": {"text": "r_a0"}}, + }, + ], + }, + { + "role": "assistant", + "tool_plan": "plan2", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": {"name": "g", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_b", + "content": [ + { + "type": "document", + "document": {"id": "res_b0", "data": {"text": "r_b0"}}, + }, + { + "type": "document", + "document": {"id": "res_b1", "data": {"text": "r_b1"}}, + }, + ], + }, + { + "role": "assistant", + "content": "final", + "citations": [ + { + "start": 0, + "end": 5, + "text": "final", + "sources": [{"type": "tool", "id": "res_b1"}], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][5] + assert cite["sources"] == [ + {"tool_call_index": 1, "tool_result_indices": [1]}, + ] + + def test_tool_citation_shifts_bucket_when_documents_present(self): + # When the request has a top-level ``documents`` array, melody + # reserves ``tool_call_index=0`` for it (see + # ``PromptRenderIds::new`` in src/templating/util.rs), so the + # first tool-call bucket becomes 1. A regression that forgets + # this shift would silently attribute tool citations to the + # documents bucket instead. + result = _convert( + _make_request( + documents=[{"id": "d0", "data": {"text": "x"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res_a0", "data": {"text": "r"}}, + }, + ], + }, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "tool", "id": "res_a0"}], + "type": "TEXT_CONTENT", + } + ], + }, + ], + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][3] + assert cite["sources"] == [ + {"tool_call_index": 1, "tool_result_indices": [0]}, + ] + + def test_multiple_sources_in_same_bucket_aggregate(self): + # Two citation sources both point at docs inside the same tool + # result. Melody's ``text`` marker packs them + # into a single ``Source`` with a list of ``tool_result_indices`` + # rather than two separate ``Source`` entries; the converter + # groups per-bucket to match that -- matching the cohere api's + # inbound citation shape. + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res0", "data": {"text": "0"}}, + }, + { + "type": "document", + "document": {"id": "res1", "data": {"text": "1"}}, + }, + ], + }, + { + "role": "assistant", + "content": "ans", + "citations": [ + { + "start": 0, + "end": 3, + "text": "ans", + "sources": [ + {"type": "tool", "id": "res0"}, + {"type": "tool", "id": "res1"}, + ], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][3] + assert cite["sources"] == [ + {"tool_call_index": 0, "tool_result_indices": [0, 1]}, + ] + + def test_tool_citation_with_unresolvable_id_drops_citation(self): + # Same policy as ``test_document_citation_with_unresolvable_id + # _drops_citation``: any unresolvable source id drops the whole + # citation rather than emitting a malformed + # ``text`` marker. + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "tool", "id": "ghost_doc"}], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations") + assert not forwarded + + def test_partial_unresolvable_source_drops_whole_citation(self): + # If any source in a citation fails to resolve, drop the entire + # citation -- mixing resolved and unresolved sources produces a + # partially-correct ```` marker, which is worse than none. + # This matches the cohere api's behavior. + result = _convert( + _make_request( + documents=[{"id": "doc_real", "data": {"text": "x"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [ + {"type": "document", "id": "doc_real"}, + {"type": "document", "id": "doc_ghost"}, + ], + "type": "TEXT_CONTENT", + } + ], + }, + ], + ) + ) + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations") + assert not forwarded + + def test_thinking_content_type_maps_to_is_thinking(self): + result = _convert( + _make_request( + documents=[{"id": "d", "data": {"text": "x"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "chain of thought"}, + {"type": "text", "text": "answer"}, + ], + "citations": [ + { + "start": 0, + "end": 5, + "text": "chain", + "sources": [{"type": "document", "id": "d"}], + "type": "THINKING_CONTENT", + } + ], + }, + ], + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][1] + assert cite["is_thinking"] is True + # And the source resolved (sanity: is_thinking is meaningful + # only if the citation actually renders). + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [0]}] + + def test_plan_type_maps_to_is_thinking(self): + # PLAN citations sit on tool-plan-style non-reasoning assistant + # turns; melody treats them like THINKING for template purposes. + result = _convert( + _make_request( + documents=[{"id": "d", "data": {"text": "x"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "citations": [ + { + "start": 0, + "end": 4, + "text": "plan", + "sources": [{"type": "document", "id": "d"}], + "type": "PLAN", + } + ], + }, + ], + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][1] + assert cite["is_thinking"] is True + + def test_tool_content_string_registers_bucket_but_has_no_doc_ids(self): + # A string tool result has no citable substructure. The bucket + # still gets registered (so a later tool call's bucket doesn't + # slide into its slot), but any citation targeting the string + # message fails to resolve and the citation is dropped. + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_a", "content": "plain text"}, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "tool", "id": "call_a"}], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations") + assert not forwarded + + def test_tool_content_all_text_blocks_registers_bucket_but_no_docs(self): + # Same as the string case, but with a structured content list + # of pure text blocks. No doc ids are exposed, so any citation + # targeting the message drops. + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + {"type": "text", "text": "part 1"}, + {"type": "text", "text": "part 2"}, + ], + }, + { + "role": "assistant", + "content": "a", + "citations": [ + { + "start": 0, + "end": 1, + "text": "a", + "sources": [{"type": "tool", "id": "part 1"}], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations") + assert not forwarded + + def test_tool_content_mixed_text_and_docs_uses_content_array_position(self): + # Melody advances ``tool_result_index`` for *every* content + # item (text items push an empty id into the row; see + # ``push_tool_message_contents`` in melody/src/templating/ + # util.rs). Docs at positions 1 and 3 in a 4-item mixed + # content must resolve to result_indices 1 and 3, not the + # "position among documents only" numbering (which would give + # 0 and 1). + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + {"type": "text", "text": "before"}, + { + "type": "document", + "document": {"id": "doc1", "data": {"text": "d1"}}, + }, + {"type": "text", "text": "between"}, + { + "type": "document", + "document": {"id": "doc2", "data": {"text": "d2"}}, + }, + ], + }, + { + "role": "assistant", + "content": "ans", + "citations": [ + { + "start": 0, + "end": 3, + "text": "ans", + "sources": [ + {"type": "tool", "id": "doc1"}, + {"type": "tool", "id": "doc2"}, + ], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][3] + assert cite["sources"] == [ + {"tool_call_index": 0, "tool_result_indices": [1, 3]}, + ] + + def test_same_tool_call_id_across_messages_offsets_result_indices(self): + # If a client splits one tool call's outputs across multiple + # tool messages with the same ``tool_call_id``, melody + # accumulates them into the same bucket (each message's + # content array is *appended* to the bucket's slot list). The + # second message's docs must therefore be offset by the first + # message's content length. Dory's melody adapter relies on + # exactly this behavior when it flattens each tool-result + # document into its own melody ``Message``. + result = _convert( + _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "d0", "data": {"text": "0"}}, + }, + { + "type": "document", + "document": {"id": "d1", "data": {"text": "1"}}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "d2", "data": {"text": "2"}}, + }, + ], + }, + { + "role": "assistant", + "content": "ans", + "citations": [ + { + "start": 0, + "end": 3, + "text": "ans", + "sources": [{"type": "tool", "id": "d2"}], + "type": "TEXT_CONTENT", + } + ], + }, + ] + ) + ) + (cite,) = result.chat_template_kwargs["_messages_citations"][4] + # ``d2`` is at position 0 within its own message's content, + # but the previous tool message contributed 2 slots, so its + # result index in bucket 0 is 2. + assert cite["sources"] == [ + {"tool_call_index": 0, "tool_result_indices": [2]}, + ] + + +# ====================================================================== +# _chat_completion_to_v2 (non-streaming response builder) +# ====================================================================== + + +class TestChatCompletionToV2: + def test_text_only(self): + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response(content="hello") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert isinstance(v2, CohereChatV2Response) + assert v2.id == "resp_1" + assert v2.finish_reason == "COMPLETE" + assert v2.message.role == "assistant" + assert len(v2.message.content) == 1 + assert v2.message.content[0].type == "text" + assert v2.message.content[0].text == "hello" + assert v2.message.tool_calls is None + assert v2.message.tool_plan is None + assert v2.usage is None + + def test_reasoning_model_keeps_thinking_with_tool_calls(self): + # Reasoning Command models: ``thinking`` block stays in + # ``message.content`` and ``tool_plan`` is left unset, even when + # tool calls are present. + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response( + content="resp text", + reasoning="thoughts", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.finish_reason == "TOOL_CALL" + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + assert v2.message.content[0].thinking == "thoughts" + assert v2.message.content[1].text == "resp text" + assert v2.message.tool_calls[0].id == "c1" + assert v2.message.tool_calls[0].function.name == "f" + + def test_non_reasoning_model_moves_reasoning_to_tool_plan(self): + # Older Command models surface reasoning as ``tool_plan`` on tool- + # call turns; the thinking block should be dropped from content. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content=None, + reasoning="plan", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan == "plan" + assert v2.message.content is None + assert v2.message.tool_calls[0].id == "c1" + + def test_non_reasoning_model_keeps_thinking_when_no_tool_calls(self): + # No tool calls => non-reasoning behavior is identical to + # reasoning behavior; the thinking block stays. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content="answer", reasoning="plan", tool_calls=None + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + + def test_id_synthesized_when_response_id_missing(self): + serving = _serving() + resp = _build_chat_completion_response(content="hi", response_id="") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.id.startswith("chat_") + + def test_kv_transfer_params_propagated(self): + serving = _serving() + resp = _build_chat_completion_response( + content="hi", kv_transfer_params={"k": 1} + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.kv_transfer_params == {"k": 1} + + +# ====================================================================== +# _build_usage +# ====================================================================== + + +class TestBuildUsage: + def test_none_passthrough(self): + resp = _build_chat_completion_response(content="hi") + # default usage is None + assert CohereServingChatV2._build_usage(resp) is None + + def test_basic_usage(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage is not None + assert usage.billed_units.input_tokens == 10 + assert usage.billed_units.output_tokens == 5 + assert usage.tokens.input_tokens == 10 + assert usage.tokens.output_tokens == 5 + assert usage.cached_tokens is None + + def test_completion_tokens_default_to_zero_when_missing(self): + resp = _build_chat_completion_response( + content="hi", + usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.billed_units.output_tokens == 0 + + def test_cached_tokens_propagated(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.cached_tokens == 3 + + +# ====================================================================== +# _extract_citations_if_any +# ====================================================================== + + +class TestExtractCitations: + """Test the parser-output -> wire-shape citation coercion. + + Sources are already fully resolved by the reasoning parser (see + ``_melody_sources_to_vllm`` in + ``vllm/reasoning/cohere_command_reasoning_parser.py``, which + consumes the position map forwarded via ``chat_template_kwargs``). + The serving layer's remaining responsibilities on the outbound + path are: coerce to :class:`cohere.types.Citation`, apply the + ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning models, + and drop citations whose ``sources`` list is entirely empty + (unattributable anchors). Id-less sources with a payload are + preserved to match the cohere api's optional-``id`` source shape. + """ + + @staticmethod + def _make_citation(**overrides: Any) -> VLLMCitation: + base = dict( + start=0, + end=11, + text="Shakespeare", + sources=[ + CitationSource( + type="document", + id="doc_hamlet", + document={"id": "doc_hamlet", "text": "Hamlet."}, + ), + ], + type="TEXT_CONTENT", + ) + base.update(overrides) + return VLLMCitation(**base) + + def test_none_or_empty_returns_none(self): + serving = _serving() + assert ( + serving._extract_citations_if_any(type("M", (), {"citations": None})()) + is None + ) + assert ( + serving._extract_citations_if_any(type("M", (), {"citations": []})()) + is None + ) + assert serving._extract_citations_if_any(type("M", (), {})()) is None + + def test_resolved_document_source_survives_to_wire(self): + # A parser-resolved document source flows through unchanged + # (aside from removing ``None`` fields). + serving = _serving() + msg = type("M", (), {})() + msg.citations = [self._make_citation()] + out = serving._extract_citations_if_any(msg) + assert out is not None and len(out) == 1 + wire = out[0].model_dump(exclude_none=True) + assert wire["type"] == "TEXT_CONTENT" + assert wire["sources"] == [ + { + "type": "document", + "id": "doc_hamlet", + "document": {"id": "doc_hamlet", "text": "Hamlet."}, + } + ] + + def test_resolved_tool_source_survives_to_wire(self): + # ``type="tool"`` sources carry ``tool_output``; validated by + # the SDK's discriminated union. + serving = _serving() + msg = type("M", (), {})() + msg.citations = [ + self._make_citation( + sources=[ + CitationSource( + type="tool", + id="res_a0", + tool_output={"id": "res_a0", "text": "r"}, + ), + ], + ) + ] + out = serving._extract_citations_if_any(msg) + assert out is not None + wire = out[0].model_dump(exclude_none=True) + assert wire["sources"] == [ + { + "type": "tool", + "id": "res_a0", + "tool_output": {"id": "res_a0", "text": "r"}, + } + ] + + def test_source_without_id_preserved_with_payload(self): + # Matches the cohere api: text-only tool results (and top-level + # docs without a client-provided id) surface on the wire + # without an ``id`` field, but the payload still rides through + # so the client can see what was cited. + serving = _serving() + msg = type("M", (), {})() + msg.citations = [ + self._make_citation( + sources=[ + CitationSource(type="tool", tool_output={"content": "hi"}), + ] + ) + ] + out = serving._extract_citations_if_any(msg) + assert out is not None + wire = out[0].model_dump(exclude_none=True) + assert wire["sources"] == [{"type": "tool", "tool_output": {"content": "hi"}}] + assert "id" not in wire["sources"][0] + + def test_citation_with_no_surviving_sources_dropped(self): + # If every source in a citation is empty (e.g. parser produced + # a stray citation with no resolvable positions), the whole + # citation is dropped so we don't emit unattributed anchors. + serving = _serving() + msg = type("M", (), {})() + msg.citations = [self._make_citation(sources=[])] + assert serving._extract_citations_if_any(msg) is None + + def test_multiple_sources_survive(self): + serving = _serving() + msg = type("M", (), {})() + msg.citations = [ + self._make_citation( + sources=[ + CitationSource(type="document", id="d0", document={"id": "d0"}), + CitationSource(type="document", id="d1", document={"id": "d1"}), + ], + ) + ] + out = serving._extract_citations_if_any(msg) + assert out is not None + wire = out[0].model_dump(exclude_none=True) + assert [s["id"] for s in wire["sources"]] == ["d0", "d1"] + + def test_plan_type_on_non_reasoning_model(self): + # ``is_thinking=True`` melody citations arrive tagged as + # ``THINKING_CONTENT``; on non-reasoning models the reasoning + # block is surfaced as ``tool_plan`` (see + # ``_chat_completion_to_v2``), so citations on it are + # ``PLAN`` on the wire. + serving = _serving(is_reasoning_model=False) + msg = type("M", (), {})() + msg.citations = [self._make_citation(type="THINKING_CONTENT")] + out = serving._extract_citations_if_any(msg) + assert out is not None + assert out[0].type == "PLAN" + + def test_thinking_type_preserved_on_reasoning_model(self): + serving = _serving(is_reasoning_model=True) + msg = type("M", (), {})() + msg.citations = [self._make_citation(type="THINKING_CONTENT")] + out = serving._extract_citations_if_any(msg) + assert out is not None + assert out[0].type == "THINKING_CONTENT" + + def test_dict_shape_from_streaming_pipeline(self): + # The streaming path receives citations as dicts (because + # ``DeltaMessage.citations`` is an untyped extras field on the + # OpenAI wire protocol). ``_to_wire_citation`` must accept + # both dict and object shapes -- unary uses objects. + serving = _serving() + msg = type("M", (), {})() + msg.citations = [ + { + "start": 0, + "end": 3, + "text": "abc", + "sources": [ + {"type": "document", "id": "d0", "document": {"id": "d0"}}, + ], + "type": "TEXT_CONTENT", + } + ] + out = serving._extract_citations_if_any(msg) + assert out is not None + assert out[0].sources[0].id == "d0" + + +# ====================================================================== +# _build_position_to_source +# ====================================================================== + + +class TestBuildPositionToSource: + """Pin the outbound numbering invariant. + + ``_build_position_to_source`` inverts the same numbering rule that + ``_build_doc_id_to_prompt_position`` follows on the inbound path + (melody's ``PromptRenderIds::from_messages`` -- see + ``melody/src/templating/util.rs``). These tests exist to catch + numbering drift between the two helpers, which would cause + outbound citations to attribute to the wrong document. + """ + + def test_top_level_documents_populate_bucket_zero(self): + request = _make_request( + documents=[ + {"id": "d0", "data": {"text": "zero"}}, + {"id": "d1", "data": {"text": "one"}}, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert set(result.keys()) == {(0, 0), (0, 1)} + assert result[(0, 0)].type == "document" + assert result[(0, 0)].id == "d0" + assert result[(0, 0)].document == {"id": "d0", "text": "zero"} + assert result[(0, 1)].id == "d1" + + def test_tool_result_documents_start_at_bucket_one_when_docs_present(self): + request = _make_request( + documents=[{"id": "d0", "data": {"text": "top"}}], + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res_a0", "data": {"text": "r"}}, + } + ], + }, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert set(result.keys()) == {(0, 0), (1, 0)} + assert result[(1, 0)].type == "tool" + assert result[(1, 0)].id == "res_a0" + assert result[(1, 0)].tool_output == {"id": "res_a0", "text": "r"} + + def test_tool_call_id_registered_from_assistant_message(self): + # Bucket assignment must match the cohere api's rule: the + # first-seen ``tool_call_id`` (whether on the assistant + # tool_calls or on the tool message) claims the next integer. + request = _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "first_call", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + }, + { + "id": "second_call", + "type": "function", + "function": {"name": "g", "arguments": "{}"}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "second_call", + "content": [ + { + "type": "document", + "document": {"id": "res_second", "data": {"text": "s"}}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "first_call", + "content": [ + { + "type": "document", + "document": {"id": "res_first", "data": {"text": "f"}}, + } + ], + }, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + # ``first_call`` was seen first -> bucket 0; ``second_call`` + # -> bucket 1. The subsequent tool messages fill their + # respective buckets, not the message-order buckets. + assert result[(0, 0)].id == "res_first" + assert result[(1, 0)].id == "res_second" + + def test_tool_content_text_and_document_blocks_both_populate_map(self): + # Mixed text + document content: text blocks now surface in + # the map as id-less ``tool`` sources (matching the cohere + # api's handling of text-only tool content) so a model + # citation pointing at the text slot still has a payload to + # attach. + request = _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + {"type": "text", "text": "prelude"}, + { + "type": "document", + "document": {"id": "res_x", "data": {"text": "x"}}, + }, + ], + }, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert set(result.keys()) == {(0, 0), (0, 1)} + assert result[(0, 0)].type == "tool" + assert result[(0, 0)].id is None + assert result[(0, 0)].tool_output == {"content": "prelude"} + assert result[(0, 1)].id == "res_x" + + def test_tool_content_json_text_parsed_as_tool_output(self): + # The cohere api tries JSON first for text tool content; parity + # here means a JSON-object text block is spread into + # ``tool_output`` directly instead of nested under ``content``. + request = _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [{"type": "text", "text": '{"title": "hi"}'}], + }, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert result[(0, 0)].tool_output == {"title": "hi"} + + def test_top_level_document_without_id_omits_id_but_keeps_payload(self): + # Client didn't provide a doc id: match cohere's optional-id + # source shape by leaving ``id=None`` on the wire, but keep + # the payload so the citation is still meaningful. + request = _make_request( + documents=[{"data": {"text": "anon"}}], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert result[(0, 0)].type == "document" + assert result[(0, 0)].id is None + assert result[(0, 0)].document == {"text": "anon"} + + def test_top_level_string_document_has_no_wire_id(self): + request = _make_request(documents=["raw text"]) + result = CohereServingChatV2._build_position_to_source(request) + assert result[(0, 0)].id is None + assert result[(0, 0)].document == {"text": "raw text"} + + def test_tool_content_string_synthesizes_idless_source(self): + # ``ToolChatMessageV2.content`` may be a bare string; the + # cohere renderer wraps it as a single text block, so we + # synthesize a matching id-less source for the slot. + request = _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_a", "content": "just text"}, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert result[(0, 0)].type == "tool" + assert result[(0, 0)].id is None + assert result[(0, 0)].tool_output == {"content": "just text"} + + def test_multiple_tool_messages_same_call_id_offset_slots(self): + # Tool results streamed across multiple messages accumulate + # into the same bucket, so later docs are offset by the sum + # of previous messages' content lengths. + request = _make_request( + messages=[ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "tool_plan": "plan", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res_first", "data": {"text": "1"}}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + { + "type": "document", + "document": {"id": "res_second", "data": {"text": "2"}}, + }, + ], + }, + ], + ) + result = CohereServingChatV2._build_position_to_source(request) + assert result[(0, 0)].id == "res_first" + assert result[(0, 1)].id == "res_second" + + +# ====================================================================== +# create_error_response +# ====================================================================== + + +class TestCreateErrorResponse: + def test_envelope_uses_400(self): + serving = _serving() + err = serving.create_error_response("oops") + assert err.error.message == "oops" + assert err.error.code == 400 + assert err.error.type == "bad_request" diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py new file mode 100644 index 000000000000..34cd2eaf0788 --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -0,0 +1,944 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 SSE stream conversion in +``vllm/entrypoints/cohere/serving.py``. + +The stream-translation entry point is +:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an +async iterable of OpenAI SSE chunks into Cohere's +``message-start → (content|tool-call|citation)* → message-end → [DONE]`` +event stream. + +We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in +isolation, plus a handful of end-to-end scenarios that exercise the +state machine. +""" + +import json +from collections.abc import AsyncGenerator +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.cohere.cohere_chat_message import ( + CitationSource, +) +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + MessageStartEvent, +) +from vllm.entrypoints.cohere.serving import ( + _DONE_FRAME, + CohereServingChatV2, + ContentBlockType, + _emit, + _sse, + _StreamState, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain (which would need a + real engine client, model registry, render service, etc.). + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_stream_to_v2`` and the per-delta handlers); the + rest is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _parse_event(frame: str) -> dict[str, Any]: + """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload.""" + assert frame.startswith("data: ") + assert frame.endswith("\n\n") + return json.loads(frame[len("data: ") : -2]) + + +def _make_chunk( + *, + chunk_id: str = "chunk_0", + role: str | None = None, + content: str | None = None, + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + omit_choices: bool = False, + citations: list[Any] | None = None, +) -> str: + """Build the ``data: {...}\\n\\n`` SSE frame the production code + consumes.""" + delta: dict[str, Any] = {} + if role is not None: + delta["role"] = role + if content is not None: + delta["content"] = content + if reasoning is not None: + delta["reasoning"] = reasoning + if tool_calls is not None: + delta["tool_calls"] = tool_calls + if citations is not None: + delta["citations"] = citations + + payload: dict[str, Any] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + } + if not omit_choices: + payload["choices"] = [ + {"index": 0, "delta": delta, "finish_reason": finish_reason} + ] + else: + payload["choices"] = [] + if usage is not None: + payload["usage"] = usage + # Use ChatCompletionStreamResponse to normalize the payload shape. + chunk = ChatCompletionStreamResponse.model_validate(payload) + return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n" + + +def _make_done() -> str: + return "data: [DONE]\n\n" + + +async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: + for item in items: + yield item + + +async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: + """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect + the emitted SSE frames.""" + request = CohereChatV2Request( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True + ) + gen = serving._chat_completion_stream_to_v2(_stream_from(items), request) + return [frame async for frame in gen] + + +# ====================================================================== +# Low-level helpers: _sse, _emit, _DONE_FRAME +# ====================================================================== + + +class TestSSEHelpers: + def test_sse_wraps_payload(self): + assert _sse("hello") == "data: hello\n\n" + + def test_done_frame_constant(self): + # cohere-python and Fern-generated clients key off this exact + # sentinel; keep it byte-for-byte stable. + assert _DONE_FRAME == "data: [DONE]\n\n" + + def test_emit_serializes_event_with_type_discriminator(self): + frame = _emit( + MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}}) + ) + payload = _parse_event(frame) + assert payload["type"] == "message-start" + assert payload["id"] == "abc" + assert payload["delta"] == {"message": {"role": "assistant"}} + + +# ====================================================================== +# _StreamState +# ====================================================================== + + +class TestStreamState: + def test_content_index_monotonic(self): + st = _StreamState() + assert st.next_content_index() == 0 + assert st.next_content_index() == 1 + assert st.next_content_index() == 2 + + def test_citation_index_separate_from_content_index(self): + st = _StreamState() + st.next_content_index() # 0 + st.next_content_index() # 1 + # Citation indexing is independent of content indexing. + assert st.next_citation_index() == 0 + assert st.next_citation_index() == 1 + + +# ====================================================================== +# _close_open_blocks +# ====================================================================== + + +class TestCloseOpenBlocks: + def test_no_block_open_emits_nothing(self): + serving = _serving() + state = _StreamState() + assert serving._close_open_blocks(state) == [] + + def test_text_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TEXT + state.active_block_index = 1 + out = serving._close_open_blocks(state) + assert len(out) == 1 + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 1} + assert state.active_block is None + assert state.active_block_index is None + + def test_thinking_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.THINKING + state.active_block_index = 3 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 3} + + def test_tool_call_block_emits_tool_call_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TOOL_CALL + state.active_tool_index = 7 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "tool-call-end", "index": 7} + assert state.active_tool_index is None + + +# ====================================================================== +# _handle_text_delta +# ====================================================================== + + +class TestHandleTextDelta: + def test_opens_block_first_time(self): + serving = _serving() + state = _StreamState() + events = serving._handle_text_delta(state, "Hi") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["index"] == 0 + assert start["delta"]["message"]["content"]["type"] == "text" + assert delta["type"] == "content-delta" + assert delta["index"] == 0 + assert delta["delta"]["message"]["content"]["text"] == "Hi" + assert state.active_block == ContentBlockType.TEXT + assert state.active_block_index == 0 + + def test_continues_block_with_just_delta(self): + serving = _serving() + state = _StreamState() + serving._handle_text_delta(state, "Hi") + events = serving._handle_text_delta(state, " there") + # Only a delta event, no new content-start. + assert len(events) == 1 + delta = _parse_event(events[0]) + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["text"] == " there" + + def test_switches_from_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + # Open a thinking block first, then switch to text. + serving._handle_thinking_delta(state, "ponder") + events = serving._handle_text_delta(state, "answer") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "content-start", "content-delta"] + # The text block gets a new index (1), distinct from thinking's 0. + assert _parse_event(events[1])["index"] == 1 + assert state.active_block == ContentBlockType.TEXT + + +# ====================================================================== +# _handle_thinking_delta +# ====================================================================== + + +class TestHandleThinkingDelta: + def test_reasoning_model_opens_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + events = serving._handle_thinking_delta(state, "thought") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["delta"]["message"]["content"]["type"] == "thinking" + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["thinking"] == "thought" + assert state.active_block == ContentBlockType.THINKING + + def test_reasoning_model_continues_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + serving._handle_thinking_delta(state, "first") + events = serving._handle_thinking_delta(state, " more") + assert len(events) == 1 + assert _parse_event(events[0])["type"] == "content-delta" + + def test_non_reasoning_model_emits_tool_plan_delta(self): + # Older Command models stream reasoning as ``tool_plan`` deltas; + # no content-start/end pair is emitted. + serving = _serving(is_reasoning_model=False) + state = _StreamState() + events = serving._handle_thinking_delta(state, "planning") + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-plan-delta" + assert payload["delta"]["message"]["tool_plan"] == "planning" + # ``tool_plan`` deltas don't claim an active content block. + assert state.active_block is None + + def test_non_reasoning_model_closes_open_text_block(self): + serving = _serving(is_reasoning_model=False) + state = _StreamState() + # Open a text block first. + serving._handle_text_delta(state, "answer") + events = serving._handle_thinking_delta(state, "rethink") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-plan-delta"] + + +# ====================================================================== +# _handle_tool_call_deltas +# ====================================================================== + + +class TestHandleToolCallDeltas: + def test_new_tool_call_opens_tool_call_start(self): + serving = _serving() + state = _StreamState() + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type( + "Fn", (), {"name": "calc", "arguments": '{"x":'} + )(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-start" + assert payload["index"] == 0 + tc = payload["delta"]["message"]["tool_calls"] + assert tc["id"] == "c1" + assert tc["function"]["name"] == "calc" + assert tc["function"]["arguments"] == '{"x":' + assert state.active_block == ContentBlockType.TOOL_CALL + assert state.active_tool_index == 0 + assert 0 in state.tool_calls_seen + + def test_subsequent_arguments_emit_delta(self): + serving = _serving() + state = _StreamState() + # First call: start. + first = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": ""})(), + }, + )() + ] + serving._handle_tool_call_deltas(state, first) + # Second call: same index, additional arguments fragment. + more = [ + type( + "Delta", + (), + { + "index": 0, + "id": None, + "function": type("Fn", (), {"name": None, "arguments": "1}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, more) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-delta" + assert payload["index"] == 0 + assert ( + payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}" + ) + + def test_new_tool_call_closes_existing_content_block(self): + serving = _serving() + state = _StreamState() + # Open a text block, then start a tool call. + serving._handle_text_delta(state, "I'll call:") + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-call-start"] + + +# ====================================================================== +# _handle_citation_deltas +# ====================================================================== + + +class TestHandleCitationDeltas: + """Test the per-delta emitter. Sources are pre-resolved by the + reasoning parser via the position map forwarded through + ``chat_template_kwargs`` (see :func:`_melody_sources_to_vllm` in + ``vllm/reasoning/cohere_command_reasoning_parser.py``), so the + handler only has to coerce to the SDK wire shape, apply the + ``THINKING_CONTENT`` -> ``PLAN`` rewrite, and drop citations with + no attributable sources. + """ + + @staticmethod + def _resolved_source(**overrides: Any) -> CitationSource: + base = dict( + type="document", + id="doc_hamlet", + document={"id": "doc_hamlet", "text": "Hamlet by Shakespeare."}, + ) + base.update(overrides) + return CitationSource(**base) + + def test_resolved_citation_emits_start_and_end_with_real_id(self): + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=11, + text="Shakespeare", + sources=[self._resolved_source()], + type="TEXT_CONTENT", + ) + ] + events = serving._handle_citation_deltas(state, citations) + assert len(events) == 2 + start = _parse_event(events[0]) + end = _parse_event(events[1]) + assert start["type"] == "citation-start" + assert start["index"] == 0 + cit_payload = start["delta"]["message"]["citations"] + assert cit_payload["start"] == 0 + assert cit_payload["end"] == 11 + assert cit_payload["text"] == "Shakespeare" + assert cit_payload["sources"] == [ + { + "type": "document", + "id": "doc_hamlet", + "document": {"id": "doc_hamlet", "text": "Hamlet by Shakespeare."}, + } + ] + assert end == {"type": "citation-end", "index": 0} + + def test_multiple_sources_survive_to_wire(self): + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="works", + sources=[ + self._resolved_source(id="d0", document={"id": "d0"}), + self._resolved_source(id="d1", document={"id": "d1"}), + ], + type="TEXT_CONTENT", + ) + ] + events = serving._handle_citation_deltas(state, citations) + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"] + ids = [s["id"] for s in cit_payload["sources"]] + assert ids == ["d0", "d1"] + + def test_source_without_id_streamed_with_payload(self): + # Match the cohere api: id-less tool sources (from text-only + # tool results) still ride through with their ``tool_output`` + # payload; only the ``id`` field is omitted. + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="tool", tool_output={"content": "hi"})], + ) + ] + events = serving._handle_citation_deltas(state, citations) + assert len(events) >= 1 + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"] + assert cit_payload["sources"] == [ + {"type": "tool", "tool_output": {"content": "hi"}} + ] + assert "id" not in cit_payload["sources"][0] + + def test_citation_with_no_sources_streamed_dropped(self): + # Citations whose source list is entirely empty are still + # dropped so we don't emit an unattributed anchor. + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[], + ) + ] + events = serving._handle_citation_deltas(state, citations) + assert events == [] + + def test_plan_rewrite_on_non_reasoning_model(self): + serving = _serving(is_reasoning_model=False) + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[self._resolved_source()], + type="THINKING_CONTENT", + ) + ] + events = serving._handle_citation_deltas(state, citations) + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"] + assert cit_payload["type"] == "PLAN" + + +# ====================================================================== +# _build_message_end_event +# ====================================================================== + + +class TestBuildMessageEndEvent: + def test_without_usage(self): + serving = _serving() + frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop") + payload = _parse_event(frame) + assert payload["type"] == "message-end" + assert payload["id"] == "abc" + assert payload["delta"]["finish_reason"] == "COMPLETE" + assert "usage" not in payload["delta"] + + def test_with_usage(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="length", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["finish_reason"] == "MAX_TOKENS" + usage = payload["delta"]["usage"] + assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5} + assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5} + assert "cached_tokens" not in usage + + def test_with_cached_tokens(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="stop", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["usage"]["cached_tokens"] == 3 + + +# ====================================================================== +# End-to-end: _chat_completion_stream_to_v2 +# ====================================================================== + + +class TestChatCompletionStreamToV2: + """End-to-end stream lifecycle tests.""" + + @pytest.mark.asyncio + async def test_text_only_happy_path_emits_full_lifecycle(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_chunk(content=" there"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 4, + "completion_tokens": 2, + "total_tokens": 6, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # message-start, content-start, content-delta, content-delta, + # content-end, message-end, then [DONE] as the last frame. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + # message-end should carry usage stats from the trailing chunk. + end_payload = _parse_event(frames[-2]) + assert end_payload["delta"]["finish_reason"] == "COMPLETE" + assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4 + + @pytest.mark.asyncio + async def test_text_then_tool_call_closes_text_first(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="planning..."), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Text block is opened+delta, then closed before the tool-call + # opens, and the final close happens on the usage chunk path. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # finish_reason captured from the prior chunk. + end = _parse_event(frames[-2]) + assert end["delta"]["finish_reason"] == "TOOL_CALL" + + @pytest.mark.asyncio + async def test_thinking_then_text_reasoning_model(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="thinking..."), + _make_chunk(content="answer"), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Thinking block opens with index 0; text block reopens with index 1. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + thinking_start = _parse_event(frames[1]) + assert thinking_start["delta"]["message"]["content"]["type"] == "thinking" + text_start = _parse_event(frames[4]) + assert text_start["delta"]["message"]["content"]["type"] == "text" + + @pytest.mark.asyncio + async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self): + serving = _serving(is_reasoning_model=False) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="planning"), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "tool-plan-delta", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # No thinking content blocks should be present. + assert "content-start" not in types + assert "content-end" not in types + + @pytest.mark.asyncio + async def test_done_marker_in_middle_closes_open_block(self): + # Some upstreams send [DONE] without a trailing usage-only chunk. + # The translator must still emit message-end before [DONE] so + # Cohere clients don't hang. + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_done(), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + + @pytest.mark.asyncio + async def test_skips_empty_and_non_data_lines(self): + serving = _serving() + items = [ + "\n", + "event: ping\n\n", + _make_chunk(role="assistant"), + "data: \n\n", # empty data + _make_chunk(content="Hi"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Should still produce a complete lifecycle. + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types[0] == "message-start" + assert types[-1] == "message-end" + + @pytest.mark.asyncio + async def test_exception_in_chunk_parsing_emits_error_message_end(self): + # An invalid JSON payload after the first chunk triggers the + # error path: a synthetic message-end with finish_reason=ERROR + # followed by [DONE]. + serving = _serving() + items = [ + _make_chunk(role="assistant"), + "data: {not valid json}\n\n", + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Find the error-shaped message-end. + error_end = _parse_event(frames[-2]) + assert error_end["type"] == "message-end" + assert error_end["delta"]["finish_reason"] == "ERROR" + assert "error" in error_end["delta"] + + @pytest.mark.asyncio + async def test_error_message_end_is_sanitized(self): + # The client-visible ``error`` field on the terminal message-end + # must be routed through ``sanitize_message`` so tracebacks, + # host filesystem paths, and Python object memory addresses + # can't leak from a mid-stream exception into the SSE payload. + serving = _serving() + + async def _raises_mid_stream() -> AsyncGenerator[str, None]: + yield _make_chunk(role="assistant") + raise RuntimeError( + "boom reading /Users/dev/vllm/x.py" + ) + + request = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + ) + frames = [ + f + async for f in serving._chat_completion_stream_to_v2( + _raises_mid_stream(), request + ) + ] + assert frames[-1] == _DONE_FRAME + error_end = _parse_event(frames[-2]) + assert error_end["delta"]["finish_reason"] == "ERROR" + err_msg = error_end["delta"]["error"] + # Memory address stripped and path replaced with ````. + assert "0x7fabc0deface" not in err_msg + assert "/Users/dev/vllm/x.py" not in err_msg + assert "" in err_msg + # The human-facing prefix should still survive sanitization. + assert "boom" in err_msg + + @pytest.mark.asyncio + async def test_citations_in_delta_emit_citation_events(self): + # End-to-end streaming citation flow. The parser hands us + # already-resolved sources (see ``_melody_sources_to_vllm`` + # in ``vllm/reasoning/cohere_command_reasoning_parser.py``); + # the streaming loop just has to coerce and emit. + serving = _serving() + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="hello"), + _make_chunk( + citations=[ + { + "start": 0, + "end": 5, + "text": "hello", + "sources": [ + { + "type": "document", + "id": "d1", + "document": {"id": "d1", "text": "cited"}, + } + ], + "type": "TEXT_CONTENT", + } + ] + ), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + parsed = [_parse_event(f) for f in frames[:-1]] + types = [p["type"] for p in parsed] + assert "citation-start" in types + assert "citation-end" in types + cit_start = next(p for p in parsed if p["type"] == "citation-start") + assert cit_start["delta"]["message"]["citations"]["sources"] == [ + { + "type": "document", + "id": "d1", + "document": {"id": "d1", "text": "cited"}, + } + ] + + @pytest.mark.asyncio + async def test_first_chunk_emits_message_start_with_chunk_id(self): + serving = _serving() + items = [ + _make_chunk(chunk_id="my-id", role="assistant"), + _make_chunk(chunk_id="my-id", content="hi"), + _make_chunk( + chunk_id="my-id", + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + ms = _parse_event(frames[0]) + assert ms["type"] == "message-start" + assert ms["id"] == "my-id" + assert ms["delta"]["message"]["role"] == "assistant" diff --git a/tests/entrypoints/conftest.py b/tests/entrypoints/conftest.py index c2e9a1de3188..27383f87ea7d 100644 --- a/tests/entrypoints/conftest.py +++ b/tests/entrypoints/conftest.py @@ -190,30 +190,32 @@ def sample_sql_statements(): @pytest.fixture(scope="session") def qwen3_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="charent/self_cognition_Alice") + return hf_api().snapshot_download(repo_id="charent/self_cognition_Alice") @pytest.fixture(scope="session") def qwen3_meowing_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") @pytest.fixture(scope="session") def qwen3_woofing_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") @pytest.fixture(scope="session") def opt125_lora_files() -> str: """Download opt-125m LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="peft-internal-testing/opt-125m-dummy-lora") + return hf_api().snapshot_download( + repo_id="peft-internal-testing/opt-125m-dummy-lora" + ) diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index 12dae6dfd523..92c606154592 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -285,7 +285,7 @@ async def test_item_ordering(self, item_first, expected): ) for i, exp in enumerate(expected): - assert engine_inputs[i]["prompt_token_ids"] == exp + assert engine_inputs[i]["prompt_token_ids"] == exp # type: ignore[typeddict-item] class TestGeneration: diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 61cdbd3eee21..cbc57b80da48 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -6,6 +6,7 @@ from vllm import LLM from vllm.distributed import cleanup_dist_env_and_memory +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams @@ -157,7 +158,7 @@ def test_chat_batch_failure_cleanup(llm_for_failure_test): batch_2 = [valid_msg, valid_msg] sampling_params = SamplingParams(temperature=0, max_tokens=10) - with pytest.raises(ValueError, match="maximum context length is"): + with pytest.raises(VLLMValidationError, match="maximum context length is"): llm.chat(batch_1, sampling_params=sampling_params) assert llm.llm_engine.get_num_unfinished_requests() == 0 diff --git a/tests/entrypoints/llm/test_prompt_validation.py b/tests/entrypoints/llm/test_prompt_validation.py index c17486d962f3..8dd55c6b10e0 100644 --- a/tests/entrypoints/llm/test_prompt_validation.py +++ b/tests/entrypoints/llm/test_prompt_validation.py @@ -5,17 +5,18 @@ import torch from vllm import LLM +from vllm.exceptions import VLLMValidationError def test_empty_prompt(): llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(ValueError, match="decoder prompt cannot be empty"): + with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"): llm.generate([""]) def test_out_of_vocab_token(): llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(ValueError, match="out of vocabulary"): + with pytest.raises(VLLMValidationError, match="out of vocabulary"): llm.generate({"prompt_token_ids": [999999]}) diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 3ece27234368..219ee7cd3875 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -212,6 +212,7 @@ class CarDescription(BaseModel): PARAMS_MODELS_BACKENDS_TOKENIZER_MODE, ) def test_structured_output( + request: pytest.FixtureRequest, backend: str, tokenizer_mode: str, model_name: str, @@ -242,6 +243,7 @@ def test_structured_output( speculative_config=speculative_config, **platform_args, ) + request.addfinalizer(llm.llm_engine.engine_core.shutdown) # # Test 1: Generate JSON output based on a provided schema diff --git a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py index 57bec9c1188a..0ea180712469 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py +++ b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py @@ -6,6 +6,7 @@ from tests.entrypoints.multimodal.conftest import managed_llm from vllm import LLM, SamplingParams from vllm.assets.image import ImageAsset +from vllm.exceptions import VLLMValidationError MODEL = "llava-hf/llava-1.5-7b-hf" PROMPT = "USER: \nDescribe this image briefly.\nASSISTANT:" @@ -42,7 +43,7 @@ def test_generate_with_embedding(llm: LLM): def test_raw_image_rejected(llm: LLM): """Raw image input is still rejected when limit=0.""" raw_image = ImageAsset("stop_sign").pil_image - with pytest.raises(ValueError, match=r"At most 0 image\(s\)"): + with pytest.raises(VLLMValidationError, match=r"At most 0 image\(s\)"): llm.generate( {"prompt": PROMPT, "multi_modal_data": {"image": raw_image}}, sampling_params=SamplingParams(max_tokens=16), diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py index fa0f141afee0..d806100e8e75 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py @@ -32,6 +32,7 @@ def server(): "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"audio": MAXIMUM_AUDIOS}), + "--no-enable-prefix-caching", ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index d005edc950cc..aac68334bf94 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -12,11 +12,12 @@ import safetensors import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoConfig, AutoTokenizer from tests.utils import RemoteOpenAIServer +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.serial_utils import tensor2base64 +from vllm.utils.torch_utils import is_torch_equal_or_newer QWEN2AUDIO_MODEL = "Qwen/Qwen2-Audio-7B-Instruct" @@ -125,11 +126,13 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]: content = "Describe this audio." tokenizer = AutoTokenizer.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True) - index_path = hf_hub_download(QWEN2AUDIO_MODEL, "model.safetensors.index.json") + index_path = hf_api().hf_hub_download( + QWEN2AUDIO_MODEL, "model.safetensors.index.json" + ) with open(index_path) as f: weight_map = json.load(f)["weight_map"] embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) - shard_path = hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key]) + shard_path = hf_api().hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key]) with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: embed_weight = f.get_tensor(embed_key) embed_layer = nn.Embedding.from_pretrained(embed_weight.to(QWEN2AUDIO_DTYPE)) @@ -142,8 +145,20 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]: @pytest.mark.asyncio @pytest.mark.parametrize( "audio_first", - [True, False], - ids=["audio_embeds-then-text", "text-then-audio_embeds"], + [ + pytest.param(True, id="audio_embeds-then-text"), + pytest.param( + False, + id="text-then-audio_embeds", + marks=pytest.mark.xfail( + condition=is_torch_equal_or_newer("2.12.0"), + reason="torch 2.12 regression: prompt_embeds output diverges " + "from raw-text when text precedes audio; " + "https://github.com/pytorch/pytorch/issues/184431", + strict=True, + ), + ), + ], ) async def test_text_content_and_prompt_embeds_match_with_audio_embeds( qwen2audio_client: openai.AsyncOpenAI, diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py index dbbed3c47127..c6ce165cf390 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -13,12 +13,12 @@ import safetensors import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoTokenizer from tests.utils import RemoteOpenAIServer from vllm.assets.image import ImageAsset from vllm.multimodal.utils import encode_image_url +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.serial_utils import tensor2base64 MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct" @@ -84,11 +84,11 @@ def aligned_content_and_embeds_b64() -> tuple[str, str]: content = "Describe this image." tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) - index_path = hf_hub_download(MODEL_NAME, "model.safetensors.index.json") + index_path = hf_api().hf_hub_download(MODEL_NAME, "model.safetensors.index.json") with open(index_path) as f: weight_map = json.load(f)["weight_map"] embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) - shard_path = hf_hub_download(MODEL_NAME, weight_map[embed_key]) + shard_path = hf_api().hf_hub_download(MODEL_NAME, weight_map[embed_key]) with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: embed_weight = f.get_tensor(embed_key) embed_layer = nn.Embedding.from_pretrained(embed_weight.to(MODEL_DTYPE)) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py b/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py index e285c8d3139e..eca23d0a48a2 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py @@ -6,17 +6,19 @@ import openai # use the official client for correctness check import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from tests.conftest import AudioTestAssets from tests.utils import RemoteOpenAIServer +from vllm.transformers_utils.repo_utils import hf_api # NOTE - the tests in this module are currently analogous to test_chat, but are # separated to avoid OOM killing due to module-scoped servers, since we # need a multimodal model for these tests. # Contains a modality specific lora alongside the base model -MULTIMODAL_MODEL_NAME = snapshot_download("microsoft/Phi-4-multimodal-instruct") +MULTIMODAL_MODEL_NAME = hf_api().snapshot_download( + "microsoft/Phi-4-multimodal-instruct" +) AUDIO_LORA_PATH = os.path.join(MULTIMODAL_MODEL_NAME, "speech-lora") ACTIVE_MM_LORA_RESPONSE = "Spoken text: The first words I spoke in the original chronograph, a little piece of practical poetry. Mary had a little lamb, it slept with quite a snow, and everywhere that Mary went, the lamb was sure to go." # noqa: E501 diff --git a/tests/entrypoints/openai/chat_completion/test_chat.py b/tests/entrypoints/openai/chat_completion/test_chat.py index dbfb48f2351e..5541b605d99e 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_chat.py @@ -18,6 +18,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams # any model with a chat template should work here @@ -27,9 +28,9 @@ @pytest.fixture(scope="module") def zephyr_lora_files(): """Download zephyr LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="typeof/zephyr-7b-beta-lora") + return hf_api().snapshot_download(repo_id="typeof/zephyr-7b-beta-lora") @pytest.fixture(scope="module") @@ -1074,7 +1075,7 @@ def test_chat_completion_request_n_parameter_exceeds_default_limit( max_tokens=10, ) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): request.to_sampling_params( max_tokens=10, default_sampling_params={}, @@ -1136,7 +1137,7 @@ def test_chat_completion_request_n_parameter_custom_limit( max_tokens=10, ) - with pytest.raises(ValueError, match="n must be at most 128"): + with pytest.raises(VLLMValidationError, match="n must be at most 128"): request_over.to_sampling_params( max_tokens=10, default_sampling_params={}, @@ -1160,7 +1161,7 @@ def test_chat_completion_request_n_parameter_massive_value( max_tokens=1, ) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): request.to_sampling_params( max_tokens=1, default_sampling_params={}, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion.py b/tests/entrypoints/openai/chat_completion/test_chat_completion.py index b5aa20448dfc..e0422a520aae 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_completion.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion.py @@ -158,3 +158,92 @@ async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> Non ], extra_body={"structured_outputs": {"grammar": ""}}, ) + + +# Decode-side token reuse for disaggregated serving. The router forwards the +# prefill stage's prompt token ids in kv_transfer_params so the decode stage +# skips re-tokenizing. + +TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}] +DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}] + + +@pytest.mark.asyncio +async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI): + """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize. + + The decode request carries different messages, so a response whose + prompt_token_ids match the forwarded ids proves the ids were used rather + than the request's own messages. Generated text is not compared across + requests because vLLM greedy decoding is not bitwise-reproducible. + """ + baseline = await client.chat.completions.create( + model=MODEL_NAME, + messages=TOKEN_IN_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={"return_token_ids": True}, + ) + reused_ids = baseline.prompt_token_ids + assert reused_ids + + decode = await client.chat.completions.create( + model=MODEL_NAME, + messages=DECODE_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={ + "kv_transfer_params": {"prompt_token_ids": reused_ids}, + "return_token_ids": True, + }, + ) + + # The engine saw the forwarded ids, not the decode request's own messages. + assert decode.prompt_token_ids == reused_ids + # text-out: reuse still yields a detokenized message. + assert decode.choices[0].message.content + + +@pytest.mark.asyncio +async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI): + """Decode-side token reuse streams chat-formatted text-out.""" + baseline = await client.chat.completions.create( + model=MODEL_NAME, + messages=TOKEN_IN_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={"return_token_ids": True}, + ) + reused_ids = baseline.prompt_token_ids + assert reused_ids + + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=DECODE_MESSAGES, + max_completion_tokens=16, + temperature=0, + stream=True, + extra_body={ + "kv_transfer_params": {"prompt_token_ids": reused_ids}, + "return_token_ids": True, + }, + ) + + content = "" + delta_token_ids: list[int] = [] + first_chunk = True + async for chunk in stream: + if first_chunk: + # prompt_token_ids arrives once, on the first chunk. + assert chunk.prompt_token_ids == reused_ids + first_chunk = False + if not chunk.choices: + continue + if chunk.choices[0].delta.content: + content += chunk.choices[0].delta.content + if tids := getattr(chunk.choices[0], "token_ids", None): + delta_token_ids.extend(tids) + + # streamed text-out, reconstructed from deltas, with generated token ids. + assert content + assert delta_token_ids diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 4b6be87ae5c2..9e805254a2f0 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -18,6 +17,7 @@ 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 +from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -479,7 +479,7 @@ def test_json_schema_response_format_missing_schema(): def test_structural_tag_response_format_invalid(format_value): """Malformed structural tags should be rejected during request validation.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): ChatCompletionRequest( @@ -493,7 +493,7 @@ def test_structural_tag_response_format_invalid(format_value): def test_batch_structural_tag_response_format_invalid(format_value): """Batch chat should reject malformed structural tags at request parsing.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): BatchChatCompletionRequest( @@ -507,7 +507,7 @@ def test_batch_structural_tag_response_format_invalid(format_value): def test_structured_outputs_structural_tag_invalid(structural_tag): """Malformed direct structured_outputs structural tags should be rejected.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid structured_outputs structural_tag", ): ChatCompletionRequest( @@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): messages=[{"role": "user", "content": "hello"}], structured_outputs={"structural_tag": structural_tag}, ) + + +@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"]) +def test_non_numeric_logprobs_rejected(field_name): + """A non-numeric logprobs value must be a clean 400 validation error, not a + TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): + ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hello"}], + **{field_name: "2"}, + ) diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py index aa04d787cccc..1eea6db48a7c 100644 --- a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -14,11 +14,11 @@ import math import pytest -from pydantic import ValidationError from tests.utils import RemoteOpenAIServer from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.exceptions import VLLMValidationError MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" @@ -87,7 +87,7 @@ def test_completion_request_decouples_top_k_from_explicit_token_ids(): def test_completion_rejects_explicit_token_ids_without_generated_tokens(): - with pytest.raises(ValidationError, match="no output tokens are generated"): + with pytest.raises(VLLMValidationError, match="no output tokens are generated"): CompletionRequest( model=MODEL_NAME, prompt="Hello", @@ -99,7 +99,7 @@ def test_completion_rejects_explicit_token_ids_without_generated_tokens(): def test_requests_reject_explicit_token_ids_with_beam_search(): - with pytest.raises(ValidationError, match="not supported with beam search"): + with pytest.raises(VLLMValidationError, match="not supported with beam search"): ChatCompletionRequest( model=MODEL_NAME, messages=[{"role": "user", "content": "Hello"}], @@ -108,7 +108,7 @@ def test_requests_reject_explicit_token_ids_with_beam_search(): use_beam_search=True, ) - with pytest.raises(ValidationError, match="not supported with beam search"): + with pytest.raises(VLLMValidationError, match="not supported with beam search"): CompletionRequest( model=MODEL_NAME, prompt="Hello", diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 2126acfe0273..1d6919e9c89f 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -2314,3 +2314,34 @@ async def result_generator(): f"Choice {choice_idx}: expected finish_reason='tool_calls', " f"got '{reasons[0]}'" ) + + +def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids(): + """The Harmony reuse branch honors ids forwarded in kv_transfer_params. + + A GPT-OSS server is impractical to stand up here, so this exercises the + branch directly on a harmony-configured renderer. + """ + engine = MockEngine() + engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss") + models = OpenAIServingModels(engine, BASE_MODEL_PATHS) + online_renderer = _build_online_renderer(engine, models.registry) + assert online_renderer.use_harmony + + request = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hi"}], + kv_transfer_params={ + "prompt_token_ids": [10, 20, 30], + "do_remote_prefill": True, + }, + ) + conversation, engine_inputs = online_renderer._make_request_with_harmony(request) + + assert conversation == [] + assert len(engine_inputs) == 1 + engine_input = engine_inputs[0] + assert engine_input["type"] == "token" + assert engine_input["prompt_token_ids"] == [10, 20, 30] + # The reuse key is consumed and other kv_transfer_params are preserved. + assert request.kv_transfer_params == {"do_remote_prefill": True} diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py index e66205b7df28..1b2b76dc0935 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py @@ -2,15 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from pydantic import ValidationError from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.exceptions import VLLMValidationError @pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5]) def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value): - with pytest.raises(ValidationError, match="thinking_token_budget"): + with pytest.raises(VLLMValidationError, match="thinking_token_budget"): ChatCompletionRequest.model_validate( { "model": "qwen", @@ -44,7 +44,7 @@ def test_chat_completion_request_accepts_minus_one_as_unlimited(): @pytest.mark.parametrize("raw_value", [0.6, 3.14, -2]) def test_completion_request_rejects_invalid_thinking_token_budget(raw_value): - with pytest.raises(ValidationError, match="thinking_token_budget"): + with pytest.raises(VLLMValidationError, match="thinking_token_budget"): CompletionRequest.model_validate( { "model": "qwen", diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index aa9e9c1d72e0..2f3db9f697d0 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.completion.protocol import CompletionRequest @@ -18,6 +17,7 @@ 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 +from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -430,7 +430,7 @@ def test_json_schema_response_format_missing_schema(): def test_structural_tag_response_format_invalid(format_value): """Malformed structural tags should be rejected during request validation.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): CompletionRequest( @@ -445,7 +445,7 @@ def test_structural_tag_response_format_invalid(format_value): def test_structured_outputs_structural_tag_invalid(structural_tag): """Malformed direct structured_outputs structural tags should be rejected.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid structured_outputs structural_tag", ): CompletionRequest( @@ -610,3 +610,16 @@ def test_bounded_prompt_embeds_list_allowed(self, monkeypatch): max_tokens=1, ) assert len(request.prompt_embeds) == 5 + + +@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"]) +def test_non_numeric_logprobs_rejected(field_name): + """A non-numeric logprobs value must be a clean 400 validation error, not a + TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + **{field_name: "2"}, + ) diff --git a/tests/entrypoints/openai/completion/test_prompt_validation.py b/tests/entrypoints/openai/completion/test_prompt_validation.py index 87c6b6e1668b..6c40037e07c9 100644 --- a/tests/entrypoints/openai/completion/test_prompt_validation.py +++ b/tests/entrypoints/openai/completion/test_prompt_validation.py @@ -13,6 +13,7 @@ from tests.utils import RemoteOpenAIServer from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.renderers.embed_utils import safe_load_prompt_embeds @@ -111,5 +112,5 @@ def test_disable_prompt_embeds(dtype: torch.dtype, seq_len: int, hidden_size: in buffer.seek(0) encoded_tensor = pybase64.b64encode(buffer.getvalue()) - with pytest.raises(ValueError, match="--enable-prompt-embeds"): + with pytest.raises(VLLMValidationError, match="--enable-prompt-embeds"): safe_load_prompt_embeds(model_config, encoded_tensor) diff --git a/tests/entrypoints/openai/responses/test_harmony.py b/tests/entrypoints/openai/responses/test_harmony.py index 2c70b06d8129..574f33a8f22b 100644 --- a/tests/entrypoints/openai/responses/test_harmony.py +++ b/tests/entrypoints/openai/responses/test_harmony.py @@ -13,7 +13,7 @@ import pytest import pytest_asyncio import requests -from openai import InternalServerError, NotFoundError, OpenAI +from openai import NotFoundError, OpenAI from openai_harmony import Message from tests.utils import RemoteOpenAIServer @@ -368,8 +368,12 @@ async def test_streaming_types( @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) async def test_function_calling_with_streaming_types( - pairs_of_event_types: dict[str, str], client: OpenAI, model_name: str + pairs_of_event_types: dict[str, str], + client: OpenAI, + model_name: str, + tool_choice: str, ): """Streaming event nesting for function-calling responses.""" @@ -382,6 +386,7 @@ def _has_function_events(evts: list) -> bool: validate_events=_has_function_events, input=[{"role": "user", "content": "What's the weather like in Paris today?"}], tools=[GET_WEATHER_SCHEMA], + tool_choice=tool_choice, temperature=0.0, ) @@ -558,7 +563,8 @@ async def test_reasoning_item(client: OpenAI, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -async def test_function_calling(client: OpenAI, model_name: str): +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +async def test_function_calling(client: OpenAI, model_name: str, tool_choice: str): tools = [GET_WEATHER_SCHEMA] response = await retry_for_tool_call( @@ -567,8 +573,9 @@ async def test_function_calling(client: OpenAI, model_name: str): expected_tool_type="function_call", input="What's the weather like in Paris today?", tools=tools, + tool_choice=tool_choice, temperature=0.0, - extra_body={"request_id": "test_function_calling_non_resp"}, + extra_body={"request_id": f"test_function_calling_non_resp_{tool_choice}"}, ) assert response.status == "completed" assert has_output_type(response, "function_call"), ( @@ -610,7 +617,10 @@ async def test_function_calling(client: OpenAI, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -async def test_function_calling_multi_turn(client: OpenAI, model_name: str): +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +async def test_function_calling_multi_turn( + client: OpenAI, model_name: str, tool_choice: str +): """Multi-tool, multi-turn function calling with retry at API level.""" tools = [ { @@ -635,6 +645,7 @@ async def test_function_calling_multi_turn(client: OpenAI, model_name: str): expected_tool_type="function_call", input="Help me plan a trip to a random place. And tell me the weather there.", tools=tools, + tool_choice=tool_choice, temperature=0.0, ) assert response.status == "completed" @@ -659,6 +670,7 @@ async def test_function_calling_multi_turn(client: OpenAI, model_name: str): } ], tools=tools, + tool_choice=tool_choice, previous_response_id=response.id, temperature=0.0, ) @@ -695,20 +707,6 @@ async def test_function_calling_multi_turn(client: OpenAI, model_name: str): ) -@pytest.mark.asyncio -@pytest.mark.parametrize("model_name", [MODEL_NAME]) -async def test_function_calling_required(client: OpenAI, model_name: str): - tools = [GET_WEATHER_SCHEMA] - - with pytest.raises(InternalServerError): - await client.responses.create( - model=model_name, - input="What's the weather like in Paris today?", - tools=tools, - tool_choice="required", - ) - - @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_system_message_with_tools(client: OpenAI, model_name: str): @@ -726,7 +724,10 @@ async def test_system_message_with_tools(client: OpenAI, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -async def test_function_calling_full_history(client: OpenAI, model_name: str): +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +async def test_function_calling_full_history( + client: OpenAI, model_name: str, tool_choice: str +): tools = [GET_WEATHER_SCHEMA] input_messages = [ @@ -739,6 +740,7 @@ async def test_function_calling_full_history(client: OpenAI, model_name: str): expected_tool_type="function_call", input=input_messages, tools=tools, + tool_choice=tool_choice, temperature=0.0, ) assert response.status == "completed" @@ -772,7 +774,10 @@ async def test_function_calling_full_history(client: OpenAI, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -async def test_function_calling_with_stream(client: OpenAI, model_name: str): +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +async def test_function_calling_with_stream( + client: OpenAI, model_name: str, tool_choice: str +): """Function calling via streaming, with retry for non-determinism.""" tools = [GET_WEATHER_SCHEMA] input_list = [ @@ -792,6 +797,7 @@ def _has_function_call(evts: list) -> bool: validate_events=_has_function_call, input=input_list, tools=tools, + tool_choice=tool_choice, temperature=0.0, ) @@ -853,8 +859,9 @@ def _has_function_call(evts: list) -> bool: @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) async def test_function_calling_no_code_interpreter_events( - client: OpenAI, model_name: str + client: OpenAI, model_name: str, tool_choice: str ): """Verify that function calls don't trigger code_interpreter events. @@ -880,6 +887,7 @@ def _has_function_call(evts: list) -> bool: validate_events=_has_function_call, input=input_list, tools=tools, + tool_choice=tool_choice, temperature=0.0, ) @@ -1048,8 +1056,9 @@ async def test_output_messages_enabled(client: OpenAI, model_name: str, server): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) async def test_function_call_with_previous_input_messages( - client: OpenAI, model_name: str + client: OpenAI, model_name: str, tool_choice: str ): """Multi-turn function calling using previous_input_messages.""" tools = [ @@ -1074,6 +1083,7 @@ async def test_function_call_with_previous_input_messages( expected_tool_type="function_call", input="What is the horoscope for Aquarius today?", tools=tools, + tool_choice=tool_choice, temperature=0.0, extra_body={"enable_response_messages": True}, max_output_tokens=1000, diff --git a/tests/entrypoints/openai/responses/test_parsable_context.py b/tests/entrypoints/openai/responses/test_parsable_context.py index 3a17da61276d..8a360fdd7105 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context.py +++ b/tests/entrypoints/openai/responses/test_parsable_context.py @@ -186,6 +186,12 @@ async def test_function_call_first_turn(client: OpenAI, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) +@pytest.mark.xfail( + reason=( + "MCP tools are not properly supported: tool name parameters " + "are not extracted for prompt rendering." + ), +) async def test_mcp_tool_call(client: OpenAI, model_name: str): """MCP tool calling with code_interpreter. diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index efbfb5c07e6a..2cda73d98a6e 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -20,7 +20,6 @@ _construct_message_from_response_item, construct_chat_messages_with_tool_call, construct_input_messages, - convert_tool_responses_to_completions_format, should_continue_final_message, ) @@ -116,27 +115,7 @@ def make_function_call_output( class TestResponsesUtils: - """Tests for convert_tool_responses_to_completions_format function.""" - - def test_convert_tool_responses_to_completions_format(self): - """Test basic conversion of a flat tool schema to nested format.""" - input_tool = { - "type": "function", - "name": "get_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string"}, - "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, - }, - "required": ["location", "unit"], - }, - } - - result = convert_tool_responses_to_completions_format(input_tool) - - assert result == {"type": "function", "function": input_tool} + """Tests for Responses API utils.""" def test_construct_chat_messages_with_tool_call(self): """Test construction of chat messages with tool calls.""" diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 5a68e3a9c0d4..6ede3f1f7f28 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -14,6 +14,7 @@ ResponsesRequest, ResponseTextConfig, ) +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import StructuredOutputsParams @@ -163,7 +164,7 @@ def test_structured_outputs_and_json_schema_conflict(self): text=text_config, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(VLLMValidationError) as exc_info: request.to_sampling_params(default_max_tokens=1000) assert "Cannot specify both structured_outputs and text.format" in str( diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index 1f382f61b797..ba55a58e7d3c 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -39,6 +39,7 @@ async def test_basic(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, input="What is 123 * 456?", + reasoning={"effort": "none"}, ) assert response is not None print("response: ", response) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 576e7ef16df4..df80053deb5e 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -201,6 +201,7 @@ def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch): monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None) monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None) monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None) + monkeypatch.setattr(dp_sup.envs, "VLLM_USE_RUST_FRONTEND", True, raising=False) monkeypatch.setattr( dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 2985c5395187..6d3fc2f44746 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -148,6 +148,7 @@ def test_openapi_stateless(case: schemathesis.Case): "/start_draft_weight_update", "/update_weights", "/finish_weight_update", + "/update_weight_version", ): return diff --git a/tests/entrypoints/openai/test_render_parity.py b/tests/entrypoints/openai/test_render_parity.py new file mode 100644 index 000000000000..57a113117fc0 --- /dev/null +++ b/tests/entrypoints/openai/test_render_parity.py @@ -0,0 +1,479 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cross-API HF render-input parity tests. + +Drives the real Chat Completions and Responses prep paths and captures the +arguments each passes to ``BaseRenderer.render_chat_async`` (via shared +``OnlineRenderer.preprocess_chat``): + + - conversation messages + - ChatParams fields that affect the HF template / media path + (template, content format, template kwargs including tools, + media_io_kwargs, mm_processor_kwargs) + +``TokenizeParams``, ``prompt_extras``, Harmony / GPT-OSS, prefill / continue, +prompt cache salt, and truncation are out of scope for this file. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import MagicMock + +import pytest +from openai.types.shared import Reasoning + +from vllm.config.multimodal import MultiModalConfig +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses +from vllm.inputs import tokens_input +from vllm.renderers.online_renderer import OnlineRenderer +from vllm.renderers.params import ChatParams + +_MODEL = "test-model" +_USER = [{"role": "user", "content": "Hello"}] + +_WEATHER_PARAMETERS = { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], +} + +_CHAT_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": _WEATHER_PARAMETERS, + }, +} + +_RESPONSES_WEATHER_TOOL = { + "type": "function", + "name": "get_weather", + "description": "Get the weather", + "parameters": _WEATHER_PARAMETERS, +} + + +@dataclass +class MockHFConfig: + model_type: str = "llama" + + +@dataclass +class MockModelConfig: + task = "generate" + runner_type = "generate" + model = _MODEL + tokenizer = _MODEL + trust_remote_code = False + tokenizer_mode = "auto" + max_model_len = 100 + tokenizer_revision = None + multimodal_config = MultiModalConfig() + hf_config = MockHFConfig() + hf_text_config = MockHFConfig() + logits_processors: list[str] | None = None + diff_sampling_param: dict | None = None + allowed_local_media_path: str = "" + allowed_media_domains: list[str] | None = None + encoder_config = None + generation_config: str = "auto" + override_generation_config: dict[str, Any] = field(default_factory=dict) + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) + skip_tokenizer_init = False + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + enable_prompt_embeds: bool = False + + def get_diff_sampling_param(self): + return self.diff_sampling_param or {} + + +@dataclass(frozen=True) +class CapturedRenderInputs: + """Args observed at ``render_chat_async`` (HF render boundary).""" + + messages: list[Any] + chat_params: ChatParams + + +class RenderCapture: + """Install a ``render_chat_async`` stub and record its HF-bound inputs.""" + + def __init__(self, online_renderer: OnlineRenderer) -> None: + self.online_renderer = online_renderer + self.captured: CapturedRenderInputs | None = None + + async def fake_render_chat_async( + conversations, + chat_params, + tok_params=None, + *, + prompt_extras=None, + skip_mm_cache=False, + ): + assert len(conversations) == 1 + self.captured = CapturedRenderInputs( + messages=list(conversations[0]), + chat_params=chat_params, + ) + return [list(conversations[0])], [ + tokens_input(prompt_token_ids=[0]), + ] + + online_renderer.renderer.render_chat_async = fake_render_chat_async + + def take(self) -> CapturedRenderInputs: + assert self.captured is not None + captured = self.captured + self.captured = None + return captured + + +async def _capture_chat( + online_renderer: OnlineRenderer, + request: ChatCompletionRequest, +) -> CapturedRenderInputs: + capture = RenderCapture(online_renderer) + result = await online_renderer.render_chat(request) + assert not isinstance(result, ErrorResponse), result + return capture.take() + + +async def _capture_responses( + serving: OpenAIServingResponses, + request: ResponsesRequest, +) -> CapturedRenderInputs: + capture = RenderCapture(serving.online_renderer) + await serving._make_request(request, prev_response=None) + return capture.take() + + +async def _assert_parity( + online_renderer: OnlineRenderer, + serving: OpenAIServingResponses, + *, + chat_kwargs: dict[str, Any], + responses_kwargs: dict[str, Any], +) -> None: + """Build paired requests, capture HF render inputs, and assert equality.""" + chat_req = ChatCompletionRequest(model=_MODEL, **chat_kwargs) + responses_req = ResponsesRequest(model=_MODEL, **responses_kwargs) + chat = await _capture_chat(online_renderer, chat_req) + responses = await _capture_responses(serving, responses_req) + + assert chat.messages == responses.messages + + chat_params = chat.chat_params + responses_params = responses.chat_params + assert chat_params.chat_template == responses_params.chat_template + assert ( + chat_params.chat_template_content_format + == responses_params.chat_template_content_format + ) + assert chat_params.media_io_kwargs == responses_params.media_io_kwargs + assert chat_params.mm_processor_kwargs == responses_params.mm_processor_kwargs + assert dict(chat_params.chat_template_kwargs) == dict( + responses_params.chat_template_kwargs + ) + + +def _weather_tools( + *, overrides: dict[str, Any] | None = None +) -> tuple[list[dict], list[dict]]: + """Return (chat_tools, responses_tools) for the shared weather function.""" + overrides = overrides or {} + chat_tool = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": _WEATHER_PARAMETERS, + **overrides, + }, + } + responses_tool = { + "type": "function", + "name": "get_weather", + "description": "Get the weather", + "parameters": _WEATHER_PARAMETERS, + **overrides, + } + return [chat_tool], [responses_tool] + + +@pytest.fixture +def model_config() -> MockModelConfig: + return MockModelConfig() + + +@pytest.fixture +def online_renderer(model_config: MockModelConfig, request) -> OnlineRenderer: + exclude_tools_when_tool_choice_none = getattr(request, "param", False) + renderer = MagicMock() + # Non-Mistral stub; only needed so render_chat / preprocess_chat can + # inspect tokenizer type before render_chat_async is mocked. + renderer.tokenizer = MagicMock() + + return OnlineRenderer( + model_config=model_config, # type: ignore[arg-type] + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + enable_auto_tools=True, + tool_parser="openai", + exclude_tools_when_tool_choice_none=exclude_tools_when_tool_choice_none, + ) + + +@pytest.fixture +def serving_responses( + model_config: MockModelConfig, + online_renderer: OnlineRenderer, +) -> OpenAIServingResponses: + engine_client = MagicMock() + engine_client.model_config = model_config + engine_client.renderer = online_renderer.renderer + engine_client.input_processor = MagicMock() + engine_client.vllm_config = MagicMock() + + return OpenAIServingResponses( + engine_client=engine_client, + models=MagicMock(), + online_renderer=online_renderer, + request_logger=None, + chat_template=online_renderer.chat_template, + chat_template_content_format=online_renderer.chat_template_content_format, + enable_auto_tools=True, + tool_parser="openai", + ) + + +@pytest.mark.asyncio +class TestConversationRenderParity: + async def test_multiturn_tool_calling(self, online_renderer, serving_responses): + """System/instructions, tool-call turn, and a follow-up user message.""" + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "Weather in NYC?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location":"NYC"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "72F", + }, + {"role": "user", "content": "Thanks"}, + ], + "tools": [_CHAT_WEATHER_TOOL], + "tool_choice": "auto", + }, + responses_kwargs={ + "instructions": "Be helpful.", + "input": [ + {"role": "user", "content": "Weather in NYC?"}, + { + "type": "function_call", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"location":"NYC"}', + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "72F", + }, + {"role": "user", "content": "Thanks"}, + ], + "tools": [_RESPONSES_WEATHER_TOOL], + "tool_choice": "auto", + }, + ) + + +@pytest.mark.asyncio +class TestToolsRenderParity: + @pytest.mark.parametrize( + "chat_tool_choice,responses_tool_choice", + [ + ("auto", "auto"), + ("required", "required"), + ( + {"type": "function", "function": {"name": "get_weather"}}, + {"type": "function", "name": "get_weather"}, + ), + ], + ids=["auto", "required", "named"], + ) + async def test_tools_with_tool_choice( + self, + online_renderer, + serving_responses, + chat_tool_choice, + responses_tool_choice, + ): + chat_tools, responses_tools = _weather_tools() + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": _USER, + "tools": chat_tools, + "tool_choice": chat_tool_choice, + }, + responses_kwargs={ + "input": _USER, + "tools": responses_tools, + "tool_choice": responses_tool_choice, + }, + ) + + @pytest.mark.parametrize( + "online_renderer", + [False, True], + indirect=True, + ids=["include_tools", "exclude_tools"], + ) + async def test_tools_with_tool_choice_none( + self, online_renderer, serving_responses + ): + chat_tools, responses_tools = _weather_tools() + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": _USER, + "tools": chat_tools, + "tool_choice": "none", + }, + responses_kwargs={ + "input": _USER, + "tools": responses_tools, + "tool_choice": "none", + }, + ) + + async def test_function_tool_optional_fields( + self, online_renderer, serving_responses + ): + """Optional FunctionDefinition fields dump the same on both APIs.""" + chat_tools, responses_tools = _weather_tools( + overrides={ + "strict": True, + "defer_loading": False, + "unrelated_extra": "should_be_ignored", + } + ) + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": _USER, + "tools": chat_tools, + "tool_choice": "auto", + }, + responses_kwargs={ + "input": _USER, + "tools": responses_tools, + "tool_choice": "auto", + }, + ) + + +@pytest.mark.asyncio +class TestReasoningRenderParity: + @pytest.mark.parametrize( + "effort", + ["none", "minimal", "low", "medium", "high", "xhigh"], + ) + async def test_reasoning_effort( + self, online_renderer, serving_responses, effort: str + ): + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": _USER, + "reasoning_effort": effort, + "tool_choice": "none", + }, + responses_kwargs={ + "input": _USER, + "reasoning": Reasoning(effort=effort), + "tool_choice": "none", + }, + ) + + async def test_explicit_enable_thinking_not_overridden( + self, online_renderer, serving_responses + ): + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": _USER, + "reasoning_effort": "high", + "chat_template_kwargs": {"enable_thinking": False}, + "tool_choice": "none", + }, + responses_kwargs={ + "input": _USER, + "reasoning": Reasoning(effort="high"), + "chat_template_kwargs": {"enable_thinking": False}, + "tool_choice": "none", + }, + ) + + +@pytest.mark.asyncio +class TestTemplateKwargsRenderParity: + async def test_passthrough_fields(self, online_renderer, serving_responses): + """chat_template_kwargs / media_io_kwargs / mm_processor_kwargs parity.""" + await _assert_parity( + online_renderer, + serving_responses, + chat_kwargs={ + "messages": [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "Hello"}, + ], + "tools": [_CHAT_WEATHER_TOOL], + "tool_choice": "auto", + "reasoning_effort": "medium", + "chat_template_kwargs": {"custom_flag": True}, + "media_io_kwargs": {"image": {"max_pixels": 512}}, + "mm_processor_kwargs": {"num_crops": 2}, + }, + responses_kwargs={ + "instructions": "Be helpful.", + "input": _USER, + "tools": [_RESPONSES_WEATHER_TOOL], + "tool_choice": "auto", + "reasoning": Reasoning(effort="medium"), + "chat_template_kwargs": {"custom_flag": True}, + "media_io_kwargs": {"image": {"max_pixels": 512}}, + "mm_processor_kwargs": {"num_crops": 2}, + }, + ) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index 5a7a8aab2a60..f06e43d143ea 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -9,8 +9,6 @@ from vllm import PoolingParams from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor from vllm.entrypoints.pooling.embed.protocol import ( - CohereEmbedContent, - CohereEmbedInput, CohereEmbedRequest, EmbeddingBatchChatInputRequest, EmbeddingBatchChatRequest, @@ -19,7 +17,10 @@ EmbeddingCompletionRequest, EmbeddingRequest, ) -from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.entrypoints.pooling.typing import ( + PoolingEngineInput, + PoolingServeContext, +) from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -410,6 +411,7 @@ class _FakeModelConfig: def _make_handler(cls): handler = object.__new__(EmbedIOProcessor) handler.model_config = cls._FakeModelConfig() + handler.enable_chunked_processing = True return handler @staticmethod @@ -421,15 +423,29 @@ def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]: } ) assert isinstance(request, EmbeddingCompletionRequest) + pooling_params = PoolingParams() return PoolingServeContext( request=request, - pooling_params=PoolingParams(), + pooling_params=pooling_params, model_name="test", request_id="embd-client-prompt-999-chunk-888", engine_inputs=[ - {"prompt_token_ids": [0, 1, 2, 3, 4]}, - {"prompt_token_ids": [10, 11]}, + PoolingEngineInput( + prompts={"prompt_token_ids": [0, 1, 2, 3, 4]}, + params=pooling_params, + lora_requests=None, + priorities=0, + ), + PoolingEngineInput( + prompts={"prompt_token_ids": [10, 11]}, + params=pooling_params, + lora_requests=None, + priorities=0, + ), ], + lora_request=None, + priorities=0, + prompt_extras=None, ) @staticmethod @@ -450,7 +466,7 @@ def test_aggregation_uses_metadata_not_request_id_parsing(self): handler = self._make_handler() ctx = self._make_context() - handler._pre_process_chunked(ctx) + handler.maybe_pre_process_chunked(ctx) assert ctx.prompt_request_ids == [ "embd-client-prompt-999-chunk-888-prompt-0-chunk-0", @@ -488,227 +504,3 @@ def test_aggregation_uses_metadata_not_request_id_parsing(self): ctx.final_res_batch[1].outputs.data, torch.tensor([9.0, 9.0]), ) - - -class TestPreProcessCohereOnline: - """Unit tests for EmbedIOProcessor._pre_process_cohere_online.""" - - @staticmethod - def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]: - return PoolingServeContext( - request=CohereEmbedRequest(model="test", **request_kwargs), - pooling_params=PoolingParams(), - model_name="test", - request_id="embd-test", - ) - - @staticmethod - def _make_handler(): - handler = object.__new__(EmbedIOProcessor) - handler._validate_input_type = lambda _input_type: None - return handler - - def test_text_only_without_task_prefix_uses_completion_path(self): - handler = self._make_handler() - ctx = self._make_context(texts=["hello"]) - calls: list[tuple[str, object]] = [] - - def preprocess_cmpl_online(request, prompt_input, prompt_embeds): - calls.append(("completion", prompt_input)) - return ["completion"] - - handler._get_task_instruction_prefix = lambda _input_type: None - handler._has_chat_template = lambda: False - handler._preprocess_cmpl_online = preprocess_cmpl_online - handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail( - "text-only request should not require chat rendering" - ) - - handler._pre_process_cohere_online(ctx) - - assert ctx.engine_inputs == ["completion"] - assert calls == [("completion", ["hello"])] - - def test_text_only_falls_back_to_prefixed_completion_without_template(self): - handler = self._make_handler() - ctx = self._make_context(texts=["hello"], input_type="query") - calls: list[tuple[str, object]] = [] - - def preprocess_cmpl(request, prompt_input, prompt_embeds): - calls.append(("completion", prompt_input)) - return ["fallback"] - - handler._get_task_instruction_prefix = lambda _input_type: "query: " - handler._has_chat_template = lambda: False - handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail( - "chat rendering should be skipped without a template" - ) - handler._preprocess_cmpl_online = preprocess_cmpl - - handler._pre_process_cohere_online(ctx) - - assert ctx.engine_inputs == ["fallback"] - assert calls == [("completion", ["query: hello"])] - - def test_text_only_with_template_uses_chat_path(self): - handler = self._make_handler() - ctx = self._make_context(texts=["hello"], input_type="query") - calls: list[tuple[str, object]] = [] - - def batch_render_chat( - request, - all_messages, - truncate_prompt_tokens, - truncation_side, - ): - calls.append( - ( - "chat", - { - "request": request, - "all_messages": all_messages, - "truncate_prompt_tokens": truncate_prompt_tokens, - "truncation_side": truncation_side, - }, - ) - ) - return ["chat"] - - handler._get_task_instruction_prefix = lambda _input_type: "query: " - handler._has_chat_template = lambda: True - handler._batch_render_chat = batch_render_chat - handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail( - "completion path should be skipped when a template exists" - ) - - handler._pre_process_cohere_online(ctx) - - assert ctx.engine_inputs == ["chat"] - assert calls == [ - ( - "chat", - { - "request": ctx.request, - "all_messages": [ - handler._mixed_input_to_messages( - CohereEmbedInput( - content=[CohereEmbedContent(type="text", text="hello")] - ), - task_prefix="query: ", - ) - ], - "truncate_prompt_tokens": -1, - "truncation_side": None, - }, - ) - ] - - -class TestPreProcessOpenAIEmbeddingChatOnline: - """Unit tests for OpenAI embedding chat preprocessing.""" - - class _FakeModelConfig: - max_model_len = 128 - encoder_config: dict[str, object] = {} - pooler_config = None - multimodal_config = None - is_encoder_decoder = False - - class _FakeRenderer: - tokenizer = object() - - def __init__(self): - self.calls = [] - - def render_chat( - self, - all_messages, - chat_params, - tok_params, - prompt_extras=None, - ): - self.calls.append( - { - "all_messages": all_messages, - "chat_params": chat_params, - "tok_params": tok_params, - "prompt_extras": prompt_extras, - } - ) - return all_messages, [ - {"prompt_token_ids": [index]} for index, _ in enumerate(all_messages) - ] - - @classmethod - def _make_handler(cls, renderer): - handler = object.__new__(EmbedIOProcessor) - handler.renderer = renderer - handler.model_config = cls._FakeModelConfig() - handler.chat_template = "template" - handler.chat_template_content_format = "auto" - handler.trust_request_chat_template = False - handler.enable_chunked_processing = False - return handler - - @staticmethod - def _make_context( - request: ( - EmbeddingChatRequest - | EmbeddingBatchChatRequest - | EmbeddingChatInputRequest - | EmbeddingBatchChatInputRequest - ), - ) -> PoolingServeContext[ - EmbeddingChatRequest - | EmbeddingBatchChatRequest - | EmbeddingChatInputRequest - | EmbeddingBatchChatInputRequest - ]: - return PoolingServeContext( - request=request, - pooling_params=PoolingParams(), - model_name="test", - request_id="embd-test", - ) - - def test_chat_template_kwargs_forwarded_for_batched_input_messages(self): - request = TypeAdapter(EmbeddingRequest).validate_python( - { - "model": "test", - "input": [ - [{"role": "user", "content": "hello"}], - [{"role": "user", "content": "goodbye"}], - ], - "add_generation_prompt": True, - "chat_template_kwargs": {"instruction": "Represent the query: "}, - "mm_processor_kwargs": {"max_pixels": 1}, - "cache_salt": "salt", - } - ) - assert isinstance(request, EmbeddingBatchChatInputRequest) - - renderer = self._FakeRenderer() - handler = self._make_handler(renderer) - ctx = self._make_context(request) - - handler.pre_process_online(ctx) - - assert ctx.engine_inputs == [ - {"prompt_token_ids": [0]}, - {"prompt_token_ids": [1]}, - ] - assert len(renderer.calls) == 1 - - call = renderer.calls[0] - assert call["all_messages"] == request.messages - assert call["prompt_extras"] == { - "mm_processor_kwargs": {"max_pixels": 1}, - "cache_salt": "salt", - } - - chat_template_kwargs = call["chat_params"].chat_template_kwargs - assert chat_template_kwargs["instruction"] == "Represent the query: " - assert chat_template_kwargs["add_generation_prompt"] is True - assert chat_template_kwargs["continue_final_message"] is False - assert "tools" not in chat_template_kwargs - assert chat_template_kwargs["tokenize"] is False diff --git a/tests/entrypoints/pooling/embed/test_online_long_text.py b/tests/entrypoints/pooling/embed/test_online_long_text.py index eaefbc02383f..c1dc5bf1ed09 100644 --- a/tests/entrypoints/pooling/embed/test_online_long_text.py +++ b/tests/entrypoints/pooling/embed/test_online_long_text.py @@ -20,7 +20,9 @@ def _generate_random_text(word_count: int) -> str: - """Generate random text with approximately the specified word count.""" + """Generate deterministic text with approximately the specified word count.""" + rng = random.Random(word_count) + # Common English words with focus on verbs and nouns for realistic text common_words = [ # Essential articles and pronouns (minimal) @@ -178,7 +180,7 @@ def _generate_random_text(word_count: int) -> str: words = [] for _ in range(word_count): - words.append(random.choice(common_words)) + words.append(rng.choice(common_words)) # Add some punctuation for more realistic text text = " ".join(words) @@ -187,7 +189,7 @@ def _generate_random_text(word_count: int) -> str: result = [] for i, word in enumerate(words_list): result.append(word) - if (i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1: + if (i + 1) % rng.randint(10, 20) == 0 and i < len(words_list) - 1: result[-1] += "." return " ".join(result) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py index df79a387afd2..56e83de3f74f 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import weakref -from types import SimpleNamespace import pytest import torch @@ -10,10 +9,7 @@ from tests.models.utils import softmax from vllm import LLM, PoolingParams from vllm.distributed import cleanup_dist_env_and_memory -from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor -from vllm.entrypoints.pooling.scoring.typing import ScoringData from vllm.platforms import current_platform -from vllm.renderers import TokenizeParams MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls" PROMPT = "The chef prepared a delicious meal." @@ -145,45 +141,6 @@ def test_max_tokens_per_doc(llm: LLM): assert with_limit_tokens < no_limit_tokens -def test_token_type_ids_follow_post_tokenization(): - processor = object.__new__(CrossEncoderIOProcessor) - processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1) - processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt) - processor.model_config = None - processor.get_score_prompt = lambda **_: ( - "", - { - "prompt_token_ids": list(range(32)), - "token_type_ids": [0] * 16 + [1] * 16, - }, - ) - - engine_inputs, pooling_params = processor._pre_process( - ScoringData(data_1=["query"], data_2=["document"]), - TokenizeParams( - max_total_tokens=None, - truncate_prompt_tokens=16, - truncation_side="left", - ), - PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}), - ) - - assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32)) - assert pooling_params[0].extra_kwargs == { - "cache_salt": "salt", - "compressed_token_type_ids": 0, - } - - engine_inputs, pooling_params = processor._pre_process( - ScoringData(data_1=["query"], data_2=["document"]), - TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40), - PoolingParams(task="classify"), - ) - - assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8 - assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16} - - def test_pooling_params(llm: LLM): def get_outputs(use_activation): outputs = llm.score( diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index 8c1d75cd762c..aae69fe217db 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -473,7 +473,7 @@ async def test_rerank_api_instruction_field( async def test_rerank_api_instruction_field_matches_chat_template_kwargs( server: tuple[RemoteOpenAIServer, str], ): - remote_server, _ = server + remote_server, backend = server doc_list = [ document, @@ -514,4 +514,6 @@ async def test_rerank_api_instruction_field_matches_chat_template_kwargs( kwargs_scores = [ r.relevance_score for r in sorted(kwargs_rerank.results, key=lambda x: x.index) ] - assert field_scores == pytest.approx(kwargs_scores) + assert field_scores == pytest.approx( + kwargs_scores, rel=get_tol(backend), abs=get_abs_tol(backend) + ) diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py new file mode 100644 index 000000000000..c2fe62f609df --- /dev/null +++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for JinaRankingIOProcessor online request building.""" + +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor +from vllm.entrypoints.pooling.scoring.protocol import RerankRequest +from vllm.entrypoints.pooling.scoring.typing import ScoringData + +pytestmark = pytest.mark.skip_global_cleanup + + +def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch): + """The proxy request handed to the base factory must carry + truncate_prompt_tokens/truncation_side from the real request. + + JinaRankingIOProcessor swaps ctx.request for a proxy + PoolingCompletionRequest before delegating to the base factory, which + reads truncation off ctx.request. Dropping the fields on the proxy + silently disables truncate_prompt_tokens for Jina rerank/score. + """ + proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + proc.valid_inputs_online = MagicMock( + return_value=ScoringData(data_1=["query"], data_2=["doc"]) + ) + proc._get_token_limits = MagicMock(return_value=(0, 0)) + proc.ensure_str = MagicMock(side_effect=lambda data: list(data)) + proc.format_docs_prompts_func = MagicMock(return_value="formatted prompt") + + captured: dict[str, object] = {} + + def _spy_base(self, ctx): + captured["truncate_prompt_tokens"] = ctx.request.truncate_prompt_tokens + captured["truncation_side"] = ctx.request.truncation_side + return [] + + monkeypatch.setattr(PoolingIOProcessor, "get_request_factory_online", _spy_base) + + request = RerankRequest( + model="m", + query="query", + documents=["doc"], + truncate_prompt_tokens=512, + truncation_side="left", + ) + ctx = MagicMock() + ctx.request = request + ctx.prompt_extras = None + + proc.get_request_factory_online(ctx) + + assert captured["truncate_prompt_tokens"] == 512 + assert captured["truncation_side"] == "left" + # The real request is restored after delegating. + assert ctx.request is request diff --git a/tests/entrypoints/pooling/scoring/util.py b/tests/entrypoints/pooling/scoring/util.py index 8aab9cd10692..dbf4b3dc3acd 100644 --- a/tests/entrypoints/pooling/scoring/util.py +++ b/tests/entrypoints/pooling/scoring/util.py @@ -6,7 +6,6 @@ import pybase64 as base64 import torch import torch.nn.functional as F -from huggingface_hub import hf_hub_download from PIL import Image from safetensors.torch import load_file from transformers import AutoModel, AutoTokenizer @@ -18,6 +17,7 @@ ) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score +from vllm.transformers_utils.repo_utils import hf_api class ColBERTScoringHfRunner(torch.nn.Module): @@ -38,7 +38,7 @@ def __init__(self, model_name, linear_weights_key): ).to(self.device) self.model.eval() - path = hf_hub_download(model_name, filename="model.safetensors") + path = hf_api().hf_hub_download(model_name, filename="model.safetensors") weights = load_file(path) self.linear_weight = weights[linear_weights_key].to(self.device).float() diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py new file mode 100644 index 000000000000..014a8d7bc6b6 --- /dev/null +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -0,0 +1,765 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for streaming derender. + +Tests are split into two layers: + +1. Unit tests (no server): covers ``_detokenize_delta`` correctness + (chunked == one-shot) and ``derender_completion_stream`` / + ``derender_chat_stream`` logic via a real tokenizer on a tiny model. + +2. Integration tests (require a running render server): covers the full + HTTP round-trip through the streaming endpoint. Marked with + ``@pytest.mark.asyncio`` and gated by the ``server`` / ``client`` + fixtures from the sibling ``test_derender.py``. +""" + +import pytest +import pytest_asyncio + +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderStreamState, + GenerateResponseStreamChoice, + GenerateStreamResponse, +) + +MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" + + +# --------------------------------------------------------------------------- +# Helpers shared across tests +# --------------------------------------------------------------------------- + + +def _make_stream_chunk( + token_ids: list[int], + index: int = 0, + finish_reason: str | None = None, + request_id: str = "test-req", + usage: dict | None = None, +) -> GenerateStreamResponse: + """Build a GenerateStreamResponse SSE chunk.""" + from vllm.entrypoints.openai.engine.protocol import UsageInfo + + return GenerateStreamResponse( + request_id=request_id, + choices=[ + GenerateResponseStreamChoice( + index=index, + token_ids=token_ids, + finish_reason=finish_reason, + ) + ], + usage=UsageInfo(**usage) if usage else None, + ) + + +def _make_usage_chunk( + completion_tokens: int, + prompt_tokens: int = 0, + request_id: str = "test-req", +) -> GenerateStreamResponse: + """Build a usage only final SSE chunk (empty choices).""" + from vllm.entrypoints.openai.engine.protocol import UsageInfo + + return GenerateStreamResponse( + request_id=request_id, + choices=[], + usage=UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +# --------------------------------------------------------------------------- +# Unit tests — no running server +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def tokenizer(): + """Load the tiny tokenizer used across unit tests.""" + from vllm.tokenizers import get_tokenizer + + return get_tokenizer(MODEL_NAME) + + +@pytest.fixture(scope="module") +def derenderer(tokenizer): + """Construct a minimal OnlineDerenderer backed by a stub renderer.""" + from unittest.mock import MagicMock + + from vllm.renderers.online_derenderer import OnlineDerenderer + + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + + model_config = MagicMock() + model_config.hf_config.model_type = "llama" + model_config.model = MODEL_NAME + + return OnlineDerenderer( + model_config=model_config, + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + trust_request_chat_template=False, + enable_auto_tools=False, + tool_parser=None, + reasoning_parser=None, + ) + + +class TestDetokenizeDelta: + """_detokenize_delta: chunked decode must equal one shot decode.""" + + def _one_shot(self, tokenizer, token_ids: list[int]) -> str: + return tokenizer.decode(token_ids, skip_special_tokens=True) + + def _chunked(self, derenderer, tokenizer, chunks: list[list[int]]) -> str: + state = DerenderStreamState() + parts: list[str] = [] + for delta in chunks: + text, state = derenderer._detokenize_delta( + tokenizer, delta, state, skip_special_tokens=True + ) + parts.append(text) + return "".join(parts) + + def test_single_chunk(self, derenderer, tokenizer): + """All tokens in one chunk == one shot decode.""" + token_ids = tokenizer.encode("Hello world")[:8] + assert self._chunked(derenderer, tokenizer, [token_ids]) == self._one_shot( + tokenizer, token_ids + ) + + def test_two_equal_chunks(self, derenderer, tokenizer): + """Split in half and reassemble == one shot.""" + token_ids = tokenizer.encode("Hello world from streaming derender")[:12] + mid = len(token_ids) // 2 + chunks = [token_ids[:mid], token_ids[mid:]] + assert self._chunked(derenderer, tokenizer, chunks) == self._one_shot( + tokenizer, token_ids + ) + + def test_single_token_per_chunk(self, derenderer, tokenizer): + """One token per chunk (most granular streaming) == one shot.""" + token_ids = tokenizer.encode("incremental detokenization test")[:10] + chunks = [[t] for t in token_ids] + assert self._chunked(derenderer, tokenizer, chunks) == self._one_shot( + tokenizer, token_ids + ) + + def test_empty_delta_passthrough(self, derenderer, tokenizer): + """An empty delta (usage only chunk) emits empty string and preserves state.""" + token_ids = tokenizer.encode("Hello")[:4] + _, state = derenderer._detokenize_delta( + tokenizer, token_ids, DerenderStreamState(), skip_special_tokens=True + ) + text, new_state = derenderer._detokenize_delta( + tokenizer, [], state, skip_special_tokens=True + ) + assert text == "" + assert new_state.prev_tokens == state.prev_tokens + assert new_state.prefix_offset == state.prefix_offset + assert new_state.read_offset == state.read_offset + + def test_multibyte_char_split_across_chunks(self, derenderer, tokenizer): + """A CJK/emoji char straddling chunk boundaries == one shot. + + Regression test for held back trailing incomplete UTF-8 byte + sequences being dropped when the rebuild window marks them as + already read (see #46159). + """ + token_ids = tokenizer.encode("Hello ✅ world 日本語")[:16] + chunks = [[t] for t in token_ids] + assert self._chunked(derenderer, tokenizer, chunks) == self._one_shot( + tokenizer, token_ids + ) + + def test_state_carries_across_calls(self, derenderer, tokenizer): + """Decode state threads across calls. Text still matches one shot.""" + t1 = tokenizer.encode("Hello")[:2] + t2 = tokenizer.encode(" world")[:2] + state = DerenderStreamState() + text1, state = derenderer._detokenize_delta(tokenizer, t1, state) + text2, state = derenderer._detokenize_delta(tokenizer, t2, state) + assert text1 + text2 == self._one_shot(tokenizer, t1 + t2) + # Offsets are rebased to the carried tail each chunk + assert state.prefix_offset == 0 + + def test_state_window_stays_bounded(self, derenderer, tokenizer): + """prev_tokens must not grow with the number of chunks (bounded transport). + + Guards that the carried decode window is a small constant + tail, so cumulative ``stream_state`` transport is O(n) and not O(n^2). + """ + token_ids = tokenizer.encode( + "a reasonably long ascii stream of tokens used to exercise the " + "window bound across many single token chunks so the carried " + "state cannot grow linearly with the generation length" + ) + assert len(token_ids) > 32 + state = DerenderStreamState() + max_window = 0 + for tok in token_ids: + _, state = derenderer._detokenize_delta(tokenizer, [tok], state) + max_window = max(max_window, len(state.prev_tokens)) + # Bounded by a small constant independent of len(token_ids) + assert max_window <= 32 + + def test_n_independent_streams_same_result(self, derenderer, tokenizer): + """N parallel streams with the same token sequence give the same text.""" + token_ids = tokenizer.encode("parallel streams")[:8] + mid = len(token_ids) // 2 + + results = [] + for _ in range(3): + state = DerenderStreamState() + text, state = derenderer._detokenize_delta( + tokenizer, token_ids[:mid], state + ) + text2, _ = derenderer._detokenize_delta(tokenizer, token_ids[mid:], state) + results.append(text + text2) + + assert len(set(results)) == 1, "All independent streams must produce same text" + assert results[0] == self._one_shot(tokenizer, token_ids) + + +class TestDerenderCompletionStream: + """derender_completion_stream: streaming output parity with one shot.""" + + @pytest.mark.asyncio + async def test_chunked_equals_oneshot(self, derenderer, tokenizer): + """Sum of streaming text chunks == one shot tokenizer.decode.""" + token_ids = tokenizer.encode("streaming completion test")[:10] + mid = len(token_ids) // 2 + + state = DerenderStreamState() + chunk1, state = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[:mid]), + state=state, + ) + chunk2, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[mid:], finish_reason="stop"), + state=state, + ) + + streamed_text = chunk1.choices[0].text + chunk2.choices[0].text + one_shot = tokenizer.decode(token_ids, skip_special_tokens=True) + assert streamed_text == one_shot + + @pytest.mark.asyncio + async def test_usage_chunk_passthrough(self, derenderer, tokenizer): + """Usage only final chunk (empty choices) is passed through correctly.""" + usage_chunk = _make_usage_chunk(completion_tokens=10, prompt_tokens=5) + chunk, state = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=usage_chunk, + ) + assert chunk.choices == [] + assert chunk.usage is not None + assert chunk.usage.completion_tokens == 10 + assert chunk.usage.prompt_tokens == 5 + + @pytest.mark.asyncio + async def test_prompt_tokens_in_usage(self, derenderer, tokenizer): + """prompt_tokens is correctly forwarded into usage on a usage chunk.""" + token_ids = tokenizer.encode("hello")[:3] + usage_chunk = _make_usage_chunk( + completion_tokens=len(token_ids), prompt_tokens=7 + ) + chunk, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=usage_chunk, + prompt_tokens=7, + ) + assert chunk.usage is not None + assert chunk.usage.prompt_tokens == 7 + + @pytest.mark.asyncio + async def test_none_state_initialises_correctly(self, derenderer, tokenizer): + """Passing state=None (first call) initialises an empty DerenderStreamState.""" + token_ids = tokenizer.encode("hello")[:4] + chunk, state = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids), + state=None, + ) + assert isinstance(state, DerenderStreamState) + assert chunk.choices[0].text == tokenizer.decode( + token_ids, skip_special_tokens=True + ) + + @pytest.mark.asyncio + async def test_skip_special_tokens_threaded(self, derenderer, tokenizer): + """completion_request.skip_special_tokens is honored (not hardcoded True).""" + from vllm.entrypoints.openai.completion.protocol import CompletionRequest + + eos = tokenizer.eos_token_id + if eos is None: + pytest.skip("tokenizer has no eos token to exercise special stripping") + token_ids = tokenizer.encode("hi")[:2] + [eos] + + async def _text(skip: bool) -> str: + req = CompletionRequest( + model=MODEL_NAME, prompt="x", skip_special_tokens=skip + ) + chunk, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids), + completion_request=req, + ) + return chunk.choices[0].text + + # skip=False must retain the special token; skip=True must strip it. + assert await _text(False) != await _text(True) + + @pytest.mark.asyncio + async def test_finish_reason_forwarded(self, derenderer, tokenizer): + """finish_reason from the generate chunk reaches the derendered choice.""" + token_ids = tokenizer.encode("done")[:2] + chunk, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids, finish_reason="length"), + ) + assert chunk.choices[0].finish_reason == "length" + + +class TestDerenderChatStream: + """derender_chat_stream: plain detok branch (no parser).""" + + @pytest.mark.asyncio + async def test_role_on_first_chunk_only(self, derenderer, tokenizer): + """role='assistant' appears in the first chunk, not subsequent ones.""" + token_ids = tokenizer.encode("hello world")[:6] + mid = len(token_ids) // 2 + + state = DerenderStreamState() + chunk1, state = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[:mid]), + state=state, + ) + chunk2, _ = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[mid:], finish_reason="stop"), + state=state, + ) + + assert chunk1.choices[0].delta.role == "assistant" + assert chunk2.choices[0].delta.role is None + + @pytest.mark.asyncio + async def test_chunked_equals_oneshot(self, derenderer, tokenizer): + """Sum of streaming content deltas == one shot decode.""" + token_ids = tokenizer.encode("streaming chat derender text")[:10] + mid = len(token_ids) // 2 + + state = DerenderStreamState() + chunk1, state = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[:mid]), + state=state, + ) + chunk2, _ = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids[mid:]), + state=state, + ) + + streamed = (chunk1.choices[0].delta.content or "") + ( + chunk2.choices[0].delta.content or "" + ) + one_shot = tokenizer.decode(token_ids, skip_special_tokens=True) + assert streamed == one_shot + + @pytest.mark.asyncio + @pytest.mark.parametrize("with_chat_request", [True, False]) + async def test_parser_raises_not_implemented(self, tokenizer, with_chat_request): + """Stream chat with a parser active must fail closed (NotImplementedError). + + Checks that a parser configured model must 501 even when ``chat_request`` is + omitted, otherwise reasoning/tool markup would leak into ``delta.content` + via the plain detok fallback. + """ + from unittest.mock import MagicMock + + from vllm.renderers.online_derenderer import OnlineDerenderer + + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + + model_config = MagicMock() + model_config.hf_config.model_type = "llama" + model_config.model = MODEL_NAME + + # Construct a derenderer WITH a tool/reasoning parser active. + # ParserManager.get_parser returns None when no parser is named, so + # we inject a mock parser class directly. + dr = OnlineDerenderer( + model_config=model_config, + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + ) + dr.parser = MagicMock() # simulate a parser being active + + chat_request = MagicMock() if with_chat_request else None + token_ids = tokenizer.encode("hello")[:3] + + with pytest.raises(NotImplementedError): + await dr.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(token_ids), + state=None, + chat_request=chat_request, + ) + + +class TestDerenderStreamStateValidation: + """DerenderStreamState rejects malformed caller supplied offsets/lengths.""" + + def test_negative_prefix_offset_rejected(self): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + DerenderStreamState(prefix_offset=-1) + + def test_negative_read_offset_rejected(self): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + DerenderStreamState(read_offset=-1) + + def test_prev_tokens_over_cap_rejected(self): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + DerenderStreamState(prev_tokens=["a"] * 1025) + + def test_prev_tokens_at_cap_accepted(self): + state = DerenderStreamState(prev_tokens=["a"] * 1024) + assert len(state.prev_tokens) == 1024 + + +class TestServingDerenderStreamErrorHandling: + """Malformed stream_state must surface as 400 and not an unhandled 500.""" + + def _make_serving(self, side_effect: Exception): + from unittest.mock import AsyncMock, MagicMock + + from vllm.entrypoints.scale_out.derender.serving import ServingDerender + + models = MagicMock() + models.is_base_model.return_value = True + models.model_config = MagicMock() + + online_derenderer = MagicMock() + online_derenderer.derender_completion_stream = AsyncMock( + side_effect=side_effect + ) + online_derenderer.derender_chat_stream = AsyncMock(side_effect=side_effect) + + return ServingDerender(models=models, online_derenderer=online_derenderer) + + @pytest.mark.asyncio + @pytest.mark.parametrize("exc", [KeyError("bad byte"), IndexError("oob")]) + async def test_completion_stream_bad_state_returns_400(self, exc): + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderCompletionStreamRequest, + ) + + serving = self._make_serving(exc) + request = DerenderCompletionStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + stream_state=DerenderStreamState(), + ) + result = await serving.derender_completion_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + + @pytest.mark.asyncio + @pytest.mark.parametrize("exc", [KeyError("bad byte"), IndexError("oob")]) + async def test_chat_stream_bad_state_returns_400(self, exc): + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(exc) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + stream_state=DerenderStreamState(), + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + + +# --------------------------------------------------------------------------- +# Integration tests — require a live render server +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def server(): + from tests.utils import RemoteLaunchRenderServer + + with RemoteLaunchRenderServer(MODEL_NAME, []) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + import httpx + + async with httpx.AsyncClient( + base_url=server.url_for(""), timeout=30.0 + ) as http_client: + yield http_client + + +async def _render_chat(client) -> dict: + """Render a minimal chat request and return the GenerateRequest dict.""" + + resp = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [{"role": "user", "content": "Hello"}], + }, + ) + assert resp.status_code == 200 + return resp.json() + + +@pytest.mark.asyncio +async def test_streaming_completion_derender_roundtrip(client): + """Streaming completions derender: chunked text == non streaming text.""" + gen_req = await _render_chat(client) + token_ids: list[int] = gen_req["token_ids"][:8] + mid = len(token_ids) // 2 + chunk1_ids, chunk2_ids = token_ids[:mid], token_ids[mid:] + + # Non streaming baseline. + non_stream_resp = await client.post( + "/v1/completions/derender", + json={ + "model": MODEL_NAME, + "generate_responses": [ + { + "request_id": "test-ns", + "choices": [ + { + "index": 0, + "token_ids": token_ids, + "finish_reason": "stop", + } + ], + } + ], + }, + ) + assert non_stream_resp.status_code == 200 + expected_text = non_stream_resp.json()["choices"][0]["text"] + + # Streaming call 1. + r1 = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "test-s", + "choices": [ + {"index": 0, "token_ids": chunk1_ids, "finish_reason": None} + ], + }, + "stream_state": None, + }, + ) + assert r1.status_code == 200 + d1 = r1.json() + text1 = d1["chunk"]["choices"][0]["text"] + state1 = d1["stream_state"] + + # Streaming call 2 (final chunk). + r2 = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "test-s", + "choices": [ + {"index": 0, "token_ids": chunk2_ids, "finish_reason": "stop"} + ], + }, + "stream_state": state1, + }, + ) + assert r2.status_code == 200 + text2 = r2.json()["chunk"]["choices"][0]["text"] + + assert text1 + text2 == expected_text + + +@pytest.mark.asyncio +async def test_streaming_chat_derender_roundtrip(client): + """Streaming chat derender (plain detok): chunked text == non streaming text.""" + gen_req = await _render_chat(client) + token_ids: list[int] = gen_req["token_ids"][:8] + mid = len(token_ids) // 2 + chunk1_ids, chunk2_ids = token_ids[:mid], token_ids[mid:] + + # Non streaming baseline. + ns = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": { + "request_id": "test-ns", + "choices": [ + { + "index": 0, + "token_ids": token_ids, + "finish_reason": "stop", + } + ], + }, + }, + ) + assert ns.status_code == 200 + expected_content = ns.json()["choices"][0]["message"]["content"] + + # Streaming call 1. + r1 = await client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "test-s", + "choices": [ + {"index": 0, "token_ids": chunk1_ids, "finish_reason": None} + ], + }, + "stream_state": None, + }, + ) + assert r1.status_code == 200 + d1 = r1.json() + text1 = d1["chunk"]["choices"][0]["delta"].get("content") or "" + state1 = d1["stream_state"] + # role=assistant on the first chunk + assert d1["chunk"]["choices"][0]["delta"].get("role") == "assistant" + + # Streaming call 2. + r2 = await client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "test-s", + "choices": [ + {"index": 0, "token_ids": chunk2_ids, "finish_reason": "stop"} + ], + }, + "stream_state": state1, + }, + ) + assert r2.status_code == 200 + d2 = r2.json() + text2 = d2["chunk"]["choices"][0]["delta"].get("content") or "" + # role must NOT be repeated on subsequent chunks + assert d2["chunk"]["choices"][0]["delta"].get("role") is None + + assert text1 + text2 == expected_content + + +@pytest.mark.asyncio +async def test_streaming_derender_invalid_body_returns_400(client): + """Missing required field in streaming request returns 400.""" + r = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + # missing required 'model' and 'generate_chunk' + }, + ) + assert r.status_code == 400 + + +@pytest.mark.asyncio +async def test_streaming_derender_non_object_body_returns_400(client): + """A non object JSON body (e.g. a list) returns 400, not a 500.""" + r = await client.post( + "/v1/completions/derender", + json=[1, 2, 3], + ) + assert r.status_code == 400 + + +@pytest.mark.asyncio +async def test_streaming_usage_chunk(client): + """Usage only final chunk is forwarded with correct token counts.""" + gen_req = await _render_chat(client) + token_ids: list[int] = gen_req["token_ids"][:6] + state: dict = {} + + # Send content chunk first. + r1 = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "usage-test", + "choices": [ + {"index": 0, "token_ids": token_ids, "finish_reason": "stop"} + ], + }, + "stream_state": None, + }, + ) + assert r1.status_code == 200 + state = r1.json()["stream_state"] + + # Send usage only final chunk. + r2 = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + "model": MODEL_NAME, + "generate_chunk": { + "request_id": "usage-test", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": len(token_ids), + "total_tokens": 10 + len(token_ids), + }, + }, + "stream_state": state, + "prompt_tokens": 10, + }, + ) + assert r2.status_code == 200 + d2 = r2.json() + assert d2["chunk"]["choices"] == [] + assert d2["chunk"]["usage"]["prompt_tokens"] == 10 + assert d2["chunk"]["usage"]["completion_tokens"] == len(token_ids) diff --git a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py index 0f96bf161d27..060b05bbaca8 100644 --- a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py @@ -8,54 +8,95 @@ """ from argparse import Namespace -from http import HTTPStatus import httpx import pytest -from fastapi import FastAPI, HTTPException, Request -from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse +from fastapi import HTTPException from prometheus_client import CollectorRegistry -from prometheus_fastapi_instrumentator import Instrumentator -from vllm.entrypoints.serve.utils.server_utils import exception_handler -from vllm.exceptions import VLLMNotFoundError, VLLMValidationError +from vllm.entrypoints.openai.api_server import build_app +from vllm.exceptions import ( + VLLMNotFoundError, + VLLMServerError, + VLLMValidationError, +) + + +@pytest.fixture(scope="module") +def should_do_global_cleanup_after_test() -> bool: + # This suite never initializes distributed/accelerator state. + return False + + +def _build_args() -> Namespace: + """Minimal args for ``build_app``; avoids ``make_arg_parser`` device probing.""" + return Namespace( + disable_fastapi_docs=True, + enable_offline_docs=False, + root_path=None, + allowed_origins=["*"], + allow_credentials=False, + allowed_methods=["*"], + allowed_headers=["*"], + api_key=None, + enable_request_id_headers=False, + enable_fault_tolerance=False, + middleware=[], + log_error_stack=False, + ) -@pytest.fixture +@pytest.fixture(scope="module") def registry(): - """Create a fresh Prometheus registry for each test.""" + """Shared Prometheus registry for the module-scoped app.""" return CollectorRegistry() -@pytest.fixture +@pytest.fixture(scope="module") def app(registry): - """Create a minimal FastAPI app that mirrors vLLM's exception handler - and Prometheus middleware setup.""" - - app = FastAPI() - - # Mock app state that exception_handler needs - app.state.args = Namespace(log_error_stack=False) - - # Register exception handlers exactly as vLLM does in build_app() - app.exception_handler(HTTPException)(_http_exception_handler) - app.exception_handler(RequestValidationError)(_validation_exception_handler) - app.exception_handler(ValueError)(exception_handler) - app.exception_handler(TypeError)(exception_handler) - app.exception_handler(OverflowError)(exception_handler) - app.exception_handler(NotImplementedError)(exception_handler) - app.exception_handler(VLLMValidationError)(exception_handler) - app.exception_handler(VLLMNotFoundError)(exception_handler) - app.exception_handler(Exception)(exception_handler) - - # Instrument with Prometheus (same as vLLM's attach_router) - Instrumentator( - excluded_handlers=["/metrics"], - registry=registry, - ).add().instrument(app) - - # Test routes that raise different exception types + """Build the real vLLM FastAPI app once and attach probe routes that raise. + + Patch the name used by ``attach_router`` (imported into the instrumentator + metrics module), not ``vllm.v1.metrics.prometheus`` alone — that binding is + captured at import time. + """ + import vllm.entrypoints.serve.instrumentator.metrics as metrics_mod + + original = metrics_mod.get_prometheus_registry + metrics_mod.get_prometheus_registry = lambda: registry + try: + app = build_app(_build_args(), supported_tasks=()) + finally: + metrics_mod.get_prometheus_registry = original + + @app.get("/raise_http_exception_400") + async def raise_http_exception_400(): + raise HTTPException(status_code=400, detail="bad request") + + @app.get("/raise_http_exception_404") + async def raise_http_exception_404(): + raise HTTPException(status_code=404, detail="not found") + + @app.get("/raise_request_validation_error") + async def raise_request_validation_error(n: int): + # Invalid ``n`` triggers FastAPI's RequestValidationError. + return {"n": n} + + @app.get("/raise_vllm_validation_error") + async def raise_vllm_validation_error(): + raise VLLMValidationError("bad parameter", parameter="temperature") + + @app.get("/raise_vllm_not_found_error") + async def raise_vllm_not_found_error(): + raise VLLMNotFoundError("model not found") + + @app.get("/raise_vllm_server_error") + async def raise_vllm_server_error(): + # Bare VLLMServerError goes through vllm_error_handler → 500. + # EngineGenerateError / EngineDeadError are not used here: they call + # terminate_if_errored and need engine/server state. + raise VLLMServerError("internal server failure") + @app.get("/raise_value_error") async def raise_value_error(): raise ValueError("invalid input value") @@ -72,22 +113,6 @@ async def raise_overflow_error(): async def raise_not_implemented_error(): raise NotImplementedError("feature not supported") - @app.get("/raise_vllm_validation_error") - async def raise_vllm_validation_error(): - raise VLLMValidationError("bad parameter", parameter="temperature") - - @app.get("/raise_vllm_not_found_error") - async def raise_vllm_not_found_error(): - raise VLLMNotFoundError("model not found") - - @app.get("/raise_http_exception_400") - async def raise_http_exception_400(): - raise HTTPException(status_code=400, detail="bad request") - - @app.get("/raise_http_exception_404") - async def raise_http_exception_404(): - raise HTTPException(status_code=404, detail="not found") - @app.get("/raise_runtime_error") async def raise_runtime_error(): raise RuntimeError("unexpected server error") @@ -99,14 +124,6 @@ async def success(): return app -async def _http_exception_handler(req: Request, exc: HTTPException): - return JSONResponse({"error": exc.detail}, status_code=exc.status_code) - - -async def _validation_exception_handler(req: Request, exc: RequestValidationError): - return JSONResponse({"error": str(exc)}, status_code=HTTPStatus.BAD_REQUEST) - - def _get_http_requests_total(registry, method: str, handler: str): """Extract the http_requests_total metric values grouped by status. @@ -128,31 +145,31 @@ def _get_http_requests_total(registry, method: str, handler: str): @pytest.mark.asyncio @pytest.mark.parametrize( - "endpoint,expected_status_group,expected_http_code", + "endpoint,expected_status_group,expected_http_code,request_kwargs", [ - # These should record as 4xx in Prometheus - ("/raise_value_error", "4xx", 400), - ("/raise_type_error", "4xx", 400), - ("/raise_overflow_error", "4xx", 400), - ("/raise_vllm_validation_error", "4xx", 400), - ("/raise_vllm_not_found_error", "4xx", 404), - ("/raise_http_exception_400", "4xx", 400), - ("/raise_http_exception_404", "4xx", 404), - # NotImplementedError returns 501 which is still 5xx group - ("/raise_not_implemented_error", "5xx", 501), - # These should record as 5xx in Prometheus (genuine server errors) - ("/raise_runtime_error", "5xx", 500), - # Successful requests should record as 2xx - ("/success", "2xx", 200), + ("/raise_http_exception_400", "4xx", 400, {}), + ("/raise_http_exception_404", "4xx", 404, {}), + ("/raise_request_validation_error", "4xx", 400, {"params": {"n": "x"}}), + ("/raise_vllm_validation_error", "4xx", 400, {}), + ("/raise_vllm_not_found_error", "4xx", 404, {}), + ("/raise_vllm_server_error", "5xx", 500, {}), + ("/raise_value_error", "4xx", 400, {}), + ("/raise_type_error", "4xx", 400, {}), + ("/raise_overflow_error", "4xx", 400, {}), + ("/raise_not_implemented_error", "5xx", 501, {}), + ("/raise_runtime_error", "5xx", 500, {}), + ("/success", "2xx", 200, {}), ], ids=[ + "HTTPException(400)->4xx", + "HTTPException(404)->4xx", + "RequestValidationError->4xx", + "VLLMValidationError->4xx", + "VLLMNotFoundError->4xx", + "VLLMServerError->5xx", "ValueError->4xx", "TypeError->4xx", "OverflowError->4xx", - "VLLMValidationError->4xx", - "VLLMNotFoundError->4xx", - "HTTPException(400)->4xx", - "HTTPException(404)->4xx", "NotImplementedError->5xx", "RuntimeError->5xx", "success->2xx", @@ -164,6 +181,7 @@ async def test_http_requests_total_records_correct_status( endpoint, expected_status_group, expected_http_code, + request_kwargs, ): """Verify that http_requests_total records the correct status group. @@ -177,7 +195,7 @@ async def test_http_requests_total_records_correct_status( async with httpx.AsyncClient( transport=transport, base_url="http://testserver" ) as client: - response = await client.get(endpoint) + response = await client.get(endpoint, **request_kwargs) # Verify the HTTP response code returned to the client is correct assert response.status_code == expected_http_code, ( diff --git a/tests/entrypoints/serve/lora/test_serving_models.py b/tests/entrypoints/serve/lora/test_serving_models.py index 658d004580a2..e75cbe908b1d 100644 --- a/tests/entrypoints/serve/lora/test_serving_models.py +++ b/tests/entrypoints/serve/lora/test_serving_models.py @@ -161,7 +161,9 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing: return serving -def _make_pooling_ctx(model_name: str) -> PoolingServeContext: +def _make_pooling_ctx( + model_name: str, serving: PoolingBaseServing +) -> PoolingServeContext: mock_request = MagicMock() mock_request.model = model_name return PoolingServeContext( @@ -169,6 +171,9 @@ def _make_pooling_ctx(model_name: str) -> PoolingServeContext: model_name=MODEL_NAME, request_id="test-id", pooling_params=PoolingParams(), + lora_request=serving._maybe_get_adapters(mock_request), + priorities=0, + prompt_extras=None, ) @@ -176,9 +181,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request(): """LoRA adapter name must populate ctx.lora_request without raising.""" lora_name = "bot-embed-lora" serving = _make_pooling_serving(lora_name) - ctx = _make_pooling_ctx(lora_name) - - ctx.lora_request = serving._maybe_get_adapters(ctx.request) + ctx = _make_pooling_ctx(lora_name, serving) assert ctx.lora_request is not None assert ctx.lora_request.lora_name == lora_name @@ -187,7 +190,6 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request(): def test_pooling_maybe_get_adapters_unknown_model_raises(): """An unrecognised model name must still raise VLLMNotFoundError.""" serving = _make_pooling_serving("some-lora") - ctx = _make_pooling_ctx("unknown-model") with pytest.raises(VLLMNotFoundError): - serving._maybe_get_adapters(ctx.request) + _make_pooling_ctx("unknown-model", serving) diff --git a/tests/entrypoints/serve/sagemaker/conftest.py b/tests/entrypoints/serve/sagemaker/conftest.py index d36c20ccd9af..0c76d25b4235 100644 --- a/tests/entrypoints/serve/sagemaker/conftest.py +++ b/tests/entrypoints/serve/sagemaker/conftest.py @@ -21,9 +21,9 @@ @pytest.fixture(scope="session") def smollm2_lora_files(): """Download LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM) + return hf_api().snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index af61ebc52648..713fa48eee48 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -17,7 +17,7 @@ import soundfile import torch from datasets import Audio, load_dataset -from evaluate import load +from jiwer import wer from transformers.models.whisper.english_normalizer import EnglishTextNormalizer from vllm.benchmarks.datasets.datasets import ASRDataset @@ -202,8 +202,7 @@ def run_evaluation( # Compute WER predictions = [res[2] for res in results] references = [res[3] for res in results] - wer = load("wer") - wer_score = 100 * wer.compute(references=references, predictions=predictions) + wer_score = 100 * wer(references, predictions) print("WER:", wer_score) return wer_score @@ -302,8 +301,7 @@ def run_longform_evaluation( predictions = [res[2] for res in results] references = [res[3] for res in results] - wer = load("wer") - wer_score = 100 * wer.compute(references=references, predictions=predictions) + wer_score = 100 * wer(references, predictions) print("WER:", wer_score) return wer_score @@ -358,7 +356,12 @@ def test_wer_correctness( print(f"Expected WER: {expected_wer}, Actual WER: {wer}") if expected_wer: - torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2) + wer_atol, wer_rtol = 1e-1, 1e-2 + max_wer = expected_wer + wer_atol + wer_rtol * abs(expected_wer) + assert wer <= max_wer, ( + f"WER {wer:.6f} exceeds maximum allowed {max_wer:.6f} " + f"(baseline {expected_wer:.6f})" + ) # 14-22mins of 6 audio samples of total ~115 mins and just 37MB. diff --git a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py index 040fc1a48ff1..45772708d92e 100644 --- a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py +++ b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py @@ -53,7 +53,13 @@ async def test_non_streaming_cancel_aborts_engine_requests( server.asr_config = SimpleNamespace(max_audio_clip_s=30) server._check_model = AsyncMock(return_value=None) server._maybe_get_adapters = Mock(return_value=None) - server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 40.0)) + server._preprocess_speech_to_text = AsyncMock( + return_value=( + engine_inputs, + 40.0, + [30.0 * i for i in range(len(engine_inputs))], + ) + ) server._log_inputs = Mock() request = SimpleNamespace( @@ -122,7 +128,9 @@ async def test_non_streaming_cancel_advances_all_chunk_generators(): server.asr_config = SimpleNamespace(max_audio_clip_s=30) server._check_model = AsyncMock(return_value=None) server._maybe_get_adapters = Mock(return_value=None) - server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 90.0)) + server._preprocess_speech_to_text = AsyncMock( + return_value=(engine_inputs, 90.0, [0.0, 29.5, 29.5 + 29.7]) + ) server._log_inputs = Mock() request = SimpleNamespace( diff --git a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py new file mode 100644 index 000000000000..22231e0d85cd --- /dev/null +++ b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for chunk timestamp offset drift in _preprocess_speech_to_text.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest + +from vllm.config.speech_to_text import SpeechToTextConfig +from vllm.entrypoints.speech_to_text.base.serving import SpeechToTextBaseServing + +SR = 16_000 +_PATCH = "vllm.entrypoints.speech_to_text.base.serving" + + +@pytest.mark.asyncio +async def test_chunk_offsets_are_cumulative_not_nominal(): + """ + chunk_start_offsets must be cumulative actual chunk lengths, not the old approach of + 'idx * max_audio_clip_s'. When split_audio places a boundary before the + nominal 30 s mark, the old formula drifts; the fixed formula stays exact. + """ + # Chunks shorter than exactly 30 s, as split_audio produces when a quiet + # region falls inside the 1 s overlap window before the nominal boundary. + chunk_lengths = [int(29.5 * SR), int(29.7 * SR), int(5.0 * SR)] + chunks = [np.zeros(n, dtype=np.float32) for n in chunk_lengths] + + duration = sum(chunk_lengths) / SR + + expected_offsets = [0.0, 29.5, 29.5 + 29.7] # cumulative seconds + wrong_offsets = [0.0, 30.0, 60.0] # what the old bug produced + + serving = SpeechToTextBaseServing.__new__(SpeechToTextBaseServing) + serving._decode_and_chunk_speech_async = AsyncMock(return_value=(chunks, duration)) + serving.asr_config = SpeechToTextConfig( + sample_rate=float(SR), + max_audio_clip_s=30, + overlap_chunk_second=1, + min_energy_split_window_size=1600, + ) + serving.max_audio_filesize_mb = 100.0 + serving.model_cls = MagicMock() + serving.model_cls.validate_language.side_effect = lambda lang: lang + serving.model_cls.supports_explicit_language_detection = False + serving.model_cls.get_generation_prompt.return_value = {} + serving.model_config = MagicMock() + serving.task_type = "transcribe" + serving.renderer = MagicMock() + serving.renderer.render_cmpl_async = AsyncMock( + return_value=[MagicMock()] * len(chunks) + ) + + request = MagicMock() + request.language = "en" + request.to_language = None + request.response_format = "json" + request.build_stt_params.return_value = MagicMock() + + with patch(f"{_PATCH}.parse_model_prompt", return_value=MagicMock()): + _, _, offsets = await serving._preprocess_speech_to_text( + request=request, + audio_data=b"\x00", + request_id="test", + ) + + assert offsets == pytest.approx(expected_offsets, abs=1e-6) + assert offsets != pytest.approx(wrong_offsets, abs=1e-6) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py index 7e51e5be44a3..8c99429af1ea 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py @@ -342,7 +342,9 @@ async def gen_world() -> AsyncGenerator[RequestOutput, None]: models.lora_requests = {} models.is_base_model.return_value = True - preprocess_mock = AsyncMock(return_value=([MagicMock(), MagicMock()], 1.0)) + preprocess_mock = AsyncMock( + return_value=([MagicMock(), MagicMock()], 1.0, [0.0, 29.5]) + ) with ( patch( diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index bbfde877b998..63357195d0f6 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]: falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER) - TRITON_ATTN: always available on ROCm - - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950 + - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250 On non-ROCm platforms, we just run with the default backend. """ @@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]: if current_platform.is_rocm(): backends: list[str | None] = [None, "TRITON_ATTN"] - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: backends.append("ROCM_AITER_UNIFIED_ATTN") return backends except Exception: diff --git a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py index ed3cff5f1c22..25f26404d0d9 100644 --- a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py +++ b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py @@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name): if "whisper" in model_name.lower(): try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_UNIFIED_ATTN"} except ImportError: logger.warning( - "Could not import _ON_MI3XX from rocm platform, " + "Could not check cdna version from rocm platform, " "falling back to TRITON_ATTN for Whisper." ) return {"backend": "TRITON_ATTN"} diff --git a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py index 5d769c0fd885..f03b6e8a3818 100644 --- a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py @@ -6,13 +6,13 @@ import openai import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from typing_extensions import TypedDict from tests.utils import RemoteOpenAIServer from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser +from vllm.transformers_utils.repo_utils import hf_api LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci" @@ -45,7 +45,7 @@ class ServerConfig(TypedDict, total=False): model: str arguments: list[str] model_arg: str - tool_parser: ToolParser + tool_parser: type[ToolParser] CONFIGS: dict[str, ServerConfig] = { @@ -91,7 +91,7 @@ def server_config(request): config = CONFIGS[request.param] # download model and tokenizer using transformers - snapshot_download(config["model"]) + hf_api().snapshot_download(config["model"]) yield CONFIGS[request.param] diff --git a/tests/entrypoints/unit_tests/_api_server_spawn_workers.py b/tests/entrypoints/unit_tests/_api_server_spawn_workers.py new file mode 100644 index 000000000000..c8ff969a8a36 --- /dev/null +++ b/tests/entrypoints/unit_tests/_api_server_spawn_workers.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Spawn worker targets kept free of heavy imports. + +``multiprocessing`` with the ``spawn`` start method re-imports the module that +defines a process target in the child. Housing these stubs in a stdlib-only +module keeps child startup fast and deterministic, instead of paying a multi- +second ``import vllm`` before the child can run. +""" + + +def exit_before_report_worker(listen_address, sock, args, client_config=None): + """Exit immediately without touching ``actual_address_pipe``.""" + return diff --git a/tests/entrypoints/unit_tests/test_api_server_process_manager.py b/tests/entrypoints/unit_tests/test_api_server_process_manager.py index ada7a8797fa2..902d0ce2d744 100644 --- a/tests/entrypoints/unit_tests/test_api_server_process_manager.py +++ b/tests/entrypoints/unit_tests/test_api_server_process_manager.py @@ -10,6 +10,9 @@ import pytest import zmq +from tests.entrypoints.unit_tests._api_server_spawn_workers import ( + exit_before_report_worker, +) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path from vllm.v1.utils import ( APIServerProcessManager, @@ -228,6 +231,7 @@ def test_normal_completion(api_server_args): def test_external_process_monitoring(api_server_args): """Test that wait_for_completion_or_failure handles additional processes.""" global WORKER_RUNTIME_SECONDS + prev_worker_runtime = WORKER_RUNTIME_SECONDS WORKER_RUNTIME_SECONDS = 100 # Create and start the external process @@ -307,6 +311,7 @@ def run_with_exception_capture(): manager.shutdown() mock_coordinator.shutdown() time.sleep(0.2) + WORKER_RUNTIME_SECONDS = prev_worker_runtime @pytest.mark.timeout(60) @@ -383,10 +388,10 @@ def test_gather_actual_addresses_child_crash_before_report(): num_servers=num_servers, input_addresses=placeholder_inputs, output_addresses=placeholder_outputs, - # mock_run_api_server_worker exits without touching + # exit_before_report_worker exits without touching # ``actual_address_pipe`` — simulates a child that dies before # reporting its bound addresses. - target_server_fn=mock_run_api_server_worker, + target_server_fn=exit_before_report_worker, ) try: # Sentinel-first vs pipe-EOF-first both produce "reporting". diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py index 82b262321d6c..2778770e373b 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/entrypoints/unit_tests/test_chat_utils.py @@ -15,12 +15,14 @@ from vllm.assets.video import VideoAsset from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ( + MEDIA_CONNECTOR_REGISTRY, AsyncMultiModalItemTracker, ConversationMessage, _postprocess_messages, parse_chat_messages, parse_chat_messages_async, ) +from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from vllm.multimodal.utils import ( encode_audio_url, @@ -708,6 +710,29 @@ def test_parse_chat_messages_empty_system( ] +@pytest.mark.asyncio +async def test_text_only_chat_does_not_initialize_media_connector( + mistral_model_config, + monkeypatch, +): + load_connector = MagicMock() + monkeypatch.setattr(MEDIA_CONNECTOR_REGISTRY, "load", load_connector) + messages = [{"role": "user", "content": "Who are you?"}] + + parse_chat_messages( + messages, + mistral_model_config, + content_format="string", + ) + await parse_chat_messages_async( + messages, + mistral_model_config, + content_format="string", + ) + + load_connector.assert_not_called() + + @pytest.mark.asyncio async def test_parse_chat_messages_single_image_async( phi3v_model_config, @@ -1504,7 +1529,7 @@ def test_parse_chat_messages_rejects_too_many_images_in_one_message( "ignore", message="coroutine 'async_get_and_parse_image' was never awaited", ) - with pytest.raises(ValueError, match="At most"): + with pytest.raises(VLLMValidationError, match="At most"): parse_chat_messages( [ { @@ -1540,7 +1565,7 @@ def test_parse_chat_messages_rejects_too_many_images_across_messages( "ignore", message="coroutine 'async_get_and_parse_image' was never awaited", ) - with pytest.raises(ValueError, match="At most"): + with pytest.raises(VLLMValidationError, match="At most"): parse_chat_messages( [ { @@ -2067,7 +2092,7 @@ def test_parse_chat_messages_multiple_images_interleave_with_placeholders( image_url, ): with pytest.raises( - ValueError, + VLLMValidationError, match=r"Found more '<|image_1|>' placeholders in input prompt " "than actual multimodal data items.", ): diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 9088b3c5e8db..5681c1ffeba6 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -234,6 +234,7 @@ def check_update_called(self): assert shapes == test_shapes llm.finish_weight_update() + assert llm.get_weight_version() == "default" @create_new_process_for_each_test() @@ -259,6 +260,8 @@ def test_full_weight_transfer_flow(): weight_transfer_config=WeightTransferConfig(backend="nccl"), ) + assert llm.get_weight_version() == "default" + # Step 1: Initialize weight transfer engine llm.init_weight_transfer_engine( WeightTransferInitRequest(init_info={"test_param": "flow_test"}) @@ -278,8 +281,15 @@ def test_full_weight_transfer_flow(): ) ) + assert llm.get_weight_version() == "default" + # Step 4: Finish weight update - llm.finish_weight_update() + llm.finish_weight_update("step-42") + + assert llm.get_weight_version() == "step-42" + + llm.update_weight_version("manual-version") + assert llm.get_weight_version() == "manual-version" # Verify the full flow completed def check_flow(self): @@ -319,4 +329,5 @@ def test_weight_transfer_config_backend(): ) config = llm.llm_engine.vllm_config.weight_transfer_config + assert config is not None assert config.backend == "nccl" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml new file mode 100644 index 000000000000..78a583888ba8 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml new file mode 100644 index 000000000000..e711ffb331e3 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend TRITON_ATTN" diff --git a/tests/evals/gpt_oss/configs/models-xpu.txt b/tests/evals/gpt_oss/configs/models-xpu.txt new file mode 100644 index 000000000000..a9de9266b143 --- /dev/null +++ b/tests/evals/gpt_oss/configs/models-xpu.txt @@ -0,0 +1,3 @@ +# Intel XPU model configurations for GPQA evaluation +gpt-oss-20b-xpu-baseline.yaml +gpt-oss-20b-xpu-triton-attn.yaml diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml index 14b4f2c0d1da..d4ddfcc99fc6 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml @@ -13,5 +13,6 @@ server_args: >- --enable-expert-parallel --kv-cache-dtype fp8 env: + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" VLLM_LOGGING_LEVEL: "DEBUG" VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/Laguna-XS.2-NVFP4.yaml b/tests/evals/gsm8k/configs/Laguna-XS.2-NVFP4.yaml new file mode 100644 index 000000000000..81bf8da00b3b --- /dev/null +++ b/tests/evals/gsm8k/configs/Laguna-XS.2-NVFP4.yaml @@ -0,0 +1,9 @@ +model_name: "poolside/Laguna-XS.2-NVFP4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1200 +server_args: >- + --enforce-eager + --max-model-len 4096 + --trust-remote-code diff --git a/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml b/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml new file mode 100644 index 000000000000..20e3948cce6c --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml @@ -0,0 +1,5 @@ +model_name: "amd/Qwen1.5-MoE-A2.7B-Chat-w-int8-a-int8-sym" +accuracy_threshold: 0.50 +num_questions: 1319 +num_fewshot: 5 +server_args: "--enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml b/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml new file mode 100644 index 000000000000..b0249f41b602 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml @@ -0,0 +1,5 @@ +model_name: "amd-quark/Qwen3-1.7B-MXFP4" +accuracy_threshold: 0.27 +num_questions: 1319 +num_fewshot: 5 +server_args: "--enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml index d247515a0f0e..921365ae686b 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml @@ -3,8 +3,9 @@ accuracy_threshold: 0.88 tolerance: 0.03 num_questions: 1319 num_fewshot: 5 +max_tokens: 12000 server_args: >- - --max-model-len 4096 + --max-model-len 16384 --data-parallel-size 2 --enable-expert-parallel --max-num-seqs 384 diff --git a/tests/evals/gsm8k/configs/gemma-4-E4B-it-qat-mobile-ct.yaml b/tests/evals/gsm8k/configs/gemma-4-E4B-it-qat-mobile-ct.yaml index 80d0c98311b7..878db590829c 100644 --- a/tests/evals/gsm8k/configs/gemma-4-E4B-it-qat-mobile-ct.yaml +++ b/tests/evals/gsm8k/configs/gemma-4-E4B-it-qat-mobile-ct.yaml @@ -2,4 +2,7 @@ model_name: "google/gemma-4-E4B-it-qat-mobile-ct" accuracy_threshold: 0.50 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 4096" +server_args: >- + --enforce-eager + --max-model-len 4096 + --speculative-config '{"method":"mtp","model":"google/gemma-4-E4B-it-assistant","num_speculative_tokens":4}' diff --git a/tests/evals/gsm8k/configs/models-blackwell.txt b/tests/evals/gsm8k/configs/models-blackwell.txt index 3c9b1084de7b..5b975aa63429 100644 --- a/tests/evals/gsm8k/configs/models-blackwell.txt +++ b/tests/evals/gsm8k/configs/models-blackwell.txt @@ -3,3 +3,4 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml Qwen1.5-MoE-W4A16-CT.yaml DeepSeek-V2-Lite-Instruct-FP8.yaml Qwen3-30B-A3B-NVFP4.yaml +Laguna-XS.2-NVFP4.yaml diff --git a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt index bcd00044bc01..66d4224665a0 100644 --- a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt +++ b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt @@ -1,6 +1,7 @@ Qwen3-0.6B-FP8.yaml Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml Qwen1.5-MoE-W4A16-CT.yaml +Qwen1.5-MoE-A2.7B-Chat-INT8.yaml DeepSeek-V2-Lite-Instruct-FP8.yaml Qwen3-Next-FP8-EP2_MI355.yaml Qwen3-30B-A3B-Thinking-2507-FP8.yaml diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index f652dcaf1bb3..7ed97f7efcd9 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -16,11 +16,13 @@ a silently skipped reload (e.g. offloading disabled) would not be flagged. Covers both KV offloading connectors (OffloadingConnector and -SimpleCPUOffloadConnector) across four architecture families: +SimpleCPUOffloadConnector) across five architecture families: - Hybrid Mamba (NemotronH: attention + Mamba) - Heterogeneous head dim (Gemma 4) - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet) - Compressed attention (DeepSeek-V4-Flash: CSA) + - Pure MLA (DeepSeek-V2-Lite: TieringOffloadingSpec, TP=2, replicated + single-slot host layout) Usage: pytest -s -v evals/gsm8k/test_gsm8k_offloading.py @@ -47,14 +49,16 @@ _OFFLOAD_SYNC_TIMEOUT = 60 -def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str: +def _kv_transfer_config( + connector: str, cpu_gib: int = 4, spec_name: str = "CPUOffloadingSpec" +) -> str: if connector == "OffloadingConnector": return json.dumps( { "kv_connector": "OffloadingConnector", "kv_role": "kv_both", "kv_connector_extra_config": { - "spec_name": "CPUOffloadingSpec", + "spec_name": spec_name, "cpu_bytes_to_use": cpu_gib << 30, "eviction_policy": "lru", }, @@ -115,8 +119,10 @@ class OffloadingModelConfig: accuracy_threshold: float tolerance: float = 0.05 extra_server_args: list[str] = field(default_factory=list) + env_dict: dict[str, str] = field(default_factory=dict) cpu_offload_gib: int = 4 startup_timeout: int = 600 + spec_name: str = "CPUOffloadingSpec" MODELS = [ @@ -165,6 +171,20 @@ class OffloadingModelConfig: cpu_offload_gib=16, startup_timeout=1200, ), + OffloadingModelConfig( + id="offloading-deepseek-v2-lite-tiering-tp2", + model="deepseek-ai/DeepSeek-V2-Lite", + connector="OffloadingConnector", + # Baseline 0.360/0.345 over two runs (measured on 2xA100). + accuracy_threshold=0.35, + extra_server_args=[ + "--tensor-parallel-size", + "2", + ], + cpu_offload_gib=8, + startup_timeout=1200, + spec_name="TieringOffloadingSpec", + ), # ── SimpleCPUOffloadConnector ──────────────────────────────────── OffloadingModelConfig( id="simple-nemotron-h-8b", @@ -229,7 +249,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig): "--enable-prefix-caching", "--no-disable-hybrid-kv-cache-manager", "--kv-transfer-config", - _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib), + _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib, cfg.spec_name), "--trust-remote-code", "--disable-uvicorn-access-log", *cfg.extra_server_args, @@ -239,7 +259,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig): cfg.model, server_args, # /reset_prefix_cache requires dev mode. - env_dict={"VLLM_SERVER_DEV_MODE": "1"}, + env_dict={"VLLM_SERVER_DEV_MODE": "1", **cfg.env_dict}, max_wait_seconds=cfg.startup_timeout, ) as server: base_url = f"http://{server.host}:{server.port}" diff --git a/tests/jit_monitor/__init__.py b/tests/jit_monitor/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/jit_monitor/conftest.py b/tests/jit_monitor/conftest.py new file mode 100644 index 000000000000..3f869995a999 --- /dev/null +++ b/tests/jit_monitor/conftest.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.utils import jit_monitor + + +@pytest.fixture(autouse=True) +def _reset_monitor(): + """Reset global monitor state between tests. + + ``activate()`` installs process-global hooks and flips module globals, so + without this every test would inherit the previous test's monitor state. + """ + + def reset(): + jit_monitor._active = False + jit_monitor._mode = "warn" + jit_monitor._verbose = False + jit_monitor._cutedsl_hook_installed = False + jit_monitor._tilelang_hook_installed = False + jit_monitor._tilelang_jitimpl_compile_depth = 0 + + reset() + yield + reset() diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py new file mode 100644 index 000000000000..1285b4c89e8d --- /dev/null +++ b/tests/jit_monitor/test_hooks.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the JIT monitor hooks. Backends are mocked, so no GPU.""" + +import inspect +import os +import sys +from contextlib import contextmanager +from types import ModuleType, SimpleNamespace +from typing import Any, cast +from unittest import mock + +import pytest + +from vllm.utils import jit_monitor + +pytestmark = pytest.mark.cpu_test + + +# ------------------------------------------------------------------ +# Helpers — lightweight stand-ins for the modules ``activate()`` patches +# ------------------------------------------------------------------ + + +def _make_fake_knobs(*, autotuning_print=False, jit_hook=None): + """Build a minimal fake ``triton.knobs`` namespace.""" + autotuning = SimpleNamespace(print=autotuning_print) + runtime = SimpleNamespace(jit_post_compile_hook=jit_hook) + return SimpleNamespace(autotuning=autotuning, runtime=runtime) + + +def _fake_cute_import_modules(compile_fn): + """Fake Python's parent package + submodule for ``import cutlass.cute``.""" + fake_cute = cast(Any, ModuleType("cutlass.cute")) + fake_cute.compile = compile_fn + fake_parent_package = cast(Any, ModuleType("cutlass")) + fake_parent_package.__path__ = [] + fake_parent_package.cute = fake_cute + return { + "cutlass": fake_parent_package, + "cutlass.cute": fake_cute, + } + + +def _fake_cute_compile(*args, **kwargs): + return "compiled" + + +def _fake_tilelang_import_modules(): + """Fake Python's TileLang modules touched by ``jit_monitor.activate``.""" + + class FakeJITKernel: + def __init__(self, *args, **kwargs): + pass + + class FakeJITImpl: + def __init__(self, func, signature): + self.func = func + self.signature = signature + self.mode = "lazy" + self._kernel_cache = {} + + def __call__(self, *args, **kwargs): + key, _ = self.func.parse_args(*args, **kwargs) + kernel = self._kernel_cache.get(key) + if kernel is None: + kernel = "compiled" + self._kernel_cache[key] = kernel + return kernel + + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) + fake_kernel.JITKernel = FakeJITKernel + + fake_jit = cast(Any, ModuleType("tilelang.jit")) + fake_jit.JITImpl = FakeJITImpl + fake_jit.kernel = fake_kernel + + fake_tilelang = cast(Any, ModuleType("tilelang")) + fake_tilelang.jit = fake_jit + + return { + "tilelang": fake_tilelang, + "tilelang.jit": fake_jit, + "tilelang.jit.kernel": fake_kernel, + } + + +@contextmanager +def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile): + """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``.""" + fake_triton = cast(Any, ModuleType("triton")) + fake_triton.knobs = fake_knobs + with ( + mock.patch.dict( + sys.modules, + { + "triton": fake_triton, + **_fake_cute_import_modules(cute_compile), + **_fake_tilelang_import_modules(), + }, + ), + mock.patch.object(jit_monitor, "HAS_TRITON", True), + ): + yield + + +def _triton_hook_kwargs(name: str): + return dict( + key="k", + repr="r", + fn=SimpleNamespace(name=name), + compile=lambda: None, + is_manual_warmup=False, + already_compiled=False, + ) + + +# ------------------------------------------------------------------ +# activate() +# ------------------------------------------------------------------ + + +def test_activate_sets_active(): + assert not jit_monitor.is_active() + with _patch_jit_modules(_make_fake_knobs()): + jit_monitor.activate() + assert jit_monitor.is_active() + + +def test_activate_is_idempotent(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate() + first_hook = fake.runtime.jit_post_compile_hook + jit_monitor.activate() + assert fake.runtime.jit_post_compile_hook is first_hook + + +def test_activate_logs_info(): + with ( + mock.patch.object(jit_monitor.logger, "info") as m, + _patch_jit_modules(_make_fake_knobs()), + ): + jit_monitor.activate() + m.assert_called_once() + assert "Kernel JIT monitor activated" in m.call_args[0][0] + + +def test_activate_rejects_unknown_mode(): + with pytest.raises(ValueError, match="Unsupported JIT monitor mode"): + jit_monitor.activate(mode="panic") # type: ignore[arg-type] + + +def test_activate_without_triton(): + with mock.patch.object(jit_monitor, "HAS_TRITON", False): + jit_monitor.activate() + assert jit_monitor.is_active() + + +# ------------------------------------------------------------------ +# Triton autotuning print +# ------------------------------------------------------------------ + + +def test_autotuning_print_is_enabled(): + fake = _make_fake_knobs(autotuning_print=False) + with _patch_jit_modules(fake): + jit_monitor.activate() + assert fake.autotuning.print is True + + +def test_autotuning_print_respects_user_opt_out(): + fake = _make_fake_knobs(autotuning_print=False) + with ( + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}), + _patch_jit_modules(fake), + ): + jit_monitor.activate() + assert fake.autotuning.print is False + + +def test_autotuning_print_noop_when_user_already_enabled(): + fake = _make_fake_knobs(autotuning_print=True) + with ( + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}), + _patch_jit_modules(fake), + ): + jit_monitor.activate() + assert fake.autotuning.print is True + + +# ------------------------------------------------------------------ +# Triton JIT hook +# ------------------------------------------------------------------ + + +def test_triton_hook_is_registered(): + fake = _make_fake_knobs() + assert fake.runtime.jit_post_compile_hook is None + with _patch_jit_modules(fake): + jit_monitor.activate() + assert fake.runtime.jit_post_compile_hook is not None + + +def test_triton_hook_logs_warning(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + + with ( + mock.patch.object(jit_monitor.logger, "warning_once") as m, + mock.patch.object(jit_monitor.logger, "warning") as warning, + ): + hook(**_triton_hook_kwargs("test_kernel")) + + m.assert_called_once() + warning.assert_not_called() + msg = m.call_args[0][0] % m.call_args[0][1:] + assert "Triton kernel JIT compilation during inference" in msg + assert "test_kernel" in msg + + +def test_triton_hook_chains_existing_hook(): + existing = mock.MagicMock(return_value="existing_result") + fake = _make_fake_knobs(jit_hook=existing) + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + result = hook(**_triton_hook_kwargs("chained_kernel")) + + existing.assert_called_once() + assert result == "existing_result" + + +def test_triton_hook_works_without_existing_hook(): + fake = _make_fake_knobs(jit_hook=None) + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + assert hook(**_triton_hook_kwargs("solo_kernel")) is None + + +def test_triton_hook_error_mode_raises(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate(mode="error") + + hook = fake.runtime.jit_post_compile_hook + with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"): + hook(**_triton_hook_kwargs("error_kernel")) + + +# ------------------------------------------------------------------ +# CuTeDSL hook +# ------------------------------------------------------------------ + + +def test_cutedsl_compile_logs_warning(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile(lambda: None, "arg", option=True) + + assert result == "compiled" + warning_once.assert_called_once() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "CuTeDSL JIT compilation during inference" in msg + + +def test_cutedsl_compile_logs_verbose_warning(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate(verbose=True) + with mock.patch.object(jit_monitor.logger, "warning") as warning: + result = cute.compile(lambda: None, "arg", option=True) + + assert result == "compiled" + warning.assert_called_once() + msg = warning.call_args[0][0] % warning.call_args[0][1:] + assert "CuTeDSL JIT compilation during inference" in msg + + +def test_cutedsl_error_mode_raises(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate(mode="error") + with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): + cute.compile(lambda: None, "arg", option=True) + + +def test_cutedsl_subscripted_compile_is_monitored(): + """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" + + class FakeCompileCallable: + def __getitem__(self, options): + return self + + def __call__(self, *args, **kwargs): + return "compiled" + + with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile[("opt_level", 3)](lambda: None, "arg") + + assert result == "compiled" + warning_once.assert_called_once() + + +# ------------------------------------------------------------------ +# TileLang hook +# ------------------------------------------------------------------ + + +def test_tilelang_jit_kernel_logs_warning(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"}) + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi") + + warning_once.assert_called_once() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "TileLang JIT compilation during inference" in msg + assert "tl_kernel" in msg + + +def test_tilelang_jit_impl_logs_warning(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit import JITImpl + + def tilelang_fn( + gemm_out_mul, + hidden_size: int, + n_splits: int = 1, + hc_mult: int = 4, + ): + return None + + class FakeFunc: + orig_func = tilelang_fn + + def parse_args(self, *args, **kwargs): + return ( + ( + "tilelang_key", + kwargs["hidden_size"], + kwargs.get("n_splits", 1), + ), + {}, + ) + + def set_mode(self, mode): + self.mode = mode + + tensor = SimpleNamespace( + shape=(2, 16, 24), + dtype="float32", + device="cuda:0", + ) + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) + + jit_monitor.activate() + with ( + mock.patch.object(jit_monitor.logger, "warning_once") as warning_once, + mock.patch.object(jit_monitor.logger, "warning") as warning, + ): + impl(tensor, hidden_size=7168, n_splits=2) + + warning_once.assert_called_once() + warning.assert_not_called() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "TileLang JIT compilation during inference" in msg + assert "tilelang_fn" in msg + + +def test_tilelang_jit_impl_does_not_log_on_cache_hit(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit import JITImpl + + def tilelang_fn(gemm_out_mul, n_splits: int = 1): + return None + + class FakeFunc: + orig_func = tilelang_fn + + def parse_args(self, *args, **kwargs): + return (("tilelang_key", kwargs.get("n_splits", 1)), {}) + + def set_mode(self, mode): + self.mode = mode + + tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32") + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + impl(tensor, n_splits=2) + impl(tensor, n_splits=2) + + warning_once.assert_called_once() + + +def test_tilelang_from_database_does_not_log(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"}) + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + JITKernel(func=func, from_database=True) + + warning_once.assert_not_called() + + +def test_tilelang_error_mode_raises(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"}) + jit_monitor.activate(mode="error") + with pytest.raises(RuntimeError, match="TileLang JIT compilation"): + JITKernel(func=func) diff --git a/tests/jit_monitor/test_hooks_gpu.py b/tests/jit_monitor/test_hooks_gpu.py new file mode 100644 index 000000000000..7c4b93abcd3a --- /dev/null +++ b/tests/jit_monitor/test_hooks_gpu.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end JIT monitor tests: real Triton kernel, real GPU, real hook.""" + +from unittest import mock + +import pytest + +from vllm.utils import jit_monitor + +try: + import torch + + _HAS_CUDA = torch.cuda.is_available() +except ImportError: + _HAS_CUDA = False + +try: + import triton + import triton.language as tl + + _HAS_TRITON = True +except ImportError: + _HAS_TRITON = False + +pytestmark = pytest.mark.skipif( + not (_HAS_CUDA and _HAS_TRITON), + reason="Requires CUDA GPU and Triton", +) + + +if _HAS_TRITON: + + @triton.jit + def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + y = tl.load(y_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x + y, mask=mask) + + +def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None: + """Launch ``_add_kernel`` with vectors of length *n*.""" + x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment + y = torch.randn(n, device="cuda") + out = torch.empty(n, device="cuda") + grid = ((n + block - 1) // block,) + _add_kernel[grid](x, y, out, n, BLOCK=block) + torch.accelerator.synchronize() + + +def test_no_warning_on_cached_shape(): + _run_add_kernel(1024) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as w: + _run_add_kernel(1024) + w.assert_not_called() + + +def test_warning_on_new_constexpr(): + _run_add_kernel(1024, block=256) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as w: + # Different BLOCK (a tl.constexpr) forces recompilation. + _run_add_kernel(1024, block=512) + w.assert_called() + msg = w.call_args[0][0] % w.call_args[0][1:] + assert "_add_kernel" in msg + + +def test_verbose_warning_on_each_new_pointer_alignment(): + _run_add_kernel(1024) + + jit_monitor.activate(verbose=True) + with ( + mock.patch.object(jit_monitor.logger, "warning") as w, + mock.patch.object(jit_monitor.logger, "warning_once") as w_once, + ): + _run_add_kernel(1024, offset=1) + assert w.called + w_once.assert_not_called() diff --git a/tests/jit_monitor/test_no_runtime_jit.py b/tests/jit_monitor/test_no_runtime_jit.py new file mode 100644 index 000000000000..86afd626b34a --- /dev/null +++ b/tests/jit_monitor/test_no_runtime_jit.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Catch runtime (post-warmup) JIT compilations for JIT-heavy backends running +e2e tests on popular models, for which we only load a few blocks for performance. + +NOTE(NickLucche) With cuda graphs on, kernels fully covered by graphs captured during +warmup do not re-trigger the Python JIT hooks. The targeted paths (prefill +MoE/MLA/SSM and the sampler) run mixed, so they are unaffected. +""" + +from dataclasses import dataclass + +import pytest + +from vllm import LLM, SamplingParams +from vllm.inputs import TokensPrompt + +from ..models.utils import dummy_hf_overrides +from ..utils import create_new_process_for_each_test + +# Warmup coverage is still incomplete for these backends, so the monitor fires +# during inference. Tracked in https://github.com/vllm-project/vllm/issues/49349; +# drop this once the warmup contract migrations land. +pytestmark = pytest.mark.skip(reason="Kernel warmup coverage is still incomplete") + + +@dataclass(frozen=True) +class JitModel: + model: str + draft: str | None = None + trust_remote_code: bool = False + + +JIT_MONITOR_MODELS = [ + JitModel("Qwen/Qwen3-0.6B"), + JitModel("deepseek-ai/DeepSeek-V2-Lite-Chat", trust_remote_code=True), + JitModel("deepseek-ai/DeepSeek-V3", trust_remote_code=True), + JitModel("ibm-granite/granite-4.0-tiny-preview"), + JitModel( + "luccafong/deepseek_mtp_main_random", + draft="luccafong/deepseek_mtp_draft_random", + trust_remote_code=True, + ), + JitModel( + "eagle618/deepseek-v3-random", + draft="eagle618/eagle-deepseek-v3-random", + trust_remote_code=True, + ), +] + + +def _run_shape_battery(llm: LLM) -> None: + """Exercise diverse compile keys so missing warmup keys surface. + + Token-id prompts keep shapes exact and avoid depending on a tokenizer. + Outputs are meaningless under dummy weights; we assert only that no JIT + fired. + """ + short = TokensPrompt(prompt_token_ids=[1, 2, 3, 4]) + medium = TokensPrompt(prompt_token_ids=list(range(1, 33))) + long = TokensPrompt(prompt_token_ids=list(range(1, 129))) + + # Greedy single-sequence multi-step decode: prefill + autoregressive decode + # + greedy sampler. + llm.generate(medium, SamplingParams(temperature=0.0, max_tokens=16)) + + # Batched prefill with mixed lengths: varlen prefill + padded decode. + llm.generate([short, medium, long], SamplingParams(temperature=0.0, max_tokens=8)) + + # Triton sampler kernels: top_k / top_p / min_p each specialize. + for sampling_params in ( + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0), + SamplingParams( + temperature=0.8, top_k=20, top_p=0.9, min_p=0.1, max_tokens=8, seed=0 + ), + ): + llm.generate(medium, sampling_params) + + # Heterogeneous SamplingParams in one step, where missing sampler warmup + # keys most often hide. + llm.generate( + [medium] * 4, + [ + SamplingParams(temperature=0.0, max_tokens=8), + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0), + ], + ) + + +@create_new_process_for_each_test("spawn") +def can_run_without_jit(spec: JitModel): + """Boot ``spec`` with the monitor armed and run the shape battery. + + A subprocess per model is required: the monitor's hooks are process-global + and, once armed in ``error`` mode, stay armed. It must be spawned rather + than forked, since forking a pytest process that already initialized CUDA + poisons the child. + """ + llm = LLM( + spec.model, + trust_remote_code=spec.trust_remote_code, + max_model_len=2048, + max_num_seqs=8, + gpu_memory_utilization=0.80, + load_format="dummy", + hf_overrides=dummy_hf_overrides, + # cuda graphs cover captured decode shapes, run eager. + enforce_eager=False, + jit_monitor_mode="error", + speculative_config={ + "model": spec.draft, + "num_speculative_tokens": 2, + } + if spec.draft + else None, + ) + + try: + _run_shape_battery(llm) + except Exception as e: + # The monitor's message contains "during inference"; distinguish a real + # JIT miss from an unrelated crash. + if "during inference" in str(e): + pytest.fail( + f"{spec.model}: post-warmup JIT compilation detected - a warmup " + f"key is missing for a shape in the battery.\n{e}" + ) + raise + + +@pytest.mark.parametrize("spec", JIT_MONITOR_MODELS, ids=lambda s: s.model) +def test_no_runtime_jit(spec: JitModel, monkeypatch: pytest.MonkeyPatch): + """Assert JIT-heavy backends do not JIT-compile during inference.""" + # Set here rather than in the child so the spawned process inherits it: + # the engine core must not be forked once the test process has CUDA up. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + can_run_without_jit(spec) diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 1e85f76b64c3..587f75a38d6d 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -546,8 +546,11 @@ def test_flash_attn_accepts_handled_fp8_variants( ): """FlashAttentionBackend must accept the two fp8 dtypes it can actually handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'.""" - import vllm.v1.attention.backends.flash_attn as fa_mod + import vllm.v1.attention.backends.fa_utils as fa_utils_mod from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True) + # The fp8 decision is made in fa_utils, using its own current_platform + # binding, so patch is_xpu there (not on flash_attn's) to stay robust to + # import order across earlier tests that patch vllm.platforms.current_platform. + monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True) assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype) diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index e296c226d709..ff8e6db92b97 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -534,6 +534,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=False, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_without_split = torch.empty_like(query) @@ -569,6 +570,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=True, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_with_split = torch.empty_like(query) @@ -685,6 +687,42 @@ def test_varlen_encoder_attention_vec( ) +@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize( + "block_size", + [ + 128, + ], +) +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("isa", ["neon"]) +@pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.ARM, + reason="Not an Arm CPU.", +) +def test_varlen_encoder_attention_neon( + seq_lens: list[int], + num_heads: tuple[int, int], + head_size: int, + sliding_window: int | None, + dtype: torch.dtype, + block_size: int, + isa: str, +) -> None: + varlen_encoder_attention( + seq_lens=seq_lens, + num_heads=num_heads, + head_size=head_size, + sliding_window=sliding_window, + dtype=dtype, + block_size=block_size, + isa=isa, + ) + + @pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES) @@ -803,6 +841,24 @@ def test_varlen_with_paged_kv_normal_amx( ) +@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.") +def test_varlen_with_paged_kv_fp8_large_prefill_amx() -> None: + varlen_with_paged_kv( + seq_lens=[(1024, 1024)] * 4, + num_heads=(16, 2), + head_size=256, + sliding_window=None, + dtype=torch.bfloat16, + block_size=2176, + soft_cap=None, + num_blocks=4, + use_alibi=False, + use_sink=False, + isa="amx", + kv_cache_dtype="fp8_e4m3", + ) + + @pytest.mark.parametrize("seq_lens", SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16) diff --git a/tests/kernels/attention/test_flashmla.py b/tests/kernels/attention/test_flashmla.py index 657b256f4687..84744a81e797 100644 --- a/tests/kernels/attention/test_flashmla.py +++ b/tests/kernels/attention/test_flashmla.py @@ -10,7 +10,9 @@ from vllm.triton_utils import triton from vllm.v1.attention.ops.flashmla import ( flash_mla_with_kvcache, + flash_mla_with_kvcache_fp8, get_mla_metadata, + get_mla_metadata_dense_fp8, is_flashmla_dense_supported, ) @@ -86,9 +88,14 @@ def test_flash_mla( ) blocked_v = blocked_k[..., :dv] - tile_scheduler_metadata, num_splits = get_mla_metadata( - cache_seqlens, s_q * h_q // h_kv, h_kv - ) + if use_fp8: + tile_scheduler_metadata, num_splits = get_mla_metadata_dense_fp8( + cache_seqlens, s_q * h_q // h_kv, h_kv + ) + else: + tile_scheduler_metadata, num_splits = get_mla_metadata( + cache_seqlens, s_q * h_q // h_kv, h_kv + ) init_dtype = q.dtype if use_fp8: @@ -104,6 +111,19 @@ def test_flash_mla( descale_k = None def flash_mla(): + if use_fp8: + return flash_mla_with_kvcache_fp8( + q, + blocked_k, + block_table, + cache_seqlens, + dv, + tile_scheduler_metadata, + num_splits, + causal=causal, + descale_q=descale_q, + descale_k=descale_k, + ) return flash_mla_with_kvcache( q, blocked_k, @@ -113,8 +133,6 @@ def flash_mla(): tile_scheduler_metadata, num_splits, causal=causal, - descale_q=descale_q, - descale_k=descale_k, ) def scaled_dot_product_attention(query, key, value, is_causal=False): diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 010c44797665..040ad50a5bbf 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -4,6 +4,43 @@ import torch +def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): + from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata + + device = torch.device("cuda") + capacity_width = 256 + active_width = 128 + global_decode_buffer = torch.empty( + (2, capacity_width), dtype=torch.int32, device=device + ) + decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device) + prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device) + + global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( + positions=torch.tensor([255, 511], dtype=torch.int64, device=device), + compress_ratio=128, + num_decode_tokens=1, + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device), + block_table=torch.tensor([[3]], dtype=torch.int32, device=device), + block_size=capacity_width, + slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device), + global_decode_buffer=global_decode_buffer, + decode_lens_buffer=decode_lens_buffer, + prefill_buffer=prefill_buffer, + max_compressed_tokens=active_width, + ) + + assert global_decode.shape == (1, active_width) + assert prefill_local.shape == (1, active_width) + assert global_decode.stride() == (active_width, 1) + assert prefill_local.stride() == (active_width, 1) + assert global_decode[0, :2].cpu().tolist() == [768, 769] + assert decode_lens.cpu().tolist() == [2] + assert prefill_local[0, :4].cpu().tolist() == list(range(4)) + assert torch.all(global_decode[0, 2:] == -1) + assert torch.all(prefill_local[0, 4:] == -1) + + def test_sparse_flashmla_metadata_smoke(): import vllm.v1.attention.ops.flashmla as fm @@ -98,7 +135,8 @@ def test_sparse_flashmla_decode_smoke(): assert lse.shape[0] == batch_size -def test_sparse_flashmla_prefill_smoke(): +@pytest.mark.parametrize("h_q", [64, 128]) +def test_sparse_flashmla_prefill_smoke(h_q: int): import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_sparse_supported() @@ -106,22 +144,25 @@ def test_sparse_flashmla_prefill_smoke(): pytest.skip(reason) device = torch.device("cuda") + torch.manual_seed(0) s_q = 1 - s_kv = 1 - h_q = 64 # kernel expects multiple of 64 + s_kv = 8 h_kv = 1 d_qk = 576 d_v = 512 topk = 128 - - q = torch.zeros((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device) - kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device) - indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device) - - out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v) - assert out.shape == (s_q, h_q, d_v) - assert max_logits.shape == (s_q, h_q) - assert lse.shape == (s_q, h_q) + q = torch.randn((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device) + kv = torch.randn((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device) + indices = torch.randint(s_kv, (s_q, h_kv, topk), dtype=torch.int32, device=device) + reference_indices = indices.clone() + reference_indices[..., 1:] = -1 + kwargs = {"topk_length": torch.ones(1, dtype=torch.int32, device=device)} + reference = fm.flash_mla_sparse_fwd(q, kv, reference_indices, 1.0, d_v, **kwargs) + actual = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v, **kwargs) + + for actual_tensor, reference_tensor in zip(actual, reference): + torch.testing.assert_close(actual_tensor, reference_tensor, rtol=0, atol=0) + assert actual[0].shape == (s_q, h_q, d_v) def test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences(): diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py new file mode 100644 index 000000000000..0e66e0da1e0d --- /dev/null +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RoPE equivalence tests for the fused Kimi-K3 MLA epilogues.""" + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA" +) + +_DTYPE = torch.bfloat16 +_NUM_TOKENS = 3 +_NUM_HEADS = 4 +_BLOCK_SIZE = 8 +_POSITIONS = (1, 7, 13) +_SLOTS = (0, 3, 9) + + +def _randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2 + + +def _rope_cache(max_position: int = 32) -> torch.Tensor: + inv_freq = 1.0 / ( + 50000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64) + ) + positions = torch.arange(max_position, dtype=torch.float32, device="cuda") + freqs = torch.outer(positions, inv_freq) + # The fused epilogue reads the cos/sin table in fp32 (RoPE math runs in fp32). + return torch.cat((freqs.cos(), freqs.sin()), dim=-1) + + +def _apply_gptj_rope( + x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor +) -> torch.Tensor: + cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1) + for _ in range(x.ndim - 2): + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + x1 = x[..., ::2].float() + x2 = x[..., 1::2].float() + out1 = x1 * cos.float() - x2 * sin.float() + out2 = x2 * cos.float() + x1 * sin.float() + return torch.stack((out1, out2), dim=-1).flatten(-2).to(x.dtype) + + +def _cache_rows(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + return cache.reshape(-1, cache.shape[-1]).index_select(0, slots) + + +def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close( + actual.float(), + expected.to(torch.float8_e4m3fn).float(), + atol=0.03125, + rtol=0.15, + ) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(0) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + q = _randn(_NUM_TOKENS, _NUM_HEADS, 192) + k_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + k_pe = _randn(_NUM_TOKENS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + v = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + + q_expected = q.clone() + q_expected[..., 128:] = _apply_gptj_rope( + q_expected[..., 128:], positions, cos_sin_cache + ) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + k_expected = torch.cat( + (k_nope, k_pe_expected[:, None, :].expand(-1, _NUM_HEADS, -1)), dim=-1 + ) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = q.clone() + k_actual = fused_mla_key_concat_kv_cache_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros( + 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn + ) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual, k_actual, v_actual = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c, + v, + cache, + slots, + one, + one, + one, + one, + positions, + cos_sin_cache, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(k_actual, k_expected) + _assert_fp8_close(v_actual, v) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = q.clone() + k_actual = fused_mla_key_concat_ds_mla_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_decode_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(1) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + + q_pe_expected = _apply_gptj_rope(q_pe, positions, cos_sin_cache) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + q_expected = torch.cat((ql_nope, q_pe_expected), dim=-1) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + kwargs = {"positions": positions, "cos_sin_cache": cos_sin_cache} + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots, **kwargs + ) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros( + 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn + ) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c, + k_pe, + cache, + slots, + q_scale_inv=one, + cache_scale_inv=one, + **kwargs, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots, ds_mla=True, **kwargs + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@torch.inference_mode() +def test_decode_epilogue_preserves_nope_path() -> None: + torch.manual_seed(2) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots + ) + + torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1)) + torch.testing.assert_close( + _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1) + ) diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index 858d9504a184..d73acfc0ee9c 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -20,7 +20,7 @@ from vllm.platforms.cpu import CpuPlatform from vllm.platforms.cuda import CudaPlatform from vllm.platforms.interface import DeviceCapability -from vllm.platforms.rocm import RocmPlatform +from vllm.platforms.rocm import RocmPlatform, on_mi3xx from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import _cached_get_attn_backend @@ -346,3 +346,137 @@ def test_mha_attn_varlen_forward_flashinfer( torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) finally: vllm_config.model_config = old_model_config + + +@pytest.mark.skipif( + not current_platform.is_rocm() or not on_mi3xx(), + reason="AITER FP8 attention requires MI300/MI350", +) +@pytest.mark.parametrize("var_seq_len", [[128, 193], [256, 384, 511]]) +@pytest.mark.parametrize("head_size", [64, 72, 80]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half]) +def test_mha_attn_varlen_forward_aiter_fp8( + default_vllm_config, + var_seq_len: list[int], + head_size: int, + dtype: torch.dtype, +): + """Compare packed AITER FP8 ViT attention with a BF16/FP16 reference.""" + aiter = pytest.importorskip("aiter") + if not hasattr(aiter, "flash_attn_varlen_fp8_pertensor_func"): + pytest.skip("Installed AITER does not provide varlen FP8 attention") + + num_heads = 4 + set_random_seed(0) + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + + vllm_config = get_current_vllm_config() + old_model_config = getattr(vllm_config, "model_config", None) + minimal_model_config = type( + "MinimalModelConfig", + (), + { + "multimodal_config": MultiModalConfig( + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA, + mm_encoder_attn_dtype="fp8", + ), + }, + )() + vllm_config.model_config = minimal_model_config + try: + total_len = sum(var_seq_len) + # Keep Q/K/V as non-contiguous views, matching interleaved ViT QKV. + qkv = torch.randn(1, total_len, 3, num_heads, head_size) + q, k, v = qkv.unbind(dim=2) + cu_seqlens = torch.tensor( + [0] + list(itertools.accumulate(var_seq_len)), dtype=torch.int32 + ) + max_seqlen = torch.tensor(max(var_seq_len), dtype=torch.int32) + scale = 1.0 / head_size**0.5 + + attn = MMEncoderAttention(num_heads, head_size, scale=scale) + assert attn.attn_backend == AttentionBackendEnum.ROCM_AITER_FA + assert attn.fp8_enabled + output = attn( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + ref_output = torch.cat( + [ + ref_attention(q_i, k_i, v_i, scale=scale) + for q_i, k_i, v_i in zip( + torch.split(q, var_seq_len, dim=1), + torch.split(k, var_seq_len, dim=1), + torch.split(v, var_seq_len, dim=1), + ) + ], + dim=1, + ) + diff = (output.float() - ref_output.float()).abs() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), + ref_output.float().flatten(), + dim=0, + ) + + assert output.shape == ref_output.shape + assert output.dtype == dtype + assert cosine.item() >= 0.99 + assert diff.mean().item() <= 0.03 + assert diff.max().item() <= 0.30 + finally: + vllm_config.model_config = old_model_config + + +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER FP8 attention requires ROCm" +) +def test_mha_attn_aiter_fp8_rejects_unsupported_arch(default_vllm_config): + vllm_config = get_current_vllm_config() + old_model_config = getattr(vllm_config, "model_config", None) + vllm_config.model_config = type( + "MinimalModelConfig", + (), + { + "multimodal_config": MultiModalConfig( + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA, + mm_encoder_attn_dtype="fp8", + ), + }, + )() + try: + with ( + patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + pytest.raises(ValueError, match="gfx942 or gfx950"), + ): + MMEncoderAttention(4, 64) + finally: + vllm_config.model_config = old_model_config + + +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER FP8 attention requires ROCm" +) +def test_mha_attn_fp8_rejects_wrong_backend(default_vllm_config): + vllm_config = get_current_vllm_config() + old_model_config = getattr(vllm_config, "model_config", None) + vllm_config.model_config = type( + "MinimalModelConfig", + (), + { + "multimodal_config": MultiModalConfig( + mm_encoder_attn_backend=AttentionBackendEnum.FLASH_ATTN, + mm_encoder_attn_dtype="fp8", + ), + }, + )() + try: + with pytest.raises(ValueError, match="requires either"): + MMEncoderAttention(4, 64) + finally: + vllm_config.model_config = old_model_config diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 8a2bd5892a2d..8fe2342808fc 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -431,7 +431,7 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype): not current_platform.is_device_capability_family(100), reason="fmha_sm100 indexer requires SM100 (Blackwell).", ) -@pytest.mark.parametrize("topk", [8, 16]) +@pytest.mark.parametrize("topk", [16]) @pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn]) def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch): import vllm.models.minimax_m3.common.indexer as indexer_mod @@ -471,6 +471,14 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch): batch, BLOCK_SIZE, device, arange_block_indices=True ) num_tokens = batch.compute_num_tokens() + # Absolute token positions; the MSA builder derives per-token causal page + # counts from them. + common.positions = torch.cat( + [ + torch.arange(s - q, s, device=device, dtype=torch.int64) + for s, q in zip(batch.seq_lens, batch.query_lens) + ] + ) # Deterministic index cache: distinct, monotonic per-logical-block values so # the top-k is unambiguous (both kernels pick the same blocks, no fp ties). @@ -515,14 +523,14 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch): triton_impl.index_cache.kv_cache = index_cache # Exercise the shared persistent top-k buffer for BOTH impls: each must write - # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views. - # Separate buffers so the two forwards don't clobber each other. + # decode ([:nd]) and prefill ([nd:]) into its token-major buffer and (Triton + # only) return views. Separate buffers so the two forwards don't clobber. nd = sum(q for q in batch.query_lens if q <= 1) msa_impl.topk_indices_buffer = torch.full( - (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device + (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device ) triton_impl.topk_indices_buffer = torch.full( - (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device + (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device ) attn_metadata = { @@ -533,18 +541,17 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch): msa_decode, msa_prefill = msa_impl(index_q) tri_decode, tri_prefill = triton_impl(index_q) - assert msa_decode is not None and tri_decode is not None - assert msa_prefill is not None and tri_prefill is not None - _assert_topk_indices_equal_unordered(msa_decode, tri_decode) - _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill) - # decode/prefill outputs are views into each impl's persistent buffer. - for impl, dec, pre in ( - (msa_impl, msa_decode, msa_prefill), - (triton_impl, tri_decode, tri_prefill), - ): - buf = impl.topk_indices_buffer - assert dec.data_ptr() == buf[:, :nd, :].data_ptr() - assert pre.data_ptr() == buf[:, nd:, :].data_ptr() + # MSA's return is vestigial; the attend reads its buffer directly. + assert msa_decode is None and msa_prefill is None + assert tri_decode is not None and tri_prefill is not None + _assert_topk_indices_equal_unordered( + msa_impl.topk_indices_buffer[:num_tokens], + triton_impl.topk_indices_buffer[:num_tokens], + ) + # Triton's decode/prefill outputs are views into its persistent buffer. + buf_htk = triton_impl.topk_indices_buffer.transpose(0, 1) + assert tri_decode.data_ptr() == buf_htk[:, :nd, :].data_ptr() + assert tri_prefill.data_ptr() == buf_htk[:, nd:, :].data_ptr() @pytest.mark.parametrize( diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py new file mode 100644 index 000000000000..e2b59dc33462 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -0,0 +1,561 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for MiniMax M3 CUTLASS sparse decode.""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.config import AttentionConfig +from vllm.models.minimax_m3.common.ops.sparse_attn import ( + minimax_m3_sparse_attn_decode, +) +from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseBackend, + MiniMaxM3SparseMetadataBuilder, + MiniMaxM3SparseTritonImpl, + select_main_backend_and_impl_cls, +) +from vllm.models.minimax_m3.nvidia import ( + sparse_attention_msa as sparse_attention_msa_module, +) +from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( + MSACutlassDecodePlanCache, + msa_cutlass_sparse_decode, + prepare_decode_metadata, + should_prepare_decode_metadata, +) +from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( + MiniMaxM3SparseMSABackend, + MiniMaxM3SparseMSADecodeMetadata, + MiniMaxM3SparseMSAImpl, + MiniMaxM3SparseMSAMetadataBuilder, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +if not current_platform.is_device_capability_family(100): + pytest.skip( + "fmha_sm100 sparse decode requires SM100 (Blackwell).", + allow_module_level=True, + ) + + +HEAD_DIM = 128 +BLOCK_SIZE = 128 +TOPK = 16 +DEFAULT_QUERY_LEN = 4 +SM_SCALE = HEAD_DIM**-0.5 + + +@pytest.mark.parametrize( + ( + "batch_size", + "decode_query_len", + "num_q_heads", + "num_kv_heads", + "expected", + ), + [ + pytest.param(8, 4, 64, 4, False, id="tp1-below-min-batch"), + pytest.param(16, 4, 64, 4, True, id="tp1-supported"), + pytest.param(16, 4, 16, 1, True, id="tp4-min-batch"), + pytest.param(24, 4, 16, 1, True, id="tp4-intermediate-batch"), + pytest.param(32, 4, 16, 1, True, id="tp4-supported"), + pytest.param(16, 1, 64, 4, True, id="tp1-query-len-1"), + pytest.param(16, 1, 16, 1, True, id="tp4-query-len-1"), + pytest.param(16, 2, 64, 4, True, id="tp1-query-len-2"), + pytest.param(16, 2, 16, 1, True, id="tp4-query-len-2"), + pytest.param(16, 32, 64, 4, True, id="query-len-upper-bound"), + pytest.param(16, 0, 64, 4, False, id="query-len-zero"), + pytest.param(16, 33, 64, 4, False, id="query-len-above-bound"), + ], +) +def test_msa_cutlass_decode_static_dispatch( + batch_size: int, + decode_query_len: int, + expected: bool, + num_q_heads: int, + num_kv_heads: int, +) -> None: + assert ( + should_prepare_decode_metadata( + batch_size, + decode_query_len, + decode_backend="cutlass", + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + kv_cache_dtype="fp8_e4m3", + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + is expected + ) + + +def test_msa_cutlass_decode_static_dispatch_requires_opt_in() -> None: + assert not should_prepare_decode_metadata( + 32, + DEFAULT_QUERY_LEN, + decode_backend="triton", + num_q_heads=16, + num_kv_heads=1, + kv_cache_dtype="fp8_e4m3", + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + + +def test_msa_cutlass_decode_static_dispatch_accepts_fp8_alias() -> None: + assert should_prepare_decode_metadata( + 32, + DEFAULT_QUERY_LEN, + decode_backend="cutlass", + num_q_heads=16, + num_kv_heads=1, + kv_cache_dtype="fp8", + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + + +@pytest.mark.parametrize( + "kv_cache_dtype", + ["auto", "bfloat16", "float16", "fp8_e5m2"], +) +def test_msa_cutlass_decode_static_dispatch_requires_fp8_e4m3( + kv_cache_dtype: str, +) -> None: + assert not should_prepare_decode_metadata( + 32, + DEFAULT_QUERY_LEN, + decode_backend="cutlass", + num_q_heads=16, + num_kv_heads=1, + kv_cache_dtype=kv_cache_dtype, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + + +def test_msa_cutlass_decode_static_dispatch_requires_sm100( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda _: False, + ) + assert not should_prepare_decode_metadata( + 32, + DEFAULT_QUERY_LEN, + decode_backend="cutlass", + num_q_heads=16, + num_kv_heads=1, + kv_cache_dtype="fp8", + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + + +@pytest.mark.parametrize( + ("backend", "expected"), + [ + (AttentionBackendEnum.CUTLASS_MSA, "cutlass"), + (AttentionBackendEnum.TRITON_MSA, "triton"), + ], +) +def test_msa_attention_backend_alias( + backend: AttentionBackendEnum, + expected: str, +) -> None: + config = AttentionConfig(backend=backend) + assert config.backend is None + assert config.minimax_m3_msa_decode_backend == expected + + +def test_msa_backend_owns_msa_metadata_builder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend_cls, impl_cls = select_main_backend_and_impl_cls( + topk_blocks=TOPK, + kv_cache_dtype="fp8_e4m3", + num_kv_heads=1, + ) + assert backend_cls is MiniMaxM3SparseMSABackend + assert impl_cls is MiniMaxM3SparseMSAImpl + + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda _: False, + ) + backend_cls, impl_cls = select_main_backend_and_impl_cls( + topk_blocks=TOPK, + kv_cache_dtype="fp8_e4m3", + num_kv_heads=1, + ) + assert backend_cls is MiniMaxM3SparseBackend + assert impl_cls is MiniMaxM3SparseTritonImpl + + +def test_msa_metadata_builder_prepares_cutlass_for_regular_decode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + batch = 16 + block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda") + seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda") + base_metadata = SimpleNamespace( + num_decodes=batch, + decode=SimpleNamespace( + block_table=block_table, + seq_lens=seq_lens, + decode_query_len=1, + ), + ) + monkeypatch.setattr( + MiniMaxM3SparseMetadataBuilder, + "build", + lambda *args, **kwargs: base_metadata, + ) + expected_metadata = object() + monkeypatch.setattr( + sparse_attention_msa_module, + "prepare_decode_metadata", + lambda *args, **kwargs: expected_metadata, + ) + + builder = object.__new__(MiniMaxM3SparseMSAMetadataBuilder) + builder.num_q_heads = 64 + builder.num_kv_heads = 4 + builder.topk_blocks = TOPK + builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4) + builder.kv_cache_dtype = "fp8_e4m3" + builder.decode_backend = "cutlass" + builder.msa_cutlass_plan_cache = object() + + metadata = builder.build( + 0, + SimpleNamespace( + seq_lens_cpu_upper_bound=torch.full((batch,), 257, dtype=torch.int32) + ), + ) + + assert isinstance(metadata.decode, MiniMaxM3SparseMSADecodeMetadata) + assert metadata.decode.decode_query_len == 1 + assert metadata.decode.msa_cutlass is expected_metadata + + +def test_msa_cutlass_plan_cache_keys_query_len( + monkeypatch: pytest.MonkeyPatch, +) -> None: + batch = 16 + block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda") + seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda") + seq_lens_cpu = torch.full((batch,), 257, dtype=torch.int32) + plan_cache = MSACutlassDecodePlanCache() + built_query_lens: list[int] = [] + + def fake_build_plan(**kwargs): + query_len = kwargs["decode_query_len"] + built_query_lens.append(query_len) + num_rows = batch * query_len + return ( + None, + None, + None, + { + "kv_segment_lens": torch.empty( + num_rows, dtype=torch.int32, device="cuda" + ), + "qo_offset": torch.empty(num_rows, dtype=torch.int32, device="cuda"), + }, + ) + + monkeypatch.setattr(plan_cache, "_build_plan", fake_build_plan) + first = prepare_decode_metadata( + block_table, + seq_lens, + seq_lens_cpu, + 1, + num_q_heads=64, + num_kv_heads=4, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + repeated = prepare_decode_metadata( + block_table, + seq_lens, + seq_lens_cpu, + 1, + num_q_heads=64, + num_kv_heads=4, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + different = prepare_decode_metadata( + block_table, + seq_lens, + seq_lens_cpu, + 2, + num_q_heads=64, + num_kv_heads=4, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + current_platform.synchronize() + + assert first.plan is repeated.plan + assert different.plan is not first.plan + assert built_query_lens == [1, 2] + + +def _make_topk( + seq_lens: list[int], + num_kv_heads: int, + query_len: int, +) -> torch.Tensor: + topk = torch.full( + (len(seq_lens) * query_len, num_kv_heads, TOPK), + -1, + dtype=torch.int32, + device="cuda", + ) + for request, seq_len in enumerate(seq_lens): + for local_query in range(query_len): + token = request * query_len + local_query + visible_tokens = seq_len - query_len + local_query + 1 + visible_pages = math.ceil(visible_tokens / BLOCK_SIZE) + topk[token, :, :visible_pages] = torch.arange( + visible_pages, dtype=torch.int32, device="cuda" + ) + return topk + + +@pytest.mark.parametrize( + ( + "num_q_heads", + "num_kv_heads", + "num_request_pairs", + "query_len", + "capture_graph", + ), + [ + pytest.param(64, 4, 8, 1, True, id="tp1-query-len-1"), + pytest.param(64, 4, 8, 2, False, id="tp1-query-len-2"), + pytest.param(64, 4, 8, 3, False, id="tp1-query-len-3"), + pytest.param(64, 4, 8, 4, True, id="tp1-query-len-4"), + pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"), + pytest.param(16, 1, 8, 1, True, id="tp4-query-len-1"), + pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"), + ], +) +def test_msa_cutlass_decode_matches_triton_with_interleaved_cache( + num_q_heads: int, + num_kv_heads: int, + num_request_pairs: int, + query_len: int, + capture_graph: bool, +) -> None: + torch.manual_seed(0) + seq_lens_list = [257, 513] * num_request_pairs + seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32) + seq_lens = seq_lens_cpu.cuda() + pages_per_request = [math.ceil(seq_len / BLOCK_SIZE) for seq_len in seq_lens_list] + num_pages = sum(pages_per_request) + max_pages = max(pages_per_request) + + block_table = torch.zeros( + len(seq_lens_list), max_pages, dtype=torch.int32, device="cuda" + ) + physical_pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda") + offset = 0 + for request, request_pages in enumerate(pages_per_request): + block_table[request, :request_pages] = physical_pages[ + offset : offset + request_pages + ] + offset += request_pages + + key = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.25 + ).to(torch.float8_e4m3fn) + value = (torch.randn_like(key, dtype=torch.bfloat16) * 0.25).to(torch.float8_e4m3fn) + kv_cache = torch.cat((key, value), dim=-1) + assert kv_cache.stride() == ( + num_kv_heads * BLOCK_SIZE * 2 * HEAD_DIM, + BLOCK_SIZE * 2 * HEAD_DIM, + 2 * HEAD_DIM, + 1, + ) + + num_query_tokens = len(seq_lens_list) * query_len + query = torch.randn( + num_query_tokens, + num_q_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + query_fp8 = torch.empty_like(query, dtype=torch.float8_e4m3fn) + ops.scaled_fp8_quant( + query.view(num_query_tokens, -1), + scale=q_scale, + output=query_fp8.view(num_query_tokens, -1), + ) + query_dequantized = query_fp8.to(torch.bfloat16) * q_scale + + topk_token_major = _make_topk(seq_lens_list, num_kv_heads, query_len) + expected = torch.empty_like(query) + minimax_m3_sparse_attn_decode( + query_dequantized, + kv_cache, + topk_token_major.transpose(0, 1), + block_table, + seq_lens, + num_kv_heads, + SM_SCALE, + expected, + query_len, + k_scale=None, + v_scale=None, + ) + + plan_cache = MSACutlassDecodePlanCache() + metadata = prepare_decode_metadata( + block_table, + seq_lens, + seq_lens_cpu, + query_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + assert metadata.page_table.data_ptr() == block_table.data_ptr() + actual = torch.empty_like(query) + msa_cutlass_sparse_decode( + query_fp8, + kv_cache, + topk_token_major, + actual, + metadata, + scale=SM_SCALE, + q_scale_float=1.0, + k_scale_float=1.0, + v_scale_float=1.0, + ) + + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + + # The same captured plan must remain correct as ragged lengths change. + updated_seq_lens_list = [129, 385] * num_request_pairs + seq_lens.copy_( + torch.tensor(updated_seq_lens_list, dtype=torch.int32, device="cuda") + ) + updated_seq_lens_cpu = torch.tensor(updated_seq_lens_list, dtype=torch.int32) + topk_token_major.copy_(_make_topk(updated_seq_lens_list, num_kv_heads, query_len)) + updated_metadata = prepare_decode_metadata( + block_table, + seq_lens, + updated_seq_lens_cpu, + query_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + assert updated_metadata.plan is metadata.plan + assert updated_metadata.page_table.data_ptr() == metadata.page_table.data_ptr() + + minimax_m3_sparse_attn_decode( + query_dequantized, + kv_cache, + topk_token_major.transpose(0, 1), + block_table, + seq_lens, + num_kv_heads, + SM_SCALE, + expected, + query_len, + k_scale=None, + v_scale=None, + ) + msa_cutlass_sparse_decode( + query_fp8, + kv_cache, + topk_token_major, + actual, + updated_metadata, + scale=SM_SCALE, + q_scale_float=1.0, + k_scale_float=1.0, + v_scale_float=1.0, + ) + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + + if capture_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + msa_cutlass_sparse_decode( + query_fp8, + kv_cache, + topk_token_major, + actual, + updated_metadata, + scale=SM_SCALE, + q_scale_float=1.0, + k_scale_float=1.0, + v_scale_float=1.0, + ) + + replay_seq_lens_list = [257, 513] * num_request_pairs + seq_lens.copy_( + torch.tensor(replay_seq_lens_list, dtype=torch.int32, device="cuda") + ) + topk_token_major.copy_( + _make_topk(replay_seq_lens_list, num_kv_heads, query_len) + ) + prepare_decode_metadata( + block_table, + seq_lens, + torch.tensor(replay_seq_lens_list, dtype=torch.int32), + query_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + plan_cache=plan_cache, + ) + minimax_m3_sparse_attn_decode( + query_dequantized, + kv_cache, + topk_token_major.transpose(0, 1), + block_table, + seq_lens, + num_kv_heads, + SM_SCALE, + expected, + query_len, + k_scale=None, + v_scale=None, + ) + graph.replay() + current_platform.synchronize() + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index 48a236a01579..ac4607355999 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -442,7 +442,9 @@ def test_flashmla_dense_fp8_decode_unified_slot_view(): n_layers = 3 layer = 1 - q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1 + # With a quantized KV cache the decode query is quantized to fp8 as well; + # the fp8 dense MLA kernel requires q already in fp8. + q = (torch.randn(bs, 1, h_q, head_dim, device=dev) * 0.1).to(torch.float8_e4m3fn) kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to( torch.float8_e4m3fn ) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index f1c591fb671b..e63d2c0adf1f 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -323,6 +323,342 @@ def test_contexted_kv_attention( torch.testing.assert_close(output, output_ref, atol=atol, rtol=0) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_contexted_kv_attention_cached_kv( + num_heads: int, + num_queries_per_kv: int, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: str, + device: str, + block_size: int = 32, +) -> None: + # Exercises the KV_FROM_CACHE path of context_attention_fwd: the current + # chunk K/V are not passed as dense tensors (k=v=None); they are read back + # from the paged KV cache, as done by layers that re-attend an already + # cached sequence with the query only (e.g. IQuest LoopCoder's + # attn(q, None, None)). The whole sequence therefore lives in the cache and + # the result must still match a dense causal SDPA reference. + if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89): + pytest.skip( + "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89" + ) + + set_random_seed(0) + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + MAX_SEQ_LEN = 1024 + MAX_CTX_LEN = 1024 + BS = 10 + cache_size = 640 + max_block_per_request = 64 + query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)] + # ensure one sequence in batch is a decode + query_lens[-1] = 1 + ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)] + seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)] + num_kv_heads = num_heads // num_queries_per_kv + + num_tokens = sum(query_lens) + query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype) + query.uniform_(-1e-3, 1e-3) + output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype) + + kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype) + kv.uniform_(-1e-3, 1e-3) + key, value = kv.unbind(dim=1) + + if kv_cache_dtype == "auto": + cache_dtype = dtype + else: + cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype] + k_cache = torch.zeros( + cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype + ) + v_cache = torch.zeros( + cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype + ) + values = torch.arange(0, cache_size, dtype=torch.int32) + values = values[torch.randperm(cache_size)] + block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request) + b_seq_len = torch.tensor(seq_lens, dtype=torch.int32) + b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32) + max_input_len = MAX_SEQ_LEN + b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to( + torch.int32 + ) + # Unlike the dense test, write the WHOLE sequence (context + current chunk) + # into the paged cache, since the current chunk is read back from the cache. + for i in range(BS): + cur = 0 + block_id = 0 + while cur < seq_lens[i]: + start_loc = b_seq_start_loc[i] + cur + if cur + block_size > seq_lens[i]: + end_loc = b_seq_start_loc[i] + seq_lens[i] + else: + end_loc = start_loc + block_size + start_slot = block_table[i, block_id] * block_size + end_slot = start_slot + end_loc - start_loc + k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_( + key[start_loc:end_loc] + ) + v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_( + value[start_loc:end_loc] + ) + cur += block_size + block_id += 1 + # transpose to the paged cache layouts the kernel expects: + # K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8] + # V_cache[num_blocks, num_kv_heads, head_size, block_size] + k_cache = ( + k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8) + .permute(0, 2, 3, 1, 4) + .contiguous() + ) + v_cache = ( + v_cache.view(-1, block_size, num_kv_heads, head_size) + .permute(0, 2, 3, 1) + .contiguous() + ) + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + + # Cached-K/V path: current-chunk k and v are None. + context_attention_fwd( + query, + None, + None, + output, + kv_cache_dtype, + k_cache, + v_cache, + block_table, + b_start_loc, + b_seq_len, + MAX_CTX_LEN, + max_input_len, + k_scale, + v_scale, + sliding_window=0, + ) + torch.accelerator.synchronize() + + scale = float(1.0 / (head_size**0.5)) + + query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size) + query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape( + 1, num_heads, num_tokens, head_size + ) + key_sdpa = key[:, :, None, :].expand( + key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1] + ) + key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape( + 1, num_heads, sum(seq_lens), head_size + ) + value_sdpa = value[:, :, None, :].expand( + value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1] + ) + value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape( + 1, num_heads, sum(seq_lens), head_size + ) + + attn_mask = create_causal_attention_mask_for_sdpa( + query_lens, seq_lens, 0, device=device, dtype=dtype + ) + output_ref = F.scaled_dot_product_attention( + query_sdpa, + key_sdpa, + value_sdpa, + attn_mask=attn_mask, + dropout_p=0.0, + scale=scale, + ) + output_ref = output_ref.view(num_heads, num_tokens, head_size) + output_ref = output_ref.permute(1, 0, 2).contiguous() + atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4 + torch.testing.assert_close(output, output_ref, atol=atol, rtol=0) + + +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None: + # Boundary guard for the KV_FROM_CACHE block-table load. With an + # exact-sized block table (row length == number of blocks for the + # sequence) and a query that ends the sequence, the last K/V tile has + # padded lanes whose absolute positions step past the sequence end. Those + # lanes must not read a block-table entry past this batch's row. + # + # seq_len is a whole number of blocks, so bn_logical for a padded lane + # lands exactly on num_blocks (one past the last valid row entry), and + # query_len is not tile-aligned so the overshoot is actually exercised. + # The result must still match a dense causal SDPA reference. + set_random_seed(0) + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + dtype = torch.float16 + kv_cache_dtype = "auto" + num_heads = 4 + num_queries_per_kv = 1 + num_kv_heads = num_heads // num_queries_per_kv + head_size = 32 + block_size = 32 + num_blocks = 4 + seq_len = num_blocks * block_size # 128 == exact multiple of block_size + query_len = 99 # < seq_len and not a multiple of the kernel tile + + query = torch.empty(query_len, num_heads, head_size, dtype=dtype) + query.uniform_(-1e-3, 1e-3) + output = torch.empty(query_len, num_heads, head_size, dtype=dtype) + + kv = torch.empty(seq_len, 2, num_kv_heads, head_size, dtype=dtype) + kv.uniform_(-1e-3, 1e-3) + key, value = kv.unbind(dim=1) + + k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype) + v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype) + # Exact-sized block table: row length == num_blocks (identity mapping). + block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks) + b_seq_len = torch.tensor([seq_len], dtype=torch.int32) + b_start_loc = torch.tensor([0, query_len], dtype=torch.int32) + + # Write the whole sequence (context + current chunk) into the paged cache. + for cur in range(0, seq_len, block_size): + block_id = cur // block_size + end = min(cur + block_size, seq_len) + start_slot = block_table[0, block_id] * block_size + end_slot = start_slot + (end - cur) + k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_( + key[cur:end] + ) + v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_( + value[cur:end] + ) + + k_cache = ( + k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8) + .permute(0, 2, 3, 1, 4) + .contiguous() + ) + v_cache = ( + v_cache.view(-1, block_size, num_kv_heads, head_size) + .permute(0, 2, 3, 1) + .contiguous() + ) + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + + # Cached-K/V path: current-chunk k and v are None. + context_attention_fwd( + query, + None, + None, + output, + kv_cache_dtype, + k_cache, + v_cache, + block_table, + b_start_loc, + b_seq_len, + seq_len, + query_len, + k_scale, + v_scale, + sliding_window=0, + ) + torch.accelerator.synchronize() + + scale = float(1.0 / (head_size**0.5)) + query_sdpa = query.view(query_len, num_kv_heads, num_queries_per_kv, head_size) + query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape( + 1, num_heads, query_len, head_size + ) + key_sdpa = key[:, :, None, :].expand( + seq_len, num_kv_heads, num_queries_per_kv, head_size + ) + key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size) + value_sdpa = value[:, :, None, :].expand( + seq_len, num_kv_heads, num_queries_per_kv, head_size + ) + value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape( + 1, num_heads, seq_len, head_size + ) + + attn_mask = create_causal_attention_mask_for_sdpa( + [query_len], [seq_len], 0, device=device, dtype=dtype + ) + output_ref = F.scaled_dot_product_attention( + query_sdpa, + key_sdpa, + value_sdpa, + attn_mask=attn_mask, + dropout_p=0.0, + scale=scale, + ) + output_ref = output_ref.view(num_heads, query_len, head_size) + output_ref = output_ref.permute(1, 0, 2).contiguous() + torch.testing.assert_close(output, output_ref, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None: + # The cached-K/V (k=None) path is not supported together with ALiBi; the + # entry point must reject it up-front with a clear NotImplementedError + # rather than launching the kernel with an unsupported combination. + set_random_seed(0) + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + num_heads = 4 + num_kv_heads = 4 + head_size = 16 + x = 8 + block_size = 16 + num_blocks = 4 + query_len = 8 + + query = torch.empty(query_len, num_heads, head_size, dtype=torch.float16) + query.uniform_(-1e-3, 1e-3) + output = torch.empty_like(query) + k_cache = torch.zeros( + num_blocks, num_kv_heads, head_size // x, block_size, x, dtype=torch.float16 + ) + v_cache = torch.zeros( + num_blocks, num_kv_heads, head_size, block_size, dtype=torch.float16 + ) + block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks) + b_seq_len = torch.tensor([query_len], dtype=torch.int32) + b_start_loc = torch.tensor([0, query_len], dtype=torch.int32) + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + alibi_slopes = torch.ones(num_heads, dtype=torch.float32, device=device) + + with pytest.raises(NotImplementedError): + context_attention_fwd( + query, + None, + None, + output, + "auto", + k_cache, + v_cache, + block_table, + b_start_loc, + b_seq_len, + query_len, + query_len, + k_scale, + v_scale, + alibi_slopes=alibi_slopes, + ) + + @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py new file mode 100644 index 000000000000..7401203629de --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for causal masking in the AITER MLA verify flatten. + +When ``num_heads < 16`` the ROCm AITER MLA backend cannot use the ASM decode +kernel for a multi-token block, so ``AiterMLAImpl.forward_mqa`` flattens a +``1 + num_speculative_tokens`` verify block into that many single-token Gluon +decodes. Every resulting row used to be handed its request's entire paged-KV +range. ``paged_kv_indptr`` is the cumulative sum of ``seq_lens``, which counts +the tokens scheduled in the current step, so that range already spans the +verify block: each verify position could attend to the draft tokens after it, +which are the tokens it is supposed to be checking. Nothing crashes and nothing +warns, so the only way to catch it is to look at the KV window per row. + +The test drives the real builder and the real ``forward_mqa`` over a +multi-token decode batch, intercepts the ``(page_table, seq_info)`` actually +handed to the Gluon kernel, and pins verify row ``t`` of request ``r`` to the +``context_r + t + 1`` entries of that request's ascending slice. +""" + +import types +from unittest.mock import patch + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found +from vllm.platforms import current_platform + + +def _on_rocm_with_aiter() -> bool: + return current_platform.is_rocm() and is_aiter_found() + + +# Unlike the fp8 persistent-decode metadata this file's neighbour guards, the +# verify flatten is not gfx950-only: it is selected by head count alone and is +# reached on any ROCm AITER MLA deployment with fewer than 16 heads per rank. +# The Gluon kernel is replaced by a spy here, so nothing below needs it to be +# runnable either. +pytestmark = pytest.mark.skipif( + not _on_rocm_with_aiter(), + reason="ROCM_AITER_MLA verify flatten requires ROCm and AITER", +) + +# Below the 16-head threshold that selects the flatten. DeepSeek-V3 / R1 at +# TP=16 and Kimi-K3 at TP=8 both land here. +NUM_QUERY_HEADS = 12 +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 128 +QK_ROPE_HEAD_DIM = 64 +V_HEAD_DIM = 128 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM + +# 1 + num_speculative_tokens, i.e. the verify block length. +QLEN = 8 +# Committed context per request, including a fresh request at zero context +# where the whole KV range is the verify block itself. +CONTEXT_LENS = [0, 1, 37, 512] +# A cudagraph padding request: seq_len 0, so every one of its rows must clamp +# to an empty window. +PADDING_ROWS = 1 + +# The Gluon path flattens the KV cache to one page per token. +PAGE_SIZE = 1 +MAX_MODEL_LEN = 1024 +MAX_NUM_SEQS = 8 + + +def _seq_lens() -> list[int]: + return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS + + +def _expected_row_lens() -> list[int]: + """Causal row lengths: row r*QLEN + t sees seq_len_r - (QLEN - 1) + t.""" + return [max(0, s - (QLEN - 1) + t) for s in _seq_lens() for t in range(QLEN)] + + +def _run_verify_block(): + """Drive the real builder + forward_mqa over one multi-token verify block. + + Returns ``(metadata, captured)`` where ``captured`` holds the ``page_table``, + ``seq_info`` and ``min_kv_seq_len`` the backend passed to the Gluon kernel. + """ + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import SpeculativeConfig + from vllm.config.vllm import set_current_vllm_config + from vllm.v1.attention.backends.registry import AttentionBackendEnum + from vllm.v1.kv_cache_interface import MLAAttentionSpec + from vllm.v1.worker.workspace import init_workspace_manager + + device = torch.device("cuda:0") + seq_lens = _seq_lens() + num_reqs = len(seq_lens) + # One flat page per token, arange block indices, plus room to spare. + num_gpu_blocks = num_reqs * max(seq_lens) + 256 + + vllm_config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-R1", + max_model_len=MAX_MODEL_LEN, + num_gpu_blocks=num_gpu_blocks, + block_size=PAGE_SIZE, + max_num_seqs=MAX_NUM_SEQS, + max_num_batched_tokens=8192, + hf_config_override={"num_attention_heads": NUM_QUERY_HEADS}, + ) + # What raises reorder_batch_threshold to QLEN, so a QLEN-token block is + # classified as a decode instead of a prefill. ngram needs no draft model. + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=QLEN - 1 + ) + + spec = MLAAttentionSpec( + block_size=PAGE_SIZE, + num_kv_heads=1, + head_size=vllm_config.model_config.get_head_size(), + dtype=vllm_config.model_config.dtype, + cache_dtype_str="auto", + ) + + backend_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class() + builder_cls = backend_cls.get_builder_cls() + impl_cls = backend_cls.get_impl_cls() + + # The builder reads layer.prefill_backend from static_forward_context; a + # stub with the attribute is enough for metadata construction. + layer_name = "placeholder" + vllm_config.compilation_config.static_forward_context[layer_name] = ( + types.SimpleNamespace(prefill_backend=torch.empty((1,))) + ) + + init_workspace_manager(device) + + batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=[QLEN] * num_reqs) + + captured: dict = {} + + def spy(**kwargs): + captured["page_table"] = kwargs["page_table"].detach().clone() + captured["seq_info"] = kwargs["seq_info"].detach().clone() + captured["min_kv_seq_len"] = kwargs["min_kv_seq_len"] + + with set_current_vllm_config(vllm_config): + builder = builder_cls(spec, [layer_name], vllm_config, device) + common_attn_metadata = create_common_attn_metadata( + batch_spec, PAGE_SIZE, device, arange_block_indices=True + ) + metadata = builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + + impl = impl_cls( + num_heads=NUM_QUERY_HEADS, + head_size=HEAD_SIZE, + scale=(QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=KV_LORA_RANK, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM, + v_head_dim=V_HEAD_DIM, + kv_b_proj=None, + ) + + num_rows = num_reqs * QLEN + dtype = torch.bfloat16 + q_nope = torch.zeros( + num_rows, NUM_QUERY_HEADS, KV_LORA_RANK, dtype=dtype, device=device + ) + q_pe = torch.zeros( + num_rows, NUM_QUERY_HEADS, QK_ROPE_HEAD_DIM, dtype=dtype, device=device + ) + kv_cache = torch.zeros( + num_gpu_blocks, PAGE_SIZE, HEAD_SIZE, dtype=dtype, device=device + ) + + # The Gluon kernel only reads the metadata this test is about, so a spy in + # its place keeps the assertions independent of the AITER build. + with patch( + "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon", + lambda: spy, + ): + impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None) + + return metadata, captured + + +def test_verify_flatten_rows_are_causal(): + """Verify row t must see the committed prefix plus tokens 0..t, and no more. + + Regression guard: before the fix every row of a request got that request's + whole paged-KV range, so verify position t attended to the draft tokens + after it. Fails on unmodified upstream. + """ + metadata, captured = _run_verify_block() + + decode = metadata.decode + assert decode is not None, "batch was not classified as a decode" + assert decode.max_qo_len == QLEN, ( + f"expected a {QLEN}-token verify block, got max_qo_len=" + f"{decode.max_qo_len}; the flatten under test was not reached" + ) + assert captured, "forward_mqa did not reach the Gluon kernel" + + seq_lens = _seq_lens() + indptr = captured["seq_info"].tolist() + page_table = captured["page_table"] + got_row_lens = [indptr[i + 1] - indptr[i] for i in range(len(indptr) - 1)] + want_row_lens = _expected_row_lens() + + assert got_row_lens == want_row_lens, ( + "AITER MLA verify flatten is not causal: verify row r*qlen+t must get " + f"seq_len_r - {QLEN - 1} + t KV entries.\n" + f" seq_lens {seq_lens}\n" + f" got {got_row_lens}\n" + f" expected {want_row_lens}\n" + "Rows longer than expected let a verify position attend to the draft " + "tokens it is supposed to be checking." + ) + + # arange_block_indices lays request r's pages out ascending from + # r * max_blocks, so each causal window must be that slice's prefix. + max_blocks = max(seq_lens) + for r, seq_len in enumerate(seq_lens): + for t in range(QLEN): + row = r * QLEN + t + window = page_table[indptr[row] : indptr[row + 1]].tolist() + n = max(0, seq_len - (QLEN - 1) + t) + want = [r * max_blocks + j for j in range(n)] + assert window == want, ( + f"request {r} (seq_len {seq_len}) verify token {t} reads KV " + f"pages {window}, expected the ascending causal prefix {want}" + ) + + padding_rows = got_row_lens[len(CONTEXT_LENS) * QLEN :] + assert padding_rows == [0] * (PADDING_ROWS * QLEN), ( + f"cudagraph padding requests (seq_len 0) must clamp to empty windows, " + f"got {padding_rows}" + ) + + # min_kv_seq_len tells Gluon how short the shortest row is, so it must be + # the minimum over the causal rows actually submitted, not over the + # per-request lengths they were cut from. + assert captured["min_kv_seq_len"] == min(want_row_lens), ( + f"min_kv_seq_len={captured['min_kv_seq_len']} does not match the " + f"shortest submitted row ({min(want_row_lens)})" + ) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py index 2b9a4e823c08..a384af6f1317 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -40,10 +40,11 @@ def _on_gfx950() -> bool: CONTEXT_LEN = 8192 PAGE_SIZE = 1 -# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8 -# KV-cache -> fp8_e4m3 kv. -EXPECTED_Q_DTYPE = torch.bfloat16 -EXPECTED_KV_DTYPE = torch.float8_e4m3fn +# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter +# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype() +# instead of hardcoding a literal (see #47276). +EXPECTED_Q_DTYPE = current_platform.fp8_dtype() +EXPECTED_KV_DTYPE = current_platform.fp8_dtype() # The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data # is excluded: it holds raw device pointers, never equal across allocations. diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index cc9ac6b8d715..aa603626414c 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -88,6 +88,81 @@ def _make_common_metadata(): ) +def _make_mixed_common_metadata(): + # req0: query_len 1 (decode), req1: query_len 4 (prefill) -> decode-first + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32, device="cpu") + seq_lens = torch.tensor([16, 10], dtype=torch.int32, device="cpu") + return CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + _seq_lens_cpu=seq_lens, + num_reqs=2, + num_actual_tokens=5, + max_query_len=4, + max_seq_len=16, + block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + slot_mapping=torch.arange(5, dtype=torch.int64, device="cpu"), + ) + + +def _patch_build_deps(monkeypatch, events=None): + """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs + on CPU.""" + + def fake_generate_sparse_seqlen_triton( + query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len + ): + return torch.zeros(num_tokens, dtype=torch.int32, device="cpu") + + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = Mock(side_effect=lambda *a, **k: None) + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton + ) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace( + synchronize=lambda: events.append("sync") if events is not None else None + ), + ) + + +def test_build_populates_decode_only_split_fields(monkeypatch): + """Decode-only batch: all reqs count as decodes, prefill fields default.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_common_metadata() + ) + + assert md.num_decodes == 2 + assert md.num_prefills == 0 + assert md.num_decode_tokens == 2 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + +def test_build_populates_mixed_split_fields(monkeypatch): + """Mixed decode+prefill batch: split is reported, prefill fields stay + default because this impl always runs the MQA path.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_mixed_common_metadata() + ) + + assert md.num_decodes == 1 + assert md.num_prefills == 1 + assert md.num_decode_tokens == 1 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch): builder = _make_builder() common_metadata = _make_common_metadata() diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index 9e33f24ea280..a03c00e55287 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -15,15 +15,15 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -_SKIP_NON_MI3XX = True +_SKIP_NON_CDNA_2_PLUS = True if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - _SKIP_NON_MI3XX = not on_mi3xx() + _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2 pytestmark = [ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), - pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"), + pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"), ] NUM_Q_HEADS = 8 diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 77e068ab1717..6fe2a3e77587 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.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 @@ -196,6 +198,68 @@ def _ragged_from_rows( ) +@torch.inference_mode() +def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + + class FakeWorkspaceManager: + def get_simultaneous(self, *shapes_and_dtypes): + return [ + torch.empty(shape, dtype=dtype, device=device) + for shape, dtype in shapes_and_dtypes + ] + + def fake_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_logits, + context_lens, + block_tables, + max_seq_len, + **kwargs, + ): + del ( + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_seq_len, + kwargs, + ) + out_logits.fill_(float("nan")) + + monkeypatch.setattr(mod, "_ON_GFX942", False) + monkeypatch.setattr(mod, "_ON_GFX950", True) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + mod, + "paged_mqa_logits_module", + lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits), + ) + monkeypatch.setattr( + mod, "current_workspace_manager", lambda: FakeWorkspaceManager() + ) + + q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device) + kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device) + logits = mod.rocm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + torch.empty((1, 1), dtype=torch.float32, device=device), + torch.ones(1, dtype=torch.int32, device=device), + torch.zeros((1, 1), dtype=torch.int32, device=device), + torch.empty(0, dtype=torch.int32, device=device), + 1, + ) + + assert not torch.isnan(logits).any() + + @torch.inference_mode() def test_compute_global_topk_ragged_indices_and_indptr() -> None: from vllm.models.deepseek_v4.amd.rocm import ( diff --git a/tests/kernels/attention/test_xpu_mla_sparse.py b/tests/kernels/attention/test_xpu_mla_sparse.py index 419644923ec4..965c6f735f0f 100644 --- a/tests/kernels/attention/test_xpu_mla_sparse.py +++ b/tests/kernels/attention/test_xpu_mla_sparse.py @@ -116,3 +116,55 @@ def test_bf16_triton_sparse_mla(device_str, dtype): assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2) assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3) assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize("device_str", ["xpu"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.skipif( + not torch.xpu.is_available(), + reason="XPU is required", +) +def test_bf16_triton_sparse_mla_masked_chunks(device_str, dtype): + """Rows whose leading BLOCK_N index entries are all masked must not NaN. + + Regression test: with an -inf running max, a fully-masked leading chunk + produced re_scale = exp2(-inf - -inf) = NaN, permanently poisoning the + accumulator even though valid keys followed in later chunks. + """ + device = torch.device(device_str) + s_q = 3 + s_kv = 256 + h_q = 64 + h_kv = 1 + d_qk = 576 + d_v = 512 + topk = 128 # 8 chunks of BLOCK_N=16 + + torch.random.manual_seed(1234) + + q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device) + kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device) + indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device) + # row 0: valid keys only in chunks 1-2 -> leading AND trailing masked chunks + indices[0, 0, 16:48] = torch.arange(32, dtype=torch.int32, device=device) + # row 1: fully valid + indices[1, 0, :] = torch.arange(topk, dtype=torch.int32, device=device) + # row 2: no valid key at all + + sm_scale = d_qk**-0.5 + + out, max_logits, lse = triton_bf16_mla_sparse_interface( + q, kv, indices, sm_scale, d_v + ) + assert out.isfinite().all() + + ref_out, _, ref_max_logits, ref_lse = reference_mla_sparse_prefill( + q, kv, indices, sm_scale, d_v + ) + assert torch.allclose(out[:2], ref_out[:2], atol=1e-2, rtol=1e-2) + assert torch.allclose(max_logits[:2], ref_max_logits[:2], atol=1e-3, rtol=1e-3) + assert torch.allclose(lse[:2], ref_lse[:2], atol=1e-3, rtol=1e-3) + # A row with no valid key yields zeros (the reference's convention); its + # lse/max_logits are large-negative finite rather than the reference's + # +inf/-inf placeholders, so only the output is compared here. + assert torch.allclose(out[2], torch.zeros_like(out[2])) diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py new file mode 100644 index 000000000000..290d155342bf --- /dev/null +++ b/tests/kernels/conftest.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + + +@pytest.fixture(autouse=True) +def reset_default_torch_device(): + """Several kernel tests call torch.set_default_device without restoring + it, which poisons subsequent tests in the same pytest run (e.g. CPU + tensors silently created on CUDA). Restore the factory default after + every test. + """ + yield + torch.set_default_device(None) diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index f698c385fee7..f6fc0ab26702 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -197,6 +197,46 @@ def test_silu_and_mul_with_clamp( opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) +@pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) +@torch.inference_mode() +def test_masked_situ_and_mul( + default_vllm_config, + linear_beta: float, + dtype: torch.dtype, +) -> None: + """Masked SITU computes valid expert rows and preserves padded zeros.""" + device = CUDA_DEVICES[0] + num_experts, max_num_tokens, d = 4, 7, 512 + beta = 1.5 + input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device) + expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device) + output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device) + + torch.ops._C.masked_situ_and_mul( + output, input, expert_num_tokens, beta, linear_beta + ) + + gate, up = input.float().chunk(2, dim=-1) + expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if linear_beta > 0: + up = linear_beta * torch.tanh(up / linear_beta) + expected = (expected * up).to(dtype) + for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()): + torch.testing.assert_close( + output[expert, :num_tokens], + expected[expert, :num_tokens], + atol=get_default_atol(output), + rtol=get_default_rtol(output), + ) + assert torch.count_nonzero(output[expert, num_tokens:]) == 0 + + opcheck( + torch.ops._C.masked_situ_and_mul, + (output, input, expert_num_tokens, beta, linear_beta), + ) + + @pytest.mark.parametrize( "activation", [ diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index b6a70b19b03d..e5f2c878ff29 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -13,7 +13,7 @@ import pytest import torch -from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm +from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( diff --git a/tests/kernels/core/test_fused_qk_norm_rope.py b/tests/kernels/core/test_fused_qk_norm_rope.py index 43737f4f23b1..b91db79c6cf8 100644 --- a/tests/kernels/core/test_fused_qk_norm_rope.py +++ b/tests/kernels/core/test_fused_qk_norm_rope.py @@ -35,10 +35,12 @@ def _apply_qk_norm_rope( q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim) q_by_head = q_norm.forward_native(q_by_head) + assert isinstance(q_by_head, torch.Tensor) q = q_by_head.view(q.shape) k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim) k_by_head = k_norm.forward_native(k_by_head) + assert isinstance(k_by_head, torch.Tensor) k = k_by_head.view(k.shape) q, k = rope.forward_native(positions, q, k) diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index 255833c48dc8..06df13d48c42 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/tests/kernels/core/test_fused_quant_layernorm.py @@ -8,7 +8,7 @@ import torch import vllm._custom_ops as ops -from tests.kernels.utils import fp8_ulp_distance, opcheck +from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, @@ -19,6 +19,12 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +ON_GFX950 = False +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 + + ON_GFX950 = on_gfx950() + DTYPES = [torch.bfloat16, torch.float] QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()] VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029] @@ -31,7 +37,7 @@ ADD_RESIDUAL = [False, True] SCALE_UBS = [True, False] -GROUP_SIZES = [None, [1, 64], [1, 128]] +GROUP_SIZES = [[1, 64], [1, 128]] TMA_ALIGNMENTS = [0, 4] SEEDS = [0] CUDA_DEVICES = [ @@ -40,6 +46,86 @@ EPS = 1e-6 + +def _is_valid_config( + hidden_size: int, + has_scale_ub: bool, + quant_dtype: torch.dtype, + group_size: list[int] | None, + tma_alignment: int, +) -> bool: + if group_size is not None and hidden_size % group_size[1] != 0: + return False + if group_size is not None and has_scale_ub: + return False + if ( + group_size is None or quant_dtype != current_platform.fp8_dtype() + ) and tma_alignment != 0: + return False + if ( + group_size is not None + and tma_alignment != 0 + and hidden_size // group_size[1] % tma_alignment == 0 + ): + return False + return not (has_scale_ub and quant_dtype != current_platform.fp8_dtype()) + + +def _config_id( + num_tokens: int, + hidden_size: int, + has_scale_ub: bool, + quant_dtype: torch.dtype, + group_size: list[int] | None, + tma_alignment: int, +) -> str: + quant = str(quant_dtype).removeprefix("torch.") + group = "per-token" if group_size is None else f"group-{group_size[1]}" + return ( + f"{num_tokens}x{hidden_size}-{quant}-{group}-" + f"tma-{tma_alignment}-scale-ub-{has_scale_ub}" + ) + + +# Filter unsupported combinations during collection. Letting each case call +# pytest.skip still runs the global per-test teardown and dominates this suite. +RMS_NORM_CONFIGS = [ + pytest.param( + num_tokens, + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + id=_config_id( + num_tokens, + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + ), + ) + for ( + (num_tokens, hidden_size), + has_scale_ub, + quant_dtype, + (group_size, tma_alignment), + ) in itertools.product( + NUM_TOKENS_HIDDEN_SIZES, + SCALE_UBS, + QUANT_DTYPES, + [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)], + ) + if _is_valid_config( + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + ) +] + ## Helpers @@ -154,15 +240,19 @@ def ops_impl( ) -@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES) -@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) -@pytest.mark.parametrize("has_scale_ub", SCALE_UBS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES) @pytest.mark.parametrize( - "group_size, tma_alignment", - [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)], + ( + "num_tokens", + "hidden_size", + "has_scale_ub", + "quant_dtype", + "group_size", + "tma_alignment", + ), + RMS_NORM_CONFIGS, ) +@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) +@pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("strided_input", [False, True]) @@ -185,32 +275,6 @@ def test_rms_norm( torch.set_default_device(device) torch.accelerator.set_device_index(device) - if group_size is not None and hidden_size % group_size[1] != 0: - # skip - pytest.skip("Skip non-divisible group sizes") - - if group_size is not None and has_scale_ub: - # blockwise baseline doesn't support scale_ub - pytest.skip("scale_ub not supported for blockwise/group quantization") - - if ( - group_size is None or quant_dtype != current_platform.fp8_dtype() - ) and tma_alignment != 0: - # TMA alignment is only supported for groupwise fp8 kernels - pytest.skip("tma alignment not supported for per-token or int8 quantization") - - if ( - group_size is not None - and tma_alignment != 0 - and hidden_size // group_size[1] % tma_alignment == 0 - ): - # Skip tests where TMA alignment doesn't create extra padding to save time - pytest.skip("Skip TMA alignment cases where no extra padding is added") - - if has_scale_ub and quant_dtype != current_platform.fp8_dtype(): - # skip - pytest.skip("scale_ub only supported for fp8 quantization") - layer = RMSNorm(hidden_size, EPS).to(dtype=dtype) # Make weights @@ -259,6 +323,13 @@ def test_rms_norm( and dtype == torch.bfloat16 and current_platform.is_rocm() ) + use_gfx950_fp8_allclose = ( + current_platform.is_rocm() + and ON_GFX950 + and group_size is None + and dtype == torch.bfloat16 + and quant_dtype == current_platform.fp8_dtype() + ) def scales_close(rtol: float, atol: float) -> bool: if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol): @@ -283,6 +354,10 @@ def scales_close(rtol: float, atol: float) -> bool: ulp = fp8_ulp_distance(ref_out, ops_out) max_outliers = ulp.numel() // 100_000 + 8 ok = int((ulp > 0).sum().item()) <= max_outliers + elif use_gfx950_fp8_allclose: + # Valid gfx950 reduction trees can straddle an E4M3 boundary. + ok = fp8_allclose(ops_out, ref_out, rtol=0.125, atol=2e-3) + ok = ok and int(fp8_ulp_distance(ops_out, ref_out).max()) <= 1 else: # CUDA (& non-bf16): compare dequantized values with relaxed tolerance. if group_size is None: diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 69788e37721c..f411f563ac77 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -7,7 +7,9 @@ import pytest import torch -from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import ( + FusedRMSNormGated, +) from vllm.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16] @@ -47,6 +49,11 @@ def test_compiled_vs_eager( device=device, dtype=dtype, ) + # Model parameters use torch.empty because checkpoint loading overwrites + # them. Initialize the standalone test module so allocator contents cannot + # introduce NaNs and make this comparison flaky. + if module.weight is not None: + module.weight.uniform_(-1, 1) x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) g = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) @@ -92,6 +99,8 @@ def test_compiled_vs_eager_multidim( device=device, dtype=dtype, ) + if module.weight is not None: + module.weight.uniform_(-1, 1) x = torch.randn(*shape, dtype=dtype, device=device) g = torch.randn(*shape, dtype=dtype, device=device) diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index 6e546f154c2e..a1ca0c09cf66 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -5,7 +5,7 @@ import torch from tests.kernels.quant_utils import FP8_DTYPE -from tests.kernels.utils import fp8_ulp_distance, opcheck +from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck from vllm import ir from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.platforms import current_platform @@ -60,7 +60,9 @@ def test_rms_norm( x = x[..., :hidden_size] assert x.is_contiguous() != strided_input x *= scale - residual = torch.randn_like(x) * scale if add_residual else None + residual = x.new_empty_strided(x.size(), x.stride()) if add_residual else None + if residual is not None: + residual.normal_(std=scale) # NOTE(woosuk): The reference implementation should be executed first # because the custom kernel is in-place. @@ -205,14 +207,22 @@ def test_fused_rms_norm_quant( ) if current_platform.is_rocm(): - # Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie; - # tolerate a tiny number of isolated fp8 outliers on ROCm. - ulp = fp8_ulp_distance(out_quant, out_quant_fused) - max_outliers = ulp.numel() // 100_000 + 8 - num_outliers = int((ulp > 0).sum().item()) - assert num_outliers <= max_outliers, ( - f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})" - ) + from vllm.platforms.rocm import on_gfx950 + + if on_gfx950() and dtype == torch.float16 and not add_residual: + # Fusion may round normalized FP16 across an E4M3 boundary on gfx950. + assert fp8_allclose(out_quant_fused, out_quant, rtol=0.125, atol=2e-3) + assert int(fp8_ulp_distance(out_quant_fused, out_quant).max()) <= 1 + else: + # Fused and unfused FP8 paths can land on opposite sides of an E4M3 + # tie; tolerate a tiny number of isolated fp8 outliers on ROCm. + ulp = fp8_ulp_distance(out_quant, out_quant_fused) + max_outliers = ulp.numel() // 100_000 + 8 + num_outliers = int((ulp > 0).sum().item()) + assert num_outliers <= max_outliers, ( + f"FP8 quant mismatch: {num_outliers} fp8 outliers " + f"(allowed {max_outliers})" + ) else: torch.testing.assert_close( out_quant.to(dtype=torch.float32), diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 29051b4a00cc..6f64fabfe9b9 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -38,6 +38,7 @@ def generate_test_data( class MRoPETestInfo(NamedTuple): model_name: str + is_neox_style: bool = True # https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317 atol: float = 1e-2 rtol: float = 1.6e-2 @@ -45,7 +46,10 @@ class MRoPETestInfo(NamedTuple): MODELS_TO_TEST = [ - MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"), + MRoPETestInfo( + model_name="zai-org/GLM-4.1V-9B-Thinking", + is_neox_style=False, + ), MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"), @@ -92,7 +96,7 @@ def test_mrope( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings @@ -162,7 +166,7 @@ def test_mrope_torch_compile_tracing( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings mrope_helper_class = get_rope( diff --git a/tests/kernels/core/test_vit_fp8_attn.py b/tests/kernels/core/test_vit_fp8_attn.py index ef1c44cada29..a677fc942ab7 100644 --- a/tests/kernels/core/test_vit_fp8_attn.py +++ b/tests/kernels/core/test_vit_fp8_attn.py @@ -34,10 +34,9 @@ def _has_flashinfer_cudnn() -> bool: @pytest.fixture def _fp8_attention(): """Create FP8-enabled MMEncoderAttention via config.""" - from types import SimpleNamespace - from unittest.mock import patch + from unittest.mock import MagicMock, patch - from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.config.multimodal import MultiModalConfig if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): @@ -45,7 +44,7 @@ def _fp8_attention(): mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8") vllm_config = VllmConfig() - vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + vllm_config.model_config = MagicMock(spec=ModelConfig, multimodal_config=mm_config) # MMEncoderAttention reads torch.get_default_dtype() during init # to determine the output dtype. In real model loading this is bf16. diff --git a/tests/kernels/core/test_vit_fp8_scaling.py b/tests/kernels/core/test_vit_fp8_scaling.py index a197439237fb..d756bfedbb22 100644 --- a/tests/kernels/core/test_vit_fp8_scaling.py +++ b/tests/kernels/core/test_vit_fp8_scaling.py @@ -4,8 +4,7 @@ import contextlib import json -from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -14,18 +13,33 @@ _FP8_AMAX_HISTORY_LEN, _FP8_MAX, ) +from vllm.platforms import current_platform from vllm.utils.flashinfer import ( is_flashinfer_cudnn_fp8_prefill_attn_supported, ) +pytestmark = pytest.mark.skipif( + not is_flashinfer_cudnn_fp8_prefill_attn_supported(), + reason="FlashInfer cuDNN FP8 prefill attention not supported", +) + LAYER_0 = "visual.blocks.0.attn.attn" LAYER_1 = "visual.blocks.1.attn.attn" NUM_HEADS = 16 HEAD_DIM = 72 +def _is_mi3xx() -> bool: + if not current_platform.is_rocm(): + return False + + from vllm.platforms.rocm import on_mi3xx + + return on_mi3xx() + + @contextlib.contextmanager -def _build_attention(mm_config): +def _build_attention(mm_config, attn_backend=None): """Yield an MMEncoderAttention with the given multimodal config. The VllmConfig context stays active while the test runs so that @@ -33,25 +47,31 @@ def _build_attention(mm_config): invokes ``process_weights_after_loading`` to simulate the model loader's auto-scan. Yields ``None`` if FlashInfer cuDNN is not available. """ - from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.attention.mm_encoder_attention import ( MMEncoderAttention, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum - if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + if attn_backend is None: + attn_backend = AttentionBackendEnum.FLASHINFER + + if ( + attn_backend == AttentionBackendEnum.FLASHINFER + and not is_flashinfer_cudnn_fp8_prefill_attn_supported() + ): yield None return vllm_config = VllmConfig() - vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + vllm_config.model_config = MagicMock(spec=ModelConfig, multimodal_config=mm_config) with ( set_current_vllm_config(vllm_config), patch( "vllm.model_executor.layers.attention.mm_encoder_attention" ".get_vit_attn_backend", - return_value=AttentionBackendEnum.FLASHINFER, + return_value=attn_backend, ), ): attn = MMEncoderAttention( @@ -164,9 +184,34 @@ def test_static_scales_loaded(_make_static_attention) -> None: assert not hasattr(attn, "_fp8_q_amax") +@pytest.mark.skipif( + not _is_mi3xx(), + reason="AITER FP8 attention requires MI300/MI350", +) +def test_aiter_static_scales_loaded(tmp_path) -> None: + """Verify AITER reuses the existing static FP8 scale loading path.""" + from vllm.config.multimodal import MultiModalConfig + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + scale_file = tmp_path / "aiter_scales.json" + scale_file.write_text(json.dumps({LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0}})) + mm_config = MultiModalConfig( + mm_encoder_attn_dtype="fp8", + mm_encoder_fp8_scale_path=str(scale_file), + ) + + with _build_attention(mm_config, AttentionBackendEnum.ROCM_AITER_FA) as attn: + assert attn is not None + assert attn.fp8_enabled + assert not attn._fp8_dynamic_scale + assert attn._fp8_q_scale.item() == 224.0 + assert attn._fp8_k_scale.item() == 198.0 + assert attn._fp8_v_scale.item() == 210.0 + + def test_static_scales_missing_layer(tmp_path) -> None: """Verify error when requested layer is not in the scale file.""" - from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.config.multimodal import MultiModalConfig from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -182,7 +227,7 @@ def test_static_scales_missing_layer(tmp_path) -> None: mm_encoder_fp8_scale_path=str(scale_file), ) vllm_config = VllmConfig() - vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + vllm_config.model_config = MagicMock(spec=ModelConfig, multimodal_config=mm_config) from vllm.model_executor.layers.attention.mm_encoder_attention import ( MMEncoderAttention, diff --git a/tests/kernels/helion/test_benchmark_script.py b/tests/kernels/helion/test_benchmark_script.py new file mode 100644 index 000000000000..a42eea00c43b --- /dev/null +++ b/tests/kernels/helion/test_benchmark_script.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.utils.import_utils import has_helion + +if not has_helion(): + pytest.skip("Helion is not installed", allow_module_level=True) + +from scripts import benchmark_helion_kernels + +check_correctness = benchmark_helion_kernels.check_correctness + + +class _FakeKernel: + helion_settings = None + + def __init__(self, offset: float): + self.calls = 0 + self.offset = offset + + def __call__(self, output: torch.Tensor, input: torch.Tensor) -> torch.Tensor: + self.calls += 1 + output.copy_(input + self.offset) + return input * 2 + + +def test_check_correctness_runs_once_without_mutating_benchmark_inputs(): + kernel = _FakeKernel(offset=1) + baseline_calls = 0 + + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor: + nonlocal baseline_calls + baseline_calls += 1 + output.copy_(input + 1) + return input * 2 + + output = torch.zeros(4) + inputs = (output, torch.arange(4)) + + check_correctness(kernel, baseline, inputs, "matching") + + assert kernel.calls == 1 + assert baseline_calls == 1 + torch.testing.assert_close(output, torch.zeros(4)) + + +def test_check_correctness_reports_mutated_input_mismatch(): + kernel = _FakeKernel(offset=2) + + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor: + output.copy_(input + 1) + return input * 2 + + inputs = (torch.zeros(4), torch.arange(4)) + + with pytest.raises(AssertionError, match="Numerics check failed for case bad"): + check_correctness(kernel, baseline, inputs, "bad") + + +def test_check_correctness_reports_return_value_mismatch(): + kernel = _FakeKernel(offset=1) + + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor: + output.copy_(input + 1) + return input * 3 + + inputs = (torch.zeros(4), torch.arange(4)) + + with pytest.raises(AssertionError, match="Numerics check failed for case bad"): + check_correctness(kernel, baseline, inputs, "bad") + + +def test_log_versions(monkeypatch): + messages = [] + + def record(message: str, *args): + messages.append(message % args) + + monkeypatch.setattr(benchmark_helion_kernels.logger, "info", record) + + benchmark_helion_kernels.log_versions() + + assert [message.split(":", 1)[0] for message in messages] == [ + "torch", + "helion", + "triton", + ] diff --git a/tests/kernels/helion/test_rms_norm_per_block_quant.py b/tests/kernels/helion/test_rms_norm_per_block_quant.py index 4cb18d775985..dacbe7203f04 100644 --- a/tests/kernels/helion/test_rms_norm_per_block_quant.py +++ b/tests/kernels/helion/test_rms_norm_per_block_quant.py @@ -136,6 +136,19 @@ def test_config_picker_fallback_to_largest(self): {"hidden_size": 4096, "group_size": 128, "num_tokens": 32} ) + def test_b200_configs_disable_reduction_warp_specialization(self): + config_set = ConfigManager.get_instance().load_config_set( + "rms_norm_per_block_quant" + ) + configs = config_set.to_dict()["nvidia_b200"].values() + + # TODO: Remove once the Triton pin includes + # https://github.com/triton-lang/triton/pull/9716. Tracked by + # https://github.com/triton-lang/triton/issues/10901. + assert all( + config["range_warp_specializes"][1] is not True for config in configs + ) + DTYPES = [torch.bfloat16, torch.float] QUANT_DTYPES = [torch.int8, FP8_DTYPE] diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index 6fc9d7543d15..c87fb72845de 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -18,20 +18,21 @@ rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn +IS_GPGPU_DEVICE = current_platform.is_cuda_alike() or current_platform.is_xpu() + @pytest.mark.skipif( - not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + not IS_GPGPU_DEVICE, reason="Currently only kernels on CUDA, ROCm and XPU", ) def test_rms_norm_registration(): expected = { "native": True, - "vllm_c": current_platform.is_cuda_alike(), + "vllm_c": IS_GPGPU_DEVICE, "aiter": current_platform.is_rocm(), "oink": current_platform.has_device_capability(100) and hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, "rmsnorm"), - "xpu_kernels": current_platform.is_xpu(), } actual = { @@ -46,17 +47,17 @@ def test_rms_norm_registration(): @pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES) @pytest.mark.parametrize("epsilon", [1e-6, 1e-5]) @pytest.mark.skipif( - not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + not IS_GPGPU_DEVICE, reason="Currently only kernels on CUDA, ROCm and XPU", ) class TestRMSNorm: - @classmethod - def setup_class(cls, **kwargs): - torch.set_default_device(current_platform.device_type) - def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): x, weight, epsilon = ir.ops.rms_norm.generate_inputs( - num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon + num_tokens=4, + hidden_size=8, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) out = rms_norm_native(x, weight, epsilon=epsilon) @@ -87,7 +88,11 @@ def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): impl = ir.ops.rms_norm.impls[provider] x, weight, eps = ir.ops.rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = (x, weight, eps) @@ -113,13 +118,17 @@ def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps) assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight) - @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"]) + @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "native"]) def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider): if not ir.ops.rms_norm.impls[provider].supported: pytest.skip(f"{provider} impl not supported on this platform") args = ir.ops.rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) # When checking the torch op, we have to set priority and use dispatch @@ -132,11 +141,14 @@ def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider): reason="aiter is only supported on ROCm", ) def test_aiter_rejects_unsupported_dtypes(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.rms_norm.impls["aiter"] for dtype in [torch.float32, torch.float64]: args = ir.ops.rms_norm.generate_inputs( - num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5 + num_tokens=8, + hidden_size=4096, + dtype=dtype, + epsilon=1e-5, + device=current_platform.device_type, ) assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}" @@ -146,15 +158,39 @@ def test_aiter_rejects_unsupported_dtypes(): reason="ROCm vllm_c RMSNorm needs explicit ND input handling", ) def test_vllm_c_rms_norm_accepts_nd_input(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.rms_norm.impls["vllm_c"] if not impl.supported: pytest.skip("vllm_c impl not supported on this platform") - base = torch.randn(3, 8, 192, dtype=torch.float16) + base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) x = base.split(64, dim=-1)[0].view(3, 8, 4, 16) assert not x.is_contiguous() - weight = torch.randn(16, dtype=torch.float16) + weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type) + epsilon = 1e-5 + + output = impl.impl_fn(x, weight, epsilon) + ref_output = rms_norm_native(x, weight, epsilon) + + assert output.shape == x.shape + assert_close(ir.ops.rms_norm, output, ref_output) + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm vllm_c RMSNorm needs a contiguous output for strided inputs", +) +def test_vllm_c_rms_norm_accepts_transposed_input(): + impl = ir.ops.rms_norm.impls["vllm_c"] + if not impl.supported: + pytest.skip("vllm_c impl not supported on this platform") + + x = torch.randn( + 1, 320, 120, dtype=torch.float16, device=current_platform.device_type + ).transpose(1, 2) + assert x.reshape(-1, x.shape[-1]).stride(-1) != 1 + weight = torch.randn(320, dtype=torch.float16, device=current_platform.device_type) epsilon = 1e-5 output = impl.impl_fn(x, weight, epsilon) @@ -168,18 +204,17 @@ def test_vllm_c_rms_norm_accepts_nd_input(): @pytest.mark.skipif( - not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + not IS_GPGPU_DEVICE, reason="Currently only kernels on CUDA, ROCm and XPU", ) def test_fused_add_rms_norm_registration(): expected = { "native": True, - "vllm_c": current_platform.is_cuda_alike(), + "vllm_c": IS_GPGPU_DEVICE, "aiter": current_platform.is_rocm(), "oink": current_platform.has_device_capability(100) and hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, "fused_add_rms_norm"), - "xpu_kernels": current_platform.is_xpu(), } actual = { @@ -195,18 +230,21 @@ def test_fused_add_rms_norm_registration(): reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling", ) def test_vllm_c_fused_add_rms_norm_accepts_nd_input(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.fused_add_rms_norm.impls["vllm_c"] if not impl.supported: pytest.skip("vllm_c impl not supported on this platform") - base = torch.randn(3, 8, 192, dtype=torch.float16) - residual_base = torch.randn(3, 8, 192, dtype=torch.float16) + base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) + residual_base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) x = base.split(64, dim=-1)[0].view(3, 8, 4, 16) x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16) assert not x.is_contiguous() assert not x_residual.is_contiguous() - weight = torch.randn(16, dtype=torch.float16) + weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type) epsilon = 1e-5 output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon) @@ -223,17 +261,17 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input(): @pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES) @pytest.mark.parametrize("epsilon", [1e-6, 1e-5]) @pytest.mark.skipif( - not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + not IS_GPGPU_DEVICE, reason="Currently only kernels on CUDA, ROCm and XPU", ) class TestFusedAddRMSNorm: - @classmethod - def setup_class(cls, **kwargs): - torch.set_default_device(current_platform.device_type) - def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon + num_tokens=4, + hidden_size=8, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps) @@ -278,7 +316,11 @@ def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): impl = ir.ops.fused_add_rms_norm.impls[provider] x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = (x, x_residual, weight, eps, None) @@ -324,7 +366,11 @@ def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider pytest.skip(f"{provider} impl not supported on this platform") x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) # Test default overload - should NOT modify inputs even with inplace impl @@ -368,7 +414,11 @@ def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm)) def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider): args = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = args + (None,) # Add variance_size parameter diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index bd30bc4f1ce0..f43a9b4a58c4 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -2,14 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import types import pytest import torch import torch.nn.functional as F import vllm._custom_ops as ops +from vllm.model_executor.layers.mamba.ops.cpu import gdn_attention from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata if not current_platform.is_cpu(): pytest.skip("skipping CPU-only tests", allow_module_level=True) @@ -24,6 +27,13 @@ (32, 32), (64, 32), ] +# chunk_gated_delta_rule_cpu (the chunked-prefill kernel) only supports +# head_dim == head_dim_v in {64, 128}; the decode-path update kernel above +# has no such restriction and keeps using the wider HEAD_DIMS list. +CHUNK_HEAD_DIMS = [ + (64, 64), + (128, 128), +] CHUNK_SIZE = 64 CONV_DIM = 128 CONV_KERNEL = 4 @@ -256,7 +266,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu( # prefill path @pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_dims", HEAD_DIMS) +@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS) @torch.inference_mode() def test_chunk_gated_delta_rule_cpu( seq_lens: list[int], @@ -305,6 +315,7 @@ def test_chunk_gated_delta_rule_cpu( cu_seqlens=cu_seqlens, head_first=False, use_qk_l2norm_in_kernel=True, + initial_state_indices=torch.arange(len(seq_lens), dtype=torch.int32), ) torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) @@ -329,7 +340,7 @@ def test_chunk_gated_delta_rule_cpu( @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_dims", HEAD_DIMS) +@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS) @torch.inference_mode() def test_chunk_gated_delta_rule_cpu_two_call_split( total_tokens: int, @@ -368,6 +379,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split( cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32), head_first=False, use_qk_l2norm_in_kernel=True, + initial_state_indices=torch.zeros(1, dtype=torch.int32), ) # Call 1: tokens [0:split], no initial state, capture final state. @@ -382,6 +394,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split( cu_seqlens=torch.tensor([0, split], dtype=torch.int32), head_first=False, use_qk_l2norm_in_kernel=True, + initial_state_indices=torch.zeros(1, dtype=torch.int32), ) # Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens # rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk. @@ -397,6 +410,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split( cu_seqlens=torch.tensor([0, tail], dtype=torch.int32), head_first=False, use_qk_l2norm_in_kernel=True, + initial_state_indices=torch.zeros(1, dtype=torch.int32), ) out_split = torch.cat([out1, out2], dim=1) @@ -417,6 +431,221 @@ def _conv_inputs(total_tokens: int): return x, weight, bias +def _sd_conv_states( + num_slots: int, state_len: int, dim: int = CONV_DIM +) -> torch.Tensor: + storage = torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16) + return storage.transpose(1, 2) + + +def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor: + return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight + + +@torch.inference_mode() +def test_spec_aware_mixed_routing_preserves_token_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + num_tokens = 4 + projection = torch.arange(num_tokens * 16, dtype=torch.float32).view(num_tokens, 16) + mixed_qkv, b, a = projection[:, :4], projection[:, 4:6], projection[:, 6:8] + assert all(not tensor.is_contiguous() for tensor in (mixed_qkv, b, a)) + + spec_indices = torch.tensor([0, 2]) + nonspec_indices = torch.tensor([1, 3]) + metadata = GDNAttentionMetadata( + num_prefills=1, + num_prefill_tokens=0, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_actual_tokens=0, + spec_sequence_masks=torch.ones(1, dtype=torch.bool), + spec_token_indx=spec_indices, + non_spec_token_indx=nonspec_indices, + ) + routed = [] + + def record(*args): + routed.append(args[2:5]) + return args[2] + + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) + monkeypatch.setattr(gdn_attention, "_spec_forward", record) + monkeypatch.setattr(gdn_attention, "_spec_aware_nonspec_subset", record) + + layer = types.SimpleNamespace( + kv_cache=[torch.empty(1, 4, 6), torch.empty(1, 1, 1, 1)] + ) + core_attn_out = torch.empty_like(mixed_qkv) + gdn_attention._cpu_gdn_attention_spec_aware( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + width=CONV_KERNEL, + state_len=6, + ) + + expected_inputs = (mixed_qkv, b, a) + assert len(routed) == 2 + for actual_inputs, indices in zip(routed, (spec_indices, nonspec_indices)): + for actual, expected in zip(actual_inputs, expected_inputs): + assert actual.is_contiguous() + torch.testing.assert_close(actual, expected.index_select(0, indices)) + torch.testing.assert_close(core_attn_out, mixed_qkv) + + +@torch.inference_mode() +def test_spec_aware_nonspec_materializes_state_indices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + block_table = torch.arange(8, dtype=torch.int32).view(2, 4) + state_indices = block_table[:, 0] + assert not state_indices.is_contiguous() + + metadata = GDNAttentionMetadata( + num_prefills=2, + num_prefill_tokens=4, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_actual_tokens=4, + non_spec_state_indices_tensor=state_indices, + non_spec_query_start_loc=torch.tensor([0, 2, 4], dtype=torch.int32), + has_initial_state=torch.tensor([False, False]), + ) + + recorded_indices = None + + def causal_conv1d_fwd_cpu(**kwargs): + nonlocal recorded_indices + recorded_indices = kwargs["cache_indices"] + return kwargs["x"] + + def fused_gdn_gating_cpu(**kwargs): + return kwargs["a"], kwargs["b"] + + def chunk_gated_delta_rule_cpu(**kwargs): + out = torch.zeros(1, 4, 1, 1) + return out, kwargs["initial_state"] + + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + monkeypatch.setattr( + gdn_attention.ops, "causal_conv1d_fwd_cpu", causal_conv1d_fwd_cpu + ) + monkeypatch.setattr(gdn_attention.ops, "fused_gdn_gating_cpu", fused_gdn_gating_cpu) + monkeypatch.setattr( + gdn_attention.ops, + "chunk_gated_delta_rule_cpu", + chunk_gated_delta_rule_cpu, + ) + + layer = types.SimpleNamespace( + activation="silu", + conv1d=types.SimpleNamespace(weight=torch.empty(0), bias=None), + A_log=torch.empty(0), + dt_bias=torch.empty(0), + rearrange_mixed_qkv=lambda x: ( + x[:, :1].view(1, 4, 1, 1), + x[:, :1].view(1, 4, 1, 1), + x[:, :1].view(1, 4, 1, 1), + ), + ) + gdn_attention._spec_aware_nonspec( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv=torch.zeros(4, 4), + b=torch.zeros(4, 1), + a=torch.zeros(4, 1), + core_attn_out=torch.zeros(4, 1, 1), + conv_buf=torch.zeros(8, 1, 6), + ssm_state=torch.zeros(8, 1, 1, 1), + width=4, + ) + + assert recorded_indices is not None + assert recorded_indices.is_contiguous() + torch.testing.assert_close( + recorded_indices, torch.tensor([0, 4], dtype=torch.int32) + ) + + +@torch.inference_mode() +def test_spec_forward_prepares_native_conv_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + block_table = torch.tensor([[0, 1], [4, 5]], dtype=torch.int32) + state_indices = block_table[:, 0] + accepted_counts = torch.tensor([1, 4], dtype=torch.int32) + assert not state_indices.is_contiguous() + metadata = GDNAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=2, + num_spec_decode_tokens=8, + num_actual_tokens=8, + spec_state_indices_tensor=block_table, + spec_query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32), + num_accepted_tokens=accepted_counts, + ) + forwarded_indices = None + forwarded_counts = None + + def causal_conv1d_update_cpu(**kwargs): + nonlocal forwarded_counts, forwarded_indices + forwarded_indices = kwargs["conv_state_indices"] + forwarded_counts = kwargs["num_accepted_tokens"] + return kwargs["x"] + + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + monkeypatch.setattr( + gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu + ) + monkeypatch.setattr( + gdn_attention.ops, + "fused_sigmoid_gating_delta_rule_update_spec_cpu", + lambda **kwargs: kwargs["q"], + ) + + layer = types.SimpleNamespace( + activation="silu", + conv1d=types.SimpleNamespace(weight=torch.empty(1, CONV_KERNEL), bias=None), + A_log=None, + dt_bias=None, + rearrange_mixed_qkv=lambda x: (x.unsqueeze(0),) * 3, + ) + gdn_attention._spec_forward( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv_spec=torch.zeros(8, 1, dtype=torch.bfloat16), + b_spec=torch.empty(0), + a_spec=torch.empty(0), + conv_buf=torch.empty(0), + ssm_state=torch.empty(0), + width=CONV_KERNEL, + state_len=0, + ) + + expected = ( + (forwarded_indices, torch.tensor([0, 4], dtype=torch.int32)), + (forwarded_counts, torch.tensor([1, 4], dtype=torch.int32)), + ) + for actual, reference in expected: + assert actual is not None + assert actual.is_contiguous() + assert actual.dtype == torch.int32 + torch.testing.assert_close(actual, reference) + + @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @torch.inference_mode() def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None: @@ -425,7 +654,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No match the single-call result. """ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( - causal_conv1d_torch, + causal_conv1d_fn_cpu as causal_conv1d_torch, ) x, weight, bias = _conv_inputs(total_tokens) @@ -475,7 +704,182 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No @pytest.mark.skipif( not torch.cpu._is_amx_tile_supported(), - reason="causal_conv1d_fwd_cpu requires AMX/AVX512", + reason="requires AMX support", +) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_accepts_wide_state() -> None: + state_len = CONV_KERNEL - 1 + wide_state_len = state_len + 5 + batch_size = 3 + is_vnni = True + x, weight, bias = _conv_inputs(batch_size) + conv_state_indices = torch.tensor([2, 0, 1], dtype=torch.int32) + + narrow_state = _sd_conv_states(batch_size, state_len) + narrow_state.copy_( + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) + ) + wide_state = _sd_conv_states(batch_size, wide_state_len) + wide_state[:, :, :state_len].copy_(narrow_state) + wide_state[:, :, state_len:].fill_(7) + wide_tail = wide_state[:, :, state_len:].clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out_narrow = ops.causal_conv1d_update_cpu( + x=x, + conv_states=narrow_state, + weight=conv_weight, + bias=bias, + silu_activation=True, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + ) + out_wide = ops.causal_conv1d_update_cpu( + x=x, + conv_states=wide_state, + weight=conv_weight, + bias=bias, + silu_activation=True, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + ) + + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 + ) + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) + + +def _ref_causal_conv1d_update_cpu_multi( + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> torch.Tensor: + batch_size, seq_len, dim = x.shape + state_len = conv_states.size(2) + conv_out = torch.empty_like(x) + conv_weight = weight.unsqueeze(1) + + for i in range(batch_size): + slot = int(conv_state_indices[i].item()) + offset = int(num_accepted_tokens[i].item()) - 1 + state = conv_states[slot] + x_seq = x[i].transpose(0, 1).to(state.dtype) + prior = state[:, offset : offset + CONV_KERNEL - 1] + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) + out = F.conv1d(conv_in, conv_weight, bias, groups=dim)[0] + if silu_activation: + out = F.silu(out) + conv_out[i] = out.transpose(0, 1).to(conv_out.dtype) + keep = state[:, offset + 1 : offset + 1 + (state_len - seq_len)] + state.copy_(torch.cat([keep, x_seq], dim=-1)) + + return conv_out + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@pytest.mark.parametrize( + ("batch_size, seq_len, accepted_counts, has_bias, silu_activation, is_vnni"), + [ + (1, 1, [1], False, False, False), + (1, 1, [1], True, True, True), + (4, 4, [1, 2, 3, 4], False, True, False), + (4, 4, [4, 3, 2, 1], True, False, True), + (4, 16, [1, 5, 10, 16], False, False, True), + (4, 16, [16, 10, 5, 1], True, True, False), + ], +) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_multi_token_matches_python( + batch_size: int, + seq_len: int, + accepted_counts: list[int], + has_bias: bool, + silu_activation: bool, + is_vnni: bool, +) -> None: + dim = 96 + state_len = seq_len + 2 + x = tensor_cache(batch_size * seq_len * dim, torch.bfloat16).view( + batch_size, seq_len, dim + ) + weight = tensor_cache(dim * CONV_KERNEL, torch.bfloat16).view(dim, CONV_KERNEL) + bias = tensor_cache(dim, torch.bfloat16) if has_bias else None + conv_state_indices = torch.arange(batch_size - 1, -1, -1, dtype=torch.int32) + num_accepted_tokens = torch.tensor(accepted_counts, dtype=torch.int32) + + conv_states_ref = _sd_conv_states(batch_size, state_len, dim) + conv_states_ref.copy_( + tensor_cache(conv_states_ref.numel(), torch.bfloat16).view_as(conv_states_ref) + ) + conv_states = conv_states_ref.clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out = ops.causal_conv1d_update_cpu( + x=x, + conv_states=conv_states, + weight=conv_weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + num_accepted_tokens=num_accepted_tokens, + ) + ref_out = _ref_causal_conv1d_update_cpu_multi( + x=x, + conv_states=conv_states_ref, + weight=weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + ) + + torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(conv_states, conv_states_ref, atol=0, rtol=0) + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@pytest.mark.parametrize("num_accepted", [0, 17]) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count( + num_accepted: int, +) -> None: + batch_size = 1 + seq_len = 16 + dim = 96 + state_len = seq_len + 2 + x = torch.zeros(batch_size, seq_len, dim, dtype=torch.bfloat16) + weight = torch.zeros(dim, CONV_KERNEL, dtype=torch.bfloat16) + conv_states = _sd_conv_states(batch_size, state_len, dim) + + with pytest.raises(RuntimeError, match="num_accepted_tokens must be in.*seqlen"): + ops.causal_conv1d_update_cpu( + x=x, + conv_states=conv_states, + weight=weight, + bias=None, + silu_activation=True, + conv_state_indices=torch.tensor([0], dtype=torch.int32), + is_vnni=False, + num_accepted_tokens=torch.tensor([num_accepted], dtype=torch.int32), + ) + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", ) @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @torch.inference_mode() @@ -515,6 +919,62 @@ def amx(x_seg, conv_states, has_init): torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@torch.inference_mode() +def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None: + state_len = CONV_KERNEL - 1 + wide_state_len = state_len + 5 + is_vnni = True + seq_lens = [CHUNK_SIZE - 1, CHUNK_SIZE + 5] + total_tokens = sum(seq_lens) + x, weight, bias = _conv_inputs(total_tokens) + query_start_loc = torch.tensor([0, seq_lens[0], total_tokens], dtype=torch.int32) + cache_indices = torch.tensor([2, 0], dtype=torch.int32) + has_initial_state = torch.tensor([True, False]) + + narrow_state = _sd_conv_states(3, state_len) + narrow_state.copy_( + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) + ) + wide_state = _sd_conv_states(3, wide_state_len) + wide_state[:, :, :state_len].copy_(narrow_state) + wide_state[:, :, state_len:].fill_(7) + wide_tail = wide_state[:, :, state_len:].clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out_narrow = ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=conv_weight, + bias=bias, + conv_states=narrow_state, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=is_vnni, + ) + out_wide = ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=conv_weight, + bias=bias, + conv_states=wide_state, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=is_vnni, + ) + + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 + ) + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) + + @torch.inference_mode() def test_batch_memcpy_cpu_fallback() -> None: """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index c6554f131fed..d2a2981edf2c 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -18,8 +18,12 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu()), - reason="causal_conv1d Triton kernels require CUDA-alike or XPU", + not ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() + or current_platform.is_cpu() + ), + reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU", ) @@ -284,7 +288,8 @@ def test_causal_conv1d_varlen( batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype ): device = DEVICE - torch.accelerator.empty_cache() + if not current_platform.is_cpu(): + torch.accelerator.empty_cache() rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) if itype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py index fa28f4701c61..b1c30303805f 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_split.py +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -174,7 +174,9 @@ def test_forward_core_split_matches_unified( batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) builder = GDNAttentionMetadataBuilder( - kv_cache_spec=MambaSpec( + # GDNAttentionMetadataBuilder declares kv_cache_spec as AttentionSpec but + # asserts isinstance(kv_cache_spec, MambaSpec). MambaSpec is correct here. + kv_cache_spec=MambaSpec( # type: ignore[arg-type] block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,) ), layer_names=[PREFIX], @@ -198,6 +200,7 @@ def test_forward_core_split_matches_unified( # Full-batch chunk metadata for the unified reference path, built the same # way the builder would for a non-split batch (backend-matched). cu_full = meta_split.non_spec_query_start_loc + assert cu_full is not None if builder.gdn_prefill_backend == "cutedsl": from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( prepare_metadata_cutedsl, @@ -227,6 +230,7 @@ def test_forward_core_split_matches_unified( ) # Size the state pools from the indices the builder actually produced. + assert meta_split.non_spec_state_indices_tensor is not None pool_size = int(meta_split.non_spec_state_indices_tensor.max().item()) + 1 conv_state_shape, temporal_state_shape = ( MambaStateShapeCalculator.gated_delta_net_state_shape( diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py index 1f5a24fd81ff..861b686e32da 100644 --- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py +++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py @@ -33,12 +33,10 @@ @pytest.mark.parametrize("num_seqs", [1, 5, 257]) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): - seq_lens = torch.randint( - 1, - 130, - (num_seqs,), - dtype=torch.int32, - ) + rng_cpu = torch.Generator("cpu").manual_seed(1234) + rng = torch.Generator("cuda").manual_seed(2345) + + seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu) cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32) cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0) total_tokens = int(cu_seqlens[-1].item()) @@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_k_dim, device="cuda", dtype=dtype, + generator=rng, ) - k = torch.randn_like(q) + k = torch.randn_like(q, generator=rng) v = torch.randn( 1, total_tokens, @@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_v_dim, device="cuda", dtype=dtype, + generator=rng, ) q = F.normalize(q.float(), p=2, dim=-1).to(dtype) k = F.normalize(k.float(), p=2, dim=-1).to(dtype) a = torch.randn( - 1, - total_tokens, - num_v_heads, - device="cuda", - dtype=dtype, + 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng ) b = torch.randn( - 1, - total_tokens, - num_v_heads, - device="cuda", - dtype=dtype, + 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng ) # Match upstream FLA GatedDeltaNet synthetic initialization: # https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py - A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16) + A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_( + 0, 16, generator=rng + ) A_log = torch.log(A) dt = torch.exp( - torch.rand(num_v_heads, device="cuda", dtype=torch.float32) + torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) ) @@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_k_dim, device="cuda", dtype=state_dtype, + generator=rng, ) * 0.05 ) diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index 81b57d4b16bd..f8af3db56be6 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -20,8 +20,12 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu()), - reason="mamba_ssm kernels require CUDA-alike or XPU", + not ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() + or current_platform.is_cpu() + ), + reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU", ) # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op, @@ -342,6 +346,13 @@ def test_selective_scan( @pytest.mark.parametrize("has_z", [False, True]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update(dim, dstate, has_z, itype): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) @pytest.mark.parametrize("max_seq_len", [1, 2, 4]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -697,6 +715,13 @@ def test_selective_scan_varlen( @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) # tests correctness in case subset of the sequences are padded @pytest.mark.parametrize("with_padding", [True, False]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update_with_batch_indices( with_padding, dim, dstate, has_z, itype ): @@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices( @pytest.mark.parametrize("ngroups", [1, 4]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update_with_heads_with_batch_indices( dim, dstate, ngroups, has_z, tie_hdim, itype ): @@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update_with_num_accepted_tokens( dim, dstate, has_z, itype, max_seq_len ): @@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), +) def test_selective_state_update_varlen_with_num_accepted( dim, dstate, has_z, itype, max_seq_len ): diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py index be1e45594860..b7a3da37d594 100644 --- a/tests/kernels/mamba/test_precopy_mamba_align.py +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Equivalence test for ``precopy_mamba_align_fused_kernel``. -The V2 "align" pre-copy must migrate mamba state across block boundaries with -byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` / -``get_temporal_copy_spec``): +The fused "align" pre-copy must migrate mamba state across block boundaries +with byte-identical semantics to the scalar V1 copy specs +(``get_conv_copy_spec`` / ``get_temporal_copy_spec``): * conv state (SD layout, conv_width > 0): shift the sliding window by ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` -> @@ -14,33 +14,48 @@ ``state[bt[dst_col]]``. The kernel must also no-op when ``src_col < 0`` (fresh request) or -``src_col == dst_col`` (no boundary crossed). +``src_col == dst_col`` (no boundary crossed). V2 callers pass an explicit +``idx_mapping``; V1 align preprocessing launches in batch order with +``idx_mapping=None``. """ from __future__ import annotations +from collections.abc import Callable +from types import SimpleNamespace +from typing import Any + +import numpy as np import torch +from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils from vllm.platforms import current_platform +from vllm.v1.worker import mamba_utils as worker_mamba_utils from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel +_parametrize: Callable[..., Callable[[Any], Any]] + try: import pytest - pytestmark = pytest.mark.skipif( + _cuda_required = pytest.mark.skipif( not current_platform.is_cuda(), reason="precopy_mamba_align_fused_kernel needs CUDA/Triton", ) _parametrize = pytest.mark.parametrize except ModuleNotFoundError: # allow running directly as ``python `` - pytest = None - def _parametrize(_name, _values): + def _cuda_required(fn): + return fn + + def _no_parametrize(_name, _values): def _deco(fn): return fn return _deco + _parametrize = _no_parametrize + NUM_LAYERS = 3 CONV_WIDTH = 4 # conv_kernel - 1 + num_spec @@ -49,13 +64,20 @@ def _deco(fn): MAX_COLS = 8 -def _build_state(num_blocks, device): +def _build_state(num_blocks, device, conv_state_dim_first): """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools.""" convs, ssms = [], [] for _ in range(NUM_LAYERS): + conv_shape = ( + (num_blocks, CONV_DIM, CONV_WIDTH) + if conv_state_dim_first + else (num_blocks, CONV_WIDTH, CONV_DIM) + ) convs.append( torch.randn( - num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device + *conv_shape, + dtype=torch.bfloat16, + device=device, ) ) ssms.append( @@ -64,7 +86,7 @@ def _build_state(num_blocks, device): return convs, ssms -def _build_meta(convs, ssms, device): +def _build_meta(convs, ssms, device, conv_state_dim_first): """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer.""" n = NUM_LAYERS * 2 base = torch.zeros(n, dtype=torch.int64, device=device) @@ -82,8 +104,14 @@ def _build_meta(convs, ssms, device): base[i] = conv.data_ptr() blk_stride[i] = conv.stride(0) * conv.element_size() elem[i] = conv.element_size() - width[i] = conv.size(1) - inner[i] = conv.stride(1) + if conv_state_dim_first: + width[i] = conv.size(2) + inner[i] = 1 + drc[i] = conv.size(1) + drs[i] = conv.stride(1) * conv.element_size() + else: + width[i] = conv.size(1) + inner[i] = conv.stride(1) i += 1 # ssm (temporal): width = 0, inner = elems per block base[i] = ssm.data_ptr() @@ -95,7 +123,7 @@ def _build_meta(convs, ssms, device): return base, blk_stride, elem, inner, width, group, drc, drs -def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs): +def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first): """Apply the V1 copy semantics on clones, reading from the pre-copy state.""" conv_pre = [c.clone() for c in convs] ssm_pre = [s.clone() for s in ssms] @@ -108,14 +136,24 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs): sblk, dblk = int(bt[r, sc]), int(bt[r, dc]) tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias for layer in range(NUM_LAYERS): - conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:] + if conv_dim_first: + conv_ref[layer][dblk, :, : CONV_WIDTH - tb] = conv_pre[layer][ + sblk, :, tb: + ] + else: + conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:] ssm_ref[layer][dblk] = ssm_pre[layer][tblk] return conv_ref, ssm_ref +@_parametrize("conv_state_dim_first", [False, True]) @_parametrize("num_reqs", [1, 4, 16]) @_parametrize("token_bias", [0, 1, 2]) -def test_precopy_matches_v1_copy_specs(num_reqs, token_bias): +@_parametrize("has_idx_mapping", [True, False]) +@_cuda_required +def test_precopy_matches_v1_copy_specs( + num_reqs, token_bias, has_idx_mapping, conv_state_dim_first +): device = torch.device("cuda") torch.manual_seed(0) # Distinct physical block per (req, col) so copies never alias. @@ -136,13 +174,20 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias): if num_reqs >= 2: dst_col[1] = 1 # src_col == dst_col -> no copy - convs, ssms = _build_state(num_blocks, device) + convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) conv_ref, ssm_ref = _reference( - convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs + convs, + ssms, + bt.cpu(), + src_col.cpu(), + dst_col.cpu(), + bias.cpu(), + num_reqs, + conv_state_dim_first, ) base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( - convs, ssms, device + convs, ssms, device, conv_state_dim_first ) bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) @@ -161,10 +206,11 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias): group, drc, drs, - idx_mapping, + idx_mapping if has_idx_mapping else None, num_reqs, COPY_BLOCK_SIZE=1024, - CONV_STATE_DIM_FIRST=False, + CONV_STATE_DIM_FIRST=conv_state_dim_first, + HAS_IDX_MAPPING=has_idx_mapping, ) torch.accelerator.synchronize() @@ -173,8 +219,206 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias): torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) +def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch): + monkeypatch.setattr( + layer_mamba_utils, + "is_conv_state_dim_first", + lambda: True, + ) + state = torch.empty((2, CONV_DIM, CONV_WIDTH), dtype=torch.bfloat16) + + with pytest.raises(AssertionError, match="num_accepted_tokens > 1"): + layer_mamba_utils.get_conv_copy_spec( + state=state, + block_ids=[0, 1], + cur_block_idx=0, + num_accepted_tokens=3, + ) + + +class _FakeCpuGpuBuffer: + def __init__(self, n): + self.np = np.zeros(n, dtype=np.int32) + self.gpu = object() + self.copy_sizes = [] + + def copy_to_gpu(self, n=None): + self.copy_sizes.append(n) + return self.gpu + + +class _FakePrecopyContext: + def __init__(self, n): + self.is_initialized = True + self.mamba_group_ids = [0] + self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n) + self.precopy_src_col_buf = _FakeCpuGpuBuffer(n) + self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n) + self.calls = [] + + def initialize_from_forward_context(self, *args, **kwargs): + raise AssertionError("test context is pre-initialized") + + def run_fused_precopy( + self, + *, + num_reqs, + state_idx_gpu, + src_col_gpu, + token_bias_gpu, + idx_mapping, + ): + self.calls.append( + { + "num_reqs": num_reqs, + "state_idx": self.mamba_state_idx_buf.np[:num_reqs].copy(), + "src_col": self.precopy_src_col_buf.np[:num_reqs].copy(), + "token_bias": self.precopy_token_bias_buf.np[:num_reqs].copy(), + "idx_mapping": idx_mapping, + } + ) + + +def _make_preprocess_case(token_bias): + req_ids = ["fresh", "same", "cross_a", "cross_b"] + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_cached_reqs=SimpleNamespace(resumed_req_ids=set()), + num_scheduled_tokens={ + "fresh": 1, + "same": 1, + "cross_a": 1, + "cross_b": 2, + }, + ) + input_batch = SimpleNamespace( + req_ids=req_ids, + num_accepted_tokens_cpu=np.array( + [token_bias + 1, token_bias + 1, token_bias + 1, 2], + dtype=np.int32, + ), + ) + requests = { + "fresh": SimpleNamespace(req_id="fresh", num_computed_tokens=0), + "same": SimpleNamespace(req_id="same", num_computed_tokens=5), + "cross_a": SimpleNamespace(req_id="cross_a", num_computed_tokens=8), + "cross_b": SimpleNamespace(req_id="cross_b", num_computed_tokens=7), + } + mamba_state_idx = {"same": 1, "cross_a": 0, "cross_b": 1} + return scheduler_output, input_batch, requests, mamba_state_idx + + +@_parametrize("token_bias", [1, 2]) +def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bias): + block_size = 4 + mamba_spec = SimpleNamespace(block_size=block_size, num_speculative_blocks=1) + # Fakes stand in for real config/context objects. The fused path tests + # below only touch the attributes preprocess_mamba actually reads. + cache_config: Any = SimpleNamespace(enable_prefix_caching=True) + kv_cache_config: Any = SimpleNamespace() + scalar_copy_calls = [] + + def fake_collect( + copy_bufs, + kv_cache_config, + mamba_state_copy_funcs, + mamba_group_ids, + src_block_idx, + dest_block_idx, + accept_token_bias, + req_state, + forward_context, + ): + scalar_copy_calls.append( + ( + req_state.req_id, + src_block_idx, + dest_block_idx, + accept_token_bias, + ) + ) + + monkeypatch.setattr(worker_mamba_utils, "collect_mamba_copy_meta", fake_collect) + monkeypatch.setattr( + worker_mamba_utils, "do_mamba_copy_block", lambda copy_bufs: None + ) + + scalar_case = _make_preprocess_case(token_bias) + fused_case = _make_preprocess_case(token_bias) + + scalar_copy_bufs: Any = SimpleNamespace( + mamba_group_ids=[0], + mamba_spec=mamba_spec, + offset=0, + ) + worker_mamba_utils.preprocess_mamba( + scheduler_output=scalar_case[0], + kv_cache_config=kv_cache_config, + cache_config=cache_config, + mamba_state_idx=scalar_case[3], + input_batch=scalar_case[1], + requests=scalar_case[2], + forward_context={}, + mamba_state_copy_funcs=(), + copy_bufs=scalar_copy_bufs, + ) + + ctx: Any = _FakePrecopyContext(len(fused_case[1].req_ids)) + fused_copy_bufs: Any = SimpleNamespace( + mamba_group_ids=[0], + mamba_spec=mamba_spec, + offset=0, + ) + worker_mamba_utils.preprocess_mamba( + scheduler_output=fused_case[0], + kv_cache_config=kv_cache_config, + cache_config=cache_config, + mamba_state_idx=fused_case[3], + input_batch=fused_case[1], + requests=fused_case[2], + forward_context={}, + mamba_state_copy_funcs=(), + copy_bufs=fused_copy_bufs, + align_ctx=ctx, + ) + + assert fused_case[3] == scalar_case[3] + np.testing.assert_array_equal( + fused_case[1].num_accepted_tokens_cpu, + scalar_case[1].num_accepted_tokens_cpu, + ) + assert scalar_copy_calls == [ + ("cross_a", 0, 2, token_bias), + ("cross_b", 1, 2, 1), + ] + assert len(ctx.calls) == 1 + call = ctx.calls[0] + assert call["num_reqs"] == len(fused_case[1].req_ids) + assert call["idx_mapping"] is None + np.testing.assert_array_equal(call["state_idx"], np.array([0, 1, 2, 2])) + np.testing.assert_array_equal(call["src_col"], np.array([-1, -1, 0, 1])) + np.testing.assert_array_equal(call["token_bias"], np.array([0, 0, token_bias, 1])) + fused_copy_calls = [ + (req_id, int(src), int(dst), int(bias)) + for req_id, src, dst, bias in zip( + fused_case[1].req_ids, + call["src_col"], + call["state_idx"], + call["token_bias"], + ) + if int(src) != -1 and int(src) != int(dst) + ] + assert fused_copy_calls == scalar_copy_calls + + if __name__ == "__main__": for nr in (1, 4, 16): for tb in (0, 1, 2): - test_precopy_matches_v1_copy_specs(nr, tb) - print(f"OK num_reqs={nr} token_bias={tb}") + for mapping in (True, False): + for dim_first in (False, True): + test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first) + print( + f"OK num_reqs={nr} token_bias={tb} " + f"has_idx_mapping={mapping} conv_dim_first={dim_first}" + ) diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py new file mode 100644 index 000000000000..5319d8964d34 --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prefill == decode equivalence for the Mamba2 ReplaySSM kernels. + +The SSM recurrence is path-independent: vLLM's chunked prefill (the SSD kernel +``mamba_chunk_scan_combined_varlen``) and step-by-step decode must produce the +same per-position outputs and final state over a sequence. This file feeds one +set of inputs through the production dt flow (raw dt + softplus + a per-head +dt_bias, applied inside each kernel) to: + + * the exact fp32 step recurrence (``selective_state_update_ref``) -- the + ground truth, + * the chunked prefill kernel, + * the baseline decode kernel, + * the ReplaySSM output_only decode kernel, + +and checks all of them agree. Prefill (a chunked scan) and decode (a step +recurrence) are different code paths, so they differ numerically: the chunked +scan carries ~2e-2 (fp32) / ~4e-2 (bf16) vs the exact recurrence, far above the +near-exact decode. We therefore anchor every path on the exact recurrence at +SSD-level tolerances (the same regime as ``test_mamba_ssm_ssd.py``), which the +chunked scan sets, keyed off the activation dtype. + +State and activation/buffer precision are swept independently, including the +fp32-state + bf16-activation production config. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import selective_state_update_ref +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.ssd_combined import ( + mamba_chunk_scan_combined_varlen, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata + + +def _prefill_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # The chunked prefill scan, not the decode, sets these: it carries ~2e-2 + # (fp32) / ~6e-2 (bf16) vs the exact recurrence, while ReplaySSM decode is + # near-exact (~2e-6 fp32). Keyed off the activation dtype (the outputs are + # in act_dtype). Same regime as test_mamba_ssm_ssd.py. + if act_dtype == torch.float32: + return 1e-2, 3e-2 + return 6e-2, 1e-1 + + +def _run_prefill_decode_equivalence( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + seqlen: int, + chunk_size: int, + max_cache_len: int, + seed: int = 0, +) -> None: + """Prefill the whole sequence and decode it step by step; check both match + the exact fp32 recurrence (and each other). All paths use the production dt + flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill + and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` + is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" + device = "cuda" + rtol, atol = _prefill_tolerances(act_dtype) + set_random_seed(seed) + + # Production dt flow: raw dt + a per-head dt_bias (~-4 keeps softplus(dt + + # bias) small and well-conditioned). dt_bias is (nheads,) for prefill and + # (nheads, headdim) for the decode kernels/reference. + A = -torch.exp(torch.rand(nheads, device=device, dtype=act_dtype)) + dt = torch.randn(seqlen, nheads, device=device, dtype=act_dtype) + dt_bias = torch.rand(nheads, device=device, dtype=act_dtype) - 4 + dt_bias_hd = dt_bias.view(nheads, 1).expand(nheads, headdim) + X = torch.randn(seqlen, nheads, headdim, device=device, dtype=act_dtype) + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + A_bcast = A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + # Chunked prefill over the whole sequence (implicit batch=1, varlen). The + # kernel always returns the final state in fp32, so no state_dtype plumbing. + cu_seqlens = torch.tensor((0, seqlen), device=device).cumsum(0).to(torch.int32) + cu_chunk_seqlens, last_chunk_indices, seq_idx = compute_varlen_chunk_metadata( + cu_seqlens, chunk_size + ) + y_prefill = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + final_state_prefill = mamba_chunk_scan_combined_varlen( + X, + dt, + A, + B, + C, + chunk_size, + cu_seqlens=cu_seqlens, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx, + out=y_prefill, + D=None, + dt_bias=dt_bias, + dt_softplus=True, + ) + + # Step paths: exact fp32 recurrence (ground truth), baseline, ReplaySSM. + # State follows state_dtype; caches follow act_dtype (dt_cache is fp32). + state_ref = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=torch.float32 + ) + state_base = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype + ) + state_dec = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype + ) + x_cache = torch.zeros( + 1, nheads, max_cache_len, headdim, device=device, dtype=act_dtype + ) + dt_cache = torch.zeros(1, nheads, max_cache_len, device=device, dtype=torch.float32) + B_cache = torch.zeros( + 1, ngroups, max_cache_len, dstate, device=device, dtype=act_dtype + ) + bc_pre = torch.empty(1, ngroups, max_cache_len, device=device, dtype=torch.float32) + write_pos = torch.zeros(1, dtype=torch.int32, device=device) + # No skip connection (D=0) on any path here; the D!=0 path is covered by the + # standard-decode suite. The baseline kernel needs a D tensor, not None. + D_zero = torch.zeros(nheads, headdim, device=device) + + y_ref = torch.empty(seqlen, nheads, headdim, device=device, dtype=torch.float32) + y_base = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + y_dec = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + for t in range(seqlen): + dt_t = dt[t].view(1, nheads, 1).expand(1, nheads, headdim) + is_flush = write_pos == max_cache_len - 1 + + y_ref[t] = selective_state_update_ref( + state_ref, + X[t : t + 1].float(), + dt_t.float(), + A_bcast.float(), + B[t : t + 1].float(), + C[t : t + 1].float(), + dt_bias=dt_bias_hd.float(), + dt_softplus=True, + )[0] + + out_b = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + selective_state_update( + state_base, + X[t : t + 1], + dt_t, + A_bcast, + B[t : t + 1], + C[t : t + 1], + D=D_zero, + dt_bias=dt_bias_hd, + dt_softplus=True, + out=out_b, + ) + y_base[t] = out_b[0] + + out_d = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + common = dict( + dt_bias=dt_bias_hd, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_d, + ) + selective_state_update_replayssm_output_only( + state_dec, + X[t : t + 1], + dt_t, + A_bcast, + B[t : t + 1], + C[t : t + 1], + bc_pre=bc_pre, + **common, + ) + y_dec[t] = out_d[0] + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + # Every path computes the same recurrence; anchor each on the fp32 truth. + torch.testing.assert_close(y_prefill.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_base.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_dec.float(), y_ref, rtol=rtol, atol=atol) + # Headline: ReplaySSM decode matches the chunked prefill directly. + torch.testing.assert_close(y_dec.float(), y_prefill.float(), rtol=rtol, atol=atol) + # Final state too (the recurrence ends in state_ref after the loop). + torch.testing.assert_close( + final_state_prefill[0].float(), state_ref[0], rtol=rtol, atol=atol + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize( + "precision", + # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16 + # appears as an activation dtype (s32_afp16, sfp16_afp16) and as a finer- + # mantissa state under bf16 activations (sfp16_a16). + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), + ], +) +@pytest.mark.parametrize( + "geometry", # (nheads, headdim, dstate, ngroups) + [ + pytest.param((8, 64, 64, 2), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), + ], +) +def test_replayssm_prefill_decode_equivalence( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_prefill_decode_equivalence( + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + seqlen=16, + chunk_size=8, + max_cache_len=4, + ) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py new file mode 100644 index 000000000000..588ce52dbeb6 --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -0,0 +1,673 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Standard (autoregressive) decode correctness for the Mamba2 ReplaySSM kernels. + +ReplaySSM caches the recent SSM inputs ``(x, dt, B)`` in a small ring buffer and +reconstructs / reads out the recurrent state on the fly, writing the full state +back to HBM only when the buffer flushes. This file checks that, over a +multi-step decode, the ReplaySSM output_only kernel reproduces the exact SSM +recurrence. + +For every step we assert, against trusted oracles driven one token at a time: + + * the ReplaySSM output matches its pure-PyTorch reference, which models the + kernel's exact arithmetic (including its bf16 reconstruction), at every + precision, + * when the state is fp32, the output also matches the baseline decode kernel + (``selective_state_update``); at bf16 state the baseline legitimately + differs -- it downcasts the fp32 state to bf16 every step, while ReplaySSM + accumulates a whole buffer in fp32 and is the more accurate path, + * when state and activations are both fp32, the output also matches the exact + elementwise ``selective_state_update_ref``, + * the cached inputs match the reference cache management, + * the checkpoint state matches the reference (and the baseline at fp32 state). + +State and activation/buffer precision are swept independently. Nemotron-3 +defaults to fp32 SSM state (``mamba_ssm_cache_dtype=float32``), but bf16 state is +also supported; the buffer dtype follows the activation dtype and ``dt_cache`` +is always fp32. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import ( + allocate_update_caches, + selective_state_update_ref, + selective_state_update_replayssm_output_only_ref, +) +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _tolerances(dtype: torch.dtype) -> tuple[float, float]: + # fp32 demands true fp32 parity. A correct fp32 reconstruction (tl.dot with + # input_precision="tf32x3"/"ieee") matches the elementwise baseline to + # ~1e-5, while a TF32 reconstruction drifts to ~1e-2. atol=1e-3 sits between, + # so it flags TF32 degradation yet passes a correct fp32 kernel. bf16 stays + # loose (bf16 rounding dominates and is unaffected by the matmul precision). + if dtype == torch.float32: + return 1e-4, 1e-3 + return 6e-2, 2e-1 + + +def _tied_A(nheads: int, headdim: int, dstate: int, device: str) -> torch.Tensor: + A = -torch.rand(nheads, device=device) - 1.0 + return A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + +def _tied_dt( + batch: int, + nheads: int, + headdim: int, + device: str, + dtype: torch.dtype, +) -> torch.Tensor: + dt = torch.randn(batch, nheads, device=device, dtype=dtype) + return dt.unsqueeze(-1).expand(batch, nheads, headdim) + + +def _tied_dt_bias(nheads: int, headdim: int, device: str) -> torch.Tensor: + dt_bias = torch.rand(nheads, device=device) - 4.0 + return dt_bias.view(nheads, 1).expand(nheads, headdim) + + +def _run_standard_decode( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + batch: int, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + max_cache_len: int, + num_steps: int, + has_z: bool, + dt_softplus: bool, + use_dt_bias: bool, + desync_write_pos: bool = False, + seed: int = 0, +) -> None: + """Drive ``num_steps`` decode steps and check every path against the + trusted recurrence. ``state_dtype`` is the recurrent-state precision; + ``act_dtype`` is the activation/buffer precision.""" + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(seed) + + # One state copy per path; all start identical at the same precision. + state0 = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + state_anchor = state0.clone() + state_baseline = state0.clone() + state_cached = state0.clone() + state_ref = state0.clone() + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) if use_dt_bias else None + D = torch.randn(nheads, headdim, device=device) + + # Caches follow the activation dtype (dt_cache is forced to fp32 inside). + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + x_cache_ref, dt_cache_ref, B_cache_ref, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + + if desync_write_pos: + # Rows start at different ring positions so they flush on different + # steps, exercising per-row write-position handling. + write_pos = ( + torch.arange(batch, device=device, dtype=torch.int32) % max_cache_len + ) + else: + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) if has_z else None + is_flush = write_pos == max_cache_len - 1 + + # Trusted recurrence (mutates state_anchor in place, returns output). + out_anchor = selective_state_update_ref( + state_anchor, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + ) + + # Upstream baseline decode kernel. + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + out=out_baseline, + ) + + # ReplaySSM kernel under test + its pure-PyTorch reference. + out_cached = torch.empty_like(x) + common = dict( + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_cached, + ) + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common + ) + out_ref = selective_state_update_replayssm_output_only_ref( + state_ref, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + x_cache=x_cache_ref, + dt_cache=dt_cache_ref, + B_cache=B_cache_ref, + write_pos=write_pos, + max_cache_len=max_cache_len, + ) + + # The reference models the kernel's exact arithmetic (including its bf16 + # reconstruction), so the kernel must match it tightly at every + # precision. At fp32 this also flags any TF32 reconstruction drift. + torch.testing.assert_close(out_cached, out_ref, rtol=rtol, atol=atol) + # When the STATE is fp32 the baseline decode kernel is a valid oracle (it + # does not downcast the state per step), so ReplaySSM must match it. At + # bf16 state the baseline legitimately differs: it downcasts the fp32 + # state to bf16 every step, while ReplaySSM accumulates a whole buffer in + # fp32 and is the MORE accurate path -- so it is not a tight oracle there. + if state_dtype == torch.float32: + torch.testing.assert_close(out_cached, out_baseline, rtol=rtol, atol=atol) + # The exact elementwise reference is valid only when state AND + # activations are fp32; otherwise it downcasts the state (at readout for + # bf16 activations, or per step for bf16 state). + if both_fp32: + torch.testing.assert_close(out_cached, out_anchor, rtol=rtol, atol=atol) + + # Cached inputs match the reference cache management. + torch.testing.assert_close(x_cache, x_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(dt_cache, dt_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(B_cache, B_cache_ref, rtol=rtol, atol=atol) + + # Checkpoint state at flush matches the reference (and, when the state is + # fp32, the baseline kernel). + if bool(is_flush.any()): + torch.testing.assert_close( + state_cached[is_flush], state_ref[is_flush], rtol=rtol, atol=atol + ) + if state_dtype == torch.float32: + torch.testing.assert_close( + state_cached[is_flush], + state_baseline[is_flush], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + +# State/activation precisions. fp32 state is the default; bf16/fp16 are the +# reduced-footprint configs. fp16 appears both as an activation dtype (fully-fp16 +# model sfp16_afp16, or fp16 act over fp32 state s32_afp16) and as a state dtype +# under bf16 activations (sfp16_a16): fp16 has a finer mantissa than bf16 at the +# same 2 bytes, so it is a more accurate state at no extra footprint. We still +# skip fp16 state under fp32 activations (the unused low-state/high-act mix). +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), +] +# Small synthetic shapes for the full axis sweep (compile fast). +_SMALL_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((4, 64, 16, 1), id="tiny"), +] +# Production Mamba2 shapes (nheads, headdim, dstate, ngroups), TP=1. +_REAL_GEOMETRIES = [ + pytest.param((96, 80, 128, 8), id="nano4b"), + pytest.param((128, 64, 128, 8), id="super120b"), + pytest.param((256, 64, 128, 8), id="ultra550b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("max_cache_len", [1, 4, 16]) +@pytest.mark.parametrize("geometry", _SMALL_GEOMETRIES) +@pytest.mark.parametrize("has_z", [False, True]) +def test_replayssm_standard_decode_matches_reference( + precision: tuple[torch.dtype, torch.dtype], + max_cache_len: int, + geometry: tuple[int, int, int, int], + has_z: bool, +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=max_cache_len, + num_steps=2 * max_cache_len + 1, + has_z=has_z, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +def test_replayssm_standard_decode_real_geometry( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + # Production Mamba2 shapes for the Nemotron-3 family (Nano-4B / Super-120B / + # Ultra-550B), at the production buffer length (8). All three precisions, + # including the bf16 state case. + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=8, + num_steps=17, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize( + "precision", + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + ], +) +def test_replayssm_standard_decode_desync_write_pos( + precision: tuple[torch.dtype, torch.dtype], +): + # Rows start at staggered ring positions, so they flush on different steps + # and hold genuinely different cached histories. + state_dtype, act_dtype = precision + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=8, + headdim=64, + ngroups=4, + dstate=64, + max_cache_len=4, + num_steps=12, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + desync_write_pos=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_replayssm_standard_decode_with_batch_indices( + precision: tuple[torch.dtype, torch.dtype], + with_padding: bool, +): + # Sparse state allocation via state_batch_indices, with NULL_BLOCK_ID + # padding rows. The pure-PyTorch references do not model the sparse + # gather, so the anchor here is the upstream baseline decode kernel. + state_dtype, act_dtype = precision + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 3 + padding = 2 if with_padding else 0 + padded_batch = batch + padding + total_state_slots = 16 + nheads = 4 + ngroups = 2 + headdim = 64 + dstate = 16 + max_cache_len = 4 + num_steps = 2 * max_cache_len + + state = torch.randn( + total_state_slots, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + state_baseline = state.clone() + state_cached = state.clone() + state_before = state.clone() + + state_indices = ( + torch.randperm(total_state_slots - 1, device=device)[:batch] + 1 + ).to(torch.int32) + state_batch_indices = torch.cat( + [ + state_indices, + torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=device), + ] + ) + unused_states = torch.ones(total_state_slots, dtype=torch.bool, device=device) + unused_states[state_indices] = False + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + x_cache = torch.zeros( + total_state_slots, + nheads, + max_cache_len, + headdim, + device=device, + dtype=act_dtype, + ) + dt_cache = torch.zeros( + total_state_slots, nheads, max_cache_len, device=device, dtype=torch.float32 + ) + B_cache = torch.zeros( + total_state_slots, + ngroups, + max_cache_len, + dstate, + device=device, + dtype=act_dtype, + ) + bc_pre = torch.empty( + padded_batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + write_pos = torch.zeros(padded_batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(padded_batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(padded_batch, nheads, headdim, device, act_dtype) + B = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + out=out_baseline, + ) + + out_cached = torch.full_like(x, 42) + common = dict( + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + state_batch_indices=state_batch_indices, + out=out_cached, + ) + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common + ) + + torch.testing.assert_close( + out_cached[:batch], out_baseline[:batch], rtol=rtol, atol=atol + ) + if with_padding: + assert torch.equal( + out_cached[batch:], torch.full_like(out_cached[batch:], 42) + ) + + if bool(is_flush[:batch].all()): + torch.testing.assert_close( + state_cached[state_indices], + state_baseline[state_indices], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + assert torch.equal(state_cached[unused_states], state_before[unused_states]) + assert torch.equal(state_baseline[unused_states], state_before[unused_states]) + + +# Geometries with (nheads, ngroups) both divisible by the tp below. +_TP_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _TP_GEOMETRIES) +@pytest.mark.parametrize("tp", [2]) +def test_replayssm_standard_decode_tp_head_shard_equivalence( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], + tp: int, +): + """Tensor-parallel correctness at the kernel boundary (single GPU). + + The kernel has no cross-rank communication, so head sharding must be exactly + separable: one run over all heads must equal concatenating ``tp`` independent + per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This + guards the per-rank divisors and group->head mapping the state-shape wiring + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + assert nheads % tp == 0 and ngroups % tp == 0 + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 4 + max_cache_len = 4 + num_steps = 2 * max_cache_len + 1 + nh_s = nheads // tp + ng_s = ngroups // tp + + # Shards slice the tied params; never .contiguous() -- that would drop the + # stride-0 broadcast the kernel's TIE_HDIM asserts require. + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + + state_full = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + + # Per-rank shards, each seeded from the matching head slice so all start equal. + shard_state = [] + shard_caches = [] + for r in range(tp): + h0 = r * nh_s + shard_state.append(state_full[:, h0 : h0 + nh_s].contiguous()) + xc, dtc, bc, _ = allocate_update_caches( + batch, + nh_s, + ng_s, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bcp = torch.empty( + batch, ng_s, max_cache_len, device=device, dtype=torch.float32 + ) + shard_caches.append((xc, dtc, bc, bcp)) + + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_full = torch.empty_like(x) + selective_state_update_replayssm_output_only( + state_full, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=bc_pre, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_full, + ) + + for r in range(tp): + h0 = r * nh_s + g0 = r * ng_s + xc, dtc, bc, bcp = shard_caches[r] + out_shard = torch.empty( + batch, nh_s, headdim, device=device, dtype=act_dtype + ) + selective_state_update_replayssm_output_only( + shard_state[r], + x[:, h0 : h0 + nh_s].contiguous(), + dt[:, h0 : h0 + nh_s], + A[h0 : h0 + nh_s], + B[:, g0 : g0 + ng_s].contiguous(), + C[:, g0 : g0 + ng_s].contiguous(), + D=D[h0 : h0 + nh_s].contiguous(), + z=z[:, h0 : h0 + nh_s].contiguous(), + dt_bias=dt_bias[h0 : h0 + nh_s], + dt_softplus=True, + x_cache=xc, + dt_cache=dtc, + B_cache=bc, + bc_pre=bcp, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_shard, + ) + + torch.testing.assert_close( + out_full[:, h0 : h0 + nh_s], out_shard, rtol=rtol, atol=atol + ) + if bool(is_flush.any()): + torch.testing.assert_close( + state_full[:, h0 : h0 + nh_s][is_flush], + shard_state[r][is_flush], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 703a5df163e9..c37f74eb5079 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import Mock + import pytest import torch -from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( FlashInferSSUBackend, TritonSSUBackend, @@ -69,6 +71,48 @@ def test_flashinfer_backend_init(): assert backend.name == "flashinfer" +@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed") +@pytest.mark.parametrize( + ("algorithm", "expected"), + [ + (None, "auto"), + ("auto", "auto"), + ("simple", "simple"), + ("vertical", "vertical"), + ("horizontal", "horizontal"), + ], +) +def test_flashinfer_forwards_ssu_algorithm( + algorithm: MambaSSUAlgorithm | None, + expected: MambaSSUAlgorithm, + monkeypatch, +): + import flashinfer.mamba + + kernel = Mock() + monkeypatch.setattr(flashinfer.mamba, "selective_state_update", kernel) + backend = FlashInferSSUBackend( + MambaConfig( + backend=MambaBackendEnum.FLASHINFER, + ssu_algorithm=algorithm, + ) + ) + + tensor = torch.empty(1) + backend( + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + ) + + assert kernel.call_args.kwargs["algorithm"] == expected + + def test_uninitialized_backend_raises(): import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index fb8a4b0a28ec..b82c4d82c33e 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -76,3 +76,163 @@ def selective_state_update_ref( if not has_heads: out = out.squeeze(1) return out + + +def selective_state_update_replayssm_output_only_ref( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + max_cache_len: int = 16, +) -> torch.Tensor: + """Pure-PyTorch cached-bc reference for validation.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + + batch, nheads, dim, dstate = state.shape + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + + ratio = nheads // ngroups + + dt_val = dt[:, :, 0].float() + if dt_bias is not None: + dt_val = dt_val + dt_bias[:, 0].float() + if dt_softplus: + dt_val = F.softplus(dt_val) + A_val = A[:, 0, 0].float() + C_heads = C.repeat_interleave(ratio, dim=1) + out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert write_pos is not None + + for b in range(batch): + cache_len = int(write_pos[b].item()) + is_flush = cache_len == max_cache_len - 1 + n_steps = cache_len + 1 + + dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32) + if cache_len > 0: + dt_all[:, :cache_len] = dt_cache[b, :, :cache_len] + dt_all[:, cache_len] = dt_val[b] + + cumsum = torch.cumsum(dt_all, dim=-1) + total = cumsum[:, -1] + dA_cumsum = A_val[:, None] * cumsum + dA_total = A_val * total + total_decay = torch.exp(dA_total) + scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum) + + x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype) + if cache_len > 0: + x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1) + x_all[..., cache_len] = x[b] + + B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype) + if cache_len > 0: + B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :] + B_all[:, cache_len, :] = B[b] + + B_heads = B_all.repeat_interleave(ratio, dim=0) + C_heads_b = C_heads[b] + + if is_flush: + B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype) + delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float()) + state_new = state[b].float() * total_decay[:, None, None] + delta + state[b].copy_(state_new.to(state.dtype)) + out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads_b.float()) + else: + checkpoint_out = torch.einsum( + "hdn,hn->hd", state[b].float(), C_heads_b.float() + ) + checkpoint_out = checkpoint_out * total_decay[:, None] + BC = torch.einsum("hkn,hn->hk", B_heads.float(), C_heads_b.float()) + cache_out = torch.einsum("hdk,hk->hd", x_all.float(), scale * BC) + out[b] = checkpoint_out + cache_out + x_cache[b, :, cache_len, :] = x[b] + dt_cache[b, :, cache_len] = dt_val[b] + B_cache[b, :, cache_len, :] = B[b] + + if D is not None: + out = out + (x.float() * D[None]).to(out.dtype) + if z is not None: + out = out * F.silu(z.float()) + out = out.to(x.dtype) + if not has_heads: + out = out.squeeze(1) + return out + + +def allocate_update_caches( + batch: int, + nheads: int, + ngroups: int, + dim: int, + dstate: int, + max_cache_len: int, + device: torch.device, + x_dtype: torch.dtype, + B_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Allocate dense reference caches for standalone validation.""" + x_cache = torch.zeros( + batch, + nheads, + max_cache_len, + dim, + device=device, + dtype=x_dtype, + ) + dt_cache = torch.zeros( + batch, + nheads, + max_cache_len, + device=device, + dtype=torch.float32, + ) + B_cache = torch.zeros( + batch, + ngroups, + max_cache_len, + dstate, + device=device, + dtype=B_dtype, + ) + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + return x_cache, dt_cache, B_cache, write_pos diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index bb2f9efc7c48..35a2f198f5d2 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -36,6 +36,10 @@ P = ParamSpec("P") +class GINNotAvailableError(RuntimeError): + pass + + @dataclasses.dataclass class ProcessGroupInfo: world_size: int @@ -96,19 +100,27 @@ def parallel_launch( **kwargs: P.kwargs, ) -> None: assert not kwargs - spawn( - _worker_parallel_launch, - args=( - world_size, - world_size, - 0, - f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}", - worker, + try: + spawn( + _worker_parallel_launch, + args=( + world_size, + world_size, + 0, + f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}", + worker, + ) + + args, + nprocs=world_size, + join=True, ) - + args, - nprocs=world_size, - join=True, - ) + except Exception as exc: + # pytest.skip cannot propagate directly through torch.multiprocessing. + if "GINNotAvailableError" in str(exc): + import pytest + + pytest.skip("NCCL GIN not available (no IBGDA-capable hardware)") + raise ## DeepEP specific utils @@ -225,6 +237,18 @@ def make_deepep_v2_a2a( ): import deep_ep + from vllm.utils.nccl import query_nccl_gin_type + + # ElasticBuffer can segfault when GIN is unavailable. Initialize the + # lazy communicator and reject unsupported systems before entering DeepEP. + probe = torch.zeros(1, device=pgi.device) + torch.distributed.all_reduce(probe, group=pg) + gin_type = query_nccl_gin_type(pg) + if gin_type is None: + raise RuntimeError("Failed to determine NCCL GIN support") + if gin_type == 0: + raise GINNotAvailableError("NCCL GIN not available") + buffer = deep_ep.ElasticBuffer( group=pg, num_max_tokens_per_rank=v2_args.max_tokens_per_rank, diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index b9fe8ceafcdd..3c605ced688c 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -351,3 +351,187 @@ def test_fused_moe_batched_experts( torch.testing.assert_close(batched_output, baseline_output, atol=3e-2, rtol=2e-2) torch.testing.assert_close(triton_output, batched_output, atol=2e-2, rtol=2e-2) + + +# USE_TD in moe_mmk + +# K % 8 == 0 (bf16, 16-byte alignment) +_TD_SHAPES = [ + (4, 1, 128, 256), + (8, 64, 512, 512), + (4, 256, 1024, 2048), +] + + +def _td_supported() -> bool: + return current_platform.is_xpu() or ( + current_platform.is_cuda() and current_platform.has_device_capability(90) + ) + + +def _run_td(A, B, num_expert_tokens, use_td: bool): + out = torch.zeros( + A.shape[0], + A.shape[1], + B.shape[1], + dtype=torch.bfloat16, + device=A.device, + ) + import triton + + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( + batched_triton_kernel, + ) + + if use_td: + from vllm.triton_utils.allocation import set_triton_allocator + + set_triton_allocator(A.device) + + E = A.shape[0] + max_tokens = A.shape[1] + K = A.shape[2] + N = B.shape[1] + BM = BN = BK = 64 + grid = (E, triton.cdiv(max_tokens, BM) * triton.cdiv(N, BN)) + batched_triton_kernel[grid]( + A, + B, + out, + num_expert_tokens, + tl.bfloat16, + max_tokens, + K, + N, + None, + None, + None, + A.stride(0), + A.stride(1), + A.stride(2), + B.stride(0), + B.stride(2), + B.stride(1), + out.stride(0), + out.stride(1), + out.stride(2), + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + False, + False, + False, + BLOCK_M=BM, + BLOCK_N=BN, + BLOCK_K=BK, + USE_TD=use_td, + ) + return out + + +@pytest.mark.parametrize("num_experts,max_tokens_per_expert,K,N", _TD_SHAPES) +def test_batched_mm_td_matches_plain(num_experts, max_tokens_per_expert, K, N): + if not _td_supported(): + pytest.skip("TD requires XPU or CUDA sm_90+") + set_random_seed(42) + device = current_platform.device_type + A = ( + torch.randn( + num_experts, max_tokens_per_expert, K, device=device, dtype=torch.bfloat16 + ) + / 10 + ) + B = torch.randn(num_experts, N, K, device=device, dtype=torch.bfloat16) + num_expert_tokens = torch.randint( + 1, + max_tokens_per_expert + 1, + size=(num_experts,), + device=device, + dtype=torch.int32, + ) + + out_plain = _run_td(A, B, num_expert_tokens, False) + out_td = _run_td(A, B, num_expert_tokens, True) + + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2) + + +def test_batched_mm_td_zero_expert_tokens(): + if not _td_supported(): + pytest.skip("TD requires XPU or CUDA sm_90+") + set_random_seed(42) + device = current_platform.device_type + E, M, K, N = 8, 32, 256, 256 + A = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) / 10 + B = torch.randn(E, N, K, device=device, dtype=torch.bfloat16) + num_expert_tokens = torch.zeros(E, device=device, dtype=torch.int32) + num_expert_tokens[::2] = M + + out_plain = _run_td(A, B, num_expert_tokens, False) + out_td = _run_td(A, B, num_expert_tokens, True) + + for e in range(E): + if num_expert_tokens[e].item() == 0: + assert out_plain[e].abs().max().item() == 0.0 + assert out_td[e].abs().max().item() == 0.0 + + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2) + + +# BatchedTritonExperts device enablement (XPU) +def test_batched_triton_experts_supports_current_device(): + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( + BatchedTritonExperts, + ) + + if current_platform.is_xpu() or current_platform.is_cuda_alike(): + assert BatchedTritonExperts._supports_current_device() + else: + pytest.skip("No GPU device available") + + +@pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)]) +def test_batched_experts_end_to_end(m, n, k, e, topk): + """End-to-end BatchedTritonExperts via the reference (no-comms) + BatchedPrepareAndFinalize, validated against a torch reference. Exercises + the device-enablement path.""" + if not (current_platform.is_xpu() or current_platform.is_cuda_alike()): + pytest.skip("No GPU device available") + + from vllm.v1.worker.workspace import init_workspace_manager + + set_random_seed(7) + device = current_platform.device_type + init_workspace_manager(torch.device(f"{device}:0")) + + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10 + score = torch.randn((m, e), device=device, dtype=torch.bfloat16) + # w1 is the fused gate+up projection [E, 2N, K]; w2 is [E, K, N]. + w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15 + w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15 + + with set_current_vllm_config(vllm_config): + topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) + + baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids) + triton_output = batched_moe(a, w1, w2, topk_weight, topk_ids) + + torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2) + + +def test_batched_triton_backend_mapping(): + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + map_unquantized_backend, + ) + + assert ( + map_unquantized_backend("batched_triton") + == UnquantizedMoeBackend.BATCHED_TRITON + ) + assert map_unquantized_backend("triton") == UnquantizedMoeBackend.TRITON diff --git a/tests/kernels/moe/test_block_int8.py b/tests/kernels/moe/test_block_int8.py index e35ca4caa9db..415faafb3ab0 100644 --- a/tests/kernels/moe/test_block_int8.py +++ b/tests/kernels/moe/test_block_int8.py @@ -82,9 +82,9 @@ def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape): ).sum(dim=1) -@pytest.fixture(autouse=True, scope="module") +@pytest.fixture(autouse=True) def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" + """Sets the default CUDA device before each test in this module.""" torch.set_default_device("cuda") diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 41ae9be51730..9e99a274fb36 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -7,10 +7,15 @@ from tests.kernels.allclose_default import get_default_atol, get_default_rtol from vllm._custom_ops import ( cpu_fused_moe, + cpu_fused_moe_int8, cpu_prepack_moe_weight, + cpu_prepack_moe_weight_int8, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + _CPU_MOE_ACT_FN, + CPUFusedMOE, +) from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed @@ -101,6 +106,90 @@ def ref_fused_moe( return final_out +def quantize_per_channel( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetrically quantize each weight output channel.""" + weight_fp32 = weight.float() + scale = weight_fp32.abs().amax(dim=-1).clamp_min(1e-12) / 127.0 + quantized = ( + (weight_fp32 / scale.unsqueeze(-1)).round().clamp(-127, 127).to(torch.int8) + ) + return quantized, scale + + +def quantize_per_token( + input: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetrically quantize each input token.""" + input_fp32 = input.float() + scale = input_fp32.abs().amax(dim=-1, keepdim=True).clamp_min(1e-7) / 127.0 + quantized = (input_fp32 / scale).round().clamp(-127, 127).to(torch.int8) + return quantized, scale + + +def ref_fused_moe_int8( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, +) -> torch.Tensor: + """Reference the two dynamically quantized INT8 GEMMs.""" + input_int8, input_scale = quantize_per_token(input) + expert_num = w13.size(0) + + counts = topk_ids.new_zeros((topk_ids.size(0), expert_num)) + counts.scatter_(1, topk_ids.to(torch.int64), 1) + tokens_per_expert = counts.sum(dim=0).tolist() + sorted_route_ids = topk_ids.view(-1).argsort() + sorted_token_ids = sorted_route_ids // topk_ids.size(1) + + outputs = [] + start_idx = 0 + for expert_idx, token_count in enumerate(tokens_per_expert): + end_idx = start_idx + token_count + if token_count == 0: + continue + + token_ids = sorted_token_ids[start_idx:end_idx] + gate_up = torch.matmul( + input_int8[token_ids].float(), + w13[expert_idx].float().T, + ) + gate_up *= input_scale[token_ids] * w13_scale[expert_idx] + if w13_bias is not None: + gate_up += w13_bias[expert_idx].float() + + intermediate = _CPU_MOE_ACT_FN[activation](gate_up).to(input.dtype) + intermediate_int8, intermediate_scale = quantize_per_token(intermediate) + output = torch.matmul( + intermediate_int8.float(), + w2[expert_idx].float().T, + ) + output *= intermediate_scale * w2_scale[expert_idx] + if w2_bias is not None: + output += w2_bias[expert_idx].float() + + outputs.append(output) + start_idx = end_idx + + sorted_output = torch.cat(outputs, dim=0) + routed_output = torch.empty_like(sorted_output) + routed_output[sorted_route_ids] = sorted_output + return ( + routed_output.view(*topk_ids.shape, -1) + .mul_(topk_weights.unsqueeze(-1)) + .sum(dim=1) + .to(input.dtype) + ) + + @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("expert_num", EXPERT_NUM) @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) @@ -176,3 +265,214 @@ def test_cpu_fused_moe( torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), f"{torch.max(torch.abs(output - ref_output))}", ) + + +@pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.ARM, + reason="Requires Arm CPU", +) +@pytest.mark.parametrize("batch_size", BATCH_SIZE) +@pytest.mark.parametrize("expert_num", EXPERT_NUM) +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM) +@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) +@pytest.mark.parametrize("use_bias", USE_BIAS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("act", ACT) +@pytest.mark.parametrize("isa", ["neon"]) +def test_cpu_fused_moe_int8( + batch_size: int, + expert_num: int, + hidden_size: int, + intermediate_size: int, + use_bias: bool, + dtype: torch.dtype, + act: MoEActivation, + isa: str, +): + set_random_seed(0) + topk_num = max(expert_num // 2, 1) + up_dim = 2 * intermediate_size + + input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w13_fp = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w2_fp = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( + 0.5 * intermediate_size**0.5 + ) + w13, w13_scale = quantize_per_channel(w13_fp) + w2, w2_scale = quantize_per_channel(w2_fp) + + w13_bias = None + w2_bias = None + if use_bias: + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + + router_logits = torch.randn((batch_size, expert_num), dtype=dtype) + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + topk_weights, topk_ids = torch.topk(score, topk_num) + topk_ids = topk_ids.to(torch.int32) + + ref_output = ref_fused_moe_int8( + input, + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + act, + ) + packed_w13 = cpu_prepack_moe_weight_int8(w13, isa) + packed_w2 = cpu_prepack_moe_weight_int8(w2, isa) + output = cpu_fused_moe_int8( + input, + packed_w13, + packed_w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + act.value, + isa, + ) + + torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) + + +# moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel +# sharding of moe_intermediate_size=704 at tp=4 produces (704 // 4 == 176). +UNALIGNED_INTERMEDIATE_DIM = 176 + + +class _StubMoELayer(torch.nn.Module): + """Minimal stand-in for the real MoE layer module, exposing just what + CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and + optionally w13_bias/w2_bias).""" + + def __init__( + self, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + activation: MoEActivation, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + ): + super().__init__() + self.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + self.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + self.activation = activation + if w13_bias is not None: + self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) + if w2_bias is not None: + self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) + + +@pytest.mark.parametrize("expert_num", EXPERT_NUM) +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM) +@pytest.mark.parametrize("use_bias", USE_BIAS) +@pytest.mark.parametrize("dtype", DTYPE) +@pytest.mark.parametrize( + "act", + [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH], +) +def test_cpu_fused_moe_unaligned_intermediate_size( + default_vllm_config, + expert_num: int, + hidden_size: int, + use_bias: bool, + dtype: torch.dtype, + act: MoEActivation, +): + """An unaligned per-partition moe_intermediate_size must still hit the + grouped-gemm fast path via automatic zero-padding, with numerically + correct output, instead of silently falling back to the much slower + per-expert torch loop.""" + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + pytest.skip("padding is only applied on the x86 AMX/vector kernels") + + set_random_seed(0) + batch_size = 64 + intermediate_size = UNALIGNED_INTERMEDIATE_DIM + topk_num = max(expert_num // 2, 1) + up_dim = 2 * intermediate_size + + input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( + 0.5 * intermediate_size**0.5 + ) + router_logits = torch.randn((batch_size, expert_num), dtype=dtype) + w13_bias = None + w2_bias = None + if use_bias: + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk_num) + topk_ids = topk_ids.to(torch.int32) + + ref_output = ref_fused_moe( + input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act + ) + + layer = _StubMoELayer( + w13.clone(), + w2.clone(), + act, + w13_bias.clone() if w13_bias is not None else None, + w2_bias.clone() if w2_bias is not None else None, + ) + + cpu_moe = CPUFusedMOE(layer) + assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, ( + "expected the padded intermediate size to hit the grouped-gemm fast path" + ) + + output = cpu_moe.forward_method( + layer, input, topk_weight, topk_ids, act, expert_num, False + ) + + atol, rtol = get_default_atol(output), get_default_rtol(output) + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) + + +def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config): + """swigluoai's interleaved gate/up layout isn't padded automatically. On + AMX-capable CPUs this must raise rather than silently falling back to + the (correct but much slower) per-expert torch loop; elsewhere it should + fall back exactly as before.""" + set_random_seed(0) + expert_num = 8 + hidden_size = 128 + intermediate_size = UNALIGNED_INTERMEDIATE_DIM + dtype = torch.bfloat16 + up_dim = 2 * intermediate_size + + layer = _StubMoELayer( + torch.randn((expert_num, up_dim, hidden_size), dtype=dtype), + torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype), + MoEActivation.SWIGLUOAI, + ) + + if torch.cpu._is_amx_tile_supported(): + with pytest.raises(RuntimeError): + CPUFusedMOE(layer) + else: + cpu_moe = CPUFusedMOE(layer) + assert cpu_moe.forward_method == cpu_moe.forward_torch diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 8d12e2888d0d..ce9e4f343d62 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -11,7 +11,7 @@ from torch.distributed import ProcessGroup import vllm.envs as envs -from tests.kernels.moe.utils import check_accuracy, make_dummy_moe_config +from tests.kernels.moe.utils import make_dummy_moe_config from vllm import _custom_ops as ops from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.activation import SiluAndMul @@ -66,8 +66,14 @@ def make_weights( # per-out-channel weight quantization assert dtype == current_platform.fp8_dtype() - w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16) - w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16) + # Keep FP8 inputs finite and bounded. torch.empty made this test depend on + # allocator contents, while larger values exceed its fixed W8A8 tolerance. + w1 = torch.randn( + (e, 2 * n, k), device=current_platform.device_type, dtype=torch.float16 + ).div_(100) + w2 = torch.randn( + (e, k, n), device=current_platform.device_type, dtype=torch.float16 + ).div_(100) n_b_scales = 2 * n k_b_scales = k @@ -361,14 +367,6 @@ def assert_deepep_close( k: int, use_fp8_dispatch: bool, ) -> None: - if use_fp8_dispatch and current_platform.is_fp8_fnuz(): - # ROCm e4m3fnuz rounds differently than the reference quant, - # so DeepEP's fp8 dispatch can yield a few outliers even with - # a correct kernel; allow a small fraction of mismatches here. - atol = rtol = 1.5e-1 - check_accuracy(expected, actual, atol=atol, rtol=rtol, percent=0.95) - return - torch.testing.assert_close( expected, actual, diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index 93b7c136605b..a2ff56be5126 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -14,6 +14,7 @@ from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe import TritonExperts from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -36,6 +37,14 @@ ) +def assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + close = torch.isclose(actual, expected, atol=2e-1, rtol=2e-1) + close_fraction = close.float().mean().item() + assert close_fraction > 0.99, ( + f"Only {close_fraction:.1%} of FP8 outputs are within tolerance" + ) + + @dataclasses.dataclass class TestConfig: dtype: torch.dtype @@ -50,6 +59,7 @@ class TestConfig: class TestTensors: rank_tokens: torch.Tensor rank_token_scales: torch.Tensor | None + intermediate_scales: torch.Tensor | None topk: torch.Tensor topk_weights: torch.Tensor config: TestConfig @@ -63,6 +73,12 @@ def make(config: TestConfig) -> "TestTensors": rank_tokens = ( torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 ) + if config.dtype == torch.float8_e4m3fn: + rank_token_scales = torch.tensor(1 / 448, device="cuda") + intermediate_scales = torch.tensor(8 / 448, device="cuda") + else: + rank_token_scales = None + intermediate_scales = None topk = torch.stack( [ @@ -73,7 +89,8 @@ def make(config: TestConfig) -> "TestTensors": topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda") return TestTensors( rank_tokens=rank_tokens, - rank_token_scales=None, + rank_token_scales=rank_token_scales, + intermediate_scales=intermediate_scales, topk=topk, topk_weights=topk_weights, config=config, @@ -124,7 +141,6 @@ def make_modular_kernel( mk = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) return mk @@ -162,6 +178,7 @@ def build_expert_map(): w2_scale=w2_scale, per_act_token_quant=per_act_token_quant, a1_scale=test_tensors.rank_token_scales, + a2_scale=test_tensors.intermediate_scales, ) hidden_size = test_tensors.rank_tokens.size(1) @@ -231,6 +248,8 @@ def _deep_ep_v2_moe( test_tensors.topk, w1_scale=w1_scale, w2_scale=w2_scale, + a1_scale=test_tensors.rank_token_scales, + a2_scale=test_tensors.intermediate_scales, quant_dtype=q_dtype, per_act_token_quant=per_act_token_quant, ) @@ -262,12 +281,15 @@ def _deep_ep_v2_moe( per_act_token_quant, ) - torch.testing.assert_close( - torch_combined, - deepep_combined, - atol=6e-2, - rtol=6e-2, - ) + if is_quantized: + assert_fp8_close(torch_combined, deepep_combined) + else: + torch.testing.assert_close( + torch_combined, + deepep_combined, + atol=6e-2, + rtol=6e-2, + ) MNKs = [ @@ -355,17 +377,29 @@ def _deep_ep_v2_moe_cudagraph( num_local_experts = config.num_experts // pgi.world_size hidden_size = config.k - # Create FP8 weights directly, then dequantize for bf16 reference. - w1_fp8 = torch.randn( - (config.num_experts, 2 * config.n, config.k), - device="cuda", - dtype=torch.bfloat16, - ).to(torch.float8_e4m3fn) - w2_fp8 = torch.randn( - (config.num_experts, config.k, config.n), - device="cuda", - dtype=torch.bfloat16, - ).to(torch.float8_e4m3fn) + # All ranks must use the same global weights before taking their EP slice. + w1_bf16 = ( + torch.randn( + (config.num_experts, 2 * config.n, config.k), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + w2_bf16 = ( + torch.randn( + (config.num_experts, config.k, config.n), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + torch.distributed.broadcast(w1_bf16, src=0, group=pg) + torch.distributed.broadcast(w2_bf16, src=0, group=pg) + + # Round-trip through FP8 before constructing the reference and kernel weights. + w1_fp8 = w1_bf16.to(torch.float8_e4m3fn) + w2_fp8 = w2_bf16.to(torch.float8_e4m3fn) w1_ref = w1_fp8.to(torch.bfloat16) w2_ref = w2_fp8.to(torch.bfloat16) @@ -375,7 +409,7 @@ def _deep_ep_v2_moe_cudagraph( vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm") with set_current_vllm_config(vllm_cfg): - # Initialize vLLM parallel state (needed by FusedMoE layer) + # Initialize vLLM parallel state (needed by MoERunner layer) temp_file = tempfile.mktemp() init_distributed_environment( world_size=pgi.world_size, @@ -385,21 +419,8 @@ def _deep_ep_v2_moe_cudagraph( backend="nccl", ) initialize_model_parallel(tensor_model_parallel_size=1) - # Reference MoE using dequantized bf16 weights - torch_combined = torch_experts( - test_tensors.rank_tokens, - w1_ref, - w2_ref, - test_tensors.topk_weights, - test_tensors.topk, - ) - - # Use the production pipeline: make_fused_moe_layer creates - # a FusedMoE layer, quantizes weights, runs - # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK - # shuffle), and selects the kernel. - # Quantize weights using production helper, EP-slice, then - # convert to TrtLLM format. + # Mirror production weight processing: quantize, EP-slice, then + # convert to the TrtLLM BlockMajorK format. from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( TrtLlmFp8ExpertsModular, @@ -411,14 +432,32 @@ def _deep_ep_v2_moe_cudagraph( block_shape = [128, 128] qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) + assert qw.w13_weight_scale is not None + assert qw.w2_weight_scale is not None + + # Reference MoE using the same blockwise FP8 quantization scheme as + # the production kernel. torch_experts quantizes activations before + # both GEMMs and dequantizes the operands for the reference matmuls. + reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to( + torch.float32 + ) + torch_combined = torch_experts( + test_tensors.rank_tokens, + qw.w13_weight, + qw.w2_weight, + reference_topk_weights, + test_tensors.topk, + w1_scale=qw.w13_weight_scale, + w2_scale=qw.w2_weight_scale, + quant_dtype=torch.float8_e4m3fn, + block_shape=block_shape, + ) # EP-slice before format conversion e_start = num_local_experts * pgi.rank e_end = e_start + num_local_experts w1_ep = qw.w13_weight[e_start:e_end] w2_ep = qw.w2_weight[e_start:e_end] - assert qw.w13_weight_scale is not None - assert qw.w2_weight_scale is not None w1_scale_ep = qw.w13_weight_scale[e_start:e_end] w2_scale_ep = qw.w2_weight_scale[e_start:e_end] @@ -452,11 +491,23 @@ class activation: w2_scale=w2_scale_ep, ) moe_config = make_dummy_moe_config( - num_experts=num_local_experts, + num_experts=config.num_experts, + num_local_experts=num_local_experts, experts_per_token=config.topk, hidden_dim=hidden_size, intermediate_size=config.n, ) + moe_parallel_config = dataclasses.replace( + moe_config.moe_parallel_config, + ep_size=pgi.world_size, + ep_rank=pgi.rank, + use_ep=True, + all2all_backend="deepep_v2", + ) + moe_config = dataclasses.replace( + moe_config, + moe_parallel_config=moe_parallel_config, + ) fused_experts = TrtLlmFp8ExpertsModular( moe_config=moe_config, quant_config=quant_config, @@ -480,21 +531,21 @@ class activation: mk_kernel = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) - for _ in range(3): - out = mk_kernel.apply( - hidden_states=test_tensors.rank_tokens, - w1=w1_ep, - w2=w2_ep, - topk_weights=test_tensors.topk_weights, - topk_ids=test_tensors.topk, - activation=MoEActivation.SILU, - global_num_experts=config.num_experts, - expert_map=None, - apply_router_weight_on_input=False, - ) + with set_forward_context(None, vllm_cfg): + for _ in range(3): + out = mk_kernel.apply( + hidden_states=test_tensors.rank_tokens, + w1=w1_ep, + w2=w2_ep, + topk_weights=test_tensors.topk_weights, + topk_ids=test_tensors.topk, + activation=MoEActivation.SILU, + global_num_experts=config.num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) torch.testing.assert_close( torch_combined, diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index a9c6e9f5b193..ae6c9685a459 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -25,6 +25,9 @@ FusedMoEQuantDesc, fp8_w8a8_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( + deepgemm_moe_permute, +) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) @@ -41,6 +44,35 @@ BLOCK_SIZE = [128, 128] +@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") +def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init): + hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16) + activations, scales = per_token_group_quant_fp8( + hidden_states, + group_size=128, + use_ue8m0=True, + ) + topk_ids = torch.tensor([[0], [1]], device="cuda", dtype=torch.int64) + + _, permuted_scales, expert_ids, _, _ = deepgemm_moe_permute( + aq=activations, + aq_scale=scales, + topk_ids=topk_ids, + local_num_experts=2, + expert_map=None, + expert_tokens_meta=None, + ) + + padding = expert_ids < 0 + assert padding.any() + torch.testing.assert_close( + permuted_scales[padding], + torch.zeros_like(permuted_scales[padding]), + rtol=0, + atol=0, + ) + + def make_block_quant_fp8_weights( e: int, n: int, @@ -349,6 +381,7 @@ def run_single_fp4_case(m, n, k, topk, num_experts): FP4_MNKs = [ (128, 4096, 4096), # DeepSeek V4 shape (256, 2048, 2048), # Half-size variant + (128, 384, 3584), # Kimi-K3 TP8 latent-MoE shape ] FP4_TOPKS = [2] diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py new file mode 100644 index 000000000000..a7a7c5251bc3 --- /dev/null +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for FlashInfer CuTeDSL NVFP4 MoE.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.quantization.nvfp4_utils import ( + FLOAT4_E2M1_MAX, + FLOAT8_E4M3_MAX, + break_fp4_bytes, +) +from tests.kernels.utils import torch_moe +from vllm import _custom_ops as ops +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import fused_topk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( + FlashInferCuteDSLExperts, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4 +from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import set_random_seed + +if not has_flashinfer_cutedsl_moe_nvfp4() or not ( + current_platform.is_device_capability_family(100) +): + pytest.skip( + "Requires FlashInfer CuTeDSL NVFP4 MoE on SM100", + allow_module_level=True, + ) + + +def _quantize_nvfp4_linear( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + weights_q = [] + scales = [] + global_scales = [] + for expert_weight in weight: + global_scale = ( + FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max() + ).to(torch.float32) + weight_q, scale = ops.scaled_fp4_quant( + expert_weight, + global_scale, + is_sf_swizzled_layout=False, + ) + weights_q.append(weight_q) + scales.append(scale) + global_scales.append(global_scale) + return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales) + + +def _dequantize_nvfp4_linear( + tensor_fp4: torch.Tensor, + tensor_sf: torch.Tensor, + global_scale: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + assert tensor_fp4.dtype == torch.uint8 + m, packed_k = tensor_fp4.shape + k = packed_k * 2 + tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32) + tensor_f32 = tensor_f32.reshape(m, k // 16, 16) + tensor_sf = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32) + tensor_sf = tensor_sf[:, : k // 16] / global_scale + return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype) + + +@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.inference_mode() +def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + + w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) + w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + + activation = MoEActivation.RELU2_NO_MUL + fake_layer = SimpleNamespace(activation=activation) + a1_scale = torch.ones(1, device="cuda", dtype=torch.float32) + a2_scale = torch.ones(1, device="cuda", dtype=torch.float32) + ( + w1_cutedsl, + w1_scale_cutedsl, + w1_alpha, + a1_scale, + w2_cutedsl, + w2_scale_cutedsl, + w2_alpha, + a2_scale, + ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( + layer=fake_layer, + w13=w1_q, + w13_scale=w1_scale, + w13_scale_2=(1.0 / w1_global_scale), + a13_scale=a1_scale, + w2=w2_q, + w2_scale=w2_scale, + w2_scale_2=(1.0 / w2_global_scale), + a2_scale=a2_scale, + ) + quant_config = nvfp4_moe_quant_config( + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + a1_gscale=(1.0 / a1_scale), + a2_gscale=(1.0 / a2_scale), + w1_scale=w1_scale_cutedsl, + w2_scale=w2_scale_cutedsl, + is_scale_swizzled=False, + ) + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=dtype, + routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), + ) + + cutedsl_experts = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + FlashInferCuteDSLExperts( + moe_config=moe_config, + quant_config=quant_config, + ), + ) + + cutedsl_output = cutedsl_experts.apply( + hidden_states=hidden_states, + w1=w1_cutedsl, + w2=w2_cutedsl, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ) + + a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32) + a_q, a_scale = ops.scaled_fp4_quant( + hidden_states, + a_global_scale, + is_sf_swizzled_layout=False, + ) + a_in_dtype = _dequantize_nvfp4_linear( + a_q, + a_scale, + a_global_scale, + dtype=dtype, + ) + + w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype) + w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) + for idx in range(e): + w1_d[idx] = _dequantize_nvfp4_linear( + w1_q[idx], + w1_scale[idx], + w1_global_scale[idx], + dtype=dtype, + ) + w2_d[idx] = _dequantize_nvfp4_linear( + w2_q[idx], + w2_scale[idx], + w2_global_scale[idx], + dtype=dtype, + ) + + torch_output = torch_moe( + a_in_dtype, + w1_d, + w2_d, + score, + topk, + activation=activation, + ) + torch.testing.assert_close( + torch_output, + cutedsl_output, + atol=2e-1, + rtol=2e-1, + ) diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index c58c8474b06e..1ab8550d64e8 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -23,6 +23,53 @@ from vllm.utils.torch_utils import set_random_seed +def _run_single_group_topk( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + return fused_grouped_topk( + hidden_states=torch.empty( + (logits.shape[0], 0), dtype=logits.dtype, device=logits.device + ), + gating_output=logits, + topk=topk, + renormalize=renormalize, + e_score_correction_bias=bias, + num_expert_group=1, + topk_group=1, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ) + + +def _single_group_reference( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + if scoring_func == "sigmoid": + scores = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 + else: + scores = torch.softmax(logits, dim=-1).float() + indices = torch.argsort( + scores + bias.float(), dim=-1, descending=True, stable=True + )[:, :topk] + values = scores.gather(1, indices) + if renormalize: + values /= values.sum(dim=-1, keepdim=True) + 1e-20 + values *= routed_scaling_factor + return values, indices.to(torch.int32) + + @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) @@ -101,3 +148,222 @@ def test_grouped_topk( baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0 ) torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_grouped_topk_single_group_large_batch(): + set_random_seed(0) + logits = torch.randn((1536, 896), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((896,), dtype=torch.float32, device="cuda") + + expected_values, expected_ids = _single_group_reference( + logits, bias, 16, scoring_func="sigmoid", renormalize=True + ) + actual_values, actual_ids = _run_single_group_topk( + logits, bias, 16, scoring_func="sigmoid", renormalize=True + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize( + "num_experts,topk,input_dtype,bias_dtype", + [ + (512, 9, torch.bfloat16, torch.float32), + (512, 16, torch.float16, torch.float16), + (513, 9, torch.float32, torch.bfloat16), + (513, 16, torch.bfloat16, torch.float32), + (895, 9, torch.float16, torch.bfloat16), + (896, 16, torch.float32, torch.float16), + (897, 9, torch.bfloat16, torch.bfloat16), + (897, 16, torch.float16, torch.float32), + (1024, 9, torch.float32, torch.bfloat16), + (1024, 16, torch.bfloat16, torch.float16), + ], +) +@pytest.mark.parametrize( + "scoring_func,renormalize,routed_scaling_factor", + [ + ("sigmoid", True, 1.0), + ("sigmoid", False, 2.5), + ("softmax", True, 2.5), + ("softmax", False, 1.0), + ], +) +def test_grouped_topk_single_group_tiers( + num_experts: int, + topk: int, + input_dtype: torch.dtype, + bias_dtype: torch.dtype, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float, +): + set_random_seed(7) + logits = torch.randn((17, num_experts), dtype=input_dtype, device="cuda") + bias = torch.randn((num_experts,), dtype=bias_dtype, device="cuda") + + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize( + "num_experts,topk,scoring_func", + [ + (128, 8, "sigmoid"), + (129, 8, "sigmoid"), + (257, 8, "sigmoid"), + (385, 8, "sigmoid"), + (512, 9, "sigmoid"), + (513, 9, "sigmoid"), + (769, 9, "sigmoid"), + (897, 16, "sigmoid"), + (1024, 16, "sigmoid"), + (128, 4, "softmax"), + (128, 5, "softmax"), + (129, 8, "softmax"), + (161, 8, "softmax"), + (256, 9, "softmax"), + (257, 8, "softmax"), + (512, 9, "softmax"), + (512, 17, "softmax"), + (512, 23, "softmax"), + (513, 8, "softmax"), + (577, 9, "softmax"), + (769, 9, "softmax"), + (897, 9, "softmax"), + (1024, 16, "softmax"), + ], +) +def test_grouped_topk_single_group_capacity_tiers( + num_experts: int, + topk: int, + scoring_func: str, +): + set_random_seed(11) + logits = torch.randn((3, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +def test_grouped_topk_single_group_stable_ties(num_experts: int): + logits = torch.zeros((1, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=True, + routed_scaling_factor=2.5, + ) + + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + expected_values = torch.full((1, 16), 2.5 / 16, dtype=torch.float32, device="cuda") + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +@pytest.mark.parametrize("num_finite", [0, 15]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_grouped_topk_single_group_nonfinite_scores( + num_experts: int, num_finite: int, renormalize: bool +): + logits = torch.full( + (1, num_experts), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + if num_finite: + logits[0, :num_finite] = torch.arange( + num_finite, dtype=torch.bfloat16, device="cuda" + ) + logits[0, num_finite] = torch.inf + logits[0, num_finite + 1] = -torch.inf + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=renormalize, + routed_scaling_factor=2.5, + ) + + if num_finite == 0: + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + if renormalize: + expected_values = torch.full( + (1, 16), 1 / 16, dtype=torch.float32, device="cuda" + ) + else: + expected_values = torch.zeros((1, 16), dtype=torch.float32, device="cuda") + else: + expected_ids = torch.cat( + ( + torch.arange(num_finite - 1, -1, -1, dtype=torch.int32, device="cuda"), + torch.tensor([num_finite], dtype=torch.int32, device="cuda"), + ) + )[None] + finite_values = logits[0, :num_finite].float().sigmoid().flip(0) + if renormalize: + finite_values /= finite_values.sum() + finite_values *= 2.5 + expected_values = torch.cat( + (finite_values, torch.zeros(1, dtype=torch.float32, device="cuda")) + )[None] + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0315d8d89e56..b3a42912da7f 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -4,6 +4,8 @@ Test modular OAI Triton MoE """ +from __future__ import annotations + import pytest import torch @@ -48,6 +50,35 @@ ] +def deepseek_v4_flash_moe_topology(): + """MoE sizes for the DeepSeek-V4-Flash`. + Default weights from: ``deepseek-ai/DeepSeek-V4-Flash`. + ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. + """ + defaults = { + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + } + + return defaults + + +def scaled_deepseek_v4_flash_problem( + *, + dim_scale: int = 8, + expert_scale: int = 8, +): + """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" + t = deepseek_v4_flash_moe_topology() + k = max(128, t["hidden_size"] // dim_scale) + n = max(64, t["moe_intermediate_size"] // dim_scale) + num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) + topk = t["num_experts_per_tok"] + return k, n, num_experts, topk + + def unshuffle_weight(w: torch.Tensor): first = w[..., ::2] second = w[..., 1::2] @@ -160,8 +191,8 @@ def oai_triton_moe_impl( x: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - w1_scale: "PrecisionConfig", - w2_scale: "PrecisionConfig", + w1_scale: PrecisionConfig, + w2_scale: PrecisionConfig, w1_bias: torch.Tensor | None, w2_bias: torch.Tensor | None, num_experts: int, @@ -261,3 +292,93 @@ def test_oai_triton_moe( ) assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): + """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. + + Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, + with linear dimensions and expert count scaled down for test GPU memory. + """ + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) + set_random_seed(0) + + k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() + m = 7 + dtype = torch.bfloat16 + + ( + w1, + w2, + w1_bias, + w2_bias, + w1_tri, + w2_tri, + w1_bias_tri, + w2_bias_tri, + w1_precision_config, + w2_precision_config, + ) = make_weights(dtype, k, n, num_experts) + + x = torch.randn((m, k), dtype=dtype, device="cuda") + router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + quant_config = mxfp4_w4a16_moe_quant_config( + w1_bias=w1_bias_tri, + w2_bias=w2_bias_tri, + w1_scale=w1_precision_config, + w2_scale=w2_precision_config, + ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + ) + experts = UnfusedOAITritonExperts(moe_config, quant_config) + + if not UnfusedOAITritonExperts._supports_current_device(): + pytest.skip("UnfusedOAITritonExperts does not support this device") + + _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + assert top_k == topk + ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( + m, + N, + K, + topk, + num_experts, + num_experts, + None, + MoEActivation.SWIGLUOAI, + ) + workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") + workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") + output = torch.empty(out_shape, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) + experts.apply( + output=output, + hidden_states=x, + w1=w1_tri, + w2=w2_tri, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SWIGLUOAI, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 7fc0e777e60d..8e5f6ce79f4a 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -36,6 +36,9 @@ batched_fused_marlin_moe, fused_marlin_moe, ) +from vllm.model_executor.layers.fused_moe.utils import ( + moe_use_td_hw_supported, +) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_permute_bias, ) @@ -53,9 +56,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types +from vllm.triton_utils import tl from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + def iterative_moe( hidden_states: torch.Tensor, @@ -289,6 +295,7 @@ def run_moe_test( @pytest.mark.parametrize("ep_size", EP_SIZE) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("padding", [True, False]) +@pytest.mark.parametrize("use_td", [False, True]) def test_fused_moe( m: int, n: int, @@ -298,9 +305,19 @@ def test_fused_moe( ep_size: int, dtype: torch.dtype, padding: bool, + use_td: bool, monkeypatch, workspace_init, ): + if use_td and not hasattr(tl, "make_tensor_descriptor"): + pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor") + if use_td and not moe_use_td_hw_supported(): + pytest.skip( + "tensor_descriptor.gather requires XPU or NVIDIA Blackwell " + "(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas " + "rejects on Hopper (sm90) and earlier" + ) + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") set_random_seed(7) # @@ -311,17 +328,17 @@ def test_fused_moe( # Setup test data # - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 + a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 - score = torch.randn((m, e), device="cuda", dtype=dtype) + score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype) if ep_size > 1: local_e = e // ep_size - e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32) - e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32) - e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32) + e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32) + e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32) + e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32) w1 = w1[e_ids] w2 = w2[e_ids] else: @@ -1165,6 +1182,174 @@ def test_fused_marlin_moe_non_gated( torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0) +def _make_humming_indexed_experts(activation: MoEActivation): + pytest.importorskip("humming") + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingIndexedExperts, + ) + from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.utils import humming + + top_k, num_experts = 6, 12 + hidden_size, intermediate_size = 2688, 1856 + gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size + num_w13_stacks = 2 if activation.is_gated else 1 + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + activation=activation, + ) + + layer = torch.nn.Module() + layer.moe_config = moe_config + layer.params_dtype = torch.bfloat16 + + weight_schema = humming.ModeloptNvfp4WeightSchema() + for sublayer_name, shape_n, shape_k, stack_size in ( + ("w13", gate_up_size, hidden_size, num_w13_stacks), + ("w2", hidden_size, intermediate_size, 1), + ): + tensor_attrs = weight_schema.get_tensors_attrs( + shape_n=shape_n, + shape_k=shape_k, + param_dtype=layer.params_dtype, + num_experts=num_experts, + stack_size=stack_size, + ) + for tensor_name, attrs in tensor_attrs.items(): + layer.register_parameter( + f"{sublayer_name}_{tensor_name}", + Parameter( + torch.ones( + attrs["shape"], + dtype=attrs["dtype"], + device="cuda", + ), + requires_grad=False, + ), + ) + + humming_utils.convert_to_humming_moe_kernel_format( + layer, + weight_schema=weight_schema, + input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16), + ) + + layer.local_num_experts = layer.global_num_experts = num_experts + layer.hidden_size = hidden_size + layer.intermediate_size_per_partition = intermediate_size + quant_config = humming_utils.get_humming_moe_quant_config(layer) + experts = HummingIndexedExperts( + layer, + moe_config, + quant_config, + ) + return experts + + +@pytest.mark.parametrize( + "activation", + [ + MoEActivation.SILU, + MoEActivation.RELU2_NO_MUL, + ], + ids=["gated", "non_gated"], +) +def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): + from vllm.utils import humming + + experts = _make_humming_indexed_experts(activation) + layer = experts.layer + moe_config = experts.moe_config + top_k = moe_config.experts_per_token + num_experts = moe_config.num_experts + hidden_size = moe_config.hidden_dim + intermediate_size = moe_config.intermediate_size + gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size + + w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2")) + for meta in (w13_meta, w2_meta): + assert meta.a_dtype == humming.dtypes.bfloat16 + assert meta.b_dtype == humming.dtypes.float4e2m1 + + assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size + assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size + + buffer_metas, _ = experts.get_buffer_metas( + M=1, + topk=top_k, + activation=moe_config.activation, + ) + assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size + assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size + assert experts.moe_problem_size( + a1=torch.empty(1, hidden_size), + w1=torch.empty(num_experts, 1), + w2=torch.empty(num_experts, 1), + topk_ids=torch.empty(1, top_k, dtype=torch.long), + ) == (num_experts, 1, intermediate_size, hidden_size, top_k) + + +def test_humming_indexed_writes_supplied_output_buffer(): + from vllm.forward_context import set_forward_context + + activation = MoEActivation.SILU + experts = _make_humming_indexed_experts(activation) + moe_config = experts.moe_config + num_tokens = 1 + top_k = moe_config.experts_per_token + hidden_size = moe_config.hidden_dim + num_experts = moe_config.num_experts + workspace13_shape, workspace2_shape, _ = experts.workspace_shapes( + M=num_tokens, + N=moe_config.intermediate_size, + K=hidden_size, + topk=top_k, + global_num_experts=num_experts, + local_num_experts=num_experts, + expert_tokens_meta=None, + activation=activation, + ) + + device = torch.device("cuda") + dtype = experts.layer.params_dtype + workspace13 = torch.empty(workspace13_shape, dtype=dtype, device=device) + workspace2 = torch.empty(workspace2_shape, dtype=dtype, device=device) + hidden_states = torch.ones((num_tokens, hidden_size), dtype=dtype, device=device) + output = torch.full_like(hidden_states, torch.nan) + topk_weights = torch.full( + (num_tokens, top_k), + 1 / top_k, + dtype=dtype, + device=device, + ) + topk_ids = torch.arange(top_k, dtype=torch.int32, device=device).unsqueeze(0) + unused = torch.empty((num_experts, 0), device=device) + + with set_forward_context(None, vllm_config, num_tokens=num_tokens): + experts.apply( + output=output, + hidden_states=hidden_states, + w1=unused, + w2=unused, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert torch.isfinite(output).all() + + @pytest.mark.parametrize("ep_size", [1, 2]) def test_moe_align_block_size_opcheck(ep_size): num_experts = 4 diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index a017fa07bf90..7d07af4080c4 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -269,6 +269,7 @@ def test_moe_align_block_size_with_expert_map( if (experts[k] in local_experts) or not mask_inactive_experts else -1 ) + topk_ids[0, 0] = -1 actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( topk_ids=topk_ids, diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index cc8d9c36dc0a..1942ce96bcca 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -42,7 +42,11 @@ from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace from vllm.forward_context import set_forward_context -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, fused_experts +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, + MoERunner, + fused_experts, +) from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.router.router_factory import ( @@ -80,7 +84,7 @@ TOP_KS = [2, 6] # dp_size, tp_size, use_ep -# Note: DP+TP is not yet supported in the FusedMoE layer. +# Note: DP+TP is not yet supported in the FusedMoEFactory layer. PARALLEL_COMBOS = [ [1, 2, False], [1, 4, False], @@ -1023,7 +1027,7 @@ def make_fused_moe_layer( kwargs["routed_input_transform"] = routed_input_transform kwargs["routed_output_transform"] = routed_output_transform - layer = FusedMoE( + layer = FusedMoEFactory( num_experts=global_num_experts, top_k=top_k, hidden_size=hidden_size, @@ -1267,7 +1271,7 @@ def _test_body_eplb( ): output_before = sp_wrapper(moe_layer)(hidden_states, router_logits) - # Create a fresh FusedMoE layer with enable_eplb=True + # Create a fresh MoERunner layer with enable_eplb=True # Delete the original layer's registration so the constructor can # re-use the same "from_forward_context" prefix cc = vllm_config.compilation_config @@ -1401,7 +1405,8 @@ def _run_one_config( - When is_sequence_parallel=True (EP + sequence splitting): * ep_size: Number of expert parallel ranks (equals dp_size * tp_size) - * tp_size: Number of ranks to split sequence across (becomes sp_size in FusedMoE) + * tp_size: Number of ranks to split sequence across (becomes sp_size in + MoERunner) * Weights are chunked by ep_size (experts) but NOT by tp_size * Input sequences are chunked by tp_size (via sp_wrapper) """ @@ -1476,8 +1481,8 @@ def _run_one_config( torch.accelerator.empty_cache() with set_current_vllm_config(vllm_config): - # Chunk weights for EP BEFORE creating FusedMoE - # FusedMoE uses EP-chunked weights and handles reductions internally + # Chunk weights for EP BEFORE creating MoERunner. + # MoERunner uses EP-chunked weights and handles reductions internally. if ep_size > 1: # Split experts across ranks (dimension 0 is the expert dimension) # When EP is enabled, use EP group rank and ep_size for chunking diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index 2fd4e0fed5e1..bb377049d91e 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for FusedMoE weight loading with padded hidden dimensions. +"""Tests for FusedMoEFactory weight loading with padded hidden dimensions. When using DeepEP backends or NIXL EP with models like nemotron_h, hidden_size may be rounded up (e.g., 2688 -> 3072) for backend requirements. @@ -61,6 +61,55 @@ def test_invalid_shard_dim_raises(self): RoutedExperts._get_hidden_dim(shard_dim=0, ndim=3) +class TestOrientFusedWeight: + """Unit tests for _orient_fused_weight. + + E=8 experts, hidden=3072, intermediate=1024. + """ + + HIDDEN = 3072 + + def test_w13_standard_orientation_is_untouched(self): + weight = torch.randn(8, 2048, self.HIDDEN) + result = RoutedExperts._orient_fused_weight(weight, "w1", self.HIDDEN) + assert result.shape == (8, 2048, self.HIDDEN) + + def test_w13_transposed_checkpoint_is_normalised(self): + # e.g. Qwen3 VL MoE stores [experts, hidden, 2 * intermediate] + weight = torch.randn(8, self.HIDDEN, 2048) + result = RoutedExperts._orient_fused_weight(weight, "w3", self.HIDDEN) + assert result.shape == (8, 2048, self.HIDDEN) + + def test_w2_standard_orientation_is_untouched(self): + weight = torch.randn(8, self.HIDDEN, 1024) + result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN) + assert result.shape == (8, self.HIDDEN, 1024) + + def test_w2_transposed_checkpoint_is_normalised(self): + weight = torch.randn(8, 1024, self.HIDDEN) + result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN) + assert result.shape == (8, self.HIDDEN, 1024) + + def test_w13_per_channel_scale_is_untouched(self): + # A fused per-channel scale has no hidden dim, so transposing it would + # leave chunk()/TP sharding operating on the wrong axis. + scale = torch.randn(8, 2048, 1) + result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN) + assert result.shape == (8, 2048, 1) + assert result.chunk(2, dim=1)[0].shape == (8, 1024, 1) + + def test_w2_per_channel_scale_is_untouched(self): + scale = torch.randn(8, self.HIDDEN, 1) + result = RoutedExperts._orient_fused_weight(scale, "w2", self.HIDDEN) + assert result.shape == (8, self.HIDDEN, 1) + + def test_block_scale_is_untouched(self): + # Block scales are [experts, 2 * intermediate / block, hidden / block] + scale = torch.randn(8, 16, 24) + result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN) + assert result.shape == (8, 16, 24) + + class TestNarrowExpertDataForPadding: """Unit tests for _narrow_expert_data_for_padding.""" diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 7c2fdbabe294..830ff48734e1 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -19,9 +19,17 @@ pytest.skip("This test can only run on ROCm.", allow_module_level=True) from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, +) from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402 + _AITER_SWIGLU_ALPHA, + _AITER_SWIGLU_BETA, AiterMxfp8Experts, ) +from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( # noqa: E402 + Mxfp8EmulationTritonExperts, +) from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402 FusedMoEActivationFormat, ) @@ -33,6 +41,7 @@ _SUPPORTED_BACKENDS, _mxfp8_backend_to_kernel_cls, _select_kernel_cls, + select_mxfp8_moe_backend, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 kMxfp8Dynamic, @@ -43,7 +52,17 @@ def _config(ep_size: int = 1): - cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144) + # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and + # alpha/beta so is_supported_config doesn't reject the config on those grounds. + cfg = make_dummy_moe_config( + num_experts=128, + experts_per_token=4, + hidden_dim=6144, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + cfg = dataclasses.replace( + cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA + ) if ep_size != 1: cfg = dataclasses.replace( cfg, @@ -76,12 +95,6 @@ def test_aiter_mxfp8_registered(): ] -def test_triton_selectable(): - assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8 - # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto. - assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS - - @pytest.mark.parametrize("ep_size", [1, 2]) def test_ep_supported(ep_size): """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask.""" @@ -133,3 +146,24 @@ def test_explicit_moe_backend_aiter(): pytest.raises(ValueError, match="flydsl package"), ): _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) + + +def test_gfx950_picks_aiter(): + """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins.""" + # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require VLLM_ROCM_USE_AITER=1 + with ( + patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True), + _flydsl_installed(True), + ): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is AiterMxfp8Experts + + +def test_gfx942_picks_emulation(): + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead.""" + with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.EMULATION + assert experts_cls is Mxfp8EmulationTritonExperts diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 2f819c09aaae..7e8d4e3028ab 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1558,3 +1558,51 @@ def test_mxfp4_emulation_rounds_up_to_block_size( # The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate. assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0 assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0 + + +def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( + monkeypatch: pytest.MonkeyPatch, +): + """ + select_mxfp4_moe_backend() must raise NotImplementedError, with the + collected per-backend unsupported reasons in the message, when no + backend supports the requested deployment configuration. + """ + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle + from vllm.model_executor.layers.fused_moe import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + RoutingMethodType, + ) + + class UnsupportedExperts: + @staticmethod + def is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ): + return False, f"unsupported reason for {cls.__name__}" + + monkeypatch.setattr( + mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [UnsupportedExperts] + ) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + monkeypatch.setattr(current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + + moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cpu", + routing_method=RoutingMethodType.Renormalize, + ) + + with pytest.raises(NotImplementedError, match="Unsupported reasons"): + mxfp4_oracle.select_mxfp4_moe_backend(moe_config) diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py new file mode 100644 index 000000000000..d4a95750ae27 --- /dev/null +++ b/tests/kernels/moe/test_routed_experts_capture_monolithic.py @@ -0,0 +1,880 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for routed-expert capture on the monolithic MoE path. + +These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the +expert IDs picked by FlashInfer's fused router-and-experts kernels (the +"monolithic" path). When ``set_capture_fn`` is installed on +a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call +should: + + * allocate an int16 ``(num_tokens, top_k)`` buffer, + * pass it to FlashInfer as ``routing_replay_out``, + * invoke the callback after the kernel returns. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, + fp8_w8a8_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + TrtLlmBf16ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsMonolithic, +) +from vllm.platforms import current_platform + +try: + from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +except ImportError: + pytest.skip("flashinfer not available", allow_module_level=True) + +if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda(): + pytest.skip( + "Requires FlashInfer TRT-LLM fused MoE on CUDA", + allow_module_level=True, + ) + +if not current_platform.has_device_capability(100): + pytest.skip( + "TRT-LLM fused MoE kernels require SM100+", + allow_module_level=True, + ) + + +def _shuffle_bf16_weights_block_major_k( + w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128 +) -> torch.Tensor: + """Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by + ``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert + row shuffle. + """ + from flashinfer import shuffle_matrix_a + from flashinfer.fused_moe import convert_to_block_layout + + num_experts = w.shape[0] + shuffled = [] + for i in range(num_experts): + t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m) + shuffled.append(convert_to_block_layout(t, block_k)) + return torch.stack(shuffled).view(torch.bfloat16) + + +def _make_bf16_monolithic_experts( + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + routing_method: RoutingMethodType, + device: torch.device, +) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]: + """Construct the monolithic BF16 experts plus the BlockMajorK weights + expected by ``trtllm_bf16_moe``. + """ + parallel_cfg = FusedMoEParallelConfig.make_no_parallel() + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=parallel_cfg, + in_dtype=torch.bfloat16, + activation=MoEActivation.SILU, + device=device, + routing_method=routing_method, + max_num_tokens=16, + ) + quant_config = FusedMoEQuantConfig.make( + quant_dtype=None, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + ) + + experts = TrtLlmBf16ExpertsMonolithic( + moe_config=moe_config, quant_config=quant_config + ) + + gemm1 = ( + torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ) + gemm2 = ( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ) + w13 = _shuffle_bf16_weights_block_major_k(gemm1) + w2 = _shuffle_bf16_weights_block_major_k(gemm2) + return experts, w13, w2 + + +def _run_bf16_monolithic( + experts: TrtLlmBf16ExpertsMonolithic, + hidden_states: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + num_experts: int, + *, + n_group: int | None = None, + topk_group: int | None = None, + routed_scaling_factor: float | None = None, + e_score_correction_bias: torch.Tensor | None = None, +) -> torch.Tensor: + return experts.apply( + hidden_states=hidden_states, + w1=w13, + w2=w2, + router_logits=router_logits, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + apply_router_weight_on_input=False, + num_expert_group=n_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + +_DSV3_NUM_EXPERTS = 32 +_DSV3_N_GROUP = 4 +_DSV3_TOPK_GROUP = 2 + + +def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor: + return torch.randn(num_experts, device=device, dtype=torch.bfloat16) + + +@pytest.mark.parametrize("num_tokens", [2, 7, 16]) +@pytest.mark.parametrize("top_k", [2, 4]) +def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts( + num_tokens: int, + top_k: int, +) -> None: + """The capture callback should receive the int16 routed-expert IDs the + kernel actually used, the values should be valid expert indices, and + each token should pick ``top_k`` distinct experts.""" + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: + pytest.skip( + f"DSV3 requires top_k <= n_group * topk_group " + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" + ) + torch.manual_seed(0) + device = torch.device("cuda:0") + + num_experts = _DSV3_NUM_EXPERTS + hidden_size = 1024 + intermediate_size = 1024 + + experts, w13, w2 = _make_bf16_monolithic_experts( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + routing_method=RoutingMethodType.DeepSeekV3, + device=device, + ) + + captured: list[torch.Tensor] = [] + + def capture_fn(replay_out: torch.Tensor) -> None: + captured.append(replay_out.clone()) + + assert experts.supports_routing_replay_capture() + experts.set_capture_fn(capture_fn) + + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 + ) + router_logits = torch.rand( + num_tokens, num_experts, device=device, dtype=torch.float32 + ) + routing_bias = _make_dsv3_routing_bias(num_experts, device) + + _ = _run_bf16_monolithic( + experts, + hidden_states=hidden_states, + w13=w13, + w2=w2, + router_logits=router_logits, + num_experts=num_experts, + n_group=_DSV3_N_GROUP, + topk_group=_DSV3_TOPK_GROUP, + routed_scaling_factor=1.0, + e_score_correction_bias=routing_bias, + ) + + assert len(captured) == 1 + replay = captured[0] + assert replay.dtype == torch.int16 + assert replay.shape == (num_tokens, top_k) + assert (replay >= 0).all(), f"got out-of-range values: {replay}" + assert (replay < num_experts).all(), f"got out-of-range values: {replay}" + + for t in range(num_tokens): + unique = replay[t].unique() + assert unique.numel() == top_k, ( + f"token {t}: expected {top_k} distinct experts, " + f"got {unique.numel()} ({replay[t].tolist()})" + ) + + +@pytest.mark.parametrize("num_tokens", [2, 7, 16]) +@pytest.mark.parametrize( + "routing_method", + [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ], +) +def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( + num_tokens: int, + routing_method: RoutingMethodType, +) -> None: + """Routing replay works for non-DeepSeekV3 routing methods too. + FlashInfer's ``routing_replay_out`` is routing-method-agnostic.""" + torch.manual_seed(0) + device = torch.device("cuda:0") + + num_experts = 8 + top_k = 2 + hidden_size = 1024 + intermediate_size = 1024 + + experts, w13, w2 = _make_bf16_monolithic_experts( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + routing_method=routing_method, + device=device, + ) + + captured: list[torch.Tensor] = [] + experts.set_capture_fn(lambda r: captured.append(r.clone())) + + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 + ) + router_logits = torch.rand( + num_tokens, num_experts, device=device, dtype=torch.float32 + ) + + _ = _run_bf16_monolithic( + experts, + hidden_states=hidden_states, + w13=w13, + w2=w2, + router_logits=router_logits, + num_experts=num_experts, + ) + + assert len(captured) == 1 + replay = captured[0] + assert replay.dtype == torch.int16 + assert replay.shape == (num_tokens, top_k) + assert (replay >= 0).all(), f"got out-of-range values: {replay}" + assert (replay < num_experts).all(), f"got out-of-range values: {replay}" + for t in range(num_tokens): + unique = replay[t].unique() + assert unique.numel() == top_k, ( + f"token {t}: expected {top_k} distinct experts, " + f"got {unique.numel()} ({replay[t].tolist()})" + ) + + +def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: + """With no callback installed the kernel should not see a + ``routing_replay_out`` tensor — verify the helper short-circuits.""" + torch.manual_seed(0) + device = torch.device("cuda:0") + experts, _, _ = _make_bf16_monolithic_experts( + num_experts=_DSV3_NUM_EXPERTS, + top_k=2, + hidden_size=1024, + intermediate_size=1024, + routing_method=RoutingMethodType.DeepSeekV3, + device=device, + ) + # No callback installed. + buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device) + assert buf is None + + # Dispatch is also a no-op. + experts._maybe_dispatch_routing_replay(buf, num_tokens=4) + + +def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: + """FlashInfer's ``routing_replay_out`` is supported by all routing + methods, so ``supports_routing_replay_capture`` should be True + regardless of routing method.""" + device = torch.device("cuda:0") + for routing_method in ( + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ): + experts, _, _ = _make_bf16_monolithic_experts( + num_experts=_DSV3_NUM_EXPERTS, + top_k=2, + hidden_size=1024, + intermediate_size=1024, + routing_method=routing_method, + device=device, + ) + assert experts.supports_routing_replay_capture() is True, ( + f"{routing_method!r} should support routing replay capture" + ) + + +def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None: + """When capture is installed, the allocated buffer is int16 and shaped + ``(num_tokens, experts_per_token)``.""" + device = torch.device("cuda:0") + experts, _, _ = _make_bf16_monolithic_experts( + num_experts=_DSV3_NUM_EXPERTS, + top_k=4, + hidden_size=1024, + intermediate_size=1024, + routing_method=RoutingMethodType.DeepSeekV3, + device=device, + ) + experts.set_capture_fn(lambda r: None) + buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device) + assert buf is not None + assert buf.dtype == torch.int16 + assert buf.shape[0] >= 11 + assert buf.shape[1] == 4 + assert buf.device.type == "cuda" + + +def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None: + """End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback + on the monolithic experts and verify the captured rows land in the + capturer's device buffer at the correct layer slot. + + Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer`` + for the monolithic path: a single closure is installed on the monolithic + ``fused_experts`` (in addition to ``router.set_capture_fn`` on the + non-monolithic path) and the capturer routes per-layer based on the + closed-over ``layer_id``. + """ + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, + ) + + torch.manual_seed(7) + device = torch.device("cuda:0") + num_tokens = 4 + top_k = 2 + num_experts = _DSV3_NUM_EXPERTS + hidden_size = 1024 + intermediate_size = 1024 + + experts, w13, w2 = _make_bf16_monolithic_experts( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + routing_method=RoutingMethodType.DeepSeekV3, + device=device, + ) + + num_layers = 3 + layer_id = 1 + capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) + capturer.dp_rank = 0 + capturer.tp_size = 1 + capturer.device_buffer = torch.full( + (num_tokens + 4, num_layers, top_k), + -1, + dtype=torch.int32, + device=device, + ) + + def capture_fn(replay_out: torch.Tensor) -> None: + capturer.capture(layer_id, replay_out) + + experts.set_capture_fn(capture_fn) + + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 + ) + router_logits = torch.rand( + num_tokens, num_experts, device=device, dtype=torch.float32 + ) + routing_bias = _make_dsv3_routing_bias(num_experts, device) + + # Patch get_forward_context to return a dp_metadata=None context so the + # capturer takes the single-DP branch. + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec + + with patch.object( + rec, + "get_forward_context", + return_value=SimpleNamespace(dp_metadata=None), + ): + _ = _run_bf16_monolithic( + experts, + hidden_states=hidden_states, + w13=w13, + w2=w2, + router_logits=router_logits, + num_experts=num_experts, + n_group=_DSV3_N_GROUP, + topk_group=_DSV3_TOPK_GROUP, + routed_scaling_factor=1.0, + e_score_correction_bias=routing_bias, + ) + + captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu() + # Valid expert IDs at this layer. + assert (captured >= 0).all() + assert (captured < num_experts).all() + for t in range(num_tokens): + unique = captured[t].unique() + assert unique.numel() == top_k, ( + f"token {t}: expected {top_k} distinct experts at layer " + f"{layer_id}, got {unique.numel()}" + ) + + # Other layers / trailing token rows untouched. + for other_layer in range(num_layers): + if other_layer == layer_id: + continue + assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), ( + f"layer {other_layer} should be untouched, got writes" + ) + assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), ( + "tail rows beyond num_tokens should remain sentinel" + ) + + +# ---------------------------------------------------------------------------- +# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic`` +# ---------------------------------------------------------------------------- + + +def _make_fp8_block_scale_monolithic_experts( + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + device: torch.device, +) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]: + """Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale + code path with DSV3 routing. + + Weights are shuffled into the BlockMajorK layout the kernel expects + (same helper the vLLM weight loader uses for DeepSeek-FP8 models). + """ + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + _shuffle_deepseek_fp8_moe_weights, + ) + + block_k = 128 + parallel_cfg = FusedMoEParallelConfig.make_no_parallel() + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=parallel_cfg, + in_dtype=torch.bfloat16, + activation=MoEActivation.SILU, + device=device, + routing_method=RoutingMethodType.DeepSeekV3, + max_num_tokens=16, + ) + + # Random fp8 weights + ones-block scales (the kernel decoder only cares + # that the per-block scales are present and finite for routing/replay). + gemm1 = torch.randn( + num_experts, 2 * intermediate_size, hidden_size, device=device + ).to(torch.float8_e4m3fn) + gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to( + torch.float8_e4m3fn + ) + w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2) + + w1_scale = torch.ones( + num_experts, + 2 * intermediate_size // block_k, + hidden_size // block_k, + device=device, + dtype=torch.float32, + ) + w2_scale = torch.ones( + num_experts, + hidden_size // block_k, + intermediate_size // block_k, + device=device, + dtype=torch.float32, + ) + quant_config = fp8_w8a8_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + block_shape=[block_k, block_k], + per_act_token_quant=False, + ) + + experts = TrtLlmFp8ExpertsMonolithic( + moe_config=moe_config, quant_config=quant_config + ) + return experts, w13_shuffled, w2_shuffled + + +def _run_fp8_block_scale_monolithic( + experts: TrtLlmFp8ExpertsMonolithic, + hidden_states_fp8: torch.Tensor, + hidden_states_scale: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + num_experts: int, + routing_bias: torch.Tensor, +) -> torch.Tensor: + return experts.apply( + hidden_states=hidden_states_fp8, + w1=w13, + w2=w2, + router_logits=router_logits, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + # The block-scale apply path reads ``a1q_scale`` and transposes it + # to ``(hidden_size/128, num_tokens)`` for the kernel call. + a1q_scale=hidden_states_scale, + apply_router_weight_on_input=False, + num_expert_group=_DSV3_N_GROUP, + topk_group=_DSV3_TOPK_GROUP, + e_score_correction_bias=routing_bias, + routed_scaling_factor=1.0, + ) + + +@pytest.mark.parametrize("num_tokens", [2, 7, 16]) +@pytest.mark.parametrize("top_k", [2, 4]) +def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts( + num_tokens: int, + top_k: int, +) -> None: + """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale + path, DSV3 routing) captures valid expert IDs.""" + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: + pytest.skip( + f"DSV3 requires top_k <= n_group * topk_group " + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" + ) + torch.manual_seed(0) + device = torch.device("cuda:0") + + num_experts = _DSV3_NUM_EXPERTS + hidden_size = 1024 + intermediate_size = 1024 + + experts, w13, w2 = _make_fp8_block_scale_monolithic_experts( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + assert experts.supports_routing_replay_capture() + + captured: list[torch.Tensor] = [] + experts.set_capture_fn(lambda r: captured.append(r.clone())) + + # Per-token / per-block hidden scales (ones is fine for the routing + # path; the GEMM output isn't being asserted on). + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 + ).to(torch.float8_e4m3fn) + hidden_states_scale = torch.ones( + num_tokens, hidden_size // 128, device=device, dtype=torch.float32 + ) + router_logits = torch.rand( + num_tokens, num_experts, device=device, dtype=torch.float32 + ) + routing_bias = _make_dsv3_routing_bias(num_experts, device) + + _ = _run_fp8_block_scale_monolithic( + experts, + hidden_states_fp8=hidden_states, + hidden_states_scale=hidden_states_scale, + w13=w13, + w2=w2, + router_logits=router_logits, + num_experts=num_experts, + routing_bias=routing_bias, + ) + + assert len(captured) == 1 + replay = captured[0] + assert replay.dtype == torch.int16 + assert replay.shape == (num_tokens, top_k) + assert (replay >= 0).all(), f"got out-of-range values: {replay}" + assert (replay < num_experts).all(), f"got out-of-range values: {replay}" + for t in range(num_tokens): + unique = replay[t].unique() + assert unique.numel() == top_k, ( + f"token {t}: expected {top_k} distinct experts, " + f"got {unique.numel()} ({replay[t].tolist()})" + ) + + +# ---------------------------------------------------------------------------- +# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic`` +# ---------------------------------------------------------------------------- + + +def _make_nvfp4_monolithic_experts( + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + device: torch.device, +) -> tuple[ + TrtLlmNvFp4ExpertsMonolithic, + torch.Tensor, # w13 (packed nvfp4 uint8) + torch.Tensor, # w13 block-scale (fp8) + torch.Tensor, # w2 (packed nvfp4 uint8) + torch.Tensor, # w2 block-scale (fp8) + torch.Tensor, # input global scale (per-tensor float32) +]: + """Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights + and DSV3 routing. + + NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor + "global" scale. We follow the layout in ``test_ocp_mx_moe.py`` / + ``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``: + * weights: uint8 (packed fp4) ``(E, M, K//2)`` + * weight scales: fp8 ``(E, M, K//16)`` + * hidden states: uint8 (packed fp4) ``(N, K//2)`` + * hidden state scales: fp8 ``(N, K//16)`` + """ + from flashinfer import fp4_quantize + + block_size = 16 + parallel_cfg = FusedMoEParallelConfig.make_no_parallel() + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=parallel_cfg, + in_dtype=torch.bfloat16, + activation=MoEActivation.SILU, + device=device, + routing_method=RoutingMethodType.DeepSeekV3, + max_num_tokens=16, + ) + + gemm1 = torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) + gemm2 = torch.randn( + num_experts, + hidden_size, + intermediate_size, + device=device, + dtype=torch.bfloat16, + ) + # Per-tensor weight scaling factor (used to build ``g1_alphas`` / + # ``g2_alphas`` below). + w_global_scale = torch.tensor(1.0, device=device) + # Per-tensor input scaling factor. + a_global_scale = torch.tensor(1.0, device=device) + + w13_q, w13_scale = fp4_quantize( + gemm1, + w_global_scale, + block_size, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + ) + w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // block_size + ) + w2_q, w2_scale = fp4_quantize( + gemm2, + w_global_scale, + block_size, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + ) + w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // block_size + ) + + # NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2, + # g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these. + g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32) + a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32) + + quant_config = FusedMoEQuantConfig.make( + quant_dtype="nvfp4", + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + w1_scale=w13_scale, + w2_scale=w2_scale, + g1_alphas=g_alphas, + g2_alphas=g_alphas, + a1_gscale=a_global_scale, + a2_gscale=a2_gscale, + ) + + experts = TrtLlmNvFp4ExpertsMonolithic( + moe_config=moe_config, quant_config=quant_config + ) + return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale + + +def _run_nvfp4_monolithic( + experts: TrtLlmNvFp4ExpertsMonolithic, + hidden_states_q: torch.Tensor, + hidden_states_scale: torch.Tensor, + router_logits: torch.Tensor, + num_experts: int, + routing_bias: torch.Tensor, +) -> torch.Tensor: + """The monolithic NVFP4 apply expects packed fp4 hidden states + the + matching fp8 per-block scale stored in the ``a1q_scale`` slot.""" + # Stash the weight tensors on the experts in the locations the apply() + # implementation reads from (it pulls them from quant_config / scales + # already; w1/w2 come in as args). + return experts.apply( + hidden_states=hidden_states_q, + w1=experts._w13_packed, + w2=experts._w2_packed, + router_logits=router_logits, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=hidden_states_scale, + apply_router_weight_on_input=False, + num_expert_group=_DSV3_N_GROUP, + topk_group=_DSV3_TOPK_GROUP, + e_score_correction_bias=routing_bias, + routed_scaling_factor=1.0, + ) + + +@pytest.mark.parametrize("num_tokens", [2, 7, 16]) +@pytest.mark.parametrize("top_k", [2, 4]) +def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts( + num_tokens: int, + top_k: int, +) -> None: + """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs + on the DSV3 routing path.""" + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: + pytest.skip( + f"DSV3 requires top_k <= n_group * topk_group " + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" + ) + from flashinfer import fp4_quantize + + torch.manual_seed(0) + device = torch.device("cuda:0") + + num_experts = _DSV3_NUM_EXPERTS + hidden_size = 1024 + intermediate_size = 1024 + block_size = 16 + + experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + # The apply() reads w1/w2 from its args, but we keep them on the experts + # for convenience of the helper. + experts._w13_packed = w13_q + experts._w2_packed = w2_q + + assert experts.supports_routing_replay_capture() + captured: list[torch.Tensor] = [] + experts.set_capture_fn(lambda r: captured.append(r.clone())) + + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 + ) + hidden_states_q, hidden_states_scale = fp4_quantize( + hidden_states, + a_gs, + block_size, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + ) + # The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave + # ``hidden_states_scale`` in its native (uint8 packed) form. + router_logits = torch.rand( + num_tokens, num_experts, device=device, dtype=torch.float32 + ) + routing_bias = _make_dsv3_routing_bias(num_experts, device) + + _ = _run_nvfp4_monolithic( + experts, + hidden_states_q=hidden_states_q, + hidden_states_scale=hidden_states_scale, + router_logits=router_logits, + num_experts=num_experts, + routing_bias=routing_bias, + ) + + assert len(captured) == 1 + replay = captured[0] + assert replay.dtype == torch.int16 + assert replay.shape == (num_tokens, top_k) + assert (replay >= 0).all(), f"got out-of-range values: {replay}" + assert (replay < num_experts).all(), f"got out-of-range values: {replay}" + for t in range(num_tokens): + unique = replay[t].unique() + assert unique.numel() == top_k, ( + f"token {t}: expected {top_k} distinct experts, " + f"got {unique.numel()} ({replay[t].tolist()})" + ) diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 62a4968a0d1f..0eded09b1ae5 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -8,9 +8,19 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + FusedTopKBiasRouter, +) +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( + FusedTopKRouter, +) +from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + GroupedTopKRouter, +) from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) @@ -19,13 +29,13 @@ def _is_aiter_capable() -> bool: - """Check if the platform supports AITER (gfx942/gfx950).""" + """Check if the platform supports AITER (gfx942/gfx950/gfx1250).""" if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - return _ON_MI3XX + return get_cdna_version() > 2 except ImportError: return False @@ -36,6 +46,112 @@ def _is_aiter_capable() -> bool: NUM_EXPERTS = [8, 16, 64] +def test_degenerate_grouped_config_uses_standard_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, FusedTopKRouter) + hidden_states, router_logits = make_test_data(32, 256, 128) + + topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) + baseline_weights, baseline_ids = baseline_fused_topk( + router_logits, + top_k=4, + renormalize=True, + ) + + assert_routing_results_close( + topk_weights, + topk_ids, + baseline_weights, + baseline_ids, + ) + + +def test_multiple_expert_groups_use_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=8, + topk_group=4, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, GroupedTopKRouter) + + +def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + + +def test_degenerate_grouped_config_with_bias_keeps_routed_scale() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + assert router.routed_scaling_factor == 1.1 + + +def test_degenerate_deepseek_v3_routing_stays_grouped() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, GroupedTopKRouter) + assert router.routing_method_type == RoutingMethodType.DeepSeekV3 + + +def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + ) + + assert isinstance(router, GroupedTopKRouter) + + def setup_eplb_state( enable_eplb: bool, global_num_experts: int ) -> EplbLayerState | None: diff --git a/tests/kernels/moe/test_routing_simulator.py b/tests/kernels/moe/test_routing_simulator.py index 4ef984a3296b..35b1af618cf7 100644 --- a/tests/kernels/moe/test_routing_simulator.py +++ b/tests/kernels/moe/test_routing_simulator.py @@ -6,7 +6,7 @@ This script demonstrates how to use the routing simulator to test different routing strategies and analyze their performance, including -integration tests with FusedMoE layer. +integration tests with FusedMoEFactory layer. """ import tempfile @@ -77,11 +77,11 @@ def test_basic_functionality( def test_routing_strategy_integration(monkeypatch, device): """Test that the routing strategy environment variable works with - FusedMoE.""" + FusedMoEFactory.""" pytest.importorskip("vllm.model_executor.layers.fused_moe.layer") import vllm.envs as envs - from vllm.model_executor.layers.fused_moe.layer import FusedMoE + from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory # Test parameters num_tokens = 32 @@ -111,7 +111,7 @@ def test_routing_strategy_integration(monkeypatch, device): ) for strategy in strategies: - fused_moe = FusedMoE( + fused_moe = FusedMoEFactory( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index 4515021a4e91..af9455b96039 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Tests for FusedMoE with routed_input_transform. +Tests for FusedMoEFactory with routed_input_transform. -Verifies that applying routed_input_transform inside FusedMoE +Verifies that applying routed_input_transform inside FusedMoEFactory produces the same results as applying the transform manually outside. """ @@ -13,7 +13,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import set_forward_context -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed @@ -133,9 +133,9 @@ def test_routed_input_transform_inside_vs_outside( workspace_init, monkeypatch, ): - """Compare FusedMoE with transform inside vs manually applying outside. - Method A (inside): FusedMoE with routed_input_transform - Method B (outside): Manually transform, then FusedMoE without transform + """Compare FusedMoEFactory with transform inside vs manually applying outside. + Method A (inside): FusedMoEFactory with routed_input_transform + Method B (outside): Manually transform, then FusedMoEFactory without transform """ if current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0") @@ -157,8 +157,8 @@ def test_routed_input_transform_inside_vs_outside( routed_transform = SimpleLinear(hidden_size, latent_size, dtype) with set_current_vllm_config(vllm_config): - # Method A: FusedMoE WITH routed_input_transform - moe_with_transform = FusedMoE( + # Method A: FusedMoEFactory WITH routed_input_transform + moe_with_transform = FusedMoEFactory( shared_experts=shared_experts, routed_input_transform=routed_transform, num_experts=num_experts, @@ -173,9 +173,9 @@ def test_routed_input_transform_inside_vs_outside( prefix="moe_with_transform", ) - # Method B: FusedMoE WITHOUT routed_input_transform + # Method B: FusedMoEFactory WITHOUT routed_input_transform # Note: shared_experts=None because when transform is done outside, - moe_without_transform = FusedMoE( + moe_without_transform = FusedMoEFactory( shared_experts=None, routed_input_transform=None, num_experts=num_experts, diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 5c6691bbc32e..c5f00253e781 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -6,6 +6,7 @@ import torch import torch.nn.functional as F +import vllm._custom_ops as ops from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -231,3 +232,119 @@ def test_dsv4_fast_topk( atol=2e-5, rtol=2e-5, ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="This test is skipped on non-CUDA platform.", +) +@pytest.mark.parametrize("use_hash", [False, True]) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("use_padding_mask", [False, True]) +@pytest.mark.parametrize("pad_with_nan", [False, True]) +@pytest.mark.parametrize("num_experts", [128, 256, 384]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) +def test_fused_topk_softplus_sqrt_padding( + use_hash: bool, + use_bias: bool, + use_padding_mask: bool, + pad_with_nan: bool, + num_experts: int, + dtype: torch.dtype, +): + """Verify explicit padding and NaN-padded rows do not affect real rows.""" + torch.manual_seed(0) + num_tokens = 8 + topk = 6 + indices_dtype = torch.int32 + + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + + padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") + padding_rows[1::2] = True + if pad_with_nan: + gating_output[padding_rows] = float("nan") + is_padding = padding_rows if use_padding_mask else None + + # A negative correction bias makes explicit pad rows look selectable unless + # the kernel uses the is_padding guard. + e_score_correction_bias = None + if use_bias: + e_score_correction_bias = ( + -torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0 + ) + + input_ids = None + hash_indices_table = None + if use_hash: + vocab_size = 64 + hash_indices_table = torch.stack( + [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] + ).to(device="cuda", dtype=indices_dtype) + input_ids = torch.randint( + 0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda" + ) + + topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda") + topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda") + token_expert_indices = torch.empty( + num_tokens, topk, dtype=torch.int32, device="cuda" + ) + + ops.topk_hash_softplus_sqrt( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize=True, + routed_scaling_factor=1.0, + e_score_correction_bias=e_score_correction_bias, + input_tokens=input_ids, + hash_indices_table=hash_indices_table, + is_padding=is_padding, + ) + + if use_padding_mask: + pad_ids = topk_ids[padding_rows] + pad_weights = topk_weights[padding_rows] + assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), ( + f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}" + ) + assert (pad_weights == 0).all(), ( + "Explicit pad rows should have all-zero weights, " + f"got {pad_weights.tolist()}" + ) + + if pad_with_nan: + nan_pad_weights = topk_weights[padding_rows] + assert torch.isfinite(nan_pad_weights).all(), ( + f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}" + ) + assert (nan_pad_weights == 0).all(), ( + "NaN-padded rows should have all-zero weights, " + f"got {nan_pad_weights.tolist()}" + ) + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output, + topk=topk, + renormalize=True, + routed_scaling_factor=1.0, + e_score_correction_bias=e_score_correction_bias, + input_ids=input_ids, + hash_indices_table=hash_indices_table, + ) + + rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda") + if use_padding_mask or pad_with_nan: + rows_to_compare = ~padding_rows + + sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1) + sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1) + torch.testing.assert_close( + sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0 + ) + + sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref) + sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops) + torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py index 0ab025dceca4..01246325c803 100644 --- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py +++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py @@ -102,9 +102,9 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk): ).sum(dim=1) -@pytest.fixture(autouse=True, scope="module") +@pytest.fixture(autouse=True) def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" + """Sets the default CUDA device before each test in this module.""" torch.set_default_device("cuda") diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 71e33b7dfacc..c2084800f979 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for FusedMoE with zero experts. +"""Tests for FusedMoEFactory with zero experts. Verifies that: - The ZeroExpertRouter is properly created and used as the layer router. -- A forward pass through FusedMoE with zero experts produces correct output. +- A forward pass through FusedMoEFactory with zero experts produces correct output. - The output decomposes correctly into real expert + zero expert contributions. Note: tests generated with Claude. @@ -15,7 +15,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import get_forward_context, set_forward_context -from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) @@ -24,7 +24,7 @@ @pytest.fixture def zero_expert_moe(dist_init, default_vllm_config): - """Create a FusedMoE layer with zero experts.""" + """Create a FusedMoEFactory layer with zero experts.""" num_experts = 4 top_k = 2 # hidden_size must be >= 256 for the zero expert identity kernel to @@ -45,7 +45,7 @@ def zero_expert_moe(dist_init, default_vllm_config): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): init_workspace_manager(torch.accelerator.current_device_index()) - layer = FusedMoE( + layer = FusedMoEFactory( zero_expert_type="identity", e_score_correction_bias=e_score_correction_bias, num_experts=num_experts, @@ -66,7 +66,7 @@ def zero_expert_moe(dist_init, default_vllm_config): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens): - """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter.""" + """Verify that FusedMoEFactory with zero_expert_type creates a ZeroExpertRouter.""" layer, _ = zero_expert_moe assert isinstance(layer.router, ZeroExpertRouter), ( f"Expected ZeroExpertRouter but got {type(layer.router).__name__}." @@ -83,7 +83,10 @@ def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_token @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): - """Run a forward pass through FusedMoE with zero experts and verify output shape.""" + """ + Run a forward pass through FusedMoEFactory with zero experts + and verify output shape. + """ layer, vllm_config = zero_expert_moe hidden_size = layer.routed_experts.hidden_size @@ -118,16 +121,16 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): - """Validate that the FusedMoE output equals a plain FusedMoE + """Validate that the FusedMoEFactory output equals a plain FusedMoEFactory output (real experts only) plus the zero expert contribution. The key invariant is: zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real) + zero_expert_output - We create a plain FusedMoE layer with the same weights and real-expert-only + We create a plain FusedMoEFactory layer with the same weights and real-expert-only router logits, compute the zero expert output via the ZeroExpertRouter, and - verify the sum matches the FusedMoE output. + verify the sum matches the FusedMoEFactory output. """ layer, vllm_config = zero_expert_moe num_experts = 4 @@ -152,9 +155,9 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): get_forward_context().all_moe_layers = None - # Create a plain FusedMoE layer with the same config but no zero + # Create a plain FusedMoEFactory layer with the same config but no zero # experts. Use a separate prefix to avoid collision. - plain_layer = FusedMoE( + plain_layer = FusedMoEFactory( num_experts=num_experts, top_k=layer.routed_experts.top_k, hidden_size=layer.routed_experts.hidden_size, @@ -210,7 +213,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): expected, atol=4e-3, rtol=4e-3, - msg="FusedMoE output should equal plain FusedMoE output " + msg="FusedMoEFactory output should equal plain FusedMoEFactory output " "plus zero expert contribution", ) diff --git a/tests/kernels/quantization/test_block_int8.py b/tests/kernels/quantization/test_block_int8.py index 310091b6a554..3a13a53118be 100644 --- a/tests/kernels/quantization/test_block_int8.py +++ b/tests/kernels/quantization/test_block_int8.py @@ -28,12 +28,6 @@ SEEDS = [0] -@pytest.fixture(autouse=True, scope="module") -def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" - torch.set_default_device("cuda") - - @pytest.mark.parametrize( "M,N,K,block_size,out_dtype,seed", itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS), @@ -41,22 +35,28 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed): torch.manual_seed(seed) + device = current_platform.device_type factor_for_scale = 1e-2 int8_info = torch.iinfo(torch.int8) int8_max, int8_min = int8_info.max, int8_info.min - A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max + A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device) + A_fp32 = (A_fp32 - 0.5) * 2 * int8_max A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max + B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device) + B_fp32 = (B_fp32 - 0.5) * 2 * int8_max B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) block_n, block_k = block_size[0], block_size[1] n_tiles = (N + block_n - 1) // block_n k_tiles = (K + block_k - 1) // block_k - As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale - Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale + As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale + Bs = ( + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device) + * factor_for_scale + ) ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index 0daa4889227c..66385f04eb0b 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -82,12 +82,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids ).sum(dim=1) -@pytest.fixture(autouse=True, scope="module") -def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" - torch.set_default_device("cuda") - - DTYPES = [torch.half, torch.bfloat16] M = [1, 33] N = [128, 1024] @@ -104,6 +98,7 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed): torch.manual_seed(seed) + device = current_platform.device_type # Initialize int8 quantization parameters factor_for_scale = 1e-2 int8_max = 127 @@ -111,19 +106,26 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed): # Input tensor # M * K - a = torch.randn((M, K), dtype=dtype) / 10 + a = torch.randn((M, K), dtype=dtype, device=device) / 10 # Generate int8 weights - w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 + w1_fp32 = ( + torch.rand( + (E, 2 * N, K), + dtype=torch.float32, + device=device, + ) + - 0.5 + ) * 2 w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8) - w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 + w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2 w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8) # Generate scale for each column (per-column quantization) w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale - score = torch.randn((M, E), dtype=dtype) + score = torch.randn((M, E), dtype=dtype, device=device) score = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(score, topk) diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index be987fda6daf..981046fa5360 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -678,7 +678,7 @@ def make_layer(hidden, intermediate): layer.hidden_size = hidden layer.apply_router_weight_on_input = False layer.moe_config = SimpleNamespace( - intermediate_size_per_partition_unpadded=intermediate + hidden_dim=hidden, intermediate_size_per_partition_unpadded=intermediate ) return layer diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py new file mode 100644 index 000000000000..c3f8bc094d75 --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP4 linear kernel selection logic (CPU-only) + +Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + AiterMxfp4LinearKernel, + EmulationMxfp4LinearKernel, + FlashInferMxFp4LinearKernel, + HummingMxFp4LinearKernel, + MarlinMxFp4LinearKernel, + MxFp4LinearKernel, + MxFp4LinearLayerConfig, + XPUMxFp4LinearKernel, + init_mxfp4_linear_kernel, + register_linear_kernel, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E3M2Dynamic, +) +from vllm.platforms import PlatformEnum + +pytestmark = pytest.mark.cpu_test + +# Kernels that quantize activations themselves (true W4A4): they require an +# explicit MXFP4-dynamic activation key. +_TRUE_W4A4_KERNELS = [ + FlashInferMxFp4LinearKernel, + XPUMxFp4LinearKernel, + AiterMxfp4LinearKernel, +] + +# Weight-only (A16) kernels: they never quantize activations. They still accept +# MXFP4 activation keys as an intentional compatibility fallback. +_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel] + + +def test_can_implement_is_abstract(): + """Test that can_implement()/is_supported() are properly defined.""" + assert hasattr(MxFp4LinearKernel, "can_implement") + assert hasattr(MxFp4LinearKernel, "is_supported") + + +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_accept_dynamic_mxfp4_activation(kernel_cls): + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert can_implement, reason + + +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): + """None means weight-only/unquantized activations, not dynamic MXFP4.""" + config = MxFp4LinearLayerConfig() + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): + """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an + explicit request for a different activation format must be rejected.""" + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) +@pytest.mark.parametrize("activation_quant_key", [None, kMxfp4Dynamic]) +def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( + kernel_cls, activation_quant_key +): + """Marlin/Humming never quantize activations, so an unset activation key, + or one that already describes MXFP4-shaped data, is tolerated. When an + activation key is explicitly set, a warning must be logged noting that it + is ignored, since these kernels are weight-only (A16).""" + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) + with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: + can_implement, reason = kernel_cls.can_implement(config) + assert can_implement, reason + + if activation_quant_key is None: + warning_once.assert_not_called() + else: + warning_once.assert_called_once() + message = warning_once.call_args.args[0] + assert "the requested activation quantization" in message + assert "is ignored" in message + + +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) +def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize( + "activation_quant_key", + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], +) +def test_emulation_kernel_accepts_any_config(activation_quant_key): + """EmulationMxfp4LinearKernel is the universal fallback: it must accept + every supported activation format.""" + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) + with patch( + "vllm.model_executor.kernels.linear._get_linear_backend", + return_value="emulation", + ): + can_implement, reason = EmulationMxfp4LinearKernel.can_implement(config) + assert can_implement, reason + + +def test_emulation_kernel_derives_quant_dequant_func_from_config(): + """quant_dequant_func must be derived purely from the config's activation + QuantKey, not set externally.""" + with patch( + "vllm.model_executor.kernels.linear._get_linear_backend", + return_value="emulation", + ): + weight_only_config = MxFp4LinearLayerConfig() + kernel = EmulationMxfp4LinearKernel(weight_only_config) + x = torch.randn(4) + # identity for weight-only + assert torch.equal(kernel.quant_dequant_func(x), x) + + w4a4_config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) + kernel = EmulationMxfp4LinearKernel(w4a4_config) + assert kernel.quant_dequant_func is quant_dequant_mxfp4 + + +def test_aiter_kernel_is_supported_requires_native_mx_support(): + """AiterMxfp4LinearKernel must not be selected on platforms without + native MX compute, even if AITER itself is importable.""" + with patch( + "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", + return_value=False, + ): + is_supported, reason = AiterMxfp4LinearKernel.is_supported() + assert not is_supported + assert reason + + +class OOTMxFp4LinearKernel(MxFp4LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): + """init_mxfp4_linear_kernel should select a registered kernel that + reports itself as supported, and construct it with a fresh config.""" + platform_mock._enum = PlatformEnum.OOT + register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") + + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) + + assert isinstance(kernel, OOTMxFp4LinearKernel) + assert kernel.config == MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) + + +class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return False, "never supported" + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock): + platform_mock._enum = PlatformEnum.UNSPECIFIED + register_linear_kernel( + UnsupportedMxFp4LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp4" + ) + + with pytest.raises(ValueError, match="Failed to find a kernel"): + init_mxfp4_linear_kernel() diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py new file mode 100644 index 000000000000..2c7b565bea57 --- /dev/null +++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP6 linear kernel selection logic (CPU-only) + +Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + EmulationMxfp6LinearKernel, + MxFp6LinearKernel, + MxFp6LinearLayerConfig, + init_mxfp6_linear_kernel, + register_linear_kernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + kMxfp4Static, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, +) +from vllm.platforms import PlatformEnum + +pytestmark = pytest.mark.cpu_test + +# The only implementation available at the moment is software emulation. +_WEIGHT_QUANT_KEYS = [kMxfp6E3M2Static, kMxfp6E2M3Static] + + +def test_can_implement_is_abstract(): + """Test that can_implement()/is_supported() are properly defined.""" + assert hasattr(MxFp6LinearKernel, "can_implement") + assert hasattr(MxFp6LinearKernel, "is_supported") + + +def test_emulation_kernel_rejects_non_mxfp6_weights(): + """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight + format.""" + config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS) +@pytest.mark.parametrize( + "activation_quant_key", + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], +) +def test_emulation_kernel_accepts_any_supported_config( + weight_quant_key, activation_quant_key +): + """EmulationMxfp6LinearKernel is the only backend today: it must accept + every supported weight/activation format combination.""" + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key + ) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert can_implement, reason + + +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS) +def test_emulation_kernel_rejects_non_mxfp4_or_mxfp6_activation(weight_quant_key): + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, activation_quant_key=kMxfp4Static + ) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +class OOTMxFp6LinearKernel(MxFp6LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock): + """init_mxfp6_linear_kernel should select a registered kernel that + reports itself as supported/able to implement the given config, and + construct it with that exact config.""" + platform_mock._enum = PlatformEnum.OOT + register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6") + + kernel = init_mxfp6_linear_kernel( + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic + ) + + assert isinstance(kernel, OOTMxFp6LinearKernel) + assert kernel.config == MxFp6LinearLayerConfig( + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic + ) + + +class UnsupportedMxFp6LinearKernel(MxFp6LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return False, "never supported" + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp6_linear_kernel_raises_when_no_kernel_matches(platform_mock): + platform_mock._enum = PlatformEnum.UNSPECIFIED + register_linear_kernel( + UnsupportedMxFp6LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp6" + ) + + with pytest.raises(ValueError, match="Failed to find a kernel"): + init_mxfp6_linear_kernel(weight_quant_key=kMxfp6E3M2Static) diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index d2056fe01ebb..eb057cebd912 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast -import huggingface_hub import pytest import torch from safetensors import safe_open @@ -33,6 +32,7 @@ kNvfp4Static, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.triton_utils import triton if current_platform.is_rocm(): @@ -56,7 +56,7 @@ def on_gfx950() -> bool: @pytest.fixture(scope="module") def loaded_model_files(): return { - model_id: huggingface_hub.snapshot_download( + model_id: hf_api().snapshot_download( repo_id=model_id, allow_patterns=config["shards"] ) for model_id, config in MOE_MODEL_CONFIGS.items() diff --git a/tests/kernels/quantization/test_rocm_skinny_gemms.py b/tests/kernels/quantization/test_rocm_skinny_gemms.py index d2123db2e8da..e137231dc9ba 100644 --- a/tests/kernels/quantization/test_rocm_skinny_gemms.py +++ b/tests/kernels/quantization/test_rocm_skinny_gemms.py @@ -150,8 +150,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode xavier = ( math.sqrt(2 / k) if xnorm else 1 ) # normalize to avoid large output-bias deltas - A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier - B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier + A = torch.randn(n, k, dtype=dtype, device="cuda") * xavier + B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier if padded_a: A = pad_fp8(A) @@ -167,7 +167,12 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode out = ops.wvSplitKrc(A, B, cu_count, BIAS) if xnorm: - torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8) + # The O(1) bias lifts outputs to ~O(1), where one bf16 ULP (~3.9e-3, the + # worst measured divergence) exceeds 1e-3. Bump atol to 5e-3 only for + # biased bf16 (above that ULP, still under finfo(bf16).eps); keep 1e-3 + # otherwise. + atol = 5e-3 if (dtype == torch.bfloat16 and BIAS is not None) else 1e-3 + torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-8) else: torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py new file mode 100644 index 000000000000..0d5813dd50a3 --- /dev/null +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -0,0 +1,649 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Kimi-K3 SM103 decode GEMM selector (shape-only dispatch).""" + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import regex as re +import torch +from torch import nn + +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig +from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm +from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + +# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is +# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16). +EXPECTED_SELECTIONS = { + (1536, 128): (set(), set(range(1, 17))), + (3072, 128): (set(), set(range(1, 17))), + (1536, 7168): (set(), set(range(1, 17))), + (3072, 7168): (set(range(1, 6)), set()), + (2112, 7168): (set(), set(range(1, 17))), + (2304, 1536): (set(), set(range(1, 17))), + (4608, 1536): (set(), set(range(1, 17))), + (3584, 7168): ({1}, set(range(2, 9))), + (6288, 7168): (set(range(1, 5)), set()), + (12448, 7168): (set(range(1, 4)), set()), + (7168, 768): (set(), set(range(1, 17))), + (7168, 1536): ({1}, set()), + (7168, 3072): ({1, 2}, set()), + (7168, 3584): ({1, 2}, set()), + (7168, 4224): ({1}, set()), + (7168, 8448): (set(range(1, 4)), set()), + (8448, 7168): ({1, 2}, set()), + (16896, 7168): ({1, 2}, set()), + (20480, 7168): (set(range(1, 5)), set()), + (40960, 7168): (set(range(1, 5)), set()), + # TP16. + (3216, 7168): (set(range(1, 6)), set(range(9, 16))), + (768, 7168): (set(range(1, 5)), set(range(5, 17))), + (1152, 1536): ({1}, set(range(2, 17))), + (768, 128): (set(), set(range(1, 17))), + (7168, 384): (set(), set(range(1, 9))), + (4224, 7168): (set(range(1, 4)), set(range(4, 9))), + (10240, 7168): (set(range(1, 5)), set()), +} + +CUTE_CASES = [ + (spec.n, spec.k, num_tokens) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, _ in spec.cute_configs +] + +RESIDUAL_CUTE_CASES = [ + (spec.n, spec.k, num_tokens) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, _ in spec.residual_configs +] + +EXPECTED_CUTE_CONFIGS = { + (3072, 7168, 1): (224, 3, 4, 8), + (3072, 7168, 2): (128, 3, 2, 8), + (3072, 7168, 3): (128, 2, 1, 8), + (3072, 7168, 4): (64, 2, 2, 8), + (3072, 7168, 5): (128, 3, 1, 8), + (3584, 7168, 1): (224, 2, 4, 8), + (6288, 7168, 1): (224, 3, 4, 8), + (6288, 7168, 2): (64, 3, 2, 8), + (6288, 7168, 3): (32, 3, 4, 8), + (6288, 7168, 4): (128, 6, 1, 8), + (12448, 7168, 1): (224, 4, 2, 8), + (12448, 7168, 2): (64, 4, 2, 8), + (12448, 7168, 3): (64, 2, 2, 8), + (7168, 1536, 1): (96, 4, 2, 8), + (7168, 3072, 1): (96, 2, 4, 8), + (7168, 3072, 2): (32, 4, 4, 8), + (7168, 3584, 1): (224, 4, 2, 8), + (7168, 3584, 2): (64, 4, 2, 8), + (7168, 4224, 1): (96, 4, 2, 4), + (7168, 8448, 1): (32, 4, 4, 8), + (7168, 8448, 2): (96, 4, 1, 8), + (7168, 8448, 3): (96, 4, 1, 8), + (8448, 7168, 1): (224, 3, 4, 8), + (8448, 7168, 2): (32, 4, 4, 8), + (16896, 7168, 1): (224, 6, 4, 8), + (16896, 7168, 2): (32, 4, 4, 8), + (20480, 7168, 1): (224, 4, 2, 8), + (20480, 7168, 2): (64, 4, 2, 8), + (20480, 7168, 3): (64, 2, 2, 8), + (20480, 7168, 4): (64, 4, 1, 8), + (40960, 7168, 1): (128, 4, 2, 8), + (40960, 7168, 2): (64, 4, 2, 8), + (40960, 7168, 3): (64, 2, 2, 8), + (40960, 7168, 4): (64, 4, 1, 8), + # TP16. + (3216, 7168, 1): (224, 3, 4, 8), + (3216, 7168, 2): (128, 4, 2, 8), + (3216, 7168, 3): (128, 2, 1, 8), + (3216, 7168, 4): (64, 2, 2, 8), + (3216, 7168, 5): (128, 3, 1, 8), + (768, 7168, 1): (224, 2, 4, 8), + (768, 7168, 2): (224, 2, 2, 8), + (768, 7168, 3): (224, 2, 2, 8), + (768, 7168, 4): (224, 2, 2, 8), + (1152, 1536, 1): (192, 3, 4, 8), + (4224, 7168, 1): (224, 3, 4, 8), + (4224, 7168, 2): (128, 2, 1, 8), + (4224, 7168, 3): (64, 2, 2, 8), + (10240, 7168, 1): (224, 4, 2, 8), + (10240, 7168, 2): (32, 2, 4, 8), + (10240, 7168, 3): (64, 4, 1, 8), + (10240, 7168, 4): (64, 4, 1, 8), +} + +EXPECTED_RESIDUAL_CUTE_CONFIGS = { + (7168, 3584, 1): (64, 4, 2, 8), + (7168, 3584, 2): (64, 7, 2, 8), + (7168, 3584, 3): (64, 2, 1, 8), + (7168, 3584, 4): (64, 2, 1, 8), +} + + +def _config_tuple(config) -> tuple[int, int, int, int]: + return ( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def test_table_is_keyed_by_shape() -> None: + for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items(): + assert (spec.n, spec.k) == (n, k) + + +def test_every_dsv3_routed_shape_is_instantiated() -> None: + """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises. + + The table routes by shape while the kernel is built per shape, so a missing + instantiation only shows up at the token counts that route to dsv3. Checking + it here needs no GPU, which is the point -- a GPU-only check is exactly what + let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). + """ + source = ( + Path(__file__).resolve().parents[2] + / "csrc" + / "libtorch_stable" + / "dsv3_fused_a_gemm.cu" + ).read_text(encoding="utf-8") + # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built + # by default, so they must not count as available. + production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0] + explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split( + "#ifdef VLLM_K3_BENCH_SHAPES" + )[0] + compiled = { + (int(hd_in), int(hd_out)) + for hd_in, hd_out in re.findall( + r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros + ) + } | { + (int(hd_in), int(hd_out)) + for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit) + } + assert compiled, "failed to parse the dispatch list" + + missing = sorted( + (spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + if spec.dsv3_tokens and (spec.k, spec.n) not in compiled + ) + assert not missing, ( + f"routed to dsv3 with no instantiation: {missing}; add " + "DISPATCH_DSV3_SHAPE(K, N) for each" + ) + + +def test_packed_row_major_rejects_single_row_slice() -> None: + packed = torch.empty(1, 128) + sliced = torch.empty(1, 144)[:, :128] + + assert packed.is_contiguous() + assert sliced.is_contiguous() + assert k3_gemm._is_packed_row_major(packed) + assert not k3_gemm._is_packed_row_major(sliced) + + +def test_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.cute_configs + } + assert actual == EXPECTED_CUTE_CONFIGS + + +def test_residual_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.residual_configs + } + assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_sm103_selector_table(key: tuple[int, int]) -> None: + n, k = key + cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key] + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k) + if num_tokens in cute_tokens: + assert backend == "cute" + elif num_tokens in dsv3_tokens: + assert backend == "dsv3_fused_a" + else: + assert backend is None + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None: + n, k = key + assert k3_gemm.select_kimi_k3_backend(0, n, k) is None + assert k3_gemm.select_kimi_k3_backend(17, n, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None + + +def test_unlisted_shape_and_unselected_tokens_fall_back() -> None: + # Shape absent from the table. + assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None + # o_proj (7168,1536) is CuTe M1 only; M2+ falls back. + assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_sm103_residual_selector_table(num_tokens: int) -> None: + backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True) + assert backend == ("cute" if num_tokens <= 4 else None) + + +def test_build_plan_matches_selector() -> None: + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values(): + plan = k3_gemm._build_plan(spec) + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k) + if backend is None: + assert num_tokens not in plan + else: + assert plan[num_tokens][0] == backend + + +def test_installation_is_shape_specific_and_unquantized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeLinear(nn.Module): + def __init__(self, quant_method: object, n: int, k: int) -> None: + super().__init__() + self.quant_method = quant_method + self.weight = torch.empty(n, k) + + class FakeHead(nn.Module): + def __init__(self, n: int, k: int) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedEmbeddingMethod() + self.weight = torch.empty(n, k) + + root = nn.Module() + # dsv3-only shape (no cute warmup contribution). + root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536) + # quantized: must be left untouched. + quantized_method = object() + root.quantized = FakeLinear(quantized_method, 6288, 7168) + # cute shape. + root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168) + # cute + residual shape. + root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584) + # shape absent from the table: must be left untouched. + root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678) + root.lm_head = FakeHead(20480, 7168) + + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True) + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True) + + def request_warmup_configs(dtype, configs, *, has_residual=False): + target = residual_warmup_configs if has_residual else warmup_configs + target.update(configs) + + monkeypatch.setattr( + k3_gemm.shape_dynamic_skinny_gemm, + "request_warmup_configs", + request_warmup_configs, + ) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16) + + assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert root.quantized.quant_method is quantized_method + assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod + assert isinstance( + root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod + ) + # Warmup covers only the installed modules' local (N, K). + assert warmup_configs == { + config + for key in ((6288, 7168), (7168, 3584), (20480, 7168)) + for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs + } + assert residual_warmup_configs == { + config + for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs + } + + +@pytest.mark.parametrize( + "dtype,platform_enabled", + [(torch.float16, True), (torch.bfloat16, False)], +) +def test_installation_requires_bf16_sm103( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + platform_enabled: bool, +) -> None: + class FakeLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedLinearMethod() + self.weight = torch.empty(2304, 1536) + + root = nn.Module() + root.projection = FakeLinear() + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype) + + assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod + + +def _require_sm103_and_dsv3() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + pytest.skip("dsv3_fused_a_gemm was not built") + + +def _require_sm103_and_cute() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not k3_gemm.shape_dynamic_skinny_gemm.is_available(): + pytest.skip("CuTe DSL is not available") + + +@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES) +def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]: + """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch.""" + if not tokens: + return set() + return {min(tokens), max(tokens)} | ({8, 9} & set(tokens)) + + +# Derived from the table rather than hand-listed, so a shape routed to dsv3 +# cannot be added without being exercised here. +DSV3_CASES = sorted( + (num_tokens, spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens) +) + + +@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES) +def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None: + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + assert num_tokens in spec.dsv3_tokens + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +def test_nonpacked_single_token_dsv3_falls_back() -> None: + _require_sm103_and_dsv3() + n, k = 1536, 128 + storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda") + x = storage[:, :k] + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + method = k3_gemm.KimiK3LowLatencyLinearMethod( + k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) + ) + + assert x.is_contiguous() + assert x.stride() == (k + 16, 1) + assert not k3_gemm._runtime_ok(x, weight) # strict guard rejects the view + output = method.apply(SimpleNamespace(weight=weight), x) + + reference = torch.nn.functional.linear(x, weight) + torch.testing.assert_close(output, reference) + + +def test_selected_kernels_cuda_graph_capture() -> None: + _require_sm103_and_cute() + _require_sm103_and_dsv3() + cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)] + dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda") + cute_weight = torch.randn( + cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda" + ) + dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda") + dsv3_weight = torch.randn( + dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda" + ) + k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + graph.replay() + torch.accelerator.synchronize() + + assert cute_output is not None + assert dsv3_output is not None + for output, activation, weight in ( + (cute_output, cute_x, cute_weight), + (dsv3_output, dsv3_x, dsv3_weight), + ): + reference = torch.nn.functional.linear(activation, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", [1, 8, 9, 16]) +def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None: + """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch.""" + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + k3_gemm.try_low_latency_gemm(x, weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.try_low_latency_gemm(x, weight) + graph.replay() + torch.accelerator.synchronize() + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES) +def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42 + num_tokens) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + config = spec.residual_config(num_tokens) + assert config is not None + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None: + _require_sm103_and_cute() + from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + ShapeDynamicSkinnyGemm, + ) + + n, k = 64, 512 + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k) + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) + + +@pytest.mark.parametrize("num_tokens", range(1, 5)) +def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None: + _require_sm103_and_cute() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)] + config = spec.residual_config(num_tokens) + assert config is not None + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda") + k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + graph.replay() + torch.accelerator.synchronize() + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +class _SkinnyGemmSpy: + """Wraps the skinny-GEMM singleton to record whether CuTe was invoked.""" + + def __init__(self, real: Any) -> None: + self._real = real + self.calls: list[int] = [] + + def __call__(self, a, b, config=None, residual=None): + self.calls.append(a.shape[0]) + return self._real(a, b, config, residual) + + def is_available(self) -> bool: + return self._real.is_available() + + +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) +def test_latent_moe_production_layout_residual( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, +) -> None: + """The real Latent-MoE residual is a non-packed slice of a cat buffer. + + The strict packed-row-major guard rejects such a slice at every token count + (a size-1 leading dim reads as contiguous but its stride is not packed), so + the CuTe residual epilogue never fires for this production layout and the + method falls back to addmm. Output is correct regardless of the path. + """ + _require_sm103_and_cute() + latent_dim, shared_dim = 3584, 7168 # routed_expert_up_proj K, N + torch.manual_seed(7 + num_tokens) + buf = torch.randn( + num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda" + ) + latent = buf[:, :latent_dim] # non-contiguous view (row stride = full width) + residual = buf[:, latent_dim:] # non-contiguous view + weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda") + + spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)] + method = k3_gemm.KimiK3LowLatencyLinearMethod( + k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) + ) + spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm) + monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy) + + layer = SimpleNamespace(weight=weight) + output = method.apply_with_residual(layer, latent, residual) + + reference = latent.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 # correct regardless of the path taken + assert not spy.calls, ( + "non-packed buf-slice residual must fall back to addmm at every M" + ) + + +def test_residual_dispatch_falls_back_to_addmm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fallback = torch.randn(2, 3) + residual = torch.randn(2, 3) + x = torch.randn(2, 4) + weight = torch.randn(3, 4) + monkeypatch.setattr(torch, "addmm", lambda *args: fallback) + # CPU tensors fail the runtime check, forcing the addmm fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual) + + assert output is fallback + + +def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None: + fallback = torch.empty(2, 8) + monkeypatch.setattr( + k3_gemm.UnquantizedLinearMethod, + "apply", + lambda *args: fallback, + ) + # 1-D input fails the runtime check, forcing the base-method fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply( + SimpleNamespace(weight=torch.empty(0)), + torch.empty(0), + ) + + assert output is fallback diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py index 25402fe03ea1..eff4508b0d83 100644 --- a/tests/kernels/test_cache_kernels.py +++ b/tests/kernels/test_cache_kernels.py @@ -21,9 +21,9 @@ def test_gather_cache_oob(): seq_starts causes the block_table offset to read out of bounds. """ - batch_size = 1 block_size = 64 - entry_size = 128 + # The kernel only supports the MLA entry sizes. + entry_size = 576 block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda") @@ -34,6 +34,7 @@ def test_gather_cache_oob(): seq_len = 65 cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda") + token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda") # src_cache: [num_blocks, block_size, entry_size] num_blocks = 5 @@ -51,7 +52,8 @@ def test_gather_cache_oob(): dst, block_table, cu_seq_lens, - batch_size, + token_to_seq, + seq_len, "auto", # kv_cache_dtype scale, seq_starts, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index b8f1cb8bfa73..d8878a0aea90 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -13,12 +13,14 @@ """ import math +from types import SimpleNamespace import pytest import torch from vllm import _custom_ops as ops from vllm.models.deepseek_v4.common.ops import ( + compute_global_topk_indices_and_lens, dequantize_and_gather_k_cache, quantize_and_insert_k_cache, ) @@ -27,11 +29,29 @@ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, _launch_two_stage_sparse_attn_compressor, ) +from vllm.models.deepseek_v4.compressor import _get_c128_boundary from vllm.platforms import current_platform from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +def test_compute_global_topk_reuses_output_buffers(): + device = "cuda" + topk_indices = torch.tensor( + [[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device + ) + token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device) + block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device) + is_valid = torch.tensor([True, False], device=device) + args = (topk_indices, token_to_req, block_table, 2, is_valid) + expected = compute_global_topk_indices_and_lens(*args) + outputs = tuple(torch.empty_like(tensor) for tensor in expected) + actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs) + for result, output, reference in zip(actual, outputs, expected): + assert result.data_ptr() == output.data_ptr() + torch.testing.assert_close(result, reference) + + def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale). @@ -58,6 +78,25 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): return x_fp8, scales +@pytest.mark.parametrize( + ("starts", "query_start_loc", "expected"), + [ + ([0], [0, 127], False), + ([0], [0, 128], True), + ([127], [0, 1], True), + ([128], [0, 127], False), + ([1, 255], [0, 1, 2], True), + (None, [0, 1], None), + ], +) +def test_get_c128_boundary(starts, query_start_loc, expected): + metadata = SimpleNamespace( + _num_computed_tokens_cpu=None if starts is None else torch.tensor(starts), + query_start_loc_cpu=torch.tensor(query_start_loc), + ) + assert _get_c128_boundary(metadata) is expected + + # ── Test A: DeepseekV4 Attention path ────────────────────────────────────────────── diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index 86f26cfe8cab..a80467ecc38f 100644 --- a/tests/kernels/test_flex_attention.py +++ b/tests/kernels/test_flex_attention.py @@ -13,6 +13,7 @@ create_standard_kv_cache_spec, create_vllm_config, ) +from vllm.config import AttentionConfig from vllm.model_executor.layers.attention import Attention from vllm.v1.attention.backends.flex_attention import ( BlockSparsityHint, @@ -27,6 +28,42 @@ DIRECT_BUILD_VERSION = version.parse("2.9.dev0") +@pytest.mark.parametrize( + ("supports_small_blocks", "uses_paged_kv", "expected"), + [ + (True, True, (16, 16)), + (True, False, (128, 128)), + (False, True, (128, 128)), + (False, False, (128, 128)), + ], +) +def test_flex_attention_default_block_sizes( + supports_small_blocks: bool, + uses_paged_kv: bool, + expected: tuple[int, int], +): + block_sizes = FlexAttentionMetadataBuilder._get_block_sizes( + AttentionConfig(), + supports_small_blocks=supports_small_blocks, + cache_block_size=16, + uses_paged_kv=uses_paged_kv, + ) + assert block_sizes == expected + + +def test_flex_attention_explicit_block_sizes_override_encoder_defaults(): + block_sizes = FlexAttentionMetadataBuilder._get_block_sizes( + AttentionConfig( + flex_attn_q_block_size=64, + flex_attn_kv_block_size=32, + ), + supports_small_blocks=True, + cache_block_size=16, + uses_paged_kv=False, + ) + assert block_sizes == (64, 32) + + @pytest.mark.skipif( not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, reason="CUDA not available or PyTorch version < 2.7", @@ -213,10 +250,12 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner): the default backend for encoder models. """ model_name = "BAAI/bge-base-en-v1.5" + # Exercise packed sequence boundaries inside 128-token FlexAttention + # blocks, including sequences that span more than one block. prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", + "hello " * 120, + "world " * 130, + "attention " * 254, ] # Run with flex attention @@ -225,7 +264,7 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner): runner="pooling", dtype=torch.bfloat16, tensor_parallel_size=1, - max_model_len=100, + max_model_len=384, enforce_eager=True, attention_config={"backend": "FLEX_ATTENTION"}, ) as llm_flex: @@ -237,7 +276,7 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner): runner="pooling", dtype=torch.bfloat16, tensor_parallel_size=1, - max_model_len=100, + max_model_len=384, enforce_eager=True, ) as llm_default: default_outputs = llm_default.embed(prompts) @@ -264,7 +303,7 @@ def test_block_mask_direct_vs_slow_path(): device = torch.device("cuda") vllm_config = create_vllm_config( - model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024 + model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024 ) kv_cache_spec = create_standard_kv_cache_spec(vllm_config) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index a27e67b62459..48a1f78edbc4 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -27,7 +27,7 @@ import pytest import torch -from vllm.models.deepseek_v32.nvidia import kernels as K +from vllm.models.deepseek_v32.common import kernels as K from vllm.platforms import current_platform FP8 = torch.float8_e4m3fn diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index ed163a0472ad..5572ee89f210 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -257,8 +257,18 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device ).view(num_blocks, -1) slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device) - q_out = _call_fused( - q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs + q_out = torch.empty(num_tokens, padded_heads, HEAD_DIM, dtype=dtype, device=device) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( + q, + kv, + q_out, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + padded_heads, + eps, + bs, ) torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 6114b7efd6e7..f8bf944a98f0 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -150,6 +150,23 @@ def test_fused_indexer_q_rope_quant_matches_unfused( q_quant_ref, weights_ref = _reference( positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4 ) + output_buffers: tuple[torch.Tensor, ...] | None = None + OUTPUT_BUFFER_TEST_NUM_TOKENS = 7 + if num_tokens == OUTPUT_BUFFER_TEST_NUM_TOKENS and cache_dtype == torch.float32: + if use_fp4: + q_ref, q_scale_ref = q_quant_ref + output_buffers = ( + torch.empty_like(q_ref), + torch.empty_like(q_scale_ref) + .view(torch.uint8) + .reshape(num_tokens, N_HEAD, -1), + torch.empty_like(weights_ref), + ) + else: + output_buffers = ( + torch.empty_like(q_quant_ref), + torch.empty_like(weights_ref), + ) # use_cutedsl=False: force the triton path even when cutedsl is installed # by patching the dispatcher's has_cutedsl() binding to return False. cutedsl_patch = ( @@ -169,8 +186,17 @@ def test_fused_indexer_q_rope_quant_matches_unfused( softmax_scale, head_scale, use_fp4, + output_buffers=output_buffers, ) + if output_buffers is not None: + if use_fp4: + assert q_quant_fused[0].data_ptr() == output_buffers[0].data_ptr() + assert q_quant_fused[1].data_ptr() == output_buffers[1].data_ptr() + else: + assert q_quant_fused.data_ptr() == output_buffers[0].data_ptr() + assert weights_fused.data_ptr() == output_buffers[-1].data_ptr() + if use_fp4: q_quant_ref, q_scale_ref = q_quant_ref q_quant_fused, q_scale_fused = q_quant_fused diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index b5b81efa7727..5b71e59baa9e 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups): This catches stride/layout bugs that only manifest when the einsum kernel actually consumes the quantized activations. """ - from deep_gemm.utils.math import ceil_div - from vllm.utils.deep_gemm import ( fp8_einsum, + is_deep_gemm_supported, per_block_cast_to_fp8, transform_sf_into_required_layout, ) + if not is_deep_gemm_supported(): + pytest.skip("DeepGEMM not supported on this platform") + + def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + heads_per_group = num_heads // n_groups d = heads_per_group * HEAD_DIM o_lora_rank = 1024 @@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups): # -- Checks -- # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs # come from fp32 ordering and UE8M0 boundary shifts only. - # Use relative diff (same metric as test_fp8_einsum.py). - from deep_gemm.testing import calc_diff + # Use relative diff (same metric as deep_gemm.testing.calc_diff). + def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return 1 - sim z_diff = calc_diff(z_fused, z_ref) assert z_diff < 0.01, ( diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 5c762b63df43..0bc3c5692ab6 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps): return roped +def assert_fp8_cache_close(kv_cache, expected_kv_cache): + """Compare two e4m3 caches allowing 1 ulp. + + On CUDA the fused kernel quantizes K from its fp32 intermediate, while the + reshape_and_cache_flash reference quantizes the bf16-materialized value, so + rounding-boundary values may differ by one e4m3 code. + """ + byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs() + got = kv_cache.view(torch.float8_e4m3fn).float() + exp = expected_kv_cache.view(torch.float8_e4m3fn).float() + ok = (byte_diff <= 1) | ((got == 0) & (exp == 0)) + assert bool(ok.all()), ( + f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements" + ) + + # ── Test 1: dense mode (norm+rope only, no index, no insert) ───────────────── @@ -131,8 +147,11 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) # V is untouched. torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) @@ -189,6 +208,12 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): # index_q here (de-interleaved from the packed qkv); k/v/index_k stay in # place inside qkv and are scatter-inserted into the caches. q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device) + q_fp8 = torch.empty( + num_tokens, + qsz, + dtype=torch.float8_e4m3fn, + device=device, + ) index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device) ops.fused_minimax_m3_qknorm_rope_kv_insert( @@ -212,6 +237,8 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): q_out, index_q, kv_cache_dtype, + q_fp8_out=q_fp8, + q_fp8_scale=0.5, ) # ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are @@ -239,8 +266,18 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps ).view(num_tokens, HEAD_DIM) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + expected_q_fp8 = torch.empty_like(q_fp8) + ops.scaled_fp8_quant( + q_out, + scale=torch.tensor(0.5, dtype=torch.float32, device=device), + output=expected_q_fp8, + ) + torch.testing.assert_close(q_fp8, expected_q_fp8, rtol=0, atol=0) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2) @@ -265,7 +302,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): scale, scale, ) - torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0) + assert_fp8_cache_close(kv_cache, expected_kv_cache) else: for t in range(num_tokens): s = slot_mapping[t].item() @@ -360,8 +397,11 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) torch.testing.assert_close(index_q_out, index_q_in, rtol=0, atol=0) torch.testing.assert_close(index_k_out, index_k_in, rtol=0, atol=0) @@ -383,7 +423,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype): scale, scale, ) - torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0) + assert_fp8_cache_close(kv_cache, expected_kv_cache) else: k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM) v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index 128a02060043..928c3e5bd748 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference( A_log = torch.randn((HV,), device=device, dtype=dtype) dt_bias = torch.randn((HV,), device=device, dtype=dtype) - # Continuous batching indices (include PAD_SLOT_ID=-1 cases). - ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32) + # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is + # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1. + ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32) ssm_state_indices[-3:] = -1 - state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype) + state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype) state_ref = state0.clone() state_packed = state0.clone() @@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference( atol = 2e-2 if dtype != torch.float32 else 1e-4 rtol = 1e-2 if dtype != torch.float32 else 1e-4 - torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol) + # Output rows for PAD_SLOT_ID entries are never written (uninitialized in + # both paths), so compare only the valid rows. + valid = ssm_state_indices > 0 + torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol) torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol) diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 82a5a6f4ca91..fc923cb945f7 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py @@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec( dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype) a = torch.rand(num_tokens, num_v_heads, dtype=dtype) b = torch.rand(num_tokens, num_v_heads, dtype=dtype) + # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid + # state indices start at 1. ssm_state = torch.rand( - total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype + total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype ) - state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens] + state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens] cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32) beta = b.sigmoid() @@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec( dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype) a = torch.rand(num_tokens, num_v_heads, dtype=dtype) b = torch.rand(num_tokens, num_v_heads, dtype=dtype) + # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid + # state indices start at 1. ssm_state = torch.rand( - total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype + total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype ) - state_indices = torch.randperm( - total_entries, - dtype=torch.int32, - )[:num_tokens].view(num_reqs, num_speculative_tokens + 1) + state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[ + :num_tokens + ].view(num_reqs, num_speculative_tokens + 1) num_accepted_tokens = torch.randint( 1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32 ) diff --git a/tests/kernels/test_kda.py b/tests/kernels/test_kda.py deleted file mode 100644 index 75c553fb864b..000000000000 --- a/tests/kernels/test_kda.py +++ /dev/null @@ -1,226 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Precision tests for vllm's chunk_kda Triton operator. - -Compares chunk_kda against a naive recurrent reference (float32). -Uses torch.rand for q/k/v to match FLA's test pattern. -""" - -import pytest -import torch -import torch.nn.functional as F - -from vllm.third_party.flash_linear_attention.ops.kda import ( - chunk_kda, - chunk_kda_with_fused_gate, - fused_kda_gate, -) -from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd - -DEVICE = "cuda" - - -def naive_recurrent_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Naive recurrent KDA reference, ported from FLA's naive.py.""" - dtype = v.dtype - B, T, H, K = q.shape - V = v.shape[-1] - if scale is None: - scale = K**-0.5 - - q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) - q = q * scale - - S = k.new_zeros(B, H, K, V).to(q) - if initial_state is not None: - S += initial_state - o = torch.zeros_like(v) - for i in range(T): - q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] - S = S * g_i[..., None].exp() - S = S + torch.einsum( - "bhk,bhv->bhkv", - b_i[..., None] * k_i, - v_i - (k_i[..., None] * S).sum(-2), - ) - o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) - if not output_final_state: - S = None - return o.to(dtype), S - - -def assert_close( - name: str, - ref: torch.Tensor, - tri: torch.Tensor, - ratio: float, - err_atol: float = 1e-6, -): - """RMSE-based relative error comparison.""" - abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() - rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() - rmse_base = ref.detach().flatten().square().mean().sqrt().item() - rel_err = rmse_diff / (rmse_base + 1e-8) - print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") - if abs_err <= err_atol: - return - assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" - assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" - assert rel_err < ratio, ( - f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" - ) - - -@pytest.mark.parametrize( - ("H", "D", "cu_seqlens", "dtype"), - [ - pytest.param( - *test, - id="H{}-D{}-cu{}-{}".format(*test), - ) - for test in [ - (32, 128, [0, 64], torch.float16), - (32, 128, [0, 1024], torch.float16), - (32, 128, [0, 15], torch.float16), - (32, 128, [0, 256, 512, 768, 1024], torch.float16), - (32, 128, [0, 15, 100, 300, 1200], torch.float16), - (64, 128, [0, 256, 500, 1000], torch.float16), - (32, 128, [0, 8192], torch.float16), - (32, 128, [0, 256, 500, 1000], torch.bfloat16), - ] - ], -) -@torch.inference_mode() -def test_chunk_kda( - H: int, - D: int, - cu_seqlens: list[int], - dtype: torch.dtype, -): - T = cu_seqlens[-1] - torch.manual_seed(42) - B = 1 - cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) - N = len(cu_seqlens) - 1 - - q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( - dtype - ) - beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() - h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) - - # Naive reference with l2norm_fwd (same kernel as chunk_kda) - ref_outputs = [] - ref_states = [] - for i in range(N): - s, e = cu_seqlens[i], cu_seqlens[i + 1] - q_i = l2norm_fwd(q[:, s:e].contiguous()) - k_i = l2norm_fwd(k[:, s:e].contiguous()) - o_i, ht_i = naive_recurrent_kda( - q_i, - k_i, - v[:, s:e], - g[:, s:e], - beta[:, s:e], - initial_state=h0[i], - output_final_state=True, - ) - ref_outputs.append(o_i) - ref_states.append(ht_i) - ref_o = torch.cat(ref_outputs, dim=1) - ref_ht = torch.cat(ref_states, dim=0) - - # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) - tri_o, tri_ht = chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=g.clone(), - beta=beta.clone(), - initial_state=h0.transpose(-1, -2).contiguous().clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - - assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" - assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" - assert_close("o", ref_o, tri_o, 0.005) - assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) - - -@pytest.mark.parametrize( - ("cu_seqlens", "dtype"), - [ - ([0, 64], torch.float16), - ([0, 15, 100, 300], torch.bfloat16), - ], -) -@torch.inference_mode() -def test_chunk_kda_fused_gate_cumsum_matches_unfused( - cu_seqlens: list[int], - dtype: torch.dtype, -): - H, D = 8, 64 - T = cu_seqlens[-1] - N = len(cu_seqlens) - 1 - torch.manual_seed(123) - - cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) - q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid() - A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() - dt_bias = ( - torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 - ).contiguous() - h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) - initial_state = h0.transpose(-1, -2).contiguous() - - gate = fused_kda_gate( - raw_g.reshape(T, H * D), - A_log, - D, - g_bias=dt_bias, - ).unsqueeze(0) - old_o, old_ht = chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=gate, - beta=beta.clone(), - initial_state=initial_state.clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - new_o, new_ht = chunk_kda_with_fused_gate( - q=q.clone(), - k=k.clone(), - v=v.clone(), - raw_g=raw_g, - beta=beta.clone(), - A_log=A_log, - g_bias=dt_bias, - initial_state=initial_state.clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - - assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) - assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index dea54ed21aea..899b3129d62f 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -9,7 +9,6 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from vllm.distributed import ( cleanup_dist_env_and_memory, @@ -25,6 +24,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.interfaces import SupportsLoRA from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api @pytest.fixture() @@ -170,43 +170,43 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module: def mixtral_lora_files(): # Note: this module has incorrect adapter_config.json to test # https://github.com/vllm-project/vllm/pull/5909/files. - return snapshot_download(repo_id="SangBinCho/mixtral-lora") + return hf_api().snapshot_download(repo_id="SangBinCho/mixtral-lora") @pytest.fixture(scope="session") def chatglm3_lora_files(): - return snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider") @pytest.fixture(scope="session") def baichuan_lora_files(): - return snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider") @pytest.fixture(scope="session") def baichuan_zero_lora_files(): # all the lora_B weights are initialized to zero. - return snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init") @pytest.fixture(scope="session") def baichuan_regex_lora_files(): - return snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex") @pytest.fixture(scope="session") def ilama_lora_files(): - return snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider") @pytest.fixture(scope="session") def minicpmv_lora_files(): - return snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon") @pytest.fixture(scope="session") def qwen2vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon") @pytest.fixture(scope="session") @@ -217,74 +217,84 @@ def qwen25vl_base_huggingface_id(): @pytest.fixture(scope="session") def qwen25vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon") @pytest.fixture(scope="session") def qwen2vl_language_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-language") + return hf_api().snapshot_download( + repo_id="prashanth058/qwen2vl-flickr-lora-language" + ) @pytest.fixture(scope="session") def qwen2vl_vision_tower_connector_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector") + return hf_api().snapshot_download( + repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector" + ) @pytest.fixture(scope="session") def qwen2vl_vision_tower_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower") + return hf_api().snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower") @pytest.fixture(scope="session") def qwen25vl_vision_lora_files(): - return snapshot_download(repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector") + return hf_api().snapshot_download( + repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector" + ) @pytest.fixture(scope="session") def qwen3vl_vision_lora_files(): - return snapshot_download(repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector") + return hf_api().snapshot_download( + repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector" + ) @pytest.fixture(scope="session") def qwen3_meowing_lora_files(): """Download Qwen3 Meow LoRA files once per test session.""" - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") @pytest.fixture(scope="session") def qwen3_woofing_lora_files(): """Download Qwen3 Woof LoRA files once per test session.""" - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") @pytest.fixture(scope="session") def tinyllama_lora_files(): - return snapshot_download(repo_id="jashing/tinyllama-colorist-lora") + return hf_api().snapshot_download(repo_id="jashing/tinyllama-colorist-lora") @pytest.fixture(scope="session") def deepseekv2_lora_files(): - return snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") + return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") @pytest.fixture(scope="session") def gptoss20b_lora_files(): - return snapshot_download(repo_id="jeeejeee/gpt-oss-20b-lora-adapter-text2sql") + return hf_api().snapshot_download( + repo_id="jeeejeee/gpt-oss-20b-lora-adapter-text2sql" + ) @pytest.fixture(scope="session") def qwen3moe_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider") @pytest.fixture(scope="session") def olmoe_lora_files(): - return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider") @pytest.fixture(scope="session") def qwen3_lora_files(): - return snapshot_download(repo_id="charent/self_cognition_Alice") + return hf_api().snapshot_download(repo_id="charent/self_cognition_Alice") @pytest.fixture(scope="session") @@ -295,32 +305,40 @@ def llama32_lora_huggingface_id(): @pytest.fixture(scope="session") def llama32_lora_files(llama32_lora_huggingface_id): - return snapshot_download(repo_id=llama32_lora_huggingface_id) + return hf_api().snapshot_download(repo_id=llama32_lora_huggingface_id) @pytest.fixture(scope="session") def whisper_lora_files(): - return snapshot_download(repo_id="chengyili2005/whisper-small-mandarin-lora") + return hf_api().snapshot_download( + repo_id="chengyili2005/whisper-small-mandarin-lora" + ) @pytest.fixture(scope="session") def qwen35_text_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora") @pytest.fixture(scope="session") def qwen35_vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora" + ) @pytest.fixture(scope="session") def qwen36_moe_2d_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora" + ) @pytest.fixture(scope="session") def qwen36_moe_3d_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora" + ) @pytest.fixture diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py index 19c910e2453d..5e3c762f79b9 100644 --- a/tests/lora/test_default_mm_loras.py +++ b/tests/lora/test_default_mm_loras.py @@ -8,15 +8,15 @@ import unittest.mock as mock import pytest -from huggingface_hub import snapshot_download from vllm.lora.request import LoRARequest from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from ..conftest import AudioTestAssets, VllmRunner from ..utils import create_new_process_for_each_test -MODEL_PATH = snapshot_download("microsoft/Phi-4-multimodal-instruct") +MODEL_PATH = hf_api().snapshot_download("microsoft/Phi-4-multimodal-instruct") AUDIO_LORA_PATH = os.path.join(MODEL_PATH, "speech-lora") IMAGE_LORA_PATH = os.path.join(MODEL_PATH, "vision-lora") diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index 0f0d807441e9..0e7fe60b3642 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -95,7 +95,8 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: pytest.param( True, marks=pytest.mark.skipif( - current_platform.is_rocm(), reason="marlin not supported" + current_platform.is_rocm() or current_platform.is_xpu(), + reason="marlin not supported", ), ), ], @@ -135,7 +136,8 @@ def test_gpt_oss_lora( pytest.param( True, marks=pytest.mark.skipif( - current_platform.is_rocm(), reason="marlin not supported" + current_platform.is_rocm() or current_platform.is_xpu(), + reason="marlin not supported", ), ), ], diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 1db05b4b09d5..834273b8c5d6 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -198,7 +198,7 @@ def __init__(self, gate: nn.Module): self.gate = gate # canonical path: "moe.gate" # Inner submodule holding the SAME gate instance under another - # path. This mirrors how FusedMoE.runner.gate references the + # path. This mirrors how MoERunner.runner.gate references the # block's gate in qwen3_moe. class _Runner(nn.Module): def __init__(self, g): diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index a4a32278db01..3362aa46ed26 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -7,6 +7,7 @@ from transformers import __version__ as TRANSFORMERS_VERSION import vllm +from tests.conftest import VllmRunner from vllm.assets.image import ImageAsset from vllm.lora.request import LoRARequest from vllm.platforms import current_platform @@ -56,24 +57,29 @@ class Qwen2VLTester: def __init__(self, config: TestConfig): self.config = config - self.llm = self._initialize_llm() - - def _initialize_llm(self) -> vllm.LLM: - """Initialize the LLM with given configuration""" - return vllm.LLM( - model=self.config.model_path, - max_num_seqs=self.config.max_num_seqs, + self._runner = VllmRunner( + model_name=config.model_path, + max_num_seqs=config.max_num_seqs, enable_lora=True, - max_loras=self.config.max_loras, - max_lora_rank=self.config.max_lora_rank, - enable_tower_connector_lora=self.config.enable_tower_connector_lora, - trust_remote_code=True, - gpu_memory_utilization=self.config.gpu_memory_utilization, - mm_processor_kwargs=self.config.mm_processor_kwargs, - mm_processor_cache_gb=self.config.mm_processor_cache_gb, - max_model_len=self.config.max_model_len, + max_loras=config.max_loras, + max_lora_rank=config.max_lora_rank, + enable_tower_connector_lora=config.enable_tower_connector_lora, + gpu_memory_utilization=config.gpu_memory_utilization, + mm_processor_kwargs=config.mm_processor_kwargs, + mm_processor_cache_gb=config.mm_processor_cache_gb, + max_model_len=config.max_model_len, ) + @property + def llm(self) -> vllm.LLM: + return self._runner.get_llm() + + def __enter__(self) -> "Qwen2VLTester": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self._runner.__exit__(exc_type, exc_value, traceback) + def run_test( self, images: list[ImageAsset], @@ -180,32 +186,44 @@ def run_beam_search_test( QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct" +def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: + # These tests assert exact greedy outputs. Force the Triton LoRA shrink + # kernel to use SPLIT_K=1 so it stores the complete reduction directly + # instead of accumulating split-K partial results with atomic_add. This + # targets reduction determinism, not full batch invariance. + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + # The kernel configuration reads VLLM_BATCH_INVARIANT at import time. + # Spawn the engine process so it observes this setting even if the LoRA + # Triton utilities were already imported during test collection. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + def test_qwen2vl_lora(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id + ) def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA through beam search.""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - # NOTE currently, we only test cherry blossom since stop sign - # output is slightly different for v1; - the root cause is likely - # independent of the intent of this test, which is to ensure beam - # search passes through lora through correctly. - tester.run_beam_search_test( - [ImageAsset("cherry_blossom")], - expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + # NOTE currently, we only test cherry blossom since stop sign + # output is slightly different for v1; - the root cause is likely + # independent of the intent of this test, which is to ensure beam + # search passes through lora through correctly. + tester.run_beam_search_test( + [ImageAsset("cherry_blossom")], + expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS, + lora_id=lora_id, + ) @pytest.mark.skipif( @@ -214,11 +232,12 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): def test_qwen25vl_lora(qwen25vl_lora_files): """Test Qwen 2.5 VL model with LoRA""" config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id + ) @pytest.mark.skipif( @@ -234,16 +253,21 @@ def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files): mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS, + lora_id=lora_id, + ) + +def test_qwen3vl_vision_lora( + qwen3vl_vision_lora_files, + monkeypatch: pytest.MonkeyPatch, +): + _enable_deterministic_lora_shrink(monkeypatch) -def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): config = TestConfig( model_path=QWEN3VL_MODEL_PATH, lora_path=qwen3vl_vision_lora_files, @@ -253,19 +277,20 @@ def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS, + lora_id=lora_id, + ) def test_qwen2vl_multiple_lora_types( qwen2vl_language_lora_files, qwen2vl_vision_tower_connector_lora_files, qwen2vl_vision_tower_lora_files, + monkeypatch: pytest.MonkeyPatch, ): """ Test multiple LoRA adapter types (language, vision tower + connector, @@ -276,6 +301,8 @@ def test_qwen2vl_multiple_lora_types( the multimodal encoder cache correctly manages state transitions between language-only and vision-enabled LoRA adapters. """ + _enable_deterministic_lora_shrink(monkeypatch) + config = TestConfig( model_path=QWEN2VL_MODEL_PATH, # We'll override the lora_path for each specific test, but need to provide @@ -287,34 +314,33 @@ def test_qwen2vl_multiple_lora_types( mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) - - # Test 1: Language-only LoRA adapter - tester.config.lora_path = qwen2vl_language_lora_files - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_LANGUAGE, - lora_id=lora_id, - lora_name="language_only", - ) + with Qwen2VLTester(config) as tester: + # Test 1: Language-only LoRA adapter + tester.config.lora_path = qwen2vl_language_lora_files + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_LANGUAGE, + lora_id=lora_id, + lora_name="language_only", + ) - # Test 2: Vision tower + connector LoRA adapter - tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files - for lora_id in [3, 4]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_VISION, - lora_id=lora_id, - lora_name="vision_tower_connector", - ) + # Test 2: Vision tower + connector LoRA adapter + tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files + for lora_id in [3, 4]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_VISION, + lora_id=lora_id, + lora_name="vision_tower_connector", + ) - # Test 3: Vision tower only LoRA adapter (no connector) - tester.config.lora_path = qwen2vl_vision_tower_lora_files - for lora_id in [5, 6]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR, - lora_id=lora_id, - lora_name="vision_tower", - ) + # Test 3: Vision tower only LoRA adapter (no connector) + tester.config.lora_path = qwen2vl_vision_tower_lora_files + for lora_id in [5, 6]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR, + lora_id=lora_id, + lora_name="vision_tower", + ) diff --git a/tests/lora/test_worker.py b/tests/lora/test_worker.py index e929fcad2896..f80f6cec412c 100644 --- a/tests/lora/test_worker.py +++ b/tests/lora/test_worker.py @@ -66,7 +66,6 @@ def set_active_loras(worker: Worker, lora_requests: list[LoRARequest]): runner_type="generate", max_num_batched_tokens=32, max_num_seqs=32, - max_num_partial_prefills=32, ), device_config=DeviceConfig(DEVICE_TYPE), cache_config=CacheConfig( diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py new file mode 100644 index 000000000000..6e1e10e8b45c --- /dev/null +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +from vllm.config import CUDAGraphMode +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata + +INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" +MLA_LAYER = "model.layers.0.self_attn.attn" + + +def make_indexer_metadata( + *, + num_decodes: int = 0, + num_decode_tokens: int = 0, + num_prefills: int = 1, + num_prefill_tokens: int = 1, + slot_mapping: torch.Tensor | None = None, +) -> DeepseekV32IndexerMetadata: + if slot_mapping is None: + slot_mapping = torch.zeros(num_prefill_tokens, dtype=torch.long) + return DeepseekV32IndexerMetadata( + seq_lens=torch.empty(0, dtype=torch.int32), + max_seq_len=2048, + slot_mapping=slot_mapping, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + prefill=SimpleNamespace(chunks=[]) if num_prefills else None, + ) + + +def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0): + return SimpleNamespace( + num_decode_tokens=num_decode_tokens, + prefill=SimpleNamespace(use_dense_mha=use_dense_mha), + ) + + +@pytest.mark.parametrize( + "batch_kind", + ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"], +) +def test_short_prefill_updates_k_cache_before_scoring_decision( + monkeypatch: pytest.MonkeyPatch, + batch_kind: str, +): + slot_mapping = torch.tensor([63, 64, 127, 128, -1]) + mla_num_decode_tokens = 1 if batch_kind == "mla_decode" else 0 + runtime_mode = ( + CUDAGraphMode.FULL if batch_kind == "full" else CUDAGraphMode.PIECEWISE + ) + should_skip = batch_kind in ("short", "threshold_mismatch") + num_decodes = int(batch_kind == "threshold_mismatch") + num_decode_tokens = 3 if batch_kind == "threshold_mismatch" else 0 + num_prefills = 0 if batch_kind == "threshold_mismatch" else 2 + num_prefill_tokens = 0 if batch_kind == "threshold_mismatch" else 5 + if batch_kind == "threshold_mismatch": + # With MTP=3 the indexer threshold is four. A main MLA backend whose + # threshold is one (for example FlashMLA under DCP) still routes this + # three-token extend through dense prefill attention. + slot_mapping = slot_mapping[:3] + indexer_metadata = make_indexer_metadata( + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + slot_mapping=slot_mapping, + ) + if indexer_metadata.num_decodes: + indexer_metadata.decode = object() + mla_metadata = make_mla_metadata( + use_dense_mha=batch_kind != "force_mqa", + num_decode_tokens=mla_num_decode_tokens, + ) + + observed: dict[str, object] = {} + + monkeypatch.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={ + INDEXER_LAYER: indexer_metadata, + MLA_LAYER: mla_metadata, + }, + cudagraph_runtime_mode=runtime_mode, + ), + ) + monkeypatch.setattr( + sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16 + ) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: batch_kind == "capture", + ) + + def record_cache_update(k, kv_cache, slots, block_size, scale_fmt): + observed.update(k=k.clone(), slots=slots) + + monkeypatch.setattr( + sparse_indexer.ops, "indexer_k_quant_and_cache", record_cache_update + ) + + class ScoringReached(Exception): + pass + + def scoring_trigger(): + if should_skip: + pytest.fail("short dense-MHA prefill must not enter indexer scoring") + raise ScoringReached + + def scoring_decode(*args): + raise ScoringReached + + monkeypatch.setattr(sparse_indexer, "current_workspace_manager", scoring_trigger) + monkeypatch.setattr( + sparse_indexer, + "kv_cache_as_quant_view", + scoring_decode, + ) + + hidden_states = torch.full((7, 1), float("inf")) + k = torch.arange(28, dtype=torch.float32).reshape(7, 4) + topk_indices = torch.full((7, 2048), 17, dtype=torch.int32) + + def run_indexer(): + return sparse_indexer.sparse_attn_indexer( + hidden_states, + INDEXER_LAYER, + torch.empty(1), + torch.full((7, 1), float("inf")), + None, + k, + torch.full((7, 1), float("inf")), + 128, + "ue8m0", + 2048, + 4, + 4096, + 4096, + topk_indices, + False, + False, + MLA_LAYER, + ) + + if should_skip: + assert run_indexer() is topk_indices + assert torch.all(topk_indices == 17) + else: + with pytest.raises(ScoringReached): + run_indexer() + assert torch.all(topk_indices == -1) + + # K cache is always updated before the scoring decision. + torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()]) + assert observed["slots"] is slot_mapping diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index f4de9bc9038f..f20be7cfb618 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -52,6 +53,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -76,6 +78,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t()) diff --git a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py index 09e0a19d11b9..3d15f6f1c551 100644 --- a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob -import tempfile -import huggingface_hub.constants import pytest import torch @@ -21,29 +19,27 @@ reason="InstantTensor requires NVIDIA GPUs", ) def test_instanttensor_model_loader(): - with tempfile.TemporaryDirectory() as tmpdir: - huggingface_hub.constants.HF_HUB_OFFLINE = False - download_weights_from_hf( - "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir - ) - safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True) - assert len(safetensors) > 0 + model_dir = download_weights_from_hf( + "openai-community/gpt2", cache_dir=None, allow_patterns=["*.safetensors"] + ) + safetensors = glob.glob(f"{model_dir}/*.safetensors") + assert len(safetensors) > 0 - instanttensor_tensors = {} - hf_safetensors_tensors = {} + instanttensor_tensors = {} + hf_safetensors_tensors = {} - for name, tensor in instanttensor_weights_iterator(safetensors, True): - instanttensor_tensors[name] = tensor.to("cpu") + for name, tensor in instanttensor_weights_iterator(safetensors, True): + instanttensor_tensors[name] = tensor.to("cpu") - for name, tensor in safetensors_weights_iterator(safetensors, True): - hf_safetensors_tensors[name] = tensor + for name, tensor in safetensors_weights_iterator(safetensors, True): + hf_safetensors_tensors[name] = tensor - assert len(instanttensor_tensors) == len(hf_safetensors_tensors) + assert len(instanttensor_tensors) == len(hf_safetensors_tensors) - for name, instanttensor_tensor in instanttensor_tensors.items(): - assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype - assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape - assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name])) + for name, instanttensor_tensor in instanttensor_tensors.items(): + assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype + assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape + assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name])) if __name__ == "__main__": diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py index d60c9ba64cbd..d7abc80edbfa 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py @@ -3,10 +3,10 @@ from pathlib import Path -from huggingface_hub import snapshot_download from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher from vllm.engine.arg_utils import EngineArgs +from vllm.transformers_utils.repo_utils import hf_api from .conftest import RunaiDummyExecutor @@ -25,7 +25,7 @@ def test_runai_model_loader_download_files_s3_mocked_with_patch( # Download model from HF mock_model_dir = f"{tmp_path}/gpt2" - snapshot_download(repo_id=test_model, local_dir=mock_model_dir) + hf_api().snapshot_download(repo_id=test_model, local_dir=mock_model_dir) monkeypatch.setattr( "vllm.transformers_utils.runai_utils.runai_list_safetensors", diff --git a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py index a15a624c905d..7047ba1c17c2 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py +++ b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py @@ -9,13 +9,14 @@ import subprocess import sys from typing import Any +from unittest.mock import Mock import pytest import torch import vllm.model_executor.model_loader.tensorizer from tests.utils import VLLM_PATH, RemoteOpenAIServer -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.model_loader.tensorizer import ( TensorizerConfig, @@ -325,7 +326,7 @@ def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref): pass -def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path): +def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path, vllm_runner): serialization_params = { "limit_cpu_concurrency": 2, } @@ -334,9 +335,6 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path): config = TensorizerConfig( tensorizer_uri=str(model_path), serialization_kwargs=serialization_params ) - llm = LLM( - model=model_ref, - ) def serialization_test(self, *args, **kwargs): # This is performed in the ephemeral worker process, so monkey-patching @@ -365,7 +363,66 @@ def tensorizer_serializer_wrapper(self, *args, **kwargs): kwargs = {"tensorizer_config": config.to_serializable()} - assert assert_from_collective_rpc(llm, serialization_test, kwargs) + with vllm_runner(model_ref) as runner: + assert assert_from_collective_rpc(runner.get_llm(), serialization_test, kwargs) + + +@pytest.mark.parametrize( + ( + "serialization_error", + "renderer_shutdown_error", + "engine_shutdown_error", + "expected_error", + ), + [ + (None, None, None, None), + (RuntimeError("serialization failed"), None, None, "serialization failed"), + ( + RuntimeError("serialization failed"), + ValueError("renderer shutdown failed"), + ValueError("engine shutdown failed"), + "serialization failed", + ), + ( + None, + ValueError("renderer shutdown failed"), + None, + "renderer shutdown failed", + ), + (None, None, ValueError("engine shutdown failed"), "engine shutdown failed"), + ], +) +def test_tensorize_vllm_model_shuts_down_engine( + monkeypatch, + serialization_error, + renderer_shutdown_error, + engine_shutdown_error, + expected_error, +): + from vllm.v1.engine.llm_engine import LLMEngine + + engine_args = Mock() + config = Mock(encryption_keyfile=None) + engine = Mock() + engine.collective_rpc.side_effect = serialization_error + engine.renderer.shutdown.side_effect = renderer_shutdown_error + engine.engine_core.shutdown.side_effect = engine_shutdown_error + monkeypatch.setattr(LLMEngine, "from_vllm_config", Mock(return_value=engine)) + monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", "7") + + if expected_error is None: + tensorize_vllm_model(engine_args, config) + else: + with pytest.raises(Exception, match=expected_error) as exc_info: + tensorize_vllm_model(engine_args, config) + notes = getattr(exc_info.value, "__notes__", ()) + if serialization_error is not None and renderer_shutdown_error is not None: + assert any("renderer shutdown failed" in note for note in notes) + if serialization_error is not None and engine_shutdown_error is not None: + assert any("engine shutdown failed" in note for note in notes) + + engine.engine_core.shutdown.assert_called_once_with(timeout=17.0) + engine.renderer.shutdown.assert_called_once_with() def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd): diff --git a/tests/model_executor/model_loader/test_mtp_validation.py b/tests/model_executor/model_loader/test_mtp_validation.py new file mode 100644 index 000000000000..ecccaa0cd74b --- /dev/null +++ b/tests/model_executor/model_loader/test_mtp_validation.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.model_executor.model_loader.mtp_validation import ( + disable_mtp_completeness_check, + is_mtp_completeness_check_enabled, +) + + +def test_disable_mtp_completeness_check_is_scoped(): + assert is_mtp_completeness_check_enabled() + + with pytest.raises(RuntimeError), disable_mtp_completeness_check(): + assert not is_mtp_completeness_check_enabled() + raise RuntimeError + + assert is_mtp_completeness_check_enabled() diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 1480191083e3..c21d9b1e02f8 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -10,9 +10,11 @@ import vllm.model_executor.model_loader.reload.meta as reload_meta from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_reload, initialize_layerwise_reload, + initialize_online_processing, record_metadata_for_reloading, ) from vllm.model_executor.model_loader.reload.meta import ( @@ -165,6 +167,283 @@ def test_materialize_layer_preserves_non_meta_tensors(): assert torch.equal(layer.bias.data, bias_values) +_MARLIN_SIZE_K, _MARLIN_SIZE_N, _MARLIN_GROUP_SIZE = 128, 64, 64 + + +def _stub_marlin_ops(monkeypatch): + from vllm import _custom_ops as ops + from vllm.model_executor.layers.quantization.utils import marlin_utils + + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + monkeypatch.setattr( + ops, + "gptq_marlin_repack", + lambda w, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( + size_k // 16, size_n * 2, dtype=torch.int32 + ), + ) + + +def _make_act_order_marlin_kernel(): + from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( + MarlinLinearKernel, + ) + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, + ) + from vllm.scalar_type import scalar_types + + kernel = object.__new__(MarlinLinearKernel) + kernel.config = MPLinearLayerConfig( + full_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N), + partition_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=_MARLIN_GROUP_SIZE, + zero_points=False, + has_g_idx=True, + ) + kernel.w_q_name = "qweight" + kernel.w_s_name = "scales" + kernel.w_zp_name = None + kernel.w_gidx_name = "g_idx" + return kernel + + +def _load_marlin_checkpoint_format_weights(layer, g_idx): + from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedvLLMParameter, + RowvLLMParameter, + ) + + layer.qweight = PackedvLLMParameter( + data=torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32), + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=8, + weight_loader=default_weight_loader, + ) + layer.scales = GroupQuantScaleParameter( + data=torch.ones( + _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16 + ), + input_dim=0, + output_dim=1, + weight_loader=default_weight_loader, + ) + layer.g_idx = RowvLLMParameter( + data=g_idx.clone(), + input_dim=0, + weight_loader=default_weight_loader, + ) + + +def _random_g_idx(generator): + return torch.randint( + 0, + _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, + (_MARLIN_SIZE_K,), + dtype=torch.int32, + generator=generator, + ) + + +def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init): + """Marlin workspace and act-order sort indices must be recomputed into + the same storage when weights are reloaded (RL weight sync), so device + addresses captured by CUDA graphs remain valid.""" + from vllm.model_executor.layers.quantization.utils import marlin_utils + + _stub_marlin_ops(monkeypatch) + kernel = _make_act_order_marlin_kernel() + + generator = torch.Generator().manual_seed(0) + first_g_idx = _random_g_idx(generator) + second_g_idx = _random_g_idx(generator) + + layer = torch.nn.Module() + _load_marlin_checkpoint_format_weights(layer, first_g_idx) + kernel.process_weights_after_loading(layer) + + workspace_ptr = kernel.workspace.data_ptr() + sort_indices_ptr = layer.g_idx_sort_indices.data_ptr() + + # Reload: fresh checkpoint-format tensors with a different act-order + _load_marlin_checkpoint_format_weights(layer, second_g_idx) + kernel.process_weights_after_loading(layer) + + assert kernel.workspace.data_ptr() == workspace_ptr + assert torch.all(kernel.workspace == 0) + assert layer.g_idx_sort_indices.data_ptr() == sort_indices_ptr + expected_sort_indices = marlin_utils.marlin_sort_g_idx(second_g_idx)[1] + assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices) + # registered as a Parameter so layerwise reload copy-back preserves it + assert isinstance(layer.g_idx_sort_indices, torch.nn.Parameter) + + +@pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) +def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): + """The Marlin fallback prepare_* functions rerun on weight reload and must + reuse the workspace storage whose address captured CUDA graphs hold.""" + from vllm import _custom_ops as ops + from vllm.model_executor.layers.quantization.utils import ( + marlin_utils, + marlin_utils_fp4, + marlin_utils_fp8, + ) + + size_k, size_n = 128, 64 + + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + monkeypatch.setattr( + ops, + "gptq_marlin_repack", + lambda b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( + size_k // 16, size_n * 2, dtype=torch.int32 + ), + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = size_n + layer.input_size_per_partition = size_k + layer.orig_dtype = torch.float16 + layer.params_dtype = torch.float16 + + if variant == "fp8": + prepare = marlin_utils_fp8.prepare_fp8_layer_for_marlin + + def load_checkpoint_format_weights(): + layer.weight = torch.nn.Parameter( + torch.zeros(size_k, size_n, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.ones(1, dtype=torch.float32), requires_grad=False + ) + elif variant == "mxfp8": + prepare = marlin_utils_fp8.prepare_mxfp8_layer_for_marlin + + def load_checkpoint_format_weights(): + layer.weight = torch.nn.Parameter( + torch.zeros(size_n, size_k, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.full((size_n, size_k // 32), 127, dtype=torch.uint8), + requires_grad=False, + ) + else: + prepare = marlin_utils_fp4.prepare_fp4_layer_for_marlin + + def load_checkpoint_format_weights(): + layer.weight = torch.nn.Parameter( + torch.zeros(size_n, size_k // 2, dtype=torch.uint8), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.ones(size_n, size_k // 16, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_global_scale = torch.nn.Parameter( + torch.ones(1, dtype=torch.float32), requires_grad=False + ) + + load_checkpoint_format_weights() + prepare(layer) + workspace_ptr = layer.workspace.data_ptr() + + # Reload: fresh checkpoint-format tensors, prepare runs again + load_checkpoint_format_weights() + prepare(layer) + + assert layer.workspace.data_ptr() == workspace_ptr + assert torch.all(layer.workspace == 0) + + +def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): + """An incompatible existing workspace means the address captured by CUDA + graphs is already unusable; allocating a replacement would hide that.""" + from vllm.model_executor.layers.quantization.utils import marlin_utils + + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + device = torch.device("cpu") + + workspace = marlin_utils.marlin_make_workspace_new(device) + reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace) + assert reused is workspace + + with pytest.raises(ValueError, match="incompatible"): + marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace) + with pytest.raises(ValueError, match="incompatible"): + marlin_utils.marlin_make_workspace_new( + device, existing=workspace.to(torch.int64) + ) + + +def test_marlin_act_order_layerwise_reload_accounting(monkeypatch, dist_init): + """`g_idx_sort_indices` is generated during weight processing and never + loaded from checkpoints. Registering it as a Parameter must not count it + toward `load_numel_total`: reload restores the construction-time tensor + set before sizing, so act-order layers still process during streaming + instead of deferring (and buffering weights) until finalization.""" + from vllm.model_executor.layers.quantization.base_config import ( + QuantizeMethodBase, + ) + from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info + + _stub_marlin_ops(monkeypatch) + kernel = _make_act_order_marlin_kernel() + + class _KernelQuantMethod(QuantizeMethodBase): + def create_weights(self, layer, *args, **kwargs): + raise NotImplementedError + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + kernel.process_weights_after_loading(layer) + + generator = torch.Generator().manual_seed(0) + layer = torch.nn.Module() + layer.quant_method = _KernelQuantMethod() + _load_marlin_checkpoint_format_weights(layer, _random_g_idx(generator)) + + # Metadata is recorded at model construction, before any processing + record_metadata_for_reloading(layer) + checkpoint_numel = sum(t.numel() for t in get_layer_tensors(layer).values()) + + kernel.process_weights_after_loading(layer) + sort_indices = layer.g_idx_sort_indices + + initialize_layerwise_reload(layer) + info = get_layerwise_info(layer) + assert info.load_numel_total == checkpoint_numel + + # Stream a new checkpoint; the layer must process as soon as its last + # tensor arrives + new_g_idx = _random_g_idx(generator) + checkpoint = { + "qweight": torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32), + "scales": torch.ones( + _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16 + ), + "g_idx": new_g_idx, + } + for name, weight in checkpoint.items(): + param = getattr(layer, name) + param.weight_loader(param, weight) + + assert not info.can_load() + assert not info.loaded_weights + assert layer.g_idx_sort_indices is sort_indices + expected_sort_indices = marlin_utils.marlin_sort_g_idx(new_g_idx)[1] + assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices) + + def test_model_cleanup(dist_init, default_vllm_config): layer = QKVParallelLinear(2, 3, 4) assert layer.weight.weight_loader.__self__ is layer @@ -278,6 +557,58 @@ def materialize_with_sentinel(meta_tensor): assert torch.equal(layer.D, loaded["D"]) +class _RecordingQuantMethod(QuantizeMethodBase): + """Records the layer's bias at the moment processing runs.""" + + uses_meta_device = True + + def __init__(self): + self.bias_at_process = None + + def create_weights(self, layer, *weight_args, **extra_weight_attrs): + pass + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + self.bias_at_process = layer.bias.detach().clone() + + +class _LateBiasLayer(torch.nn.Module): + """Mimics an online-quantized linear: `weight` is created on meta by + `create_weights()`, which wraps the loaders, and the linear base registers + `bias` afterwards.""" + + def __init__(self, quant_method): + super().__init__() + self.quant_method = quant_method + weight = torch.nn.Parameter(torch.empty(4, 2, device="meta")) + weight.weight_loader = default_weight_loader + self.register_parameter("weight", weight) + initialize_online_processing(self) + bias = torch.nn.Parameter(torch.zeros(4)) + bias.weight_loader = default_weight_loader + self.register_parameter("bias", bias) + + +def test_online_processing_waits_for_late_registered_bias(): + # Regression test: `bias` is skipped by the meta device paths, but it is + # still loaded by a weight loader. Excluding it from the processing trigger + # finalized the layer one load early, so the trailing bias was written into + # an already-processed layer (e.g. over FP8 Marlin's permuted bias). + quant_method = _RecordingQuantMethod() + layer = _LateBiasLayer(quant_method) + loaded_bias = torch.full((4,), 3.0) + + layer.weight.weight_loader(layer.weight, torch.full((4, 2), 2.0)) + assert quant_method.bias_at_process is None + + layer.bias.weight_loader(layer.bias, loaded_bias) + assert quant_method.bias_at_process is not None + assert torch.equal(quant_method.bias_at_process, loaded_bias) + + def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py index a0b5a2a4aeca..48672d4314dc 100644 --- a/tests/model_executor/model_loader/test_sharded_state_loader.py +++ b/tests/model_executor/model_loader/test_sharded_state_loader.py @@ -9,11 +9,11 @@ import pytest import torch -from huggingface_hub import snapshot_download from vllm import LLM, SamplingParams from vllm.model_executor.model_loader import ShardedStateLoader from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api prompts = [ "Hello, my name is", @@ -52,7 +52,7 @@ def test_filter_subtensors(): @pytest.fixture(scope="module") def llama_3p2_1b_files(): - input_dir = snapshot_download( + input_dir = hf_api().snapshot_download( "meta-llama/Llama-3.2-1B-Instruct", ignore_patterns=["*.bin*", "original/*"] ) diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py new file mode 100644 index 000000000000..da6b21091b97 --- /dev/null +++ b/tests/model_executor/test_jit_warmup.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, + WarmupIntRange, + get_ast_full_name, + zip_inputs, +) + + +def _next_power_of_2(value: int) -> int: + return 1 << max(0, value - 1).bit_length() + + +def _round_up(value: int, *, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +def _config( + *, + bias: int = 0, + disabled: bool = False, + name: str = "base", + vectorized: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + bias=bias, + disabled=disabled, + name=name, + vectorized=vectorized, + ) + + +class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + block_size: int + work: int + vector_width: int + descriptor: tuple[object, ...] + enabled: bool + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + cfg: Any, + lanes: int = 1, + mode: str = "default", + debug: int = 0, + ) -> CompileKey: + block_size = _next_power_of_2(tokens) + work: int = block_size * lanes + cfg.bias + return self.CompileKey( + block_size=block_size, + work=work, + vector_width=4 if cfg.vectorized and block_size >= 4 else 1, + descriptor=( + cfg.name, + mode, + -block_size, + block_size % 3, + block_size**2, + ), + enabled=not cfg.disabled, + ) + + def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + tokens=WarmupIntRange(1, max_tokens + 1), + cfg=cfg, + # This argument is intentionally unused by dispatch expressions. + debug=WarmupIntRange(0, 100), + ) + + def compile(self, compile_key: CompileKey) -> None: + pass + + +class RecordingToyKernel(ToyKernel): + def __init__(self) -> None: + self.compiled: list[ToyKernel.CompileKey] = [] + super().__init__() + + def compile(self, compile_key: ToyKernel.CompileKey) -> None: + self.compiled.append(compile_key) + + +def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None: + cfg = _config() + + assert ToyKernel().get_warmup_keys(5, cfg) == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True), + ] + + +def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None: + cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True) + + assert ToyKernel().compile_key( + { + "tokens": 4, + "cfg": cfg, + "lanes": 2, + } + ) == ToyKernel.CompileKey( + block_size=4, + work=11, + vector_width=4, + descriptor=("cfg", "default", -4, 1, 16), + enabled=False, + ) + + +def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None: + cfg = _config(vectorized=True) + + keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)( + zip_inputs( + dict(tokens=1, mode="small"), + dict(tokens=4, mode="wide"), + ), + cfg=cfg, + lanes=(1, 2), + ) + + assert keys == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True), + ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True), + ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True), + ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True), + ] + + +def test_zip_inputs_validates_input_rows() -> None: + with pytest.raises(ValueError, match="requires at least one"): + zip_inputs() + with pytest.raises(ValueError, match="rows must be mappings"): + zip_inputs(cast(Any, ("tokens", 1))) + with pytest.raises(ValueError, match="at least one dispatch input name"): + zip_inputs({}) + with pytest.raises(ValueError, match="dispatch input names must be strings"): + zip_inputs(cast(Any, {1: 2})) + with pytest.raises(ValueError, match="same dispatch input names"): + zip_inputs({"tokens": 1}, {"mode": "small"}) + + +def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None: + kernel = ToyKernel() + + with pytest.raises(TypeError, match="zip_inputs"): + kernel._trace_dispatch(kernel.dispatch)( + cast(Any, {"tokens": 1}), + cfg=_config(), + ) + + with pytest.raises(ValueError, match="specified more than once"): + kernel._trace_dispatch(kernel.dispatch)( + zip_inputs(dict(tokens=1, mode="small")), + tokens=2, + cfg=_config(), + ) + + +def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None: + class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + block_size: int, + ) -> CompileKey: + return self.CompileKey(value=_round_up(tokens, multiple=block_size)) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + block_size: int, + ) -> CompileKey: + return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size})) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert HelperKernel().compile_key( + { + "tokens": 5, + "block_size": 4, + } + ) == HelperKernel.CompileKey(value=8) + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"): + StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4}) + + +def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None: + class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + if value > 0: + value = 1 + return self.CompileKey(value=value) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(**{"value": value}) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + with pytest.raises(ValueError, match="local assignments"): + BranchKernel() + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"): + KwargsReturnKernel() + + +def test_dispatch_reports_unsupported_expression_with_context() -> None: + class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: object + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value={value}) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + with pytest.raises(ValueError) as exc_info: + UnsupportedKernel().compile_key({"value": 1}) + + message = str(exc_info.value) + assert "Unsupported dispatch expression" in message + assert "{value}" in message + assert "Supported dispatch expressions" in message + + +def test_warmup_compiles_all_returned_keys_in_order() -> None: + kernel = RecordingToyKernel() + cfg = _config() + + kernel.warmup(3, cfg) + + assert kernel.compiled == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ] + + +def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None: + dotted_expr = ast.parse("foo.bar.baz").body[0] + call_expr = ast.parse("foo()").body[0] + assert isinstance(dotted_expr, ast.Expr) + assert isinstance(call_expr, ast.Expr) + + assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz" + assert get_ast_full_name(call_expr.value) is None diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 7bb7b8ba5131..b80b9039e09a 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter: ) +def _make_modular_routed_experts(): + return types.SimpleNamespace( + quant_method=types.SimpleNamespace(is_monolithic=False), + ) + + def test_base_router_capture_pre_eplb_mapping(): router = _make_router() captured = [] @@ -122,6 +128,8 @@ class DummyFusedMoE: def __init__(self): self.layer_id = 7 self.router = _make_router() + self.routed_experts = _make_modular_routed_experts() + self._quant_method = self.routed_experts.quant_method class DummyCapturer: def __init__(self): @@ -160,6 +168,8 @@ class DummyFusedMoE: def __init__(self): self.layer_id = 11 self.router = _make_router() + self.routed_experts = _make_modular_routed_experts() + self._quant_method = self.routed_experts.quant_method class DummyCapturer: def __init__(self): @@ -197,6 +207,8 @@ class DummyFusedMoE: def __init__(self, layer_id): self.layer_id = layer_id self.router = _make_router() + self.routed_experts = _make_modular_routed_experts() + self._quant_method = self.routed_experts.quant_method target_module = DummyFusedMoE(layer_id=7) draft_module = DummyFusedMoE(layer_id=0) @@ -222,6 +234,49 @@ def __init__(self, layer_id): assert draft_module.router.capture_fn is None +def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch): + from vllm.v1.worker import gpu_model_runner as gmr + + class DummyFusedMoE: + def __init__(self): + self.layer_id = 3 + self.router = _make_router() + # Use a concrete monolithic expert and override its capability + # instead of instantiating the abstract base class directly. + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsFp8, + ) + + fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8) + self.routed_experts = types.SimpleNamespace( + quant_method=types.SimpleNamespace( + is_monolithic=True, + moe_kernel=types.SimpleNamespace( + impl=types.SimpleNamespace(fused_experts=fused_experts) + ), + ) + ) + self._quant_method = self.routed_experts.quant_method + self._quant_method.moe_kernel.impl.fused_experts = fused_experts + fused_experts.supports_routing_replay_capture = lambda: False + + class DummyCapturer: + def capture(self, layer_id, topk_ids): + pass + + dummy_module = DummyFusedMoE() + import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer + + monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE) + + dummy_self = types.SimpleNamespace( + model=types.SimpleNamespace(modules=lambda: [dummy_module]) + ) + + with pytest.raises(ValueError, match="monolithic MoE kernel"): + gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer()) + + def test_routed_experts_capturer_single_dp_no_metadata(): """dp_metadata is None: capture writes the full topk_ids rows.""" capturer = _capturer_with_buffer(dp_rank=0) diff --git a/tests/models/inkling/rocm/conftest.py b/tests/models/inkling/rocm/conftest.py new file mode 100644 index 000000000000..bbf1e7e12d58 --- /dev/null +++ b/tests/models/inkling/rocm/conftest.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.platforms import current_platform + + +@pytest.fixture(autouse=True) +def require_rocm_cdna4() -> None: + if not current_platform.is_rocm(): + pytest.skip("requires ROCm") + + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + pytest.skip("requires CDNA4") diff --git a/tests/models/inkling/rocm/test_model_alignment.py b/tests/models/inkling/rocm/test_model_alignment.py new file mode 100644 index 000000000000..3bf3859de8b8 --- /dev/null +++ b/tests/models/inkling/rocm/test_model_alignment.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm guards for shared Inkling model-definition contracts.""" + +from pathlib import Path + +import torch + +import vllm.models.inkling as inkling +from vllm.lora.utils import get_supported_lora_modules +from vllm.model_executor.models.interfaces import supports_lora +from vllm.models.inkling.amd import model as amd_model +from vllm.models.inkling.amd.mtp import InklingMTP + + +def test_backend_neutral_definitions_match_nvidia_reference() -> None: + inkling_dir = Path(amd_model.__file__).parents[1] + for filename in ("logits_processor.py", "mtp.py"): + assert (inkling_dir / "amd" / filename).read_bytes() == ( + inkling_dir / "nvidia" / filename + ).read_bytes() + + +def test_rocm_exports_its_aligned_mtp_implementation() -> None: + assert inkling.InklingMTP is InklingMTP + + +def test_rocm_model_exposes_the_upstream_lora_contract() -> None: + model_cls = amd_model._TmlForCausalLMBase + + assert supports_lora(model_cls) + assert model_cls.packed_modules_mapping == { + "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"], + "w13": ["w1", "w3"], + } + assert model_cls.embedding_modules == {"lm_head": "output_embeddings"} + + model = torch.nn.Module() + model.embedding_modules = model_cls.embedding_modules + supported = get_supported_lora_modules(model) + assert "embed_tokens" not in supported + assert "lm_head" in supported + + +def test_lightseek_bundled_adapter_weights_remain_opt_in() -> None: + mapper = amd_model._TmlForCausalLMBase.hf_to_vllm_mapper + weight = torch.empty(1) + name, mapped_weight = next( + iter( + mapper.apply([("language_model.layers.3.attn.wq_du.lora_A.weight", weight)]) + ) + ) + + assert name == "model.layers.3.attn.qkvr.lora_A.weight" + assert mapped_weight.shard_id == 0 + assert amd_model._is_peft_adapter_weight(name) + + lm_head_name = mapper.apply_list(["language_model.lm_head.lora_B.weight"]) + assert lm_head_name == ["lm_head.lora_B.weight"] diff --git a/tests/models/inkling/rocm/test_mxfp4_load.py b/tests/models/inkling/rocm/test_mxfp4_load.py new file mode 100644 index 000000000000..00ff2f4018f5 --- /dev/null +++ b/tests/models/inkling/rocm/test_mxfp4_load.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Quark OCP MXFP4 weights loaded into AITER MoE layouts.""" + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.models.inkling.amd.moe import InklingMoE + + +def _fake_moe() -> tuple[InklingMoE, SimpleNamespace]: + moe = InklingMoE.__new__(InklingMoE) + nn.Module.__init__(moe) + moe.n_routed_experts = 4 + routed = SimpleNamespace( + moe_config=SimpleNamespace( + moe_parallel_config=SimpleNamespace(tp_rank=1, tp_size=2) + ), + # Logical intermediate per rank is 3; AITER pads it to 4. + w13_weight=nn.Parameter(torch.zeros(4, 8, 2, dtype=torch.uint8), False), + w2_weight=nn.Parameter(torch.zeros(4, 5, 3, dtype=torch.uint8), False), + w13_weight_scale=nn.Parameter(torch.ones(4, 8, 2, dtype=torch.uint8), False), + w2_weight_scale=nn.Parameter(torch.ones(4, 5, 3, dtype=torch.uint8), False), + ) + moe.experts = SimpleNamespace(routed_experts=routed) + moe._local_expert_slots = ( # type: ignore[method-assign] + lambda: dict(enumerate(range(4))) + ) + return moe, routed + + +def test_loads_logical_tp_shards_and_preserves_aiter_padding(): + moe, routed = _fake_moe() + + w13 = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4, 12, 2) + w2 = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4, 5, 4) + moe.load_expert_weight("experts.w13_weight", w13) + moe.load_expert_weight("experts.w2_weight", w2) + + # TP rank 1 consumes the second six interleaved gate/up rows. AITER keeps + # one padding row in each destination half and one padding w2 column. + torch.testing.assert_close(routed.w13_weight[:, :3], w13[:, 6:12:2]) + torch.testing.assert_close(routed.w13_weight[:, 4:7], w13[:, 7:12:2]) + assert torch.count_nonzero(routed.w13_weight[:, 3]) == 0 + assert torch.count_nonzero(routed.w13_weight[:, 7]) == 0 + torch.testing.assert_close(routed.w2_weight[:, :, :2], w2[:, :, 2:4]) + assert torch.count_nonzero(routed.w2_weight[:, :, 2]) == 0 + + +def test_unflattens_quark_scales_and_preserves_scale_padding(): + moe, routed = _fake_moe() + + flat_w13_scale = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4 * 12, 2) + flat_w2_scale = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4 * 5, 4) + moe.load_expert_weight("experts.w13_weight_scale", flat_w13_scale) + moe.load_expert_weight("experts.w2_weight_scale", flat_w2_scale) + + w13_scale = flat_w13_scale.view(4, 12, 2) + w2_scale = flat_w2_scale.view(4, 5, 4) + torch.testing.assert_close(routed.w13_weight_scale[:, :3], w13_scale[:, 6:12:2]) + torch.testing.assert_close(routed.w13_weight_scale[:, 4:7], w13_scale[:, 7:12:2]) + assert torch.all(routed.w13_weight_scale[:, 3] == 1) + assert torch.all(routed.w13_weight_scale[:, 7] == 1) + torch.testing.assert_close(routed.w2_weight_scale[:, :, :2], w2_scale[:, :, 2:4]) + assert torch.all(routed.w2_weight_scale[:, :, 2] == 1) diff --git a/tests/models/inkling/rocm/test_rel_attention.py b/tests/models/inkling/rocm/test_rel_attention.py new file mode 100644 index 000000000000..c154faca34fd --- /dev/null +++ b/tests/models/inkling/rocm/test_rel_attention.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for Inkling's ROCm relative-bias paged attention.""" + +import pytest +import torch + +from vllm.models.inkling.amd.ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_rel_attention, + use_gfx950_gluon_decode, + use_gfx950_gluon_extend, +) +from vllm.models.inkling.amd.ops.rel_attention_decode import ( + decode_split_count, + use_split_kv_decode, +) + +HEAD_DIM = 128 +DTYPE = torch.bfloat16 + + +def _reference( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + rel_logits: torch.Tensor, + block_table: torch.Tensor, + q_lens: list[int], + kv_lens: list[int], + rel_extent: int, + window_left: int | None, +) -> torch.Tensor: + num_heads = q.shape[1] + num_kv_heads = key_cache.shape[2] + gqa_group = num_heads // num_kv_heads + out: list[torch.Tensor] = [] + q_start = 0 + + for req, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): + page_size = key_cache.shape[1] + num_pages = (kv_len + page_size - 1) // page_size + page_ids = block_table[req, :num_pages].long() + k = key_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] + v = value_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] + k = k.float().repeat_interleave(gqa_group, dim=1) + v = v.float().repeat_interleave(gqa_group, dim=1) + + q_req = q[q_start : q_start + q_len].float() + rel_req = rel_logits[q_start : q_start + q_len].float() + scores = torch.einsum("qhd,khd->hqk", q_req, k) / HEAD_DIM + + q_pos = torch.arange(q_len, device=q.device)[:, None] + kv_len - q_len + k_pos = torch.arange(kv_len, device=q.device)[None, :] + distance = q_pos - k_pos + rel_idx = distance.clamp(0, rel_extent - 1) + bias = rel_req.permute(1, 0, 2).gather( + 2, rel_idx[None].expand(num_heads, -1, -1) + ) + in_rel_extent = (distance >= 0) & (distance < rel_extent) + scores += torch.where(in_rel_extent[None], bias, 0.0) + + masked = distance < 0 + if window_left is not None: + masked |= distance > window_left + scores.masked_fill_(masked[None], float("-inf")) + out.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v)) + q_start += q_len + + return torch.cat(out).to(DTYPE) + + +def test_query_length_bucket(): + assert bucket_max_seqlen_q(1) == 1 + assert bucket_max_seqlen_q(17) == 32 + assert bucket_max_seqlen_q(33) == 64 + + +def test_split_kv_dispatch_thresholds(): + assert not use_split_kv_decode( + max_query_len=4, + max_kv_len=65536, + page_size=128, + window_left=-1, + ) + assert not use_split_kv_decode( + max_query_len=1, + max_kv_len=4096, + page_size=16, + window_left=-1, + ) + assert use_split_kv_decode( + max_query_len=1, + max_kv_len=8192, + page_size=16, + window_left=-1, + ) + assert use_split_kv_decode( + max_query_len=1, + max_kv_len=512, + page_size=128, + window_left=511, + ) + assert decode_split_count(512, 511) == 4 + assert decode_split_count(65536, -1) == 32 + + +def test_gfx950_gluon_dispatch_is_strict(monkeypatch: pytest.MonkeyPatch): + import vllm.models.inkling.amd.ops.fa4_rel_attention as rel_attention + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) + assert use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=16, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=4, page_size=128, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=256) + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: False) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) + assert use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=128, + head_dim=128, + window_left=-1, + ) + assert not use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=128, + head_dim=128, + window_left=511, + ) + assert not use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=16, + head_dim=128, + window_left=-1, + ) + + +@pytest.mark.parametrize( + ("page_size", "window_left", "max_kv_len"), + [ + (16, None, None), + (16, 15, None), + (128, None, 8192), + ], +) +@torch.inference_mode() +def test_ragged_multi_page_relative_attention( + page_size: int, + window_left: int | None, + max_kv_len: int | None, +): + """Covers full/chunked prefill, decode, GQA, and the local window.""" + torch.manual_seed(19 + int(window_left is not None)) + device = "cuda" + q_lens = [17, 1] + kv_lens = [35, 80] + num_heads = 8 + num_kv_heads = 2 + rel_extent = 16 + max_pages = max((length + page_size - 1) // page_size for length in kv_lens) + + q = torch.randn(sum(q_lens), num_heads, HEAD_DIM, device=device) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + key_cache = torch.randn( + 1 + len(q_lens) * max_pages, + page_size, + num_kv_heads, + HEAD_DIM, + device=device, + ) + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) + value_cache = torch.randn_like(key_cache) + rel_logits = torch.randn( + sum(q_lens), num_heads, rel_extent, device=device, dtype=DTYPE + ) + + block_table = torch.stack( + [ + torch.arange(1 + req * max_pages, 1 + (req + 1) * max_pages) + for req in range(len(q_lens)) + ] + ).to(device=device, dtype=torch.int32) + cu_seqlens_q = torch.tensor( + [0, *torch.tensor(q_lens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + window_size = (-1, -1) if window_left is None else (window_left, 0) + + preallocated = torch.empty_like(q) + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=bucket_max_seqlen_q(max(q_lens)), + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + out=preallocated, + ) + assert actual.data_ptr() == preallocated.data_ptr() + + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) + + +@pytest.mark.parametrize( + ("page_size", "kv_lens", "num_kv_heads", "rel_extent", "window_left"), + [ + (16, [257, 513], 1, 64, None), + (128, [257, 512], 2, 512, 511), + ], +) +@torch.inference_mode() +def test_split_kv_decode_matches_reference( + page_size: int, + kv_lens: list[int], + num_kv_heads: int, + rel_extent: int, + window_left: int | None, +): + """Covers page-16 split-KV and page-128 gfx950 Gluon decode.""" + torch.manual_seed(41 + page_size) + device = "cuda" + batch_size = len(kv_lens) + q_lens = [1] * batch_size + num_heads = 8 + max_pages = max((length + page_size - 1) // page_size for length in kv_lens) + + q = torch.randn(batch_size, num_heads, HEAD_DIM, device=device) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + key_cache = torch.randn( + batch_size * max_pages, + page_size, + num_kv_heads, + HEAD_DIM, + device=device, + ) + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) + value_cache = torch.randn_like(key_cache) + rel_logits = torch.randn( + batch_size, + num_heads, + rel_extent, + device=device, + dtype=DTYPE, + ) + block_table = torch.arange( + batch_size * max_pages, + device=device, + dtype=torch.int32, + ).view(batch_size, max_pages) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + cu_seqlens_q = torch.arange( + batch_size + 1, + device=device, + dtype=torch.int32, + ) + window_size = (-1, -1) if window_left is None else (window_left, 0) + max_kv_len = 8192 if page_size == 16 else max(kv_lens) + + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + ) + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) + + +@pytest.mark.parametrize( + ( + "q_len", + "kv_len", + "num_kv_heads", + "rel_extent", + "window_left", + "gluon_enabled", + ), + [ + (1, 513, 2, 512, 511, True), + (1, 513, 2, 512, 511, False), + (4, 8192, 1, 1024, None, True), + ], +) +@torch.inference_mode() +def test_gfx950_page128_packed_kv_views_match_reference( + monkeypatch: pytest.MonkeyPatch, + q_len: int, + kv_len: int, + num_kv_heads: int, + rel_extent: int, + window_left: int | None, + gluon_enabled: bool, +): + """Cover both Gluon and split-KV against vLLM's packed KV allocation.""" + monkeypatch.setenv( + "INKLING_GFX950_GLUON", + "1" if gluon_enabled else "0", + ) + torch.manual_seed(71 + q_len) + device = "cuda" + page_size = 128 + batch_size = 2 + num_heads = 8 + q_lens = [q_len] * batch_size + kv_lens = [kv_len] * batch_size + pages_per_req = (kv_len + page_size - 1) // page_size + + q = torch.randn( + batch_size * q_len, + num_heads, + HEAD_DIM, + device=device, + dtype=DTYPE, + ) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + + # FlashAttentionBackend allocates logical [block, head, page, 2 * dim]. + # Inkling transposes to [block, page, head, 2 * dim] and splits K/V, + # producing non-contiguous views whose page/head strides include both. + packed_kv = torch.randn( + batch_size * pages_per_req, + num_kv_heads, + page_size, + 2 * HEAD_DIM, + device=device, + dtype=DTYPE, + ) + key_cache, value_cache = packed_kv.transpose(1, 2).split(HEAD_DIM, dim=-1) + assert not key_cache.is_contiguous() + assert key_cache.stride() == value_cache.stride() + + rel_logits = torch.randn( + batch_size * q_len, + num_heads, + rel_extent, + device=device, + dtype=DTYPE, + ) + block_table = torch.arange( + batch_size * pages_per_req, + device=device, + dtype=torch.int32, + ).view(batch_size, pages_per_req) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + cu_seqlens_q = torch.arange( + 0, + batch_size * q_len + 1, + q_len, + device=device, + dtype=torch.int32, + ) + window_size = (-1, -1) if window_left is None else (window_left, 0) + + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=q_len, + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=kv_len, + ) + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) diff --git a/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py b/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py new file mode 100644 index 000000000000..948c3a7f2a85 --- /dev/null +++ b/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm numerical-equivalence tests for the fused MTP depth-layer input kernel. + +``embed_dual_rmsnorm_cat`` must match the unfused module sequence numerically: +each rmsnorm computes in fp32 and rounds to bf16 at the same points as the +vendored ``rmsnorm`` kernel (including the bf16 round-trip between the +chained backbone embed_norm and the depth embed_norm), and the fused row +gather matches ``F.embedding``. +""" + +from typing import cast + +import pytest +import torch + +from vllm.models.inkling.amd.ops.norm import ( + embed_dual_rmsnorm_cat, + embed_rmsnorm, + rmsnorm, +) + +EPS = 1e-6 +VOCAB = 4096 +BF16_ATOL = torch.finfo(torch.bfloat16).eps + + +def _assert_bf16_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + # Different Triton reduction schedules can move a small number of outputs + # by one bf16 epsilon while remaining numerically equivalent. + torch.testing.assert_close(actual, expected, rtol=0, atol=BF16_ATOL) + + +def _ref(hidden, w_h, w_e, emb, w_pre=None): + if w_pre is not None: + emb = rmsnorm(emb, w_pre, EPS) + return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w_h = torch.randn(n, device=dev).to(torch.bfloat16) + w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) + w_pre = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) + emb = table[ids.long()] + + # Fused gather + chained backbone pre-norm (the decode draft-step path). + out = embed_dual_rmsnorm_cat( + hidden, + w_h, + w_e, + EPS, + input_ids=ids, + embed_table=table, + pre_norm_weight=w_pre, + ) + assert out.shape == (t, 2 * n) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb, w_pre)) + + # Precomputed embeds, no pre-norm (draft prefill with target-merged MM + # embeddings, already backbone-normed). + out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) + + # Fused gather, no pre-norm (use_embed_norm=False). + out = embed_dual_rmsnorm_cat( + hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table + ) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + ref_emb = table[ids.long()] + + # Gather + embed_norm (base model / MTP prefill embed path). + out = cast(torch.Tensor, embed_rmsnorm(ids, table, w, EPS)) + assert out.shape == (t, n) + _assert_bf16_close(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) + + # Pure gather (use_embed_norm=False / replicated module forward). + out = cast(torch.Tensor, embed_rmsnorm(ids, table, None, EPS)) + assert torch.equal(out, ref_emb) + + # Chained first-layer attn_norm (the target text-path forward): one launch + # emits both the residual and layer 0's normed attention input. + w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) + res, attn_in = cast( + tuple[torch.Tensor, torch.Tensor], + embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain), + ) + ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb + _assert_bf16_close(res, ref_res) + _assert_bf16_close(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) + + # Chained without embed_norm (use_embed_norm=False). + res, attn_in = cast( + tuple[torch.Tensor, torch.Tensor], + embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain), + ) + assert torch.equal(res, ref_emb) + _assert_bf16_close(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) diff --git a/tests/models/inkling/rocm/test_sconv_cache_layout.py b/tests/models/inkling/rocm/test_sconv_cache_layout.py new file mode 100644 index 000000000000..eeefb877d87a --- /dev/null +++ b/tests/models/inkling/rocm/test_sconv_cache_layout.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.inkling.amd.sconv_swa_attn import InklingConvState + + +def test_runtime_sconv_block_size_tracks_unified_cache_page(): + """The cache planner may enlarge W=4 blocks to match attention pages.""" + owner = InklingConvState.__new__(InklingConvState) + nn.Module.__init__(owner) + owner.block_size = 4 + + owner.kv_cache = torch.tensor([]) + assert owner.cache_block_size == 4 + + owner.kv_cache = torch.empty(2, 1, 32, 1024) + assert owner.cache_block_size == 32 diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py index 42313e7d6775..4bdb75dcc499 100644 --- a/tests/models/inkling/test_contract_validation.py +++ b/tests/models/inkling/test_contract_validation.py @@ -6,6 +6,7 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser +from vllm.models.inkling.common.towers import plan_out_scales from vllm.models.inkling.configs import ( InklingAudioConfig, InklingModelConfig, @@ -17,6 +18,22 @@ from vllm.v1.attention.backend import AttentionCGSupport +def test_vision_scale_plan_matches_released_config(): + assert plan_out_scales(2, 40, 4) == [ + (1, 1, 1, 3), + (1, 5, 5, 128), + (1, 10, 10, 320), + (1, 40, 40, 4800), + (2, 40, 40, 9600), + ] + + +def test_vision_scale_plan_breaks_assignment_ties_in_order(): + reductions = [np.prod(scale[:-1]) for scale in plan_out_scales(2, 52, 4)] + + assert reductions == sorted(set(reductions)) + + @pytest.mark.parametrize( ("config_cls", "kwargs", "missing"), [ diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py index 53a37e3e60bc..e15a573b8459 100644 --- a/tests/models/inkling/test_moe_weight_layout.py +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -112,6 +112,30 @@ def test_inkling_mapper_maps_modelopt_exclusions() -> None: assert not quant_config.is_layer_excluded("model.layers.3.mlp.experts") +@pytest.mark.parametrize("projection", ["w13", "w2"]) +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize( + "suffix", + [ + "input_global_scale", + "weight_global_scale", + "weight_packed", + "weight_scale", + ], +) +def test_inkling_mapper_maps_compressed_tensors_expert_params( + projection: str, suffix: str, nested: bool +) -> None: + projection_param = ( + f"{projection}_weight.{suffix}" if nested else f"{projection}_{suffix}" + ) + source = f"model.llm.layers.2.mlp.experts.{projection_param}" + + mapped = _TmlForCausalLMBase.hf_to_vllm_mapper.apply_list([source]) + + assert mapped == [f"model.layers.2.mlp.experts.{projection}_{suffix}"] + + @pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)]) def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None: experts = SimpleNamespace( @@ -132,6 +156,63 @@ def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None: assert loaded == [f"experts.routed_experts.{projection}_input_scale"] +@pytest.mark.parametrize("projection", ["w13", "w2"]) +@pytest.mark.parametrize("scale_kind", ["input", "weight"]) +def test_moe_loads_compressed_tensors_global_scale( + projection: str, scale_kind: str +) -> None: + param = torch.nn.Parameter(torch.empty(3, 2 if projection == "w13" else 1)) + if projection == "w2": + param = torch.nn.Parameter(param.squeeze(1)) + experts = SimpleNamespace( + **{f"{projection}_{scale_kind}_global_scale": param}, + moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=0)), + ) + layer = SimpleNamespace( + experts=SimpleNamespace(routed_experts=experts), + _local_expert_slots=lambda: {0: 0, 1: 1, 2: 2}, + ) + checkpoint_scale = torch.tensor([[1.0], [2.0], [3.0]]) + + loaded = moe.InklingMoE.load_expert_weight( + layer, + f"experts.{projection}_{scale_kind}_global_scale", + checkpoint_scale, + ) + + expected = ( + checkpoint_scale.expand_as(param) if param.ndim == 2 else checkpoint_scale[:, 0] + ) + torch.testing.assert_close(param, expected) + assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"] + + +@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)]) +def test_moe_loads_channelwise_scale_for_tp( + projection: str, checkpoint_rows: int +) -> None: + param = torch.nn.Parameter(torch.empty(2, 4, 1)) + experts = SimpleNamespace( + **{f"{projection}_weight_scale": param}, + moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)), + ) + layer = SimpleNamespace( + experts=SimpleNamespace(routed_experts=experts), + _local_expert_slots=lambda: {0: 0, 2: 1}, + ) + checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1) + + loaded = moe.InklingMoE.load_expert_weight( + layer, f"experts.{projection}_weight_scale", checkpoint_scale + ) + + expected = checkpoint_scale[[0, 2]] + if projection == "w13": + expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2) + torch.testing.assert_close(param, expected.float()) + assert loaded == [f"experts.routed_experts.{projection}_weight_scale"] + + def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None: monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py new file mode 100644 index 000000000000..f6dfaa422b3f --- /dev/null +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AMD AttnRes requires ROCm", +) + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + hidden_size = prefix.shape[-1] + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (hidden_size,), norm_weight, eps) + probs = (keys @ qk_weight).softmax(dim=-1) + return torch.matmul(probs.unsqueeze(1), values).squeeze(1) + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "block_capacity", + "hidden_size", + "row_padding", + ), + [ + pytest.param(0, 3, 5, 128, 0, id="empty"), + pytest.param(1, 1, 2, 128, 0, id="decode-single"), + pytest.param(17, 4, 6, 1024, 7, id="decode-padded"), + pytest.param(320, 8, 10, 7168, 0, id="prefill-full"), + ], +) +def test_amd_attn_res_matches_reference( + num_tokens: int, + num_blocks: int, + block_capacity: int, + hidden_size: int, + row_padding: int, +) -> None: + eps = 1e-5 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding) + blocks = _randn_with_row_padding( + num_tokens, + block_capacity, + hidden_size, + padding=row_padding, + ) + norm_weight = 1 + 0.1 * torch.randn( + hidden_size, device="cuda", dtype=torch.bfloat16 + ) + qk_weight = ( + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5 + ) + expected = _reference( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + original_prefix = prefix.clone() + original_blocks = blocks.clone() + + actual = attn_res( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.shape == prefix.shape + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py new file mode 100644 index 000000000000..38a34bc923a5 --- /dev/null +++ b/tests/models/kimi_k3/test_attn_res.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.common.mtp import fused_mtp_input +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.platforms import current_platform + +HIDDEN_SIZE = 7168 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if delta is not None: + prefix = prefix + delta + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS) + probs = (keys @ qk_weight).softmax(dim=-1) + output = torch.matmul(probs.unsqueeze(1), values).squeeze(1) + if output_norm_weight is not None: + output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS) + return output, prefix + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "row_padding", + "write_block", + "has_delta", + "backend", + ), + [ + pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"), + pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"), + pytest.param(17, 5, 7, True, False, "triton", id="triton-write"), + pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), + pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), + pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), + pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), + pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), + pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), + ], +) +def test_attn_res( + num_tokens: int, + num_blocks: int, + row_padding: int, + write_block: bool, + has_delta: bool, + backend: str, +): + if backend == "nvidia" and not current_platform.is_device_capability_family(100): + pytest.skip("NVIDIA AttnRes requires the SM100 family") + + prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + delta = ( + _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + if has_delta + else None + ) + blocks = _randn_with_row_padding( + num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding + ) + norm_weight = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + original_blocks = blocks.clone() + expected, expected_prefix = _reference( + prefix.clone(), + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + block_write_idx = num_blocks if write_block else -1 + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + block_write_idx, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + if write_block: + original_blocks[:, block_write_idx].copy_(expected_prefix) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.is_contiguous() + + +@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1)) +def test_attn_res_block_counts(num_blocks: int): + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + blocks = torch.randn( + 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = torch.ones_like(norm_weight) + expected, _ = _reference( + prefix.clone(), + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + + actual = attn_res( + prefix, + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + -1, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def test_attn_res_without_output_norm(): + prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn( + 7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + expected, _ = _reference( + prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS + ) + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + None, + MAX_BLOCKS, + -1, + EPS, + 0.0, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 17]) +def test_fused_mtp_input(num_tokens: int): + positions = torch.arange(num_tokens, device="cuda") + inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7) + previous_hidden_states = _randn_with_row_padding( + num_tokens, HIDDEN_SIZE, padding=11 + ) + enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + + masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + expected = torch.cat( + ( + F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS), + F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS), + ), + dim=-1, + ) + actual = fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE) + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py new file mode 100644 index 000000000000..61a24a83e041 --- /dev/null +++ b/tests/models/kimi_k3/test_eagle3.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from vllm.model_executor.models.interfaces import supports_eagle3 +from vllm.models.kimi_k3.nvidia import model as kimi_model +from vllm.models.kimi_k3.nvidia.model import ( + KimiK3ForConditionalGeneration, + KimiLinearModel, +) + + +def _make_kimi_linear_model() -> KimiLinearModel: + model = object.__new__(KimiLinearModel) + object.__setattr__(model, "aux_hidden_state_layers", (2,)) + object.__setattr__(model, "use_sequence_parallel", False) + return model + + +def test_kimi_k3_advertises_eagle3_support(): + assert supports_eagle3(KimiK3ForConditionalGeneration) + + +def test_kimi_k3_uses_shared_eagle3_layer_configuration(): + target = object.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(target) + model = _make_kimi_linear_model() + object.__setattr__(model, "layers", [None] * 93) + language_model = SimpleNamespace( + embed_input_ids=lambda _: None, + model=model, + ) + object.__setattr__(target, "language_model", language_model) + object.__setattr__(target, "_language_model_names", ["language_model"]) + + target.set_aux_hidden_state_layers((2, 46, 90)) + + assert model.aux_hidden_state_layers == (2, 46, 90) + assert target.get_eagle3_default_aux_hidden_state_layers() == ( + 2, + 46, + 90, + ) + + +def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + layer_residual = torch.tensor([[5.0, 6.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, None, layer_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", False) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + expected_layer_output = layer_hidden_states + layer_residual + torch.testing.assert_close(output, expected_layer_output) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) + torch.testing.assert_close(aux_hidden_states[1], expected_layer_output) + + +def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + final_hidden_states = torch.tensor([[9.0, 10.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + final_attn_res = Mock(return_value=final_hidden_states) + monkeypatch.setattr(kimi_model, "attn_res", final_attn_res) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + torch.testing.assert_close(output, final_hidden_states) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) + torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) + assert final_attn_res.call_args.args[2] is block_residual diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py new file mode 100644 index 000000000000..be916f6c0d0f --- /dev/null +++ b/tests/models/kimi_k3/test_kda.py @@ -0,0 +1,757 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Precision tests for vllm's chunk_kda Triton operator. + +Compares chunk_kda against a naive recurrent reference (float32). +Uses torch.rand for q/k/v to match FLA's test pattern. +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update +from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from vllm.models.kimi_k3.nvidia.kda import ( + is_flashkda_supported, + is_fused_kda_decode_supported, +) +from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda, + chunk_kda_with_fused_gate, + fused_kda_gate, + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd + +DEVICE = "cuda" + + +@torch.inference_mode() +def test_gather_initial_states_correctness(): + row_size = 8 * 128 * 128 + storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE) + state = storage[:, :row_size].view(5, 8, 128, 128) + assert not state.is_contiguous() + assert state[0].is_contiguous() + indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE) + has_initial_state = torch.tensor([True, False, True], device=DEVICE) + + expected = state[indices].clone() + expected[~has_initial_state] = 0 + + torch.testing.assert_close( + gather_initial_states(state, indices, has_initial_state), + expected, + ) + + +def naive_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Naive recurrent KDA reference, ported from FLA's naive.py.""" + dtype = v.dtype + B, T, H, K = q.shape + V = v.shape[-1] + if scale is None: + scale = K**-0.5 + + q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) + q = q * scale + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + for i in range(T): + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] + S = S * g_i[..., None].exp() + S = S + torch.einsum( + "bhk,bhv->bhkv", + b_i[..., None] * k_i, + v_i - (k_i[..., None] * S).sum(-2), + ) + o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) + if not output_final_state: + S = None + return o.to(dtype), S + + +def assert_close( + name: str, + ref: torch.Tensor, + tri: torch.Tensor, + ratio: float, + err_atol: float = 1e-6, +): + """RMSE-based relative error comparison.""" + abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() + rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() + rmse_base = ref.detach().flatten().square().mean().sqrt().item() + rel_err = rmse_diff / (rmse_base + 1e-8) + print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") + if abs_err <= err_atol: + return + assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" + assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" + assert rel_err < ratio, ( + f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" + ) + + +@pytest.mark.parametrize( + ("H", "D", "cu_seqlens", "dtype"), + [ + pytest.param( + *test, + id="H{}-D{}-cu{}-{}".format(*test), + ) + for test in [ + (32, 128, [0, 64], torch.float16), + (32, 128, [0, 1024], torch.float16), + (32, 128, [0, 15], torch.float16), + (32, 128, [0, 256, 512, 768, 1024], torch.float16), + (32, 128, [0, 15, 100, 300, 1200], torch.float16), + (64, 128, [0, 256, 500, 1000], torch.float16), + (32, 128, [0, 8192], torch.float16), + (32, 128, [0, 256, 500, 1000], torch.bfloat16), + ] + ], +) +@torch.inference_mode() +def test_chunk_kda( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + T = cu_seqlens[-1] + torch.manual_seed(42) + B = 1 + cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) + N = len(cu_seqlens) - 1 + + q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( + dtype + ) + beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + + # Naive reference with l2norm_fwd (same kernel as chunk_kda) + ref_outputs = [] + ref_states = [] + for i in range(N): + s, e = cu_seqlens[i], cu_seqlens[i + 1] + q_i = l2norm_fwd(q[:, s:e].contiguous()) + k_i = l2norm_fwd(k[:, s:e].contiguous()) + o_i, ht_i = naive_recurrent_kda( + q_i, + k_i, + v[:, s:e], + g[:, s:e], + beta[:, s:e], + initial_state=h0[i], + output_final_state=True, + ) + ref_outputs.append(o_i) + ref_states.append(ht_i) + ref_o = torch.cat(ref_outputs, dim=1) + ref_ht = torch.cat(ref_states, dim=0) + + # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) + tri_o, tri_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0.transpose(-1, -2).contiguous().clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" + assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" + assert_close("o", ref_o, tri_o, 0.005) + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) + + +@pytest.mark.parametrize( + ("cu_seqlens", "dtype", "lower_bound"), + [ + ([0, 64], torch.float16, None), + ([0, 15, 100, 300], torch.bfloat16, None), + ([0, 15, 100, 300], torch.bfloat16, -3.0), + ], +) +@torch.inference_mode() +def test_chunk_kda_fused_gate_cumsum_matches_unfused( + cu_seqlens: list[int], + dtype: torch.dtype, + lower_bound: float | None, +): + H, D = 8, 64 + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + torch.manual_seed(123) + + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) + q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE) + raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2] + beta = raw_beta.float().sigmoid() + A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() + dt_bias = ( + torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 + ).contiguous() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + initial_state = h0.transpose(-1, -2).contiguous() + + gate = fused_kda_gate( + raw_g.reshape(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ) + if lower_bound is not None: + expected_gate = lower_bound * torch.sigmoid( + A_log.exp()[None, :, None] + * (raw_g.float().view(T, H, D) + dt_bias.view(H, D)) + ) + torch.testing.assert_close(gate, expected_gate) + gate = gate.unsqueeze(0) + old_o, old_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=gate, + beta=beta, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + new_o, new_ht = chunk_kda_with_fused_gate( + q=q.clone(), + k=k.clone(), + v=v.clone(), + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) + assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize("num_seqs", [1, 8, 32]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@pytest.mark.parametrize("state_indices_stride", [1, 8]) +@torch.inference_mode() +def test_packed_kda_decode_correctness( + num_seqs: int, + lower_bound: float | None, + state_indices_stride: int, +): + H, D = 8, 128 + torch.manual_seed(321) + + packed_storage = torch.randn( + num_seqs, + 3 * H * D + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + mixed_qkv = packed_storage[:, : 3 * H * D] + assert mixed_qkv.stride(0) == 3 * H * D + 1 + q, k, v = ( + x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1) + ) + raw_g = torch.randn( + 1, + num_seqs, + H, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = torch.randn( + 1, + num_seqs, + H, + dtype=torch.bfloat16, + device=DEVICE, + ) + beta = raw_beta.float().sigmoid() + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + state_storage = torch.randn( + num_seqs + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D) + assert not state.is_contiguous() + assert state.stride()[1:] == (D * D, D, 1) + state_indices_storage = torch.zeros( + num_seqs, + state_indices_stride, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = state_indices_storage[:, 0] + state_indices.copy_( + torch.arange( + 1, + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ) + ) + gate = fused_kda_gate( + raw_g.reshape(num_seqs, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + dense_state = state.clone() + dense_out, _ = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=D**-0.5, + initial_state=dense_state, + inplace_final_state=True, + cu_seqlens=torch.arange( + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ), + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + ) + packed_state = state + packed_out, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=packed_state, + state_indices=state_indices, + ) + + assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3) + assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize( + ("H", "fuse_gate"), + [(12, True), (12, False), (12, None), (96, None)], +) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@torch.inference_mode() +def test_kda_spec_decode_correctness( + H: int, + fuse_gate: bool | None, + lower_bound: float | None, +): + num_seqs, query_len, D = 3, 3, 128 + T = num_seqs * query_len + torch.manual_seed(1234) + + qkv_storage = torch.randn( + 1, + T, + 3 * H * D + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + packed_qkv = qkv_storage[..., : 3 * H * D] + q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1)) + gate_storage = torch.randn( + 1, + T, + H * D + 5, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_g = gate_storage[..., : H * D].view(1, T, H, D) + beta_storage = torch.randn( + 1, + T, + H + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = beta_storage[..., :H] + A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.arange( + 0, + T + 1, + query_len, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = torch.arange( + 1, + T + 1, + dtype=torch.int32, + device=DEVICE, + ).view(num_seqs, query_len) + num_accepted_tokens = torch.tensor( + [1, 2, 3], + dtype=torch.int32, + device=DEVICE, + ) + state_storage = 0.01 * torch.randn( + T + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(T + 1, H, D, D) + output_storage = torch.full( + (1, T, H * D + 11), + torch.nan, + dtype=torch.bfloat16, + device=DEVICE, + ) + output = output_storage[..., : H * D].view(1, T, H, D) + + gate = fused_kda_gate( + raw_g.contiguous().view(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + beta = raw_beta.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + expected_state = state.clone() + expected_outputs = [] + for seq, accepted in enumerate(num_accepted_tokens.tolist()): + recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose( + -1, -2 + ) + start = seq * query_len + for token in range(query_len): + token_slice = slice(start + token, start + token + 1) + token_output, recurrent_state = naive_recurrent_kda( + q_norm[:, token_slice], + k_norm[:, token_slice], + v[:, token_slice], + gate[:, token_slice], + beta[:, token_slice], + initial_state=recurrent_state, + output_final_state=True, + ) + assert recurrent_state is not None + expected_outputs.append(token_output) + expected_state[state_indices[seq, token]] = recurrent_state.transpose( + -1, -2 + ) + expected = torch.cat(expected_outputs, dim=1) + + actual_state = state.clone() + actual, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=actual_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + out=output, + fuse_gate=fuse_gate, + ) + + assert actual.data_ptr() == output.data_ptr() + assert_close("o", expected, actual, 1e-3, err_atol=1e-3) + used_states = state_indices.flatten().long() + assert_close( + "ht", + expected_state[used_states], + actual_state[used_states], + 3e-3, + err_atol=3e-3, + ) + assert torch.isnan(output_storage[..., H * D :]).all() + + +@pytest.mark.parametrize( + ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), + [ + (12, 1, -5.0, True), + (12, 4, None, False), + (24, 4, None, False), + (48, 1, -5.0, True), + (96, 1, -5.0, True), + ], +) +@torch.inference_mode() +def test_fused_kda_decode_correctness( + num_heads: int, + num_seqs: int, + lower_bound: float | None, + fuse_output_norm: bool, +): + D, W = 128, 4 + if not is_fused_kda_decode_supported( + num_heads, + D, + W, + num_spec=0, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ): + pytest.skip("Fused KDA decode is not supported on this platform") + torch.manual_seed(967 + num_heads + num_seqs) + dim = num_heads * D + slots = num_seqs + 2 + packed_x_storage = torch.randn( + num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE + ) + packed_x = packed_x_storage[:, : 3 * dim] + weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE) + conv_seed = 0.1 * torch.randn( + slots, + W - 1, + 3 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ).transpose(1, 2) + raw_g = torch.randn( + 1, + num_seqs, + num_heads, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta_storage = torch.randn( + 1, + num_seqs, + num_heads + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = raw_beta_storage[:, :, :num_heads] + output_gate_storage = torch.randn( + num_seqs, + dim + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D) + norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE) + norm_eps = 1e-5 + A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE) + state_indices = torch.arange( + num_seqs, + 0, + -1, + dtype=torch.int32, + device=DEVICE, + ) + state_seed = 0.01 * torch.randn( + slots, + num_heads, + D, + D, + dtype=torch.float32, + device=DEVICE, + ) + + conv_ref = conv_seed.clone() + state_ref = state_seed.clone() + mixed_qkv = causal_conv1d_update( + packed_x, + conv_ref, + weight, + activation="silu", + conv_state_indices=state_indices, + validate_data=True, + out=torch.empty_like(packed_x), + ) + expected, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=state_ref, + state_indices=state_indices, + ) + if fuse_output_norm: + expected_float = expected.float() + expected = ( + expected_float + * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps) + * norm_weight + * output_gate.float().sigmoid().unsqueeze(0) + ).to(expected.dtype) + + conv_slot_elements = 3 * dim * (W - 1) + state_slot_elements = num_heads * D * D + conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize + page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize + cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE) + conv_actual = torch.as_strided( + cache_storage.view(torch.bfloat16), + size=(slots, 3 * dim, W - 1), + stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim), + ) + state_actual = torch.as_strided( + cache_storage.view(torch.float32), + size=(slots, num_heads, D, D), + stride=(page_bytes // torch.float32.itemsize, D * D, D, 1), + storage_offset=conv_slot_bytes // torch.float32.itemsize, + ) + conv_actual.copy_(conv_seed) + state_actual.copy_(state_seed) + fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous() + actual = ops.fused_kda_decode( + x=packed_x, + weight=fused_weight, + bias=None, + conv_state=conv_actual, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + state=state_actual, + lower_bound=lower_bound, + output_gate=output_gate if fuse_output_norm else None, + norm_weight=norm_weight if fuse_output_norm else None, + norm_eps=norm_eps, + ) + + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0) + torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) + + +def test_fused_kda_decode_rejects_speculative_conv_state(): + assert not is_fused_kda_decode_supported( + num_heads=12, + head_dim=128, + conv_width=4, + num_spec=2, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ) + + +@torch.inference_mode() +def test_flashkda_correctness(): + if not is_flashkda_supported(128, torch.bfloat16, -3.0): + pytest.skip("FlashKDA is not supported on this platform") + + import vllm._flashkda_C # noqa: F401 + + B, T, H, D = 1, 48, 2, 128 + torch.manual_seed(11) + q, k, v, raw_g = [ + torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) + ] + beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) + lower_bound = -3.0 + + gate = lower_bound * torch.sigmoid( + A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :]) + ) + beta = beta_logits.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + + expected_outputs = [] + expected_states = [] + for i, (start, end) in enumerate( + zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()) + ): + output, final_state = naive_recurrent_kda( + q_norm[:, start:end], + k_norm[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=initial_state[i].transpose(-1, -2), + output_final_state=True, + ) + expected_outputs.append(output) + expected_states.append(final_state) + expected_out = torch.cat(expected_outputs, dim=1) + expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous() + + actual_out = torch.empty_like(v) + actual_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1), + dtype=torch.uint8, + device=DEVICE, + ) + torch.ops._flashkda_C.fwd( + q, + k, + v, + raw_g, + beta_logits, + D**-0.5, + actual_out, + workspace, + A_log, + dt_bias, + lower_bound, + initial_state, + actual_state, + cu_seqlens, + ) + + assert_close("o", expected_out, actual_out, 0.01) + assert_close("ht", expected_state, actual_state, 0.01) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py new file mode 100644 index 000000000000..5352ef7a7a64 --- /dev/null +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import fields + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, + KimiK3KDAMetadataBuilder, + _mamba_get_block_table_tensor, + stage_spec_decode_metadata, +) +from vllm.v1.attention.backend import AttentionMetadataBuilder +from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + mamba_get_block_table_tensor, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +PRUNED_METADATA_FIELDS = { + "chunk_indices", + "chunk_offsets", + "prefill_query_start_loc", + "prefill_state_indices", + "prefill_has_initial_state", + "spec_sequence_masks", +} + + +def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata): + for field in fields(KimiK3KDAMetadata): + actual_value = getattr(actual, field.name) + expected_value = getattr(reference, field.name) + if field.name in PRUNED_METADATA_FIELDS: + assert actual_value is None + continue + if ( + field.name in {"spec_token_indx", "non_spec_token_indx"} + and actual.num_spec_decodes > 0 + and actual.num_prefills == 0 + and actual.num_decodes == 0 + ): + assert actual_value is None + continue + if isinstance(actual_value, torch.Tensor): + torch.testing.assert_close(actual_value, expected_value) + elif field.name == "nums_dict": + assert (actual_value is None) == (expected_value is None) + if actual_value is not None: + assert actual_value[8]["tot"] == expected_value[8]["tot"] + torch.testing.assert_close( + actual_value[8]["nums"], expected_value[8]["nums"] + ) + else: + assert actual_value == expected_value + + +def _make_builder( + builder_cls: type[AttentionMetadataBuilder], + num_speculative_tokens: int, + full_cuda_graph: bool, + device: torch.device = DEVICE, + mamba_cache_mode: str = "none", +) -> AttentionMetadataBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + if num_speculative_tokens: + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=num_speculative_tokens, + ) + vllm_config.compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + return builder_cls( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=num_speculative_tokens, + ), + layer_names=["layer.0"], + vllm_config=vllm_config, + device=device, + ) + + +@pytest.mark.parametrize( + ( + "batch", + "num_decode_draft_tokens", + "num_speculative_tokens", + "full_cuda_graph", + "is_prefilling", + ), + [ + pytest.param( + BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]), + [2, 2], + 2, + False, + [False, False], + id="pure-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[100, 65, 20], query_lens=[50, 1, 3]), + [-1, -1, 2], + 2, + False, + [True, False, False], + id="mixed-prefill-and-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + None, + 0, + False, + [False, False], + id="regular-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + [0, 0], + 2, + False, + [False, False], + id="no-scheduled-draft-tokens", + ), + ], +) +def test_kimi_k3_kda_metadata_matches_shared_gdn( + batch: BatchSpec, + num_decode_draft_tokens: list[int] | None, + num_speculative_tokens: int, + full_cuda_graph: bool, + is_prefilling: list[bool], +): + kwargs: dict[str, torch.Tensor] = {} + if num_decode_draft_tokens is not None: + kwargs = { + "num_decode_draft_tokens_cpu": torch.tensor( + num_decode_draft_tokens, dtype=torch.int32 + ), + "num_accepted_tokens": torch.ones( + batch.batch_size, dtype=torch.int32, device=DEVICE + ), + } + + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build( + 0, + common_attn_metadata, + **kwargs, + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build(0, common_attn_metadata, **kwargs) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata(): + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, False, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + # The K3 layer dispatches the non-spec subgroup to packed decode whenever + # it contains no prefill request. + assert actual.num_decodes == 2 + assert actual.num_decode_tokens == 2 + assert actual.num_prefills == 0 + assert actual.num_prefill_tokens == 0 + assert actual.has_initial_state is None + assert actual.nums_dict is None + assert actual.non_spec_query_start_loc is None + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4])) + torch.testing.assert_close( + actual.spec_query_start_loc, + torch.tensor([0, 3], dtype=torch.int32), + ) + + +def test_mixed_regular_and_spec_decode_excludes_request_padding(): + batch = BatchSpec(seq_lens=[16, 65, 20], query_lens=[0, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, False, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_decodes == 1 + assert actual.non_spec_state_indices_tensor is not None + assert actual.non_spec_state_indices_tensor.shape == (1,) + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) + + +@pytest.mark.parametrize( + ("seq_len", "expected_has_initial_state"), + [ + pytest.param(1, False, id="first-token-prefill"), + pytest.param(65, True, id="final-one-token-prefill-chunk"), + ], +) +def test_mixed_one_token_prefill_and_spec_decode_uses_prefill_metadata( + seq_len: int, + expected_has_initial_state: bool, +): + batch = BatchSpec(seq_lens=[seq_len, 20], query_lens=[1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([True, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_prefills == 1 + assert actual.num_prefill_tokens == 1 + assert actual.num_decodes == 0 + assert actual.num_decode_tokens == 0 + assert actual.has_initial_state is not None + assert actual.has_initial_state.tolist() == [expected_has_initial_state] + assert actual.non_spec_query_start_loc is not None + torch.testing.assert_close( + actual.non_spec_query_start_loc, + torch.tensor([0, 1], dtype=torch.int32), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([False, False])) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_kimi_k3_kda_backend_uses_private_metadata_builder(): + assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder + assert KimiK3KDAAttentionBackend.is_ssm() + assert issubclass(KimiK3KDAAttentionBackend, GDNAttentionBackend) + assert issubclass(KimiK3KDAMetadata, GDNAttentionMetadata) + assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_stage_spec_decode_metadata_matches_pytorch(): + device = torch.device("cuda") + num_spec_decodes = 33 + batch_size = 65 + num_state_slots = 3 + state_indices = torch.arange( + num_spec_decodes * 32, + dtype=torch.int32, + device=device, + ).reshape(num_spec_decodes, 32)[:, :num_state_slots] + query_start_loc = ( + torch.arange(num_spec_decodes + 1, dtype=torch.int32, device=device) + * num_state_slots + ) + num_accepted_tokens = ( + torch.arange(num_spec_decodes, dtype=torch.int32, device=device) + % num_state_slots + + 1 + ) + + staged_state_indices = torch.empty( + (batch_size, num_state_slots), dtype=torch.int32, device=device + ) + staged_query_start_loc = torch.empty( + batch_size + 1, dtype=torch.int32, device=device + ) + staged_num_accepted_tokens = torch.empty( + batch_size, dtype=torch.int32, device=device + ) + stage_spec_decode_metadata( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + num_spec_decodes=num_spec_decodes, + ) + + expected_state_indices = torch.full_like(staged_state_indices, NULL_BLOCK_ID) + expected_state_indices[:num_spec_decodes] = state_indices + expected_query_start_loc = torch.full( + (batch_size + 1,), + query_start_loc[-1], + dtype=torch.int32, + device=device, + ) + expected_query_start_loc[: num_spec_decodes + 1] = query_start_loc + expected_num_accepted_tokens = torch.ones( + batch_size, dtype=torch.int32, device=device + ) + expected_num_accepted_tokens[:num_spec_decodes] = num_accepted_tokens + + torch.testing.assert_close(staged_state_indices, expected_state_indices) + torch.testing.assert_close(staged_query_start_loc, expected_query_start_loc) + torch.testing.assert_close(staged_num_accepted_tokens, expected_num_accepted_tokens) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_aligned_block_table_matches_shared_gdn(): + device = torch.device("cuda") + seq_lens = torch.tensor( + [0, 1, 15, 16, 17, 31, 32, 33, 511, 512, 513], + dtype=torch.int32, + device=device, + ).repeat(6)[:65] + block_table_storage = torch.arange( + seq_lens.numel() * 128, + dtype=torch.int32, + device=device, + ).reshape(seq_lens.numel(), 128) + block_table = block_table_storage[:, ::2] + kv_cache_spec = MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=2, + ) + + expected = mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + actual = _mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + + torch.testing.assert_close(actual, expected) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py new file mode 100644 index 000000000000..f18f80a36c04 --- /dev/null +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -0,0 +1,145 @@ +# 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 +import torch.nn.functional as F + +from tests.utils import ( + init_test_distributed_environment, + multi_gpu_test, + multi_process_parallel, +) +from vllm.distributed import get_tp_group +from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp +from vllm.platforms import current_platform + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +EPS = 0.1 + + +@ray.remote(num_gpus=1, max_calls=1) +def _test_latent_moe_tail_worker( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + pp_size: int, + rank: int, + distributed_init_port: str, +) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + ) + + torch.manual_seed(0) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + up_weight = ( + torch.randn( + HIDDEN_SIZE, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + / LATENT_SIZE**0.5 + ) + + group = get_tp_group().device_group + op = KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=EPS, + ) + cutedsl_warmup() + + for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)): + torch.manual_seed(100 * iteration + rank + 1) + routed_output = torch.randn( + num_tokens, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ).mul_(0.01) + shared_output = torch.randn( + num_tokens, + HIDDEN_SIZE, + device=device, + dtype=torch.bfloat16, + ) + + routed_reference = routed_output.clone() + shared_reference = shared_output.clone() + dist.all_reduce(routed_reference, group=group) + dist.all_reduce(shared_reference, group=group) + expected = F.linear( + F.rms_norm( + routed_reference, + (LATENT_SIZE,), + rms_weight, + EPS, + ), + up_weight, + ) + expected.add_(shared_reference) + + actual = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + assert actual.is_contiguous() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + graph.replay() + torch.testing.assert_close(graph_output, expected, atol=8e-2, rtol=3e-2) + + +def _run_latent_moe_tail_test( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, +) -> None: + if not current_platform.is_device_capability_family(100): + pytest.skip("K3 latent-MoE tail fusion requires SM100") + multi_process_parallel( + monkeypatch, + tp_size, + 1, + _test_latent_moe_tail_worker, + ) + + +@multi_gpu_test(num_gpus=8) +def test_latent_moe_tail_tp8_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 8) + + +@multi_gpu_test(num_gpus=16) +def test_latent_moe_tail_tp16_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 16) diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py new file mode 100644 index 000000000000..9e8faea099cb --- /dev/null +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from types import MethodType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from vllm.config import ParallelConfig +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.models.common.ops import sequence_parallel as sp_ops +from vllm.models.kimi_k3.nvidia import model as kimi_model +from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp +from vllm.platforms import current_platform + + +class _IdentityNorm(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False) + self.variance_epsilon = 1e-5 + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + ): + if residual is None: + return hidden_states + return hidden_states, residual + + +class _RecordingMoE(nn.Module): + def __init__(self) -> None: + super().__init__() + self.num_tokens = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.num_tokens = hidden_states.shape[0] + return hidden_states + + +class _Projection(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.ones(1, hidden_size), + requires_grad=False, + ) + + +class _SequenceParallelMTPBlock: + use_sequence_parallel = True + + def __call__( + self, + *, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ): + assert residual is None + return hidden_states * 2, None, hidden_states * 3 + + +def _mock_sequence_parallel_collectives(monkeypatch): + monkeypatch.setattr( + kimi_model, + "sp_reduce_scatter", + lambda tensor: tensor.chunk(2, dim=0)[0], + ) + monkeypatch.setattr( + kimi_model, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_model, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + +@pytest.mark.parametrize( + ("num_tokens", "is_padding", "tp_rank", "expected"), + [ + (1, None, 0, [False]), + (1, None, 1, [True]), + (5, None, 2, [False, True]), + (5, None, 3, [True, True]), + (5, [False, True, False, False, False], 0, [False, True]), + ], +) +def test_sp_padding_mask_marks_added_rows( + monkeypatch, + num_tokens: int, + is_padding: list[bool] | None, + tp_rank: int, + expected: list[bool], +): + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) + + hidden_states = torch.empty(num_tokens, 2) + padding = torch.tensor(is_padding) if is_padding is not None else None + actual = sp_ops.sp_padding_mask(padding, hidden_states) + + torch.testing.assert_close(actual, torch.tensor(expected)) + + +@pytest.mark.parametrize( + ("data_parallel_size", "expected"), + [ + (1, False), + (2, True), + ], +) +def test_moe_sequence_parallel_requires_data_parallel( + monkeypatch, + data_parallel_size: int, + expected: bool, +): + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + parallel_config = ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=data_parallel_size, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ) + + assert parallel_config.use_sequence_parallel_moe is expected + + +def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = False + layer.use_sequence_parallel = True + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + + positions = torch.arange(3) + full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + hidden_states = kimi_model.sp_shard(full_hidden_states) + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = True + layer.use_sequence_parallel = True + layer.prev_valid_blocks = 0 + layer.block_write_idx = 0 + layer.is_block_write_layer = False + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.self_attention_res_norm = _IdentityNorm() + layer.mlp_res_norm = _IdentityNorm() + layer.self_attention_res_proj = _Projection() + layer.mlp_res_proj = _Projection() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + monkeypatch.setattr( + kimi_model, + "attn_res", + lambda prefix_sum, hidden_states, *args, **kwargs: ( + prefix_sum if hidden_states is None else prefix_sum + hidden_states + ), + ) + + prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2)) + block_residual = torch.zeros(2, 1, 2) + hidden_states, prefix_sum, block_residual = layer( + positions=torch.arange(3), + hidden_states=None, + prefix_sum=prefix_sum, + residual=block_residual, + ) + + assert hidden_states.shape == prefix_sum.shape == (2, 2) + assert block_residual.shape == (2, 1, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): + layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.enorm = _IdentityNorm() + layer.hnorm = _IdentityNorm() + layer.eh_proj = nn.Identity() + object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock()) + + final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1) + object.__setattr__( + layer, + "shared_head", + SimpleNamespace(norm=final_norm), + ) + + monkeypatch.setattr( + kimi_mtp, + "fused_mtp_input", + lambda positions, inputs_embeds, *args: inputs_embeds, + ) + monkeypatch.setattr( + kimi_mtp, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_mtp, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2) + logits_hidden_states, hidden_states = layer( + input_ids=torch.zeros(3, dtype=torch.long), + positions=torch.arange(3), + previous_hidden_states=torch.zeros_like(inputs_embeds), + inputs_embeds=inputs_embeds, + ) + + sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2] + expected_hidden_states = torch.cat( + [sharded_states * 5, sharded_states * 5], + dim=0, + )[:3] + torch.testing.assert_close(hidden_states, expected_hidden_states) + torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1) + final_norm.assert_called_once() + torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states) + + +@pytest.mark.parametrize( + ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"), + [ + (True, True, True, 8, True), + (False, True, True, 8, False), # opt-in only + (True, False, True, 8, False), # replication only exists under SP + (True, True, False, 8, False), # FusedMoE path owns the reduction + (True, True, True, 1, False), # nothing to shard + (True, True, True, 5, False), # 6144 % 5 -- would fail divide() + ], +) +def test_shard_sequence_parallel_mlp_gating( + monkeypatch, + enabled: bool, + use_sequence_parallel: bool, + eligible: bool, + tp_size: int, + expected: bool, +): + monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled) + monkeypatch.setattr( + kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + + assert ( + kimi_model.shard_sequence_parallel_mlp( + hidden_size=7168, + intermediate_size=6144, + use_sequence_parallel=use_sequence_parallel, + eligible=eligible, + ) + is expected + ) + + +def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config): + """Sharded SP MLP must reproduce the replicated result for every token. + + Each rank owns a *disjoint* token shard, so a weight shard alone cannot + finish a rank's own tokens: the ranks must gather the full token set, + compute partial sums over their intermediate shard, and reduce-scatter. + Splicing per-rank feature slices together instead silently mixes different + tokens and produces plausible-looking garbage. + """ + tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3 + torch.manual_seed(0) + num_tokens = tp_size * tokens_per_rank + x = torch.randn(num_tokens, hidden) + gate_weight = torch.randn(intermediate, hidden) + up_weight = torch.randn(intermediate, hidden) + down_weight = torch.randn(hidden, intermediate) + act_fn = SiluAndMul() + + replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T + + shard = intermediate // tp_size + # Every rank all-gathers the full token set, then computes its partial. + partials = [ + act_fn( + x + @ torch.cat( + [ + gate_weight[r * shard : (r + 1) * shard], + up_weight[r * shard : (r + 1) * shard], + ] + ).T + ) + @ down_weight[:, r * shard : (r + 1) * shard].T + for r in range(tp_size) + ] + reduced = torch.stack(partials).sum(0) + # Reduce-scatter: rank r keeps only its own token shard. + for r in range(tp_size): + mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank] + expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank] + torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5) + + +def test_sp_all_gather_uses_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + expected = torch.cat([hidden_states, hidden_states]) + custom_all_gather = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_all_gather=custom_all_gather, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback) + + output = sp_ops.sp_all_gather(hidden_states) + + torch.testing.assert_close(output, expected) + custom_all_gather.assert_called_once_with(hidden_states) + fallback.assert_not_called() + + +def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch): + hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + custom_reduce_scatter = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_reduce_scatter=custom_reduce_scatter, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback) + + output = sp_ops.sp_reduce_scatter(hidden_states) + + torch.testing.assert_close(output, expected) + padded = custom_reduce_scatter.call_args.args[0] + assert padded.shape == (4, 2) + torch.testing.assert_close(padded[:3], hidden_states) + torch.testing.assert_close(padded[3], torch.zeros(2)) + fallback.assert_not_called() + + +@pytest.mark.parametrize("shape", [(3,), (3, 2, 2)]) +def test_sp_shard_pads_only_the_token_axis(monkeypatch, shape): + hidden_states = torch.arange(math.prod(shape), dtype=torch.float32).view(shape) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: 1) + + output = sp_ops.sp_shard(hidden_states) + + padding = hidden_states.new_zeros((1, *shape[1:])) + expected = torch.cat([hidden_states, padding])[2:] + torch.testing.assert_close(output, expected) + + +def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=None), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + all_gather = Mock(return_value=hidden_states) + reduce_scatter = Mock(return_value=hidden_states) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather) + monkeypatch.setattr( + sp_ops, + "tensor_model_parallel_reduce_scatter", + reduce_scatter, + ) + + torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states) + torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states) + all_gather.assert_called_once_with(hidden_states, 0) + reduce_scatter.assert_called_once_with(hidden_states, 0) diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 18d06fe0fd3e..db6b74ee618e 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -128,14 +128,6 @@ def test_models( if use_rocm_aiter and (model in AITER_MODEL_LIST): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") - if model == "TitanML/tiny-mixtral": - # Untrained model: near-uniform logits make argmax sensitive to - # AITER's bfloat16 rounding error. Route the plain rms_norm and the - # fused MoE (whose near-uniform router logits flip expert selection - # under ~1 ULP drift) through the native kernels for this model. - # See ROCm/aiter#3806 for the tracking issue and minimal repro. - monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0") - monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0") elif use_rocm_aiter and model not in AITER_MODEL_LIST: # Skip model that are not using AITER tests. # When more AITER kernels are added, this list will not be diff --git a/tests/models/language/generation/test_gdn_sleep_wake.py b/tests/models/language/generation/test_gdn_sleep_wake.py new file mode 100644 index 000000000000..e61c274aeca9 --- /dev/null +++ b/tests/models/language/generation/test_gdn_sleep_wake.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for hybrid GDN/Mamba models under sleep -> wake. + +Hybrid Mamba / gated-delta-net (GDN) models (e.g. Qwen3-Next) keep a +persisted conv + recurrent state cache. With sleep mode (the RLHF reuse +pattern: ``sleep()`` -> weight update -> ``wake_up()``) the state-cache tag is +discarded on sleep and its device memory is re-created on wake. If a *new* +sequence's state slot is consumed before being reset, the gated-delta-rule +kernel faithfully propagates whatever is in that (now non-zeroed) memory; when +it contains NaN/inf the output becomes NaN and ``argmax`` collapses every token +to id 0 (which decodes to ``"!"``), giving ``reward=0`` / NaN log-probs in RL +training. + +This test sleeps and wakes a small hybrid GDN model and asserts that +post-wake generation is neither degenerate (single repeated token) nor NaN. +""" + +import pytest + +from vllm import LLM, SamplingParams + +# Small Qwen3-Next (GDN) model already used by the hybrid model test-suite. +MODEL = "tiny-random/qwen3-next-moe" + +PROMPTS = [ + "The capital of France is", + "Once upon a time,", + "1, 2, 3, 4,", + "Water is made of", +] + + +@pytest.mark.hybrid_model +def test_gdn_sleep_wake_no_stale_state(): + sampling_params = SamplingParams(temperature=0.0, max_tokens=32, logprobs=1) + + # Keep the reserved fraction low. On some (notably ROCm/amdgpu) drivers the + # VRAM discarded by ``sleep()`` is not returned to the free pool before + # ``wake_up()`` re-creates it, so the woken allocation must coexist with the + # not-yet-reclaimed one (~2x peak). A high ``gpu_memory_utilization`` then + # OOMs in ``cuMemCreate`` on wake. The model is tiny, so a small fraction + # still leaves ample KV/state cache while keeping the sleep/wake cycle well + # within device memory. + llm = LLM( + model=MODEL, + enable_sleep_mode=True, + enforce_eager=True, + max_model_len=1024, + gpu_memory_utilization=0.4, + trust_remote_code=True, + ) + + # Warm generation before sleeping. + llm.generate(PROMPTS, sampling_params) + + # Default sleep offloads weights and DISCARDS the kv / GDN state cache; + # wake_up re-creates that memory (fresh, not guaranteed zeroed). + llm.sleep() + llm.wake_up() + + after = llm.generate(PROMPTS, sampling_params) + + for output in after: + completion = output.outputs[0] + token_ids = list(completion.token_ids) + assert token_ids, "empty generation after wake_up" + # The bug collapses every token to a single id (e.g. 0 -> "!"). + assert len(set(token_ids)) > 1, ( + f"degenerate single-token output after wake_up: {token_ids[:16]}" + ) + # NaN logits surface as NaN log-probs. + for step_logprobs in completion.logprobs or []: + for logprob in step_logprobs.values(): + assert logprob.logprob == logprob.logprob, ( + "NaN log-prob after wake_up (stale GDN state)" + ) diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 3fee0662bc60..768b26c86881 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -35,7 +35,6 @@ HYBRID_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", "ibm-granite/granite-4.0-tiny-preview", "tiiuae/Falcon-H1-0.5B-Base", @@ -50,7 +49,6 @@ FULL_CUDA_GRAPH_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", ] diff --git a/tests/models/language/generation/test_mistral.py b/tests/models/language/generation/test_mistral.py index bc85d6f7220d..fc774bc6e835 100644 --- a/tests/models/language/generation/test_mistral.py +++ b/tests/models/language/generation/test_mistral.py @@ -5,12 +5,10 @@ import pytest +from vllm.parser.mistral import MistralToolCall from vllm.sampling_params import SamplingParams from vllm.tokenizers.mistral import MistralTokenizer -from vllm.tool_parsers.mistral_tool_parser import ( - MistralToolCall, - MistralToolParser, -) +from vllm.tool_parsers.mistral_tool_parser import MistralToolParser from ...utils import check_logprobs_close @@ -261,7 +259,7 @@ def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None: model_output = outputs[0].outputs[0].text.strip() assert model_output.startswith(tool_parser.bot_token), model_output - parsed_message = tool_parser.extract_tool_calls(model_output, None) + parsed_message = tool_parser.extract_tool_calls(model_output, None) # type: ignore[arg-type] assert parsed_message.tools_called @@ -304,7 +302,7 @@ def get_vocab(): model_output = f"{parser.bot_token}get_current_weather{json.dumps(args_dict)}" - parsed = parser.extract_tool_calls(model_output, None) + parsed = parser.extract_tool_calls(model_output, None) # type: ignore[arg-type] # Assertions: the tool call is detected and the full nested JSON is parsed # without truncation. @@ -337,7 +335,7 @@ def get_vocab(): ] ) - parsed = parser.extract_tool_calls(model_output, None) + parsed = parser.extract_tool_calls(model_output, None) # type: ignore[arg-type] # Assertions: the tool call is detected and the full nested JSON is parsed # without truncation. diff --git a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py b/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py index 251bc8c1d8a4..c987e92caca1 100644 --- a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py +++ b/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py @@ -15,7 +15,9 @@ ["Qwen/Qwen3-Embedding-0.6B"], ) @torch.inference_mode -def test_embed_models(hf_runner, vllm_runner, model: str): +def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str): + # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") chunk_size = 10 n_prompt_tokens = [55, 56, 57] token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens] @@ -104,3 +106,44 @@ def score_with(max_num_batched_tokens: int) -> float: assert chunked == pytest.approx(unchunked, abs=5e-2), ( f"chunked score {chunked} diverged from unchunked {unchunked}" ) + + +@torch.inference_mode +def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None: + model = "Qwen/Qwen3-Embedding-0.6B" + chunk_size = 10 + token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)] + prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts] + + with hf_runner(model, auto_cls=AutoModel) as hf_model: + hf_outputs = [] + for token_prompt in token_prompts: + inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])}) + output = hf_model.model(inputs["input_ids"]) + embedding = torch.nn.functional.normalize( + output.last_hidden_state.float()[0, -1], dim=0 + ) + hf_outputs.append(embedding.cpu().tolist()) + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + pooler_config=PoolerConfig(task="embed"), + max_model_len=64, + max_num_batched_tokens=chunk_size, + max_num_seqs=2, + gpu_memory_utilization=0.25, + enforce_eager=True, + enable_chunked_prefill=True, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_outputs = vllm_model.embed(prompts) + + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + tol=1e-2, + ) diff --git a/tests/models/language/pooling/test_classification.py b/tests/models/language/pooling/test_classification.py index e7128197bfc7..e15fc7467eed 100644 --- a/tests/models/language/pooling/test_classification.py +++ b/tests/models/language/pooling/test_classification.py @@ -49,3 +49,60 @@ def test_models( vllm_output, rtol=2e-3 if dtype == "float" else 1e-2, ) + + +@pytest.mark.core_model +def test_bert_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None: + model = "cross-encoder/ms-marco-TinyBERT-L-2-v2" + score_inputs = ( + "What is the capital of France?", + [ + "Paris.", + "Paris is the capital and largest city of France.", + "William Shakespeare wrote Hamlet in the early seventeenth century.", + ], + ) + prompt_batches = [ + ["short input"], + [ + "short input", + "a longer input that exercises mixed sequence lengths", + ], + ] + + with hf_runner( + model, dtype="half", auto_cls=AutoModelForSequenceClassification + ) as hf_model: + # HfRunner uses problem_type to preserve the model's + # sbert_ce_default_activation_function=Identity raw logits. + hf_model.config.problem_type = "regression" + hf_outputs = [hf_model.classify(prompts) for prompts in prompt_batches] + + text_1, text_2 = score_inputs + text_pairs = [[text_1, document] for document in text_2] + with hf_runner(model, dtype="half", is_cross_encoder=True) as hf_model: + hf_scores = hf_model.predict(text_pairs).tolist() + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + dtype="half", + max_model_len=64, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_outputs = [vllm_model.classify(prompts) for prompts in prompt_batches] + vllm_scores = vllm_model.score(*score_inputs) + + for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs): + hf_tensor = torch.tensor(hf_batch) + vllm_tensor = torch.tensor(vllm_batch) + assert vllm_tensor.shape == hf_tensor.shape + assert torch.allclose(vllm_tensor, hf_tensor, rtol=1e-2, atol=1e-4) + + assert torch.allclose( + torch.tensor(vllm_scores), + torch.tensor(hf_scores), + rtol=1e-2, + atol=1e-4, + ) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 6c82ad8a9ca0..3057e14060c6 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -126,10 +126,11 @@ def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device): def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device): """Download and return the ColBERT linear projection weight.""" - from huggingface_hub import hf_hub_download from safetensors.torch import load_file - path = hf_hub_download(model_name, filename=hf_spec["weights_file"]) + from vllm.transformers_utils.repo_utils import hf_api + + path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]) weights = load_file(path) return weights[hf_spec["weights_key"]].to(device) @@ -378,8 +379,16 @@ def test_colbert_embed_not_supported( vllm_model.embed([TEXTS_1[0]]) -@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys())) -def test_colbert_hf_comparison(vllm_runner, backend): +@pytest.mark.parametrize( + ("backend", "use_v2"), + [ + pytest.param("bert", True, id="bert-v2"), + pytest.param("modernbert", True, id="modernbert-v2"), + pytest.param("jina", False, id="jina-v1"), + pytest.param("lfm2", False, id="lfm2-v1"), + ], +) +def test_colbert_hf_comparison(vllm_runner, monkeypatch, backend, use_v2): """Test that vLLM ColBERT embeddings match HuggingFace for each backend.""" from transformers import AutoTokenizer @@ -392,6 +401,8 @@ def test_colbert_hf_comparison(vllm_runner, backend): assert isinstance(extra_kwargs, dict) test_texts = [TEXTS_1[0], TEXTS_2[0]] + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0") + with vllm_runner( model_name, runner="pooling", @@ -400,6 +411,7 @@ def test_colbert_hf_comparison(vllm_runner, backend): enforce_eager=True, **extra_kwargs, ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner == use_v2 vllm_outputs = vllm_model.token_embed(test_texts) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index e105195afe0f..7d48880e9c06 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/tests/models/language/pooling/test_embedding.py @@ -2,7 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch +from transformers import AutoModel +from vllm import PoolingParams from vllm.config import PoolerConfig from ...utils import check_embeddings_close @@ -87,3 +90,130 @@ def test_models( name_1="vllm", tol=1e-2, ) + + +@pytest.mark.parametrize( + "model", + [ + "BAAI/bge-base-en-v1.5", + "intfloat/multilingual-e5-small", + ], +) +@torch.inference_mode() +def test_encoder_only_model_runner_v2_attention( + hf_runner, + vllm_runner, + monkeypatch, + model: str, +) -> None: + prompts = [ + "short input", + "a longer input that exercises mixed sequence lengths", + ] + + with hf_runner(model, dtype="float", auto_cls=AutoModel) as hf_model: + hf_outputs = [] + for prompt in prompts: + inputs = hf_model.tokenizer(prompt, return_tensors="pt") + output = hf_model.model(**hf_model.wrap_device(inputs)) + embedding = torch.nn.functional.normalize( + output.last_hidden_state[0, -1].float(), dim=0 + ) + hf_outputs.append(embedding.cpu().tolist()) + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + dtype="float", + max_model_len=64, + max_num_seqs=2, + gpu_memory_utilization=0.25, + pooler_config=PoolerConfig( + task="embed", seq_pooling_type="LAST", use_activation=True + ), + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_outputs = vllm_model.embed(prompts) + + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + tol=1e-2, + ) + + +@pytest.mark.core_model +def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None: + model = "sentence-transformers/all-MiniLM-L6-v2" + prompt_batches = [ + ["short input"], + [ + "short input", + "a longer input that exercises mixed sequence lengths", + ], + ] + + with hf_runner(model, is_sentence_transformer=True) as hf_model: + hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches] + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + max_model_len=64, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches] + + for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs): + check_embeddings_close( + embeddings_0_lst=hf_batch, + embeddings_1_lst=vllm_batch, + name_0="hf", + name_1="vllm", + tol=1e-2, + ) + + +@pytest.mark.core_model +def test_matryoshka_dimensions_model_runner_v2( + hf_runner, vllm_runner, monkeypatch +) -> None: + model = "Snowflake/snowflake-arctic-embed-m-v1.5" + prompts = ["short input", "a longer input for a different output width"] + dimensions = [None, 256] + + with hf_runner(model, is_sentence_transformer=True) as hf_model: + hf_outputs = hf_model.encode(prompts) + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + max_model_len=64, + gpu_memory_utilization=0.25, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_outputs = vllm_model.embed( + prompts, + pooling_params=[PoolingParams(dimensions=d) for d in dimensions], + ) + + expected_outputs = [] + for output, dimension in zip(hf_outputs, dimensions): + output = torch.as_tensor(output) + if dimension is not None: + output = torch.nn.functional.normalize(output[:dimension], dim=0) + expected_outputs.append(output.tolist()) + + assert [len(output) for output in vllm_outputs] == [768, 256] + check_embeddings_close( + embeddings_0_lst=expected_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + tol=1e-2, + ) diff --git a/tests/models/language/pooling/test_jina_embeddings_v5.py b/tests/models/language/pooling/test_jina_embeddings_v5.py new file mode 100644 index 000000000000..fbfebcc5cedb --- /dev/null +++ b/tests/models/language/pooling/test_jina_embeddings_v5.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Backbone validation for Jina Embeddings V5. + +The V5 family ships two backbones under one `architectures` entry: `-small` is +a Qwen3 decoder, while `-nano` is a bidirectional EuroBERT encoder. Upstream +ships a separate `configuration_*.py` per repository, so the only signal +distinguishing them is `is_decoder`, which the encoder variant sets to False. +`JinaEmbeddingsV5ModelConfig` uses it to enable bidirectional attention for the +encoder variant; `JinaEmbeddingsV5Model` then dispatches to the correct backbone. +""" + +from types import SimpleNamespace +from typing import cast + +import pytest +from transformers import PretrainedConfig + +from vllm.config import ModelConfig +from vllm.model_executor.models.config import ( + MODELS_CONFIG_MAP, + JinaEmbeddingsV5ModelConfig, +) + + +def _model_config(hf_config: PretrainedConfig) -> ModelConfig: + """Minimal stand-in for ModelConfig; only hf_config is read.""" + return cast(ModelConfig, SimpleNamespace(hf_config=hf_config)) + + +@pytest.mark.cpu_test +def test_registered_for_the_architecture(): + """The handler only runs if it is wired to the architecture name.""" + assert MODELS_CONFIG_MAP["JinaEmbeddingsV5Model"] is JinaEmbeddingsV5ModelConfig + + +@pytest.mark.cpu_test +def test_encoder_backbone_enables_bidirectional_attention(): + """An encoder checkpoint (is_decoder=False) is supported. + + The handler sets is_causal=False so the Llama backbone uses + EncoderOnlyAttention; JinaEmbeddingsV5Model then dispatches to the encoder + implementation. + """ + hf_config = PretrainedConfig(is_decoder=False) + + JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(_model_config(hf_config)) + + assert hf_config.is_causal is False + + +@pytest.mark.cpu_test +def test_supported_decoder_backbone_is_accepted(): + """The Qwen3-based variants must keep loading. + + `-small` omits `is_decoder` entirely, so an absent attribute has to be + treated as a decoder. The first assertion pins that assumption: if + PretrainedConfig ever gains an `is_decoder=False` default, this fails here + rather than silently rejecting a supported checkpoint. + """ + absent = PretrainedConfig() + assert not hasattr(absent, "is_decoder") + + JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(_model_config(absent)) + JinaEmbeddingsV5ModelConfig.verify_and_update_model_config( + _model_config(PretrainedConfig(is_decoder=True)) + ) diff --git a/tests/models/language/pooling/test_reward.py b/tests/models/language/pooling/test_reward.py index 1872ca4ae09b..f981266d0c0b 100644 --- a/tests/models/language/pooling/test_reward.py +++ b/tests/models/language/pooling/test_reward.py @@ -9,6 +9,7 @@ from transformers import AutoModel from vllm.platforms import current_platform +from vllm.utils.mem_constants import MiB_bytes from ....conftest import HfRunner from ....utils import VLLM_PATH @@ -106,7 +107,12 @@ def test_prm_models( if current_platform.is_cpu(): pytest.skip("CPU only supports V1") - with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model: + with vllm_runner( + model, + max_model_len=1024, + dtype=dtype, + kv_cache_memory_bytes=64 * MiB_bytes, + ) as vllm_model: vllm_outputs = vllm_model.token_classify(math_step_prompts) with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model: @@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs( if not FIXTURE_REWARD_RESULT.get(model): pytest.skip(f"No available golden outputs for {model}.") - with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model: + with vllm_runner( + model, + max_model_len=1024, + dtype=dtype, + kv_cache_memory_bytes=64 * MiB_bytes, + ) as vllm_model: vllm_outputs = vllm_model.token_classify(math_step_prompts) golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model]) diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index 38a90d07abeb..2fbeadacb91e 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/tests/models/language/pooling/test_splade_sparse_pooler.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + +import numpy as np import pytest import torch import torch.nn as nn @@ -9,8 +12,15 @@ BertMLMHead, SPLADESparsePooler, ) +from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams +from vllm.utils.torch_utils import PIN_MEMORY +from vllm.v1.pool.late_interaction_runner import LateInteractionRunner from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.model_runner import GPUModelRunner +from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner +from vllm.v1.worker.gpu.states import RequestState # --------------------------------------------------------------------- # Functional test: SPLADE formula correctness (no HF download needed) @@ -91,3 +101,191 @@ def ref_one(hs: torch.Tensor, L: int, tid_row: torch.Tensor) -> torch.Tensor: rtol=1e-4, atol=1e-4, ) + + +def test_pooling_runner_gathers_required_token_ids() -> None: + runner = PoolingRunner.__new__(PoolingRunner) + pooling_params = PoolingParams(task="embed", requires_token_ids=True) + runner.pooling_params = {1: pooling_params, 3: pooling_params} + runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()} + runner.prompt_token_ids = { + 1: torch.tensor([101, 102]), + 3: torch.tensor([101, 11, 102]), + } + + input_batch = MagicMock(spec=InputBatch) + input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32) + input_batch.num_reqs = 2 + req_states = MagicMock(spec=RequestState) + req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32)) + metadata = runner._get_pooling_metadata( + input_batch, req_states, torch.device(current_platform.device_type) + ) + + expected = torch.tensor([[101, 11, 102], [101, 102, 0]]) + assert metadata.prompt_token_ids_cpu is not None + assert metadata.prompt_token_ids is not None + assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY + torch.testing.assert_close( + metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32) + ) + torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected) + torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected) + + +def test_pooling_runner_stores_only_required_token_ids() -> None: + runner = PoolingRunner.__new__(PoolingRunner) + runner.model = MagicMock() + runner.supported_tasks = frozenset({"embed"}) + runner.pooling_params = {} + runner.pooling_states = {} + runner.prompt_token_ids = {} + runner.late_interaction_runner = MagicMock() + + runner.add_request("req-1", 1, PoolingParams(task="embed"), [101, 102]) + runner.add_request( + "req-2", + 2, + PoolingParams(task="embed", requires_token_ids=True), + [101, 11, 102], + ) + + assert 1 not in runner.prompt_token_ids + torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102])) + + +def test_pooling_runner_releases_aborted_late_interaction_doc() -> None: + runner = PoolingRunner.__new__(PoolingRunner) + runner.late_interaction_runner = LateInteractionRunner() + + query_key = "query-abort" + late_interaction_runner = runner.late_interaction_runner + late_interaction_runner._query_cache[query_key] = torch.ones(2, 4) + late_interaction_runner._query_uses[query_key] = 1 + late_interaction_runner._doc_query_keys["doc-req"] = query_key + + runner.on_requests_finished({"doc-req"}) + + assert not late_interaction_runner._query_cache + assert not late_interaction_runner._query_uses + assert not late_interaction_runner._doc_query_keys + + +def test_encoder_cache_reset_clears_late_interaction_state() -> None: + # Cached query embeddings are only invalidated by weight reloads, which + # reach the runner via reset_encoder_cache. Resetting the multi-modal + # cache is unrelated and must leave them intact. + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.encoder_cache = MagicMock() + runner.pooling_runner = MagicMock() + + runner.reset_mm_cache() + + runner.encoder_cache.reset_mm_cache.assert_called_once_with() + runner.pooling_runner.clear.assert_not_called() + + runner.reset_encoder_cache() + + runner.encoder_cache.reset_encoder_cache.assert_called_once_with() + runner.pooling_runner.clear.assert_called_once_with() + + +def test_pooling_runner_rejects_unsupported_selected_task() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"embed", "plugin"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "encoder_only" + vllm_config.model_config.get_pooling_task.return_value = "plugin" + + with ( + patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]), + pytest.raises(ValueError, match="selects 'plugin'"), + ): + PoolingRunner(model, vllm_config) + + +def test_pooling_runner_supports_encoder_token_classification() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"token_classify"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "encoder_only" + vllm_config.model_config.get_pooling_task.return_value = "token_classify" + + runner = PoolingRunner(model, vllm_config) + + assert runner.supported_tasks == {"token_classify"} + + +def test_pooling_runner_supports_encoder_token_embedding() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"token_embed"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "encoder_only" + vllm_config.model_config.get_pooling_task.return_value = "token_embed" + + runner = PoolingRunner(model, vllm_config) + + assert runner.supported_tasks == {"token_embed"} + + +def test_pooling_runner_rejects_decoder_token_classification() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"token_classify"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "decoder" + vllm_config.model_config.get_pooling_task.return_value = "token_classify" + + with pytest.raises(ValueError, match="selects 'token_classify'") as exc_info: + PoolingRunner(model, vllm_config) + + assert "Set an explicitly supported task" not in str(exc_info.value) + + +def test_pooling_runner_filters_decoder_token_classification() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"embed", "token_classify"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "decoder" + vllm_config.model_config.get_pooling_task.return_value = "embed" + + runner = PoolingRunner(model, vllm_config) + + assert runner.supported_tasks == {"embed"} + + +def test_pooling_runner_filters_decoder_token_embedding() -> None: + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"embed", "token_embed"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + vllm_config.model_config.attn_type = "decoder" + vllm_config.model_config.get_pooling_task.return_value = "embed" + + runner = PoolingRunner(model, vllm_config) + + assert runner.supported_tasks == {"embed"} + + +def test_pooling_runner_gates_embed_token_classification_to_encoder() -> None: + # BGE-M3's combined task emits per-token weights alongside the embedding, + # so it is enabled only where prefill is unchunked. + model = MagicMock() + model.pooler.get_supported_tasks.return_value = {"embed", "embed&token_classify"} + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = 2 + + vllm_config.model_config.attn_type = "encoder_only" + vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify" + assert PoolingRunner(model, vllm_config).supported_tasks == { + "embed", + "embed&token_classify", + } + + vllm_config.model_config.attn_type = "decoder" + vllm_config.model_config.get_pooling_task.return_value = "embed" + assert PoolingRunner(model, vllm_config).supported_tasks == {"embed"} diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 2e95fb8f70a4..5b18bdf1a54b 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -29,16 +29,24 @@ def seed_everything(): ) # The float32 is required for this tiny model to pass the test. @pytest.mark.parametrize("dtype", ["float"]) +@pytest.mark.core_model @torch.inference_mode -def test_bert_models( +def test_bert_model_runner_v2( hf_runner, vllm_runner, example_prompts, + monkeypatch, model: str, dtype: str, ) -> None: + prompt_batches = [[example_prompts[0]], example_prompts] + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model: - vllm_outputs = vllm_model.token_classify(example_prompts) + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + vllm_output_batches = [ + vllm_model.token_classify(prompts) for prompts in prompt_batches + ] # Use eager attention on ROCm to avoid HF Transformers flash attention # accuracy issues: https://github.com/vllm-project/vllm/issues/30167 @@ -59,12 +67,14 @@ def test_bert_models( inputs = hf_model.wrap_device(inputs) output = hf_model.model(**inputs) hf_outputs.append(softmax(output.logits[0])) + hf_output_batches = [[hf_outputs[0]], hf_outputs] # check logits difference - for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): - hf_output = hf_output.detach().clone().cpu().float() - vllm_output = vllm_output.detach().clone().cpu().float() - torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) + for hf_outputs, vllm_outputs in zip(hf_output_batches, vllm_output_batches): + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): + hf_output = hf_output.detach().clone().cpu().float() + vllm_output = vllm_output.detach().clone().cpu().float() + torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) @pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"]) diff --git a/tests/models/language/pooling/test_truncation_control.py b/tests/models/language/pooling/test_truncation_control.py index d41a3379dc0f..50e8cdbd064c 100644 --- a/tests/models/language/pooling/test_truncation_control.py +++ b/tests/models/language/pooling/test_truncation_control.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from vllm.exceptions import VLLMValidationError + MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2" max_model_len = 128 @@ -60,7 +62,7 @@ def test_bigger_truncation_size( truncate_prompt_tokens = max_model_len + 1 with ( - pytest.raises(ValueError), + pytest.raises(VLLMValidationError), vllm_runner( model_name, runner="pooling", max_model_len=max_model_len ) as vllm_model, diff --git a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py index fc575c399d04..101a0a4690ce 100644 --- a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py +++ b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py @@ -207,9 +207,18 @@ def mteb_test_embed_models( vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype head_dtype = model_config.head_dtype - # Test embedding_size, isnan and whether to use normalize + # Test embedding_size, isnan and whether to use normalize. + # Apply the same prompt_prefix used for scoring so this check compares + # identical effective inputs: HF's SentenceTransformer.encode() applies + # the model's default prompt (e.g. "Document: "), so vLLM must receive it + # too. This mirrors VllmMtebEncoder and changes no production behavior. + consistency_prompts = ( + [prompt_prefix + p for p in example_prompts] + if prompt_prefix + else example_prompts + ) vllm_outputs = vllm_model.embed( - example_prompts, + consistency_prompts, tokenization_kwargs=dict(truncate_prompt_tokens=-1), ) outputs_tensor = torch.tensor(vllm_outputs) diff --git a/tests/models/language/pooling_mteb_test/test_jina.py b/tests/models/language/pooling_mteb_test/test_jina.py index 24aa3188f8be..fd6fd96f14a5 100644 --- a/tests/models/language/pooling_mteb_test/test_jina.py +++ b/tests/models/language/pooling_mteb_test/test_jina.py @@ -38,6 +38,14 @@ is_prefix_caching_supported=True, is_chunked_prefill_supported=True, ), + EmbedModelInfo( + "jinaai/jina-embeddings-v5-text-nano", + architecture="JinaEmbeddingsV5Model", + seq_pooling_type="LAST", + attn_type="encoder_only", + is_prefix_caching_supported=False, + is_chunked_prefill_supported=False, + ), ] RERANK_MODELS = [ diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 8583c1d21369..476cc44d9c63 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -234,7 +234,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): image_size_factors=[(0.25, 0.5, 1.0)], vllm_runner_kwargs={ "model_impl": "transformers", - "default_torch_num_threads": 1, }, marks=[pytest.mark.core_model], ), @@ -908,7 +907,15 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): multi_image_prompt="Picture 1: \nPicture 2: \nDescribe these two images with one paragraph respectively.", # noqa: E501 max_model_len=4096, max_num_seqs=2, - num_logprobs=10, + # torch 2.13 accumulates CPU numerical drift in the qwen2_vl multi-image + # path: HF and vLLM agree for a long prefix (~69 tokens) then a token + # flips outside vLLM's top-N only near the end of the generation. The + # window is already at the max_logprobs=20 cap, so widening it further is + # not possible. Treat this as acceptable drift and cap max_tokens on CPU + # so the compared prefix stays before the divergence, keeping the + # multi-image path under test. See pytorch/pytorch#187735. + max_tokens=64 if current_platform.is_cpu() else 128, + num_logprobs=20 if current_platform.is_cpu() else 10, auto_cls=AutoModelForImageTextToText, vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output, image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], @@ -1234,7 +1241,7 @@ def test_custom_inputs_models( create_new_process_for_each_test=True, ), ) -@create_new_process_for_each_test() +@create_new_process_for_each_test("spawn") def test_single_image_models_heavy( tmp_path: PosixPath, model_type: str, diff --git a/tests/models/multimodal/generation/test_granite_speech.py b/tests/models/multimodal/generation/test_granite_speech.py index 3019f5f22d4b..e1fee635ff66 100644 --- a/tests/models/multimodal/generation/test_granite_speech.py +++ b/tests/models/multimodal/generation/test_granite_speech.py @@ -45,9 +45,10 @@ def vllm_to_hf_output( def granite_speech_attention_config(): """Return attention config for Granite Speech tests on ROCm.""" if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if on_mi3xx(): + # -1 (unknown arch) is truthy; gate on CDNA3+ like other call sites. + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_FA"} return {"backend": "TRITON_ATTN"} return None diff --git a/tests/models/multimodal/generation/test_phi4mm.py b/tests/models/multimodal/generation/test_phi4mm.py index 1a4fb35a28aa..5ab75e145aee 100644 --- a/tests/models/multimodal/generation/test_phi4mm.py +++ b/tests/models/multimodal/generation/test_phi4mm.py @@ -6,7 +6,6 @@ import pytest import regex as re -from huggingface_hub import snapshot_download from transformers import AutoTokenizer from vllm.assets.image import ImageAsset @@ -14,6 +13,7 @@ from vllm.lora.request import LoRARequest from vllm.multimodal.image import convert_image_mode, rescale_image_size from vllm.multimodal.media.audio import load_audio +from vllm.transformers_utils.repo_utils import hf_api from ....conftest import ( IMAGE_ASSETS, @@ -35,7 +35,7 @@ "<|user|>\n<|image_1|>\n<|image_2|>\nDescribe these images.<|end|>\n<|assistant|>\n" # noqa: E501 ) -model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct") +model_path = hf_api().snapshot_download("microsoft/Phi-4-multimodal-instruct") # Since the vision-lora and speech-lora co-exist with the base model, # we have to manually specify the path of the lora weights. vision_lora_path = os.path.join(model_path, "vision-lora") diff --git a/tests/models/multimodal/generation/test_transformers_audio.py b/tests/models/multimodal/generation/test_transformers_audio.py new file mode 100644 index 000000000000..919557987796 --- /dev/null +++ b/tests/models/multimodal/generation/test_transformers_audio.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +from transformers import AutoModelForSeq2SeqLM + +from vllm.assets.audio import AudioAsset +from vllm.envs import disable_envs_cache +from vllm.lora.request import LoRARequest +from vllm.multimodal.audio import AudioResampler + +from ....conftest import HfRunner, VllmRunner +from ...utils import check_logprobs_close + +AUDIO_ASSET = AudioAsset("mary_had_lamb") + +AUDIO_MODEL_SETTINGS: dict[str, dict[str, Any]] = { + "ibm-granite/granite-speech-3.3-2b": { + "prompt": ( + "<|start_of_role|>system<|end_of_role|>" + "You are a helpful AI assistant<|end_of_text|>\n" + "<|start_of_role|>user<|end_of_role|>" + "<|audio|>can you transcribe the speech into a written format?" + "<|end_of_text|>\n" + "<|start_of_role|>assistant<|end_of_role|>" + ), + "audio_lora_path": "ibm-granite/granite-speech-3.3-2b", + }, + "nvidia/audio-flamingo-3-hf": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "Transcribe the input speech.<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "vllm_runner_kwargs": { + "gpu_memory_utilization": 0.85, + }, + }, + "microsoft/VibeVoice-ASR-HF": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant that transcribes audio input " + "into text output in JSON format.<|im_end|>\n" + "<|im_start|>user\n" + "<|object_ref_start|><|box_start|><|object_ref_end|>\n" + "This is a 1.0 seconds audio, please transcribe it with " + "these keys: Start time, End time, Speaker ID, Content" + "<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "sampling_rate": 24000, + "vllm_runner_kwargs": { + "max_num_batched_tokens": 2048, + "gpu_memory_utilization": 0.85, + }, + }, + "zai-org/GLM-ASR-Nano-2512": { + "prompt": ( + "<|user|>\n" + "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n" + "Please transcribe this audio into text" + "<|assistant|>\n" + ), + }, +} + + +@pytest.mark.parametrize("model_id", list(AUDIO_MODEL_SETTINGS)) +def test_transformers_audio_generation( + hf_runner: type[HfRunner], + vllm_runner: type[VllmRunner], + monkeypatch, + model_id: str, +): + """Single-process workaround for V1 fork safety deadlock issue + (vllm-project/vllm/issues/17676). Running multiple audio models together + under pytest can cause (possibly flaky) hangs, so they are grouped under + the same config. Using VLLM_WORKER_MULTIPROC_METHOD=spawn avoids the + deadlock and allows worker processes to terminate cleanly, and release + GPU memory between test runs until the issue is fixed.""" + # TODO: Remove monkeypatch once + # https://github.com/vllm-project/vllm/issues/17676 is fixed. + disable_envs_cache() + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + settings = AUDIO_MODEL_SETTINGS[model_id] + audio_lora_path = settings.get("audio_lora_path") + + audio, orig_sr = AUDIO_ASSET.audio_and_sample_rate + target_sr = settings.get("sampling_rate", orig_sr) + if orig_sr != target_sr: + audio = AudioResampler(target_sr=target_sr).resample(audio, orig_sr=orig_sr) + audio = (audio, target_sr) + + with vllm_runner( + model_id, + model_impl="transformers", + dtype="bfloat16", + max_model_len=2048, + enforce_eager=True, + limit_mm_per_prompt={"audio": 1}, + enable_lora=audio_lora_path is not None, + max_lora_rank=64, + **settings.get("vllm_runner_kwargs", {}), + ) as vllm_model: + lora_request = ( + LoRARequest("audio", 1, audio_lora_path) if audio_lora_path else None + ) + vllm_outputs = vllm_model.generate_greedy_logprobs( + [settings["prompt"]], + 128, + num_logprobs=10, + audios=[audio], + lora_request=lora_request, + ) + + with hf_runner( + model_id, dtype="bfloat16", auto_cls=AutoModelForSeq2SeqLM + ) as hf_model: + hf_outputs = hf_model.generate_greedy_logprobs_limit( + [settings["prompt"]], 128, num_logprobs=10, audios=[audio] + ) + + check_logprobs_close( + outputs_0_lst=hf_outputs, + outputs_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 387046f3e742..13fbe26880df 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -249,6 +249,19 @@ def gemma3_chat_template(content: str) -> str: }, skip=True, # TODO: Re-enable this once OOM issues are resolved on CI. ), + "gemma4": VitCudagraphTestConfig( + model="google/gemma-4-E2B-it", + image_prompt=( + "user\n<|image|>\nWhat is in this image?\n" + "model\n" + ), + video_prompt=( + "user\n<|video|>\nDescribe this video in one sentence." + "\nmodel\n" + ), + needs_video_metadata=True, + marks=[pytest.mark.core_model], + ), } diff --git a/tests/models/multimodal/pooling/test_colqwen3_5.py b/tests/models/multimodal/pooling/test_colqwen3_5.py index 3513bd025b7e..43914d819b8a 100644 --- a/tests/models/multimodal/pooling/test_colqwen3_5.py +++ b/tests/models/multimodal/pooling/test_colqwen3_5.py @@ -7,6 +7,8 @@ embeddings for both text and image inputs. """ +from types import SimpleNamespace + import pytest import torch @@ -154,12 +156,14 @@ def test_colqwen3_5_relevance_ordering( _run_relevance_test(vllm_runner, model, dtype=dtype) -def test_colqwen3_5_config_enables_bidirectional_attention() -> None: - """ColQwen3.5 retrieval must be served BIDIRECTIONAL (is_causal=False) so the - full_attention layers build with AttentionType.ENCODER_ONLY. This guards the - silent-causal regression (no GPU / model load needed).""" - from types import SimpleNamespace - +@pytest.mark.parametrize( + ("contract", "expected_is_causal"), + [("causal", True), ("bidirectional", False)], +) +def test_colqwen3_5_config_applies_declared_attention_contract( + contract: str, + expected_is_causal: bool, +) -> None: from vllm.model_executor.models.config import ( MODELS_CONFIG_MAP, ColQwen3_5Config, @@ -167,6 +171,97 @@ def test_colqwen3_5_config_enables_bidirectional_attention() -> None: assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config - model_config = SimpleNamespace(hf_config=SimpleNamespace()) + hf_config = SimpleNamespace(retrieval_attention_contract=contract) + text_config = SimpleNamespace() + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) + ColQwen3_5Config.verify_and_update_model_config(model_config) + assert hf_config.is_causal is expected_is_causal + assert text_config.is_causal is expected_is_causal + + +@pytest.mark.parametrize( + "hf_config", + [ + SimpleNamespace(), + SimpleNamespace(retrieval_attention_contract="unsupported"), + SimpleNamespace( + retrieval_attention_contract="causal", + text_config=SimpleNamespace(retrieval_attention_contract="bidirectional"), + ), + ], +) +def test_colqwen3_5_config_rejects_invalid_attention_contract(hf_config) -> None: + from vllm.model_executor.models.config import ColQwen3_5Config + + text_config = getattr(hf_config, "text_config", SimpleNamespace()) + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) + with pytest.raises(ValueError, match="retrieval_attention_contract"): + ColQwen3_5Config.verify_and_update_model_config(model_config) + + +def test_colqwen3_5_bidirectional_contract_builds_encoder_only_attention( + monkeypatch, +) -> None: + from vllm.model_executor.models import qwen3_next + from vllm.model_executor.models.config import ColQwen3_5Config + from vllm.v1.attention.backend import AttentionType + + hf_config = SimpleNamespace(retrieval_attention_contract="bidirectional") + text_config = SimpleNamespace( + hidden_size=256, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + max_position_embeddings=4096, + rope_parameters={}, + rms_norm_eps=1e-6, + ) + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) ColQwen3_5Config.verify_and_update_model_config(model_config) - assert model_config.hf_config.is_causal is False + + captured = {} + + class FakeAttention(torch.nn.Module): + def __init__(self, *args, **kwargs) -> None: + super().__init__() + captured["attn_type"] = kwargs["attn_type"] + + monkeypatch.setattr(qwen3_next, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr( + qwen3_next, "QKVParallelLinear", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr( + qwen3_next, "RowParallelLinear", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr( + qwen3_next, + "get_rope", + lambda *args, **kwargs: SimpleNamespace(is_neox_style=False), + ) + monkeypatch.setattr( + qwen3_next, "Qwen3NextRMSNorm", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr(qwen3_next, "Attention", FakeAttention) + + qwen3_next.Qwen3NextAttention(text_config) + + assert captured["attn_type"] is AttentionType.ENCODER_ONLY + + +def test_colqwen3_5_encoder_only_attention_has_no_kv_cache_spec() -> None: + from vllm.model_executor.layers.attention import Attention + from vllm.v1.attention.backend import AttentionType + + attention = SimpleNamespace(attn_type=AttentionType.ENCODER_ONLY) + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + + assert Attention.get_kv_cache_spec(attention, vllm_config) is None diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index d7b67b8bdb6a..72d2a7018bf9 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -3,11 +3,11 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -31,7 +31,7 @@ def run_intern_vit_test( *, dtype: str, ): - model = snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) + model = hf_api().snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype] img_processor = CLIPImageProcessor.from_pretrained(model) diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py index fcab077fbba8..7c039ab8b9c6 100644 --- a/tests/models/multimodal/pooling/test_radio.py +++ b/tests/models/multimodal/pooling/test_radio.py @@ -3,13 +3,13 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.models.radio import RadioModel from vllm.platforms import current_platform from vllm.transformers_utils.configs.radio import RadioConfig +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -28,7 +28,7 @@ def run_radio_test( *, dtype: str, ): - model = snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) + model = hf_api().snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype] img_processor = CLIPImageProcessor.from_pretrained(model) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 1ef39cfaa7b9..6d31e87003d1 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -87,7 +87,9 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict: _XPU_EXCLUDED_MODEL_IDS = { "baidu/Unlimited-OCR", "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4", + "moonshotai/Kimi-K3", "Qwen/Qwen2.5-Omni-7B-AWQ", + "thinkingmachines/Inkling-NVFP4", } @@ -452,6 +454,14 @@ def test_processing_correctness( "audio placeholders from processed audio lengths. Its vLLM " "processor paths are covered by test_moss_audio.py." ) + # TODO: Remove when transformers 5.15.0 is released, which contains + # https://github.com/huggingface/transformers/pull/47483. + if model_id == "microsoft/VibeVoice-ASR-HF": + pytest.skip( + "VibeVoice ASR requires audio as a positional argument and hence " + "cannot pass the processing correctness test as is. Its generation " + "is covered by test_transformers_audio.py." + ) if model_id == "lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct": pytest.skip( "LLaVA-OneVision-2 video processing routes frames through custom " diff --git a/tests/models/multimodal/processing/test_cosmos3_edge.py b/tests/models/multimodal/processing/test_cosmos3_edge.py index 5ebef08c1020..226868c9e7a4 100644 --- a/tests/models/multimodal/processing/test_cosmos3_edge.py +++ b/tests/models/multimodal/processing/test_cosmos3_edge.py @@ -70,7 +70,19 @@ def _assert_video_outputs(processor, processed) -> None: merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size expected_tokens = int(grid_thw.prod()) // merge_size**2 video_token_id = processor.info.get_hf_config().video_token_id - assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens + prompt_token_ids = processed["prompt_token_ids"] + assert prompt_token_ids.count(video_token_id) == expected_tokens + + hf_processor = processor.info.get_hf_processor() + expected_frame_wrappers = int(grid_thw[:, 0].sum()) + assert ( + prompt_token_ids.count(hf_processor.vision_start_token_id) + == expected_frame_wrappers + ) + assert ( + prompt_token_ids.count(hf_processor.vision_end_token_id) + == expected_frame_wrappers + ) @pytest.mark.parametrize("num_images", [1, 2]) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index a355501fdd80..f30afe47dde2 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -7,6 +7,7 @@ import torch from PIL import Image as PILImage +from vllm.exceptions import VLLMValidationError from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration, Gemma4ImagePixelInputs, @@ -222,7 +223,7 @@ def test_limit_mm_per_prompt( mm_data = {"image": images} # Expect ValueError when exceeding limit - with pytest.raises(ValueError, match="At most 1 image"): + with pytest.raises(VLLMValidationError, match="At most 1 image"): processor( prompt, mm_items=processor.info.parse_mm_data(mm_data), diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 473ba729b85b..67a81ddb7b41 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -7,6 +7,7 @@ import torch from PIL import Image as PILImage +from vllm.exceptions import VLLMValidationError from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -197,7 +198,7 @@ def test_limit_mm_per_prompt( mm_data = {"image": images} - with pytest.raises(ValueError, match="At most 1 image"): + with pytest.raises(VLLMValidationError, match="At most 1 image"): processor( prompt, mm_items=processor.info.parse_mm_data(mm_data), diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index 0cafa261a345..e45e741c5b47 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -63,6 +63,31 @@ def test_encoder_cudagraph_uses_model_video_frame_limit(): assert Glm4vForConditionalGeneration.get_max_frames_per_video(model) == 600 +@pytest.mark.parametrize( + ("temporal_patch_size", "expected_grid_t"), + [(2, 9), (4, 5), (8, 3)], +) +def test_vision_info_rounds_up_temporal_frames( + temporal_patch_size: int, + expected_grid_t: int, +): + info = Mock(spec=Glm4vProcessingInfo) + vision_config = info.get_hf_config.return_value.vision_config + vision_config.patch_size = 14 + vision_config.spatial_merge_size = 2 + vision_config.temporal_patch_size = temporal_patch_size + + _, num_vision_tokens = Glm4vProcessingInfo._get_vision_info( + info, + image_width=28, + image_height=28, + num_frames=17, + do_resize=False, + ) + + assert num_vision_tokens == expected_grid_t + + @pytest.mark.parametrize("model_id", ["zai-org/GLM-4.1V-9B-Thinking"]) @pytest.mark.parametrize("expected_toks_per_frame", [299]) @pytest.mark.parametrize( @@ -93,7 +118,6 @@ def test_processor_override( limit_mm_per_prompt={"video": 1}, ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - tokenizer = processor.info.get_tokenizer() hf_processor_mm_kwargs = {"fps": fps} # Build the image str / prompt based on the number of images we pass @@ -112,8 +136,8 @@ def test_processor_override( # Ensure we have the right number of placeholders per num_crops size hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs) - video_token_id = tokenizer.convert_tokens_to_ids(hf_processor.video_token) - video_tok_count = processed_inputs["prompt_token_ids"].count(video_token_id) + image_token_id = hf_processor.image_token_id + video_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id) grid_t, _, _ = processed_inputs["mm_kwargs"].get_data()["video_grid_thw"][0] assert grid_t == expected_grid_t diff --git a/tests/models/multimodal/processing/test_moss_audio.py b/tests/models/multimodal/processing/test_moss_audio.py index 6a18f6364284..8a18d90f560c 100644 --- a/tests/models/multimodal/processing/test_moss_audio.py +++ b/tests/models/multimodal/processing/test_moss_audio.py @@ -50,6 +50,7 @@ def batch_decode(self, batch_token_ids, **kwargs): class _MMConfig: enable_mm_embeds = False mm_processor_cache_gb = 1 + mm_hasher_algorithm = "blake3" def merge_mm_processor_kwargs(self, kwargs): return dict(kwargs) diff --git a/tests/models/multimodal/processing/test_moss_transcribe_diarize.py b/tests/models/multimodal/processing/test_moss_transcribe_diarize.py new file mode 100644 index 000000000000..954890d10364 --- /dev/null +++ b/tests/models/multimodal/processing/test_moss_transcribe_diarize.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.models.moss_transcribe_diarize import ( + MossTranscribeDiarizeForConditionalGeneration, +) + + +def test_parse_diarized_transcript_preserves_moss_segments(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0.48][S01]Welcome[1.66][12.26][S02]Ready[13.81]" + ) + + assert [ + (segment.start, segment.end, segment.speaker, segment.text) + for segment in segments + ] == [ + (0.48, 1.66, "S01", "Welcome"), + (12.26, 13.81, "S02", "Ready"), + ] + + +def test_parse_diarized_transcript_preserves_overlapping_segments(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01]First speaker[2][1][S02]Second speaker[3]" + ) + + assert [ + (segment.start, segment.end, segment.speaker, segment.text) + for segment in segments + ] == [ + (0.0, 2.0, "S01", "First speaker"), + (1.0, 3.0, "S02", "Second speaker"), + ] + + +def test_parse_diarized_transcript_preserves_numeric_text_markers(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01]The [2024] report is ready.[4]" + ) + + assert [segment.text for segment in segments] == ["The [2024] report is ready."] + + +def test_parse_diarized_transcript_ignores_whitespace_between_segments(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01]Hello[1]\n [2][S02]Hi[3]" + ) + + assert [(segment.start, segment.end, segment.text) for segment in segments] == [ + (0.0, 1.0, "Hello"), + (2.0, 3.0, "Hi"), + ] + + +def test_parse_diarized_transcript_ignores_noise_before_a_segment(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "noise [bad][0.1][S01]Hello[0.9]" + ) + + assert [(segment.start, segment.end, segment.text) for segment in segments] == [ + (0.1, 0.9, "Hello"), + ] + + +def test_parse_diarized_transcript_preserves_timestamps_before_the_end(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[2][S01]The earlier timestamp is [1] not the end[3]" + ) + + assert [segment.text for segment in segments] == [ + "The earlier timestamp is [1] not the end", + ] + + +def test_parse_diarized_transcript_skips_empty_segments(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01][1][2][S02]Complete[3]" + ) + + assert [(segment.speaker, segment.text) for segment in segments] == [ + ("S02", "Complete"), + ] + + +def test_parse_diarized_transcript_fails_closed_for_incomplete_output(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01]Complete[1][2][S02]Incomplete" + ) + + assert segments == [] + + +def test_parse_diarized_transcript_fails_closed_for_trailing_text(): + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + "[0][S01]Complete[1] trailing text" + ) + + assert segments == [] + + +def test_parse_diarized_transcript_preserves_overlong_timestamp_markers(): + marker = f"[{'1' * 33}]" + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript( + f"[0][S01]Value {marker}[1]" + ) + + assert [segment.text for segment in segments] == [f"Value {marker}"] diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py new file mode 100644 index 000000000000..3ff3e04379fc --- /dev/null +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import pytest + +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY + +AUDIO_MODEL_SETTINGS = { + "ibm-granite/granite-speech-3.3-2b": { + "prompt": ( + "<|start_of_role|>system<|end_of_role|>" + "You are a helpful AI assistant<|end_of_text|>\n" + "<|start_of_role|>user<|end_of_role|>" + "<|audio|>can you transcribe the speech into a written format?" + "<|end_of_text|>\n" + "<|start_of_role|>assistant<|end_of_role|>" + ), + }, + "nvidia/audio-flamingo-3-hf": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "Transcribe the input speech.<|im_end|>\n" + "<|im_start|>assistant\n" + ), + }, + "mistralai/Voxtral-Mini-3B-2507": { + "prompt": ("[INST][AUDIO]What can you tell me about this audio?[/INST]"), + }, + "microsoft/VibeVoice-ASR-HF": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant that transcribes audio input " + "into text output in JSON format.<|im_end|>\n" + "<|im_start|>user\n" + "<|object_ref_start|><|box_start|><|object_ref_end|>\n" + "This is a 1.0 seconds audio, please transcribe it with " + "these keys: Start time, End time, Speaker ID, Content" + "<|im_end|>\n" + "<|im_start|>assistant\n" + ), + }, + "zai-org/GLM-ASR-Nano-2512": { + "prompt": ( + "<|user|>\n" + "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n" + "Please transcribe this audio into text" + "<|assistant|>\n" + ), + }, +} + + +@pytest.mark.parametrize( + "model_id", + [ + "ibm-granite/granite-speech-3.3-2b", + "nvidia/audio-flamingo-3-hf", + pytest.param( + "mistralai/Voxtral-Mini-3B-2507", + marks=pytest.mark.xfail( + reason="MistralCommonBackend.encode does not produce the audio " + "placeholder token (ID 24) from raw text. apply_chat_template " + "yields token IDs with placeholders, but MultiModalProcessor." + "apply() decodes the prompt back to text and re-tokenizes, at " + "which point the placeholders are lost. Fix belongs in " + "mistral_common or in the Voxtral-specific path.", + strict=False, + ), + ), + "microsoft/VibeVoice-ASR-HF", + "zai-org/GLM-ASR-Nano-2512", + ], +) +def test_audio_multimodal_processor(model_id): + settings = AUDIO_MODEL_SETTINGS[model_id] + + model_config = ModelConfig( + model=model_id, + model_impl="transformers", + ) + + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + audio = np.zeros(16000, dtype=np.float32) + mm_data = {"audio": (audio, 16000)} + + result = mm_processor( + prompt=settings["prompt"], + mm_items=mm_processor.info.parse_mm_data(mm_data), + hf_processor_mm_kwargs={}, + ) + + assert "prompt_token_ids" in result + assert len(result["prompt_token_ids"]) > 0 + + mm_placeholders = result.get("mm_placeholders", {}) + assert "audio" in mm_placeholders, f"No audio placeholders found for {model_id}" + assert len(mm_placeholders["audio"]) == 1 + + placeholder = mm_placeholders["audio"][0] + assert placeholder.length > 0 + assert placeholder.offset >= 0 + + audio_items = result.get("mm_kwargs", {}).get("audio", []) + assert len(audio_items) == 1, f"Expected 1 audio item, got {len(audio_items)}" + item_keys = list(audio_items[0].keys()) + has_features = "input_features" in item_keys or "input_values" in item_keys + assert has_features, ( + f"No audio features (input_features/input_values) in {item_keys} for {model_id}" + ) + + +def test_audio_multiple_inputs(): + """Multiple audios per prompt are each detected as a separate placeholder + and multi-modal item by the Transformers backend.""" + model_id = "ibm-granite/granite-speech-3.3-2b" + model_config = ModelConfig(model=model_id, model_impl="transformers") + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + audio_token = mm_processor.info.get_hf_processor().audio_token + # One token per audio; the processor expands each to its placeholder run. + prompt = ( + "<|start_of_role|>user<|end_of_role|>" + f"{audio_token} and {audio_token} transcribe<|end_of_text|>\n" + ) + audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)] + + result = mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + + assert len(result["mm_placeholders"]["audio"]) == 2 + assert len(result["mm_kwargs"]["audio"]) == 2 diff --git a/tests/models/multimodal/processing/test_transformers.py b/tests/models/multimodal/processing/test_transformers_image.py similarity index 63% rename from tests/models/multimodal/processing/test_transformers.py rename to tests/models/multimodal/processing/test_transformers_image.py index a556b8f10afd..2c31bcc6347d 100644 --- a/tests/models/multimodal/processing/test_transformers.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -54,3 +54,26 @@ def test_multimodal_processor(model_id): str_processed_inputs["prompt_token_ids"] == ids_processed_inputs["prompt_token_ids"] ) + + +def test_image_multiple_inputs(): + """Multiple images per prompt are each detected as a separate placeholder + and multi-modal item by the Transformers backend.""" + model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" + model_config = ModelConfig(model=model_id, model_impl="transformers") + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + image = ImageAsset("cherry_blossom").pil_image + prompt = ( + "<|im_start|>user \n and \n" + "What do these images show?<|im_end|><|im_start|>assistant\n" + ) + + result = mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"image": [image, image]}), + hf_processor_mm_kwargs={}, + ) + + assert len(result["mm_placeholders"]["image"]) == 2 + assert len(result["mm_kwargs"]["image"]) == 2 diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py index ce9b8a9d459f..21ca71c6daa9 100644 --- a/tests/models/multimodal/test_mapping.py +++ b/tests/models/multimodal/test_mapping.py @@ -74,6 +74,44 @@ def test_cosmos3_new_checkpoint_weights_mapper(): ) +def test_cosmos3_modelopt_quantizer_weights_mapper(): + """ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers + (``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable + ``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former + (it has no parameter for them) while keeping the latter.""" + from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration + + mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper + + # Native ModelOpt quantizer buffers are dropped. + assert ( + mapper.apply_list( + [ + "layers.0.self_attn.to_q.input_quantizer._amax", + "layers.0.self_attn.to_q.weight_quantizer._amax", + "layers.0.self_attn.to_q.weight_quantizer._scale", + "layers.0.mlp.down_proj.output_quantizer._amax", + ] + ) + == [] + ) + + # The FP8 scale sidecars vLLM actually consumes are kept and remapped. + assert mapper.apply_list( + [ + "layers.0.self_attn.to_q.weight", + "layers.0.self_attn.to_q.weight_scale", + "layers.0.self_attn.to_q.input_scale", + "layers.0.mlp.down_proj.input_scale", + ] + ) == [ + "language_model.model.layers.0.self_attn.q_proj.weight", + "language_model.model.layers.0.self_attn.q_proj.weight_scale", + "language_model.model.layers.0.self_attn.q_proj.input_scale", + "language_model.model.layers.0.mlp.down_proj.input_scale", + ] + + def test_cosmos3_edge_checkpoint_weights_mapper(): from vllm.model_executor.models.cosmos3_edge import ( Cosmos3EdgeForConditionalGeneration, @@ -132,6 +170,7 @@ def test_cosmos3_edge_checkpoint_weights_mapper(): "layers.0.self_attn.to_add_out.weight", "layers.0.self_attn.norm_added_q.weight", "layers.0.self_attn.norm_added_k.weight", + "layers.0.self_attn.k_norm_und_for_gen.weight", "layers.0.self_attn.q_proj_moe_gen.weight", "layers.0.mlp_moe_gen.up_proj.weight", "norm_moe_gen.weight", diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index 03c19b0bf62a..c2f1a89df206 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -21,11 +21,11 @@ from ..utils import check_embeddings_close, check_logprobs_close if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_cdna pytestmark = pytest.mark.skipif( - on_gfx9(), - reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)", + on_cdna(), + reason="bitsandbytes not supported on CDNA (warp size 64 limitation)", ) models_4bit_to_test = [ @@ -323,7 +323,7 @@ def test_bitsandbytes_passes_revision_by_name(): return_value=["/folder/model.safetensors"], ), ): - bnb.BitsAndBytesModelLoader._prepare_weights(fake_self, "org/model", "myrev") + bnb.BitsAndBytesModelLoader._prepare_weights(fake_self, "org/model", "myrev") # type: ignore[arg-type] mock_idx.assert_called_once() assert mock_idx.call_args.kwargs.get("revision") == "myrev" diff --git a/tests/models/quantization/test_fp8.py b/tests/models/quantization/test_fp8.py index 6a13794427ab..e27d52d15bcd 100644 --- a/tests/models/quantization/test_fp8.py +++ b/tests/models/quantization/test_fp8.py @@ -9,8 +9,8 @@ import pytest from tests.quantization.utils import is_quant_method_supported +from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype from vllm.platforms import current_platform -from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from ..utils import check_logprobs_close @@ -70,13 +70,7 @@ def test_models( if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda(): pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.") - if not ( - current_platform.is_xpu() - or ( - get_flash_attn_version() == 3 - and current_platform.is_device_capability_family(90) - ) - ): + if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype): pytest.skip( f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention." ) diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index 783f1773d216..f7a7b98c3ea8 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -22,6 +22,7 @@ from packaging import version from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx950 @@ -47,7 +48,7 @@ def on_gfx950() -> bool: def has_huggingface_access(repo): try: - huggingface_hub.list_repo_refs(repo) + hf_api().list_repo_refs(repo) return True except huggingface_hub.errors.RepositoryNotFoundError: return False diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index 7c250d11576e..c12a72a09c0a 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/tests/models/quantization/test_mxfp8.py @@ -17,8 +17,10 @@ """ import pytest +import torch from tests.quantization.utils import is_quant_method_supported +from vllm.platforms import current_platform from ..utils import check_logprobs_close @@ -81,6 +83,170 @@ def test_mxfp8_logprobs( ) +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_activation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_accepts_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA, + swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA, + ) + + backend, experts_cls = select_mxfp8_moe_backend(config) + + assert backend == Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts + + @pytest.mark.skipif( not is_quant_method_supported("mxfp8"), reason="mxfp8 is not supported on this GPU type (requires sm_100+).", diff --git a/tests/models/registry.py b/tests/models/registry.py index fa23e14d9654..4ea312071593 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -267,7 +267,7 @@ def check_available_online( "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True ), "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"), - "ExaoneMoEForCausalLM": _HfExamplesInfo( + "ExaoneMoeForCausalLM": _HfExamplesInfo( "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0" ), "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"), @@ -461,7 +461,6 @@ def check_available_online( "OrionForCausalLM": _HfExamplesInfo( "OrionStarAI/Orion-14B-Chat", trust_remote_code=True ), - "OuroForCausalLM": _HfExamplesInfo("ByteDance/Ouro-1.4B", trust_remote_code=True), "PanguEmbeddedForCausalLM": _HfExamplesInfo( "FreedomIntelligence/openPangu-Embedded-7B-V1.1", trust_remote_code=True ), @@ -484,17 +483,6 @@ def check_available_online( "PhiMoEForCausalLM": _HfExamplesInfo( "microsoft/Phi-3.5-MoE-instruct", trust_remote_code=True ), - "Plamo2ForCausalLM": _HfExamplesInfo( - "pfnet/plamo-2-1b", - trust_remote_code=True, - max_transformers_version="4.57", - transformers_version_reason={ - "hf": ( - "Custom model code uses `_tied_weight_keys: list[str]` but " - "Transformers v5 now expects `_tied_weight_keys: dict[str, str]`" - ) - }, - ), "Plamo3ForCausalLM": _HfExamplesInfo( "pfnet/plamo-3-nict-2b-base", trust_remote_code=True, @@ -509,6 +497,8 @@ def check_available_online( "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"), "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"), "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"), + "Qwen3_5ForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-0.8B-text-only"), + "Qwen3_5MoeForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-35B-A3B-text-only"), "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"), "Qwen3NextForCausalLM": _HfExamplesInfo( "Qwen/Qwen3-Next-80B-A3B-Instruct", @@ -569,6 +559,7 @@ def check_available_online( "TeleFLMForCausalLM": _HfExamplesInfo( "CofeAI/FLM-2-52B-Instruct-2407", trust_remote_code=True ), + "VaultGemmaForCausalLM": _HfExamplesInfo("google/vaultgemma-1b"), "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"), "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL", trust_remote_code=True), "MiMoV2FlashForCausalLM": _HfExamplesInfo( @@ -1035,6 +1026,10 @@ def check_available_online( "moonshotai/Kimi-K2.5", trust_remote_code=True, ), + "KimiK3ForConditionalGeneration": _HfExamplesInfo( + "moonshotai/Kimi-K3", + trust_remote_code=True, + ), "KimiVLForConditionalGeneration": _HfExamplesInfo( "moonshotai/Kimi-VL-A3B-Instruct", extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, @@ -1405,6 +1400,10 @@ def check_available_online( "fixie-ai/ultravox-v0_5-llama-3_2-1b", trust_remote_code=True, ), + "VibeVoiceAsrForConditionalGeneration": _HfExamplesInfo( + "microsoft/VibeVoice-ASR-HF", + min_transformers_version="5.13.0", + ), "VoxtralForConditionalGeneration": _HfExamplesInfo( "mistralai/Voxtral-Mini-3B-2507", tokenizer_mode="mistral", @@ -1450,8 +1449,8 @@ def check_available_online( max_num_seqs=32, ), "DFlashLagunaForCausalLM": _HfExamplesInfo( - "poolside/Laguna-XS-2.1", - speculative_model="poolside/Laguna-XS-2.1-DFlash", + "poolside/Laguna-XS-2.1-NVFP4", + speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", use_original_num_layers=True, max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env max_num_seqs=32, @@ -1471,6 +1470,14 @@ def check_available_online( is_available_online=False, use_original_num_layers=True, # DSpark has >1 draft block ), + "K3DSparkModel": _HfExamplesInfo( + "moonshotai/Kimi-K3", + speculative_model="Inferact/Kimi-K3-DSpark", + use_original_num_layers=True, # DSpark has >1 draft block + trust_remote_code=True, + # FIXME: Investigate the NVML failure in CI. + is_available_online=False, + ), "Qwen3DSparkModel": _HfExamplesInfo( "Qwen/Qwen3-8B", speculative_model="deepseek-ai/dspark_qwen3_8b_block7", @@ -1669,6 +1676,12 @@ def check_available_online( trust_remote_code=True, max_model_len=4096, ), + "KimiK3MTPModel": _HfExamplesInfo( + "moonshotai/Kimi-K3", + speculative_model="moonshotai/Kimi-K3", + trust_remote_code=True, + is_available_online=False, + ), "LongCatFlashMTPModel": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True, diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py new file mode 100644 index 000000000000..28dcd8ebf4c2 --- /dev/null +++ b/tests/models/test_dspark_mla.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.registry import ModelRegistry +from vllm.models.kimi_k3.nvidia import dspark_mla +from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel + + +def test_dspark_mla_uses_compile_free_model_entrypoint(): + assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM + assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper) + + +@pytest.mark.parametrize( + ("checkpoint_name", "runtime_name", "shard_id"), + [ + ( + "layers.0.self_attn.q_a_proj.weight", + "model.layers.0.self_attn.fused_qkv_a_proj.weight", + 0, + ), + ( + "layers.0.self_attn.kv_a_proj_with_mqa.weight", + "model.layers.0.self_attn.fused_qkv_a_proj.weight", + 1, + ), + ( + "layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 0, + ), + ( + "layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 1, + ), + ("context_proj.weight", "model.context_proj.weight", None), + ], +) +def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id): + assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard( + checkpoint_name + ) == (runtime_name, shard_id) + + +def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head(): + assert not K3DSparkForCausalLM.has_own_embed_tokens + assert not K3DSparkForCausalLM.has_own_lm_head + assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == { + "confidence_head", + "embed_tokens", + "lm_head", + } + + +@pytest.mark.cpu_test +def test_dspark_markov_head_is_replicated( + monkeypatch: pytest.MonkeyPatch, +): + from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding + + monkeypatch.setattr( + vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3 + ) + monkeypatch.setattr( + vocab_parallel_embedding, + "get_tensor_model_parallel_world_size", + lambda: 8, + ) + monkeypatch.setattr( + logits_processor, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=None), + ) + + head = DSparkMarkovHead(128, 128, 8, prefix="markov_head") + assert head.markov_w2.tp_size == 1 + assert head.markov_w1.weight.shape == (128, 8) + assert head.markov_w2.weight.shape == (128, 8) + + def fail_collective(*args, **kwargs): + raise AssertionError("replicated Markov head must not invoke TP collectives") + + monkeypatch.setattr( + vocab_parallel_embedding, + "tensor_model_parallel_all_reduce", + fail_collective, + ) + logits_processor = LogitsProcessor(128) + monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective) + + markov_embed = head.embed(torch.tensor([1, 2])) + bias = head.bias(markov_embed, logits_processor) + assert markov_embed.shape == (2, 8) + assert bias.shape == (2, 128) + + +@pytest.mark.cpu_test +def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): + markov_head_calls = [] + + class DummyModule(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def make_markov_head(*args, **kwargs): + markov_head_calls.append((args, kwargs)) + return DummyModule() + + monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None) + monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule) + monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule) + monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule) + monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head) + + config = SimpleNamespace( + target_hidden_size=16, + num_target_layers=2, + hidden_size=8, + rms_norm_eps=1e-6, + num_hidden_layers=1, + vocab_size=128, + draft_vocab_size=128, + markov_rank=4, + ) + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=config) + ) + ) + + K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model") + + assert len(markov_head_calls) == 1 diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 1418704036c1..45db16e87996 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -134,6 +134,22 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): ) +@create_new_process_for_each_test() +@pytest.mark.parametrize( + "model_arch,supported", + [ + # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today. + ("NemotronHForCausalLM", True), + ("Mamba2ForCausalLM", False), + ("Zamba2ForCausalLM", False), + ], +) +def test_registry_supports_replayssm(model_arch, supported): + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_replayssm is supported + + def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path): package_dir = tmp_path / "model_package" package_dir.mkdir() diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index 546dafb79211..842a1f547942 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -11,7 +11,11 @@ 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 +from vllm.model_executor.models.transformers.fusers import ( + GLUFuser, + PackedQKVFuser, + QKVFuser, +) class SiluAndMulStub(nn.Module): @@ -203,6 +207,103 @@ def forward( return self.o_proj((q + k + v).flatten(-2)), None +class ResidDropoutAttention(FakeAttention): + """GPT-style dropout after `o_proj` -> the output projection is still found.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.resid_dropout = nn.Dropout(0.0) + + 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) + 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 + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.resid_dropout(self.o_proj(attn_output)), None + + +class PackedQKVAttention(nn.Module): + """GPTBigCode-style: one packed projection split into q/k/v in the forward.""" + + is_causal = True + + def __init__( + self, + hidden: int = 32, + head_dim: int = 8, + heads: int = 4, + kv_heads: int = 1, + 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.embed_dim = heads * head_dim + self.kv_dim = kv_heads * head_dim + self.c_attn = nn.Linear(hidden, self.embed_dim + 2 * self.kv_dim, bias=bias) + self.c_proj = nn.Linear(self.embed_dim, hidden, bias=bias) + self.resid_dropout = nn.Dropout(0.0) + + 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] + q, k, v = ( + self.c_attn(hidden_states) + .unsqueeze(1) + .split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3) + ) + q = q.view(*input_shape, -1, self.head_dim).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.resid_dropout(self.c_proj(attn_output)), attn_weights + + +class PerHeadSplitAttention(nn.Module): + """A packed projection reshaped and split *per head* -> not a q/k/v split.""" + + def __init__(self, hidden: int = 32, head_dim: int = 8, heads: int = 4): + super().__init__() + self.head_dim = head_dim + self.heads = heads + self.c_attn = nn.Linear(hidden, 3 * heads * head_dim) + self.c_proj = nn.Linear(heads * head_dim, hidden) + + def forward(self, hidden_states): + shape = (*hidden_states.shape[:2], self.heads, 3 * self.head_dim) + q, k, v = ( + self.c_attn(hidden_states) + .view(shape) + .transpose(1, 2) + .split((self.head_dim, self.head_dim, self.head_dim), dim=3) + ) + return self.c_proj((q + k + v).transpose(1, 2).flatten(-2)) + + class FakeSelfAttn(nn.Module): """Stand-in for the vLLM `Attention` looked up in `attention_instances`.""" @@ -215,6 +316,14 @@ def forward(self, q, k, v): return q + 2 * k + 3 * v +class FakeMQASelfAttn(FakeSelfAttn): + """Stand-in for grouped/multi-query layouts, where `k`/`v` are narrower.""" + + def forward(self, q, k, v): + groups = q.shape[-1] // k.shape[-1] + return q + (2 * k + 3 * v).repeat(1, groups) + + @pytest.fixture(autouse=True) def _clear_fuser_cache(): get_fuser.cache_clear() @@ -267,6 +376,15 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): return module +def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser): + """Apply a fuser at `tp_size == 1`, where the rewritten split is unchanged.""" + qkv = module.get_submodule(fuser.qkv_name) + qkv.output_sizes = [fuser.q_size, fuser.kv_size, fuser.kv_size] + qkv.tp_size = 1 + 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): @@ -366,14 +484,57 @@ def test_qkv_identifies_output_projection(): # 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" + # A module between o_proj and the return is transparent. + assert get_fuser(ResidDropoutAttention()).o_name == "o_proj" + + +@pytest.mark.parametrize("kv_heads", [1, 2]) +def test_detects_and_rewrites_packed_qkv(kv_heads): + """A single projection split into q/k/v must be re-sharded, not merged. + + Only the split sizes change: `QKVParallelLinear` loads the packed + checkpoint weight as-is, and shards q by heads while replicating k/v.""" + with torch.device("meta"): + meta = PackedQKVAttention(kv_heads=kv_heads) + fuser = get_fuser(meta) + assert isinstance(fuser, PackedQKVFuser) + assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj") + assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads) + + # The hard-coded widths become the per-rank widths of the sharded linear + names = fuser.fused_forward.__code__.co_names + assert "output_sizes" in names and "tp_size" in names + assert "kv_dim" not in names and "embed_dim" not in names + + # Numerics: the rewritten forward must match the original on a real instance + real = PackedQKVAttention(kv_heads=kv_heads, layer_idx=3) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(1, 5, 32) + attention_instances = {3: FakeMQASelfAttn()} + expected, _ = real(x, attention_instances=attention_instances) + fused = _apply_packed_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 PackedQKVAttention + assert fused.layer_idx == 3 and fused.is_causal + out, _ = fused(x, attention_instances=attention_instances) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_per_head_split_is_not_packed_qkv(): + """The split must consume the whole projection, else its sizes are head + widths and re-sharding by them would be wrong.""" + with torch.device("meta"): + assert get_fuser(PerHeadSplitAttention()) is None -def test_fuser_is_cached_per_class(): +def test_fuser_is_cached_per_class_and_structure(): 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 + assert any(key[0] is GLUMLP for key in get_fuser.cache) @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) @@ -460,8 +621,7 @@ def test_unfusable_modules_are_not_fused(cls, default_vllm_config): 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) + assert fuser is None or not fuser.validate(module, default_vllm_config) def test_act_and_mul_derived_from_module(default_vllm_config): diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 04eadac3f785..7de4589eef91 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -29,8 +29,48 @@ def forward(self, hidden_states): return logits, value, index +class ScaledRouter(TopKRouter): + """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`).""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value * 16.0 + return logits, value, index + + +class Fp32Router(TopKRouter): + """Router that computes its logits in fp32 (DeepSeek/GLM style).""" + + def forward(self, hidden_states): + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type(torch.float32) + ) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + return logits, value.to(hidden_states.dtype), index + + +class GroupedRouter(TopKRouter): + """Group-limited router (DeepSeek `group_limited_greedy`), scaled weights.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + group_scores = scores.view(-1, 4, 2).max(dim=-1).values + group_idx = torch.topk(group_scores, k=2, dim=-1)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand(-1, 4, 2).reshape(-1, 8) + scores = scores.masked_fill(~score_mask.bool(), 0.0) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value * 16.0 + return logits, value, index + + class CorrectionRouter(nn.Module): - """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined.""" + """Router with a score-correction bias buffer (DeepSeek-V3 noaux) -> matched.""" def __init__(self, num_experts=8, hidden=16): super().__init__() @@ -228,6 +268,44 @@ def test_moe_fuser_detects_router(sigmoid): assert fuser.shared_name is None and fuser.shared_gate_name is None +def test_moe_fuser_matches_scaled_router(): + """Weight scaling after the top-k (DeepSeek style) does not break matching.""" + with torch.device("meta"): + block = MoEBlock(ScaledRouter) + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) + + +def test_moe_fuser_matches_grouped_router(): + """Group masking between the score and the routing top-k still matches: the + scorer is anchored on the routing (last) top-k, not the group one.""" + with torch.device("meta"): + block = MoEBlock(GroupedRouter) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.scoring_func == "softmax" + + +def test_moe_fuser_matches_correction_router(): + """A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the + rebuilt gate carries it in fp32 for FusedMoE's biased routers.""" + with torch.device("meta"): + block = MoEBlock(CorrectionRouter) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.scoring_func == "sigmoid" + + +def test_moe_fuser_reads_router_dtype_from_the_gate(): + """A router that computes its logits in fp32 must keep routing in fp32 when + rebuilt, even though no config field names the dtype. The cast back to the + activation dtype after the top-k must not be mistaken for the routing dtype.""" + with torch.device("meta"): + fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts") + default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts") + assert fp32.router_dtype == torch.float32 + assert default.router_dtype is None + + def test_moe_fuser_detects_shared_experts(): with torch.device("meta"): block = MoEBlockShared() @@ -269,8 +347,7 @@ def test_moe_fuser_detects_non_glu_shared_expert(): @pytest.mark.parametrize( "block_cls", [ - lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped) - lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param) + lambda: MoEBlock(BiasedRouter), # router with an unrecognized 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 diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index 6497b98c4ba3..131897b19241 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -148,11 +148,13 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c # `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) + vllm_config = SimpleNamespace( + 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) + built = fuser.fuse(module, "norm", vllm_config) from vllm.model_executor.models.transformers.fusers.rms_norm import ( TPAwareNormMixin, ) @@ -176,8 +178,9 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): 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) + mc = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32) + vllm_config = SimpleNamespace(model_config=mc) + built = fuser.fuse(module, "norm", vllm_config) assert isinstance(built, VLLMRMSNorm) assert built.variance_epsilon == torch.finfo(torch.float32).eps @@ -185,11 +188,13 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): 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) + vllm_config = SimpleNamespace( + 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) + built = get_fuser(module).fuse(module, "norm", vllm_config) assert built.variance_epsilon == eps diff --git a/tests/models/utils.py b/tests/models/utils.py index 86fd80cd62a0..f938a702231d 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -507,6 +507,12 @@ class DummyConfig: hf_config = _hf_config hf_text_config = text_config + # Keep architecture conversion on the default multimodal path; this + # helper only needs HF-derived fields, not deployment MM limits. + @staticmethod + def _supports_multimodal_for_mm_prefix() -> bool: + return True + model_arch_config = ModelConfig.get_model_arch_config(DummyConfig) # Only set MoE related config when the model has MoE layers. # Otherwise all models detected as MoE by _get_transformers_backend_cls. diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index bee9d50ac1ce..ddf6dbcef868 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -6,6 +6,7 @@ import os import shutil import time +from io import BytesIO from tempfile import NamedTemporaryFile, TemporaryDirectory import aiohttp @@ -111,6 +112,34 @@ async def test_fetch_image_base64( assert _image_equals(data_image_sync, data_image_async) +@pytest.mark.asyncio +async def test_fetch_image_keep_original_mode(): + """media_io_kwargs can disable the default RGB conversion.""" + # RGBA image: opaque black pixel on a fully transparent background + rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0)) + rgba_image.putpixel((2, 2), (0, 0, 0, 255)) + buffer = BytesIO() + rgba_image.save(buffer, "PNG") + data_url = ( + f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}" + ) + + # Default behavior: RGBA is composited onto a white background + default_image = MediaConnector().fetch_image(data_url) + assert default_image.mode == "RGB" + assert default_image.getpixel((0, 0)) == (255, 255, 255) + assert default_image.getpixel((2, 2)) == (0, 0, 0) + + # image_mode=None via media_io_kwargs: original mode is preserved + connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}}) + image_sync = connector.fetch_image(data_url) + image_async = await connector.fetch_image_async(data_url) + for image in (image_sync, image_async): + assert image.mode == "RGBA" + assert image.getpixel((0, 0)) == (0, 0, 0, 0) + assert image.getpixel((2, 2)) == (0, 0, 0, 255) + + @pytest.mark.asyncio @pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) async def test_fetch_image_local_files(image_url: str): @@ -196,6 +225,44 @@ async def test_fetch_image_local_files_with_space_in_name(image_url: str): assert not ImageChops.difference(image_sync, image_async).getbbox() +@pytest.mark.asyncio +async def test_fetch_image_data_url_with_params(): + """RFC 2397 allows parameters between the mediatype and the base64 + marker; they must not be rejected or leak into the media type.""" + connector = MediaConnector() + + image = Image.new("RGB", (4, 4), color=(255, 0, 0)) + with NamedTemporaryFile(suffix=".png") as f: + image.save(f.name) + base64_image = base64.b64encode(f.read()).decode("utf-8") + + data_url = f"data:image/png;charset=utf-8;base64,{base64_image}" + image_sync = connector.fetch_image(data_url) + image_async = await connector.fetch_image_async(data_url) + assert _image_equals(image_sync, image_async) + + +def test_fetch_image_data_url_malformed(): + connector = MediaConnector() + + with pytest.raises(ValueError, match="missing ','"): + connector.fetch_image("data:image/png;base64") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:text/plain,hello") + + # ";base64" requires the ";"; here "base64" is a (bogus) media type. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:base64,aGVsbG8=") + + # Strict RFC 2397 grammar: lowercase "base64", no whitespace. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png;BASE64,aGVsbG8=") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png; base64,aGVsbG8=") + + @pytest.mark.asyncio async def test_fetch_image_error_conversion(): connector = MediaConnector() diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index c84343a3786d..5bbf54b7872e 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path): assert green_numpy[0][0][2] == 0 # B +def test_image_media_io_no_mode_conversion(tmp_path): + """image_mode=None skips conversion and preserves the original mode.""" + # RGBA image: opaque black pixel on a fully transparent background + rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0)) + rgba_image.putpixel((5, 5), (0, 0, 0, 255)) + test_image_path = tmp_path / "test_rgba.png" + rgba_image.save(test_image_path) + + # Default behavior: RGBA is composited onto a white background + image_io_default = ImageMediaIO() + converted_default = image_io_default.load_file(test_image_path) + assert converted_default.media.mode == "RGB" + assert converted_default.media.getpixel((0, 0)) == (255, 255, 255) + assert converted_default.media.getpixel((5, 5)) == (0, 0, 0) + + # image_mode=None: original mode and alpha channel are preserved + image_io_keep = ImageMediaIO(image_mode=None) + converted_keep = image_io_keep.load_file(test_image_path) + assert converted_keep.media.mode == "RGBA" + assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0) + assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255) + + def test_image_media_io_rgba_background_color_validation(): """Test that invalid rgba_background_color values are properly rejected.""" diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 8be70383b8a2..7cad42955755 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -14,7 +14,7 @@ import pytest from vllm.entrypoints.serve.utils.error_response import create_error_response -from vllm.exceptions import VLLMUnprocessableEntityError +from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector @@ -35,9 +35,9 @@ def test_creation_with_parameter_and_value(self): assert "parameter=image_url" in str(exc) assert "value=https://example.com/image.jpg" in str(exc) - def test_is_value_error_subclass(self): + def test_is_client_error_subclass(self): exc = VLLMUnprocessableEntityError("Test") - assert isinstance(exc, ValueError) + assert isinstance(exc, VLLMClientError) class TestMediaConnectorErrorHandling: diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index 671abd7b0775..a5595121b2e4 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -404,6 +404,23 @@ def test_preserves_backend_pynv_when_static(self): ) assert result["backend"] == "pynvvideocodec" + def test_strips_request_level_hw_decoders_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"hw_decoders": 4}, + ) + assert "hw_decoders" not in result + + def test_prevents_request_level_hw_decoders_override(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={ + "video_backend": "pynvvideocodec", + "hw_decoders": 2, + }, + runtime_kwargs={"hw_decoders": 4}, + ) + assert result["hw_decoders"] == 2 + @pytest.mark.parametrize("backend", ["opencv", "pyav", "torchcodec"]) def test_software_video_backend_passes_through(self, backend: str): result = VideoMediaIO.merge_kwargs( diff --git a/tests/multimodal/test_audio.py b/tests/multimodal/test_audio.py index 40d0197a4126..a4a4a9476747 100644 --- a/tests/multimodal/test_audio.py +++ b/tests/multimodal/test_audio.py @@ -744,6 +744,23 @@ def test_find_split_point_silence(self): # from start_idx, so the quietest scanned window starts at 19200. assert split_idx == 19200 + def test_split_audio_rejects_multi_channel(self): + """Chunking is mono-only; stereo must fail loudly rather than silently. + + find_split_point searches axis 0, so (channels, time) input would skip + the energy search entirely and split at fixed boundaries instead. + """ + stereo = np.ones((2, 16000 * 65), dtype=np.float32) + + with pytest.raises(ValueError, match="expects mono audio"): + split_audio( + audio_data=stereo, + sample_rate=16000, + max_clip_duration_s=30.0, + overlap_duration_s=1.0, + min_energy_window_size=1600, + ) + def test_split_audio_preserves_boundaries(self): """Verify first and last samples are preserved when chunking.""" diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index bf297946f468..816d218dd6a1 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -144,7 +144,8 @@ def _compare_caches( for _ in range(int(item_capacity / hit_rate)) ] all_hashes = [ - MultiModalHasher.hash_kwargs(item=item.get_data()) for item in all_items + MultiModalHasher.hash_kwargs("blake3", item=item.get_data()) + for item in all_items ] prompt_update = PromptInsertion("dummy", "target", "insertion").resolve(0) diff --git a/tests/multimodal/test_gpu_ipc_memory.py b/tests/multimodal/test_gpu_ipc_memory.py index bc6bf031fec0..c98fd7673cbf 100644 --- a/tests/multimodal/test_gpu_ipc_memory.py +++ b/tests/multimodal/test_gpu_ipc_memory.py @@ -6,15 +6,51 @@ import pytest +import vllm.config.multimodal as multimodal_config_module +from vllm.config.multimodal import MultiModalConfig from vllm.multimodal.gpu_ipc_memory import ( MultiModalGPUMemoryPool, get_mm_gpu_ipc_pool, maybe_init_mm_gpu_ipc_pool, + reserve_mm_ipc_gpu_memory, set_mm_gpu_ipc_pool, ) +from vllm.multimodal.video import ( + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, + PYNVVIDEOCODEC_VIDEO_BACKEND, +) from vllm.utils.mem_constants import GiB_bytes +def _mm_config( + *, + mm_ipc_gpu_memory_gb: float = 0, + video_backend: str | None = None, + hw_decoders: int | None = None, +) -> MultiModalConfig: + video_kwargs: dict[str, object] = ( + {} if video_backend is None else {"video_backend": video_backend} + ) + if hw_decoders is not None: + video_kwargs["hw_decoders"] = hw_decoders + + return MultiModalConfig( + mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, + media_io_kwargs={"video": video_kwargs} if video_kwargs else {}, + ) + + +def _pynvvideocodec_decoder_budget( + api_process_count: int = 1, + hw_decoders: int = 2, +) -> int: + return api_process_count * ( + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * hw_decoders + + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES + ) + + def test_acquire_release_accounting(): pool = MultiModalGPUMemoryPool(total_bytes=100) assert pool.available_bytes == 100 @@ -143,3 +179,91 @@ def test_global_pool_splits_budget_across_api_processes(): def test_global_pool_rejects_invalid_api_process_count(): with pytest.raises(ValueError): maybe_init_mm_gpu_ipc_pool(2, api_process_count=0) + + +@pytest.mark.parametrize("video_backend", [None, "opencv"]) +def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( + monkeypatch: pytest.MonkeyPatch, + video_backend: str | None, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + mm_config = _mm_config( + mm_ipc_gpu_memory_gb=0.25, + video_backend=video_backend, + ) + + assert reserve_mm_ipc_gpu_memory(GiB_bytes, mm_config) == int(0.75 * GiB_bytes) + + +def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + mm_config = _mm_config( + mm_ipc_gpu_memory_gb=0.25, + video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + available_bytes = 4 * GiB_bytes + + assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == ( + available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget() + ) + + +def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + available_bytes = 4 * GiB_bytes + + assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == ( + available_bytes - _pynvvideocodec_decoder_budget() + ) + + +def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + available_bytes = 8 * GiB_bytes + + assert reserve_mm_ipc_gpu_memory( + available_bytes, + _mm_config(), + api_process_count=3, + ) == available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) + + +def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + available_bytes = 4 * GiB_bytes + mm_config = _mm_config( + video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + hw_decoders=3, + ) + + assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == ( + available_bytes - _pynvvideocodec_decoder_budget(hw_decoders=3) + ) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index fdedcaea27c4..4810081ad40c 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib import uuid +from io import BytesIO from pathlib import Path import numpy as np @@ -8,7 +10,11 @@ import torch from PIL import Image, ImageDraw +from vllm.config.multimodal import MMHasherAlgorithm from vllm.multimodal.hasher import MultiModalHasher +from vllm.multimodal.media.base import MediaWithBytes +from vllm.multimodal.media.image import ImageMediaIO +from vllm.multimodal.parse import MultiModalDataParser pytestmark = pytest.mark.cpu_test @@ -16,12 +22,26 @@ assert ASSETS_DIR.exists() +@pytest.mark.parametrize("algorithm", ["sha256", "sha512"]) +def test_hash_algorithm(algorithm: MMHasherAlgorithm): + hasher = getattr(hashlib, algorithm)() + for bytes_ in MultiModalHasher.iter_item_to_bytes("value", "test"): + hasher.update(bytes_) + + assert MultiModalHasher.hash_kwargs(algorithm, value="test") == hasher.hexdigest() + + +def test_hash_algorithm_required(): + with pytest.raises(TypeError, match="algorithm"): + MultiModalHasher.hash_kwargs(value="test") # type: ignore[call-arg] + + def test_hash_single_item_different_shape(): x1 = torch.zeros(()) x2 = torch.zeros((1,)) hasher = MultiModalHasher - assert hasher.hash_kwargs(x=x1) != hasher.hash_kwargs(x=x2) + assert hasher.hash_kwargs("blake3", x=x1) != hasher.hash_kwargs("blake3", x=x2) def test_hash_key_order_invariant(): @@ -29,7 +49,9 @@ def test_hash_key_order_invariant(): y = torch.ones((5, 10)) hasher = MultiModalHasher - assert hasher.hash_kwargs(x=x, y=y) == hasher.hash_kwargs(y=y, x=x) + assert hasher.hash_kwargs("blake3", x=x, y=y) == hasher.hash_kwargs( + "blake3", y=y, x=x + ) # NOTE: Images that are the same visually are allowed to have the same hash @@ -40,7 +62,9 @@ def test_hash_collision_image_mode(mode_pair): image2 = Image.new(mode2, size=(10, 10), color=1) hasher = MultiModalHasher - assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2) + assert hasher.hash_kwargs("blake3", image=image1) != hasher.hash_kwargs( + "blake3", image=image2 + ) def test_hash_collision_image_palette(): @@ -49,7 +73,9 @@ def test_hash_collision_image_palette(): image2 = Image.open(ASSETS_DIR / "image2.png") hasher = MultiModalHasher - assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2) + assert hasher.hash_kwargs("blake3", image=image1) != hasher.hash_kwargs( + "blake3", image=image2 + ) def test_hash_collision_image_transpose(): @@ -60,7 +86,9 @@ def test_hash_collision_image_transpose(): ImageDraw.Draw(image2).line([(0, 0), (0, 10)]) hasher = MultiModalHasher - assert hasher.hash_kwargs(image=image1) != hasher.hash_kwargs(image=image2) + assert hasher.hash_kwargs("blake3", image=image1) != hasher.hash_kwargs( + "blake3", image=image2 + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -70,7 +98,9 @@ def test_hash_collision_tensor_shape(dtype): arr2 = torch.zeros((10, 20, 5, 3), dtype=dtype) hasher = MultiModalHasher - assert hasher.hash_kwargs(data=arr1) != hasher.hash_kwargs(data=arr2) + assert hasher.hash_kwargs("blake3", data=arr1) != hasher.hash_kwargs( + "blake3", data=arr2 + ) def test_hash_collision_array_shape(): @@ -79,7 +109,32 @@ def test_hash_collision_array_shape(): arr2 = np.zeros((10, 20, 5, 3)) hasher = MultiModalHasher - assert hasher.hash_kwargs(data=arr1) != hasher.hash_kwargs(data=arr2) + assert hasher.hash_kwargs("blake3", data=arr1) != hasher.hash_kwargs( + "blake3", data=arr2 + ) + + +def test_hash_collision_video_num_frames(): + source = b"x" * 100 + + def item_for_hash(num_frames: int): + frames: np.ndarray = np.zeros((num_frames, 8, 8, 3), dtype=np.uint8) + metadata = { + "total_num_frames": 16, + "fps": 2.0, + "duration": 8.0, + "video_backend": "opencv", + "frames_indices": list(range(num_frames)), + "do_sample_frames": False, + } + video = MediaWithBytes((frames, metadata), source) + items = MultiModalDataParser()._parse_video_data([video]) + return items.get_all_items_for_hash()[0] + + hasher = MultiModalHasher + assert hasher.hash_kwargs("blake3", video=item_for_hash(2)) != hasher.hash_kwargs( + "blake3", video=item_for_hash(4) + ) def test_hash_non_contiguous_array(): @@ -91,7 +146,9 @@ def test_hash_non_contiguous_array(): hasher = MultiModalHasher # Both should be hashable and produce the same hashes - assert hasher.hash_kwargs(data=arr) == hasher.hash_kwargs(data=arr_c) + assert hasher.hash_kwargs("blake3", data=arr) == hasher.hash_kwargs( + "blake3", data=arr_c + ) def test_hash_image_exif_id(): @@ -106,6 +163,52 @@ def test_hash_image_exif_id(): hasher = MultiModalHasher # first image has UUID in ImageID, so it should hash to that UUID - assert hasher.hash_kwargs(image=image1) == hasher.hash_kwargs(image=id.bytes) + assert hasher.hash_kwargs("blake3", image=image1) == hasher.hash_kwargs( + "blake3", image=id.bytes + ) # second image has non-UUID in ImageID, so it should hash to the image data - assert hasher.hash_kwargs(image=image2) == hasher.hash_kwargs(image=image2a) + assert hasher.hash_kwargs("blake3", image=image2) == hasher.hash_kwargs( + "blake3", image=image2a + ) + + +def _rgba_png_bytes() -> bytes: + image = Image.new("RGBA", (8, 8), (255, 0, 0, 128)) + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +def test_hash_collision_media_io_config(): + data = _rgba_png_bytes() + white = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data) + black = ImageMediaIO(rgba_background_color=(0, 0, 0)).load_bytes(data) + white2 = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data) + keep = ImageMediaIO(image_mode=None).load_bytes(data) + + hasher = MultiModalHasher + assert hasher.hash_kwargs("blake3", image=white) != hasher.hash_kwargs( + "blake3", image=black + ) + assert hasher.hash_kwargs("blake3", image=white) != hasher.hash_kwargs( + "blake3", image=keep + ) + assert hasher.hash_kwargs("blake3", image=white) == hasher.hash_kwargs( + "blake3", image=white2 + ) + + +def test_hash_media_io_noop_config_preserves_hash(): + image = Image.new("RGB", (8, 8), (0, 128, 255)) + buf = BytesIO() + image.save(buf, format="PNG") + data = buf.getvalue() + + loaded = ImageMediaIO().load_bytes(data) + assert loaded.io_config is None + + plain = MediaWithBytes(loaded.media, data) + hasher = MultiModalHasher + assert hasher.hash_kwargs("blake3", image=loaded) == hasher.hash_kwargs( + "blake3", image=plain + ) diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index 66acdbe62fff..2153cd2c2595 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -8,6 +8,7 @@ import pytest from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.processing.context import InputProcessingContext from vllm.multimodal.processing.processor import ( @@ -931,7 +932,11 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid): else: mm_data = {"image": [image] * num_images} - exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most") + exc_ctx = ( + nullcontext() + if is_valid + else pytest.raises(VLLMValidationError, match="At most") + ) with exc_ctx: processor( diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py index 4e765ab1b8b3..ffc8ce504ccb 100644 --- a/tests/multimodal/test_utils.py +++ b/tests/multimodal/test_utils.py @@ -199,6 +199,31 @@ def test_group_and_batch_mm_items_split_by_fieldset(): assert [num_items for num_items, _ in res] == [2, 1, 1, 1] +def test_group_and_batch_mm_items_splits_shared_data_by_dtype(): + elem1 = MultiModalFieldElem( + data=torch.zeros(1, dtype=torch.int32), + field=MultiModalSharedField(batch_size=1), + ) + elem2 = MultiModalFieldElem( + data=torch.zeros(1, dtype=torch.float32), + field=MultiModalSharedField(batch_size=1), + ) + elem3 = MultiModalFieldElem( + data=[torch.zeros(1, dtype=torch.int32), torch.zeros(1, dtype=torch.float32)], + field=MultiModalSharedField(batch_size=1), + ) + + res = group_and_batch_mm_items( + [ + MultiModalKwargsItem({"x": elem1}), + MultiModalKwargsItem({"x": elem2}), + MultiModalKwargsItem({"x": elem3}), + ] + ) + + assert [num_items for num_items, _ in res] == [1, 1, 1] + + def test_group_and_batch_mm_items_split_by_shared_data(): elem1 = MultiModalFieldElem( data=torch.zeros(1, dtype=torch.uint8), diff --git a/tests/multimodal/test_vidcom2.py b/tests/multimodal/test_vidcom2.py new file mode 100644 index 000000000000..4a62bbd87cfe --- /dev/null +++ b/tests/multimodal/test_vidcom2.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count, + compute_retention_mask, +) + + +def _fake_video_embeds( + num_frames: int, + rows: int, + cols: int, + hidden: int = 64, + seed: int = 0, +) -> torch.Tensor: + """Deterministic fake ViT output with a distinct mean per frame.""" + g = torch.Generator().manual_seed(seed) + frames = [] + for f in range(num_frames): + base = torch.randn(hidden, generator=g) * (0.1 + 0.05 * f) + frames.append( + base[None, :].expand(rows * cols, hidden) + + 0.01 * torch.randn(rows * cols, hidden, generator=g) + ) + return torch.cat(frames, dim=0) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_mask_shape_and_dtype(q: float, num_frames: int) -> None: + merge = 2 + rows, cols = 6, 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + assert mask.dtype == torch.bool + assert mask.shape == (num_frames * rows * cols,) + + +def test_retained_count_floors_at_one_token_per_frame() -> None: + """The global minimum is one token per frame (not a full first frame).""" + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.999) == 4 + ) + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.0) + == 48 * 4 + ) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_total_retained_matches_target(q: float, num_frames: int) -> None: + """Mask total must equal the placeholder-sizing helper.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + expected = compute_retained_tokens_count( + tokens_per_frame=tpf, num_frames=num_frames, q=q + ) + assert int(mask.sum().item()) == expected + + +def test_per_frame_min_one_when_budget_allows() -> None: + """No frame is fully dropped when the budget allows.""" + merge = 2 + rows, cols = 6, 8 + num_frames = 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.25, + ) + per_frame = mask.view(num_frames, rows * cols).sum(dim=1) + assert (per_frame >= 1).all(), f"zero-token frame detected: {per_frame.tolist()}" + + +def test_dynamic_per_frame_budget() -> None: + """A distinctive frame gets more retained tokens than bland ones.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + hidden = 64 + torch.manual_seed(0) + bland = 0.01 * torch.randn(tpf, hidden) + frames = [torch.randn(tpf, hidden) * 1.0] + for _ in range(7): + frames.append(bland + 0.001 * torch.randn(tpf, hidden)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.5, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() > per_frame[1:].float().mean().item() + + +def test_empty_input_safe() -> None: + embeds = torch.zeros(0, 32) + mask = compute_retention_mask(embeds, (0, 0, 0), spatial_merge_size=2, q=0.25) + assert mask.numel() == 0 + + +@pytest.mark.parametrize("q", [0.0, 0.25, 0.5, 0.75]) +def test_first_frame_not_privileged(q: float) -> None: + """A bland first frame is not force-retained (unlike EVS).""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + torch.manual_seed(1) + bland = 0.01 * torch.randn(tpf, 64) + frames = [bland] + for f in range(7): + frames.append(torch.randn(tpf, 64) * (1.0 + 0.1 * f)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() <= tpf + if q > 0.0: + assert per_frame[0].item() < int(mask.sum().item()) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 6aeb7dc486ec..1a7b63ae9274 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -14,9 +14,9 @@ from transformers.video_utils import VideoMetadata from vllm.assets.base import get_vllm_public_assets +from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend from vllm.multimodal.video import ( PYNVVIDEOCODEC_DECODER_CACHE_SIZE, - PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, DynamicVideoBackend, @@ -26,6 +26,7 @@ PyNvVideoCodecVideoBackend, Qwen2VLVideoBackend, Qwen3VLVideoBackend, + VideoBackend, VideoLoader, VideoSourceMetadata, VideoTargetMetadata, @@ -199,7 +200,11 @@ def fake_decode(cls, file_path: str, frame_idx: list[int], nvc): assert metadata["frames_indices"] == [0, 9] -def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("hw_decoders", [1, 3]) +def test_pynvvideocodec_decoder_slots_are_bounded( + monkeypatch: pytest.MonkeyPatch, + hw_decoders: int, +): class FakeSlot: pass @@ -207,10 +212,13 @@ class FakeSlot: old_slots = PyNvVideoCodecVideoBackend._decoder_slots old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond + old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots try: PyNvVideoCodecVideoBackend._decoder_slots = [] PyNvVideoCodecVideoBackend._active_decoder_slots = 0 PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() + PyNvVideoCodecVideoBackend._max_decoder_slots = None + PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders) def fake_create_slot(cls): nonlocal create_count @@ -229,7 +237,7 @@ def fake_create_slot(cls): with ExitStack() as stack: retained_slots = [ stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot()) - for _ in range(PYNVVIDEOCODEC_MAX_RETAINED_DECODERS) + for _ in range(hw_decoders) ] def borrow_extra_slot(): @@ -246,11 +254,34 @@ def borrow_extra_slot(): assert not thread.is_alive() assert seen_slots[0] in retained_slots - assert create_count == PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + assert create_count == hw_decoders finally: PyNvVideoCodecVideoBackend._decoder_slots = old_slots PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond + PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + + +def test_pynvvideocodec_decoder_slots_are_configured_once( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(PyNvVideoCodecVideoBackend, "_max_decoder_slots", None) + + PyNvVideoCodecVideoBackend._configure_decoder_slots(2) + PyNvVideoCodecVideoBackend._configure_decoder_slots(2) + + with pytest.raises(RuntimeError, match="already configured as 2, got 3"): + PyNvVideoCodecVideoBackend._configure_decoder_slots(3) + + +@pytest.mark.parametrize("hw_decoders", [0, -1, 1.5, True, "2"]) +def test_pynvvideocodec_rejects_invalid_hw_decoders(hw_decoders: object): + with pytest.raises(ValueError, match="hw_decoders must be a positive integer"): + VideoBackend.load_bytes( + b"fake video", + backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + hw_decoders=hw_decoders, # type: ignore[arg-type] + ) def test_pynvvideocodec_decoder_slot_retains_simple_decoder(): @@ -304,6 +335,13 @@ class OutputColorType: # ============================================================================ +def test_cosmos3_edge_uses_qwen3_vl_video_backend(): + backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor") + + assert backend == "qwen3_vl" + assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend) + + @pytest.mark.parametrize( "model_repo, expected_loader_cls, hf_sample_kwargs", [ @@ -351,6 +389,12 @@ class OutputColorType: {"fps": 2}, id="qwen2_5_vl", ), + pytest.param( + "MiniMaxAI/MiniMax-M3", + MiniMaxM3VideoBackend, + None, + id="minimax_m3_vl", + ), ], ) def test_video_processor_from_model_repo( diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 3aa7b3b9d21a..e5e7b075bac2 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -920,3 +920,45 @@ def test_eos_drop_token_does_not_swallow_tool_calls(self): assert output.tool_calls[0]["name"] == "get_weather" args = json.loads(output.tool_calls[0]["arguments"]) assert args == {"location": "Berlin"} + + @pytest.mark.parametrize( + "chunk_size", + [1, 2, 3, 5, None], + ids=lambda c: f"chunk={c}", + ) + def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size): + """EOS must not leak into reasoning_content when the model never + emits (generation ends while still in REASONING state).""" + eos_text = "<｜end▁of▁sentence｜>" + eos_id = 128801 + vocab = { + **_DSV4_FULL_VOCAB, + eos_text: eos_id, + } + + reasoning_text = "Good morning! How can I help you today?" + tokens: list[tuple[int, str]] = [] + tid = 100 + for word in reasoning_text.split(" "): + prefix = " " if tokens else "" + tokens.append((tid, prefix + word)) + tid += 1 + tokens.append((eos_id, eos_text)) + + tokenizer = MockTokenizer(vocab=vocab, tokens=tokens) + parser = _DeepSeekV4Delegating( + tokenizer, + chat_template_kwargs={"thinking": True}, + ) + deltas = replay_streaming( + parser, + tokens, + chunk_size=chunk_size, + finished_on_last=True, + ) + output = collect_output(deltas) + + assert reasoning_text in output.reasoning + assert eos_text not in output.reasoning + assert output.content == "" + assert output.tool_calls == [] diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py index b922e71e3982..50c41a8ea6e9 100644 --- a/tests/parser/engine/test_delegating_replay.py +++ b/tests/parser/engine/test_delegating_replay.py @@ -40,6 +40,7 @@ ParserEngineReasoningAdapter, ParserEngineToolAdapter, ) +from vllm.parser.mistral import MistralParser _TOOLS_VALIDATOR = TypeAdapter(list[ChatCompletionToolsParam]) @@ -84,6 +85,11 @@ def _discover_pairings() -> list[_PairingInfo]: for engine_cls, adapters in engines.items(): if "tool" not in adapters or "reasoning" not in adapters: continue + if engine_cls is MistralParser: + # Mistral uses brace-balanced JSON tool args with no TOOL_END + # token, so it does not fit this TOOL_END-based replay harness. + # It is covered by tests/parser/mistral/ instead. + continue cfg = engine_cls(bare_tok, None).parser_engine_config if cfg.name not in _BUILDERS: missing_builders.append(f"{engine_cls.__name__} (config.name={cfg.name!r})") diff --git a/tests/parser/engine/test_inkling.py b/tests/parser/engine/test_inkling.py index c0b9877db34a..582783fd5d0b 100644 --- a/tests/parser/engine/test_inkling.py +++ b/tests/parser/engine/test_inkling.py @@ -178,6 +178,13 @@ def test_non_object_args_rejected(self): class TestNonStreaming: + @pytest.mark.parametrize("suffix", ["", END_MESSAGE, END_SAMPLING]) + def test_bare_text_after_model_opener(self, parser, mock_request, suffix): + reasoning, content, tools = parser.parse(f"hello world{suffix}", mock_request) + assert reasoning is None + assert content == "hello world" + assert tools is None + def test_plain_text(self, parser, mock_request): reasoning, content, tools = parser.parse( f"{TEXT_START}hello world{END_MESSAGE}", mock_request @@ -426,6 +433,28 @@ def test_generation_prompt_header_hides_tool_name(self, parser, mock_request): assert delta.content is None assert delta.tool_calls[0].function.name == "get_weather" + def test_generation_prompt_header_flushes_bare_text_at_finish( + self, parser, mock_request + ): + prompt_token_ids = [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]] + first = parser.parse_delta( + "plain ", + [ord(char) for char in "plain "], + mock_request, + prompt_token_ids=prompt_token_ids, + finished=False, + ) + assert first is None + + second = parser.parse_delta( + "answer", + [ord(char) for char in "answer"], + mock_request, + finished=True, + ) + assert second is not None + assert second.content == "plain answer" + class TestToolCallFiltering: """Inkling equivalents of the generic tool-call-filtering replay tests diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 36258668215b..d97aafe365ec 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -1624,19 +1624,41 @@ def test_no_special_tokens_means_no_drops(self): assert delta is not None assert "" in delta.reasoning - def test_drops_suppressed_with_skip_tool_parsing(self): - """When skip_tool_parsing is active, drop tokens are preserved - as content so a later tool-call pass can see them.""" + def test_drops_applied_with_skip_tool_parsing(self): + """Drop tokens are always dropped, even with skip_tool_parsing. + DROP_TERMINALs have no transitions by construction, so no parser + pass can use them.""" + for initial_state in (ParserState.REASONING, ParserState.CONTENT): + engine = _make_engine( + vocab=_DROP_VOCAB, + special_tokens=list(_DROP_VOCAB.keys()), + ) + engine._engine.skip_tool_parsing = True + engine._engine.reset(initial_state=initial_state) + events = engine._engine.feed("helloworld", [72, 204, 73]) + delta = engine._events_to_delta(events) + assert delta is not None + output = (delta.reasoning or "") + (delta.content or "") + assert "" not in output, f" leaked in state {initial_state}" + + def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self): + """With skip_tool_parsing in REASONING state, drop tokens are + removed but configured terminals still fire their transitions.""" engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), ) engine._engine.skip_tool_parsing = True engine._engine.reset() - events = engine._engine.feed("helloworld", [72, 204, 73]) - delta = engine._events_to_delta(events) - assert delta is not None - assert "" in delta.reasoning + events = engine._engine.feed("thoughtanswer", [72, 204, 201, 73]) + types = [e.type for e in events] + assert EventType.REASONING_CHUNK in types + assert EventType.REASONING_END in types + assert EventType.TEXT_CHUNK in types + reasoning_text = "".join( + e.value for e in events if e.type == EventType.REASONING_CHUNK + ) + assert "" not in reasoning_text def test_drops_in_tool_args_state(self): """Drop tokens in TOOL_ARGS state are silently discarded.""" diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py index 064e1b623679..5bba7f7eafd3 100644 --- a/tests/parser/engine/test_replay.py +++ b/tests/parser/engine/test_replay.py @@ -34,6 +34,7 @@ from vllm.parser.engine import registered_adapters as _adapters_mod from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ParserState +from vllm.parser.mistral import MistralParser # ── Parser discovery ───────────────────────────────────────────────── @@ -64,6 +65,11 @@ def _discover_parsers() -> list[_ParserInfo]: and obj is not ParserEngine ): continue + if obj is MistralParser: + # Mistral uses brace-balanced JSON tool args with no TOOL_END + # token, so it does not fit this TOOL_END-based replay harness. + # It is covered by tests/parser/mistral/ instead. + continue cfg = obj(bare_tok, None).parser_engine_config if cfg.name not in _BUILDERS: missing_builders.append(f"{obj.__name__} (config.name={cfg.name!r})") diff --git a/tests/parser/mistral/__init__.py b/tests/parser/mistral/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/parser/mistral/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/parser/mistral/test_reasoning.py b/tests/parser/mistral/test_reasoning.py new file mode 100644 index 000000000000..137da648538b --- /dev/null +++ b/tests/parser/mistral/test_reasoning.py @@ -0,0 +1,599 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest +from mistral_common.protocol.instruct.messages import ( + AssistantMessage, + ThinkChunk, + ToolMessage, + UserMessage, +) +from mistral_common.protocol.instruct.request import InstructRequest +from mistral_common.protocol.instruct.tool_calls import FunctionCall, ToolCall + +from tests.reasoning.utils import run_reasoning_extraction_mistral +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.parser_manager import ParserManager +from vllm.reasoning import ReasoningParser, ReasoningParserManager +from vllm.tokenizers.detokenizer_utils import detokenize_incrementally +from vllm.tokenizers.mistral import MistralTokenizer + +_PARSER_NAME = "mistral" +_MODEL_V13 = "mistralai/Magistral-Small-2509" +_MODEL_V11 = "mistralai/Magistral-Small-2506" + +_SAMPLE_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] + + +@pytest.fixture(scope="module") +def mistral_tokenizer() -> MistralTokenizer: + """v13 Magistral tokenizer with `[THINK]`/`[/THINK]` special tokens.""" + return MistralTokenizer.from_pretrained(_MODEL_V13) + + +@pytest.fixture(scope="module") +def mistral_v11_tokenizer() -> MistralTokenizer: + """v11 Magistral tokenizer using plain-text ``/`` tags.""" + return MistralTokenizer.from_pretrained(_MODEL_V11) + + +def _encode_v13(tokenizer: MistralTokenizer, output: str) -> list[int]: + """Encode output string with `[THINK]`/`[/THINK]` markers into token IDs. + + `[THINK]` and `[/THINK]` placeholders in `output` are replaced by their + special-token IDs; all surrounding text is encoded with the base tokenizer. + Mirrors the encoding helper from the deleted + `tests/reasoning/test_mistral_reasoning_parser.py`. + """ + think_start = "[THINK]" + think_end = "[/THINK]" + len_start = len(think_start) + len_end = len(think_end) + + index_start = output.find(think_start) + index_end = output.find(think_end) + + out: list[int] = [] + if index_start != -1: + out += tokenizer.tokenizer.encode(output[:index_start], False, False) + out += [tokenizer.instruct.BEGIN_THINK] + + if index_end != -1: + middle = output[index_start + len_start : index_end] + suffix = output[index_end + len_end :] + out += tokenizer.tokenizer.encode(middle, False, False) + out += [tokenizer.instruct.END_THINK] + out += tokenizer.tokenizer.encode(suffix, False, False) + else: + out += tokenizer.tokenizer.encode( + output[index_start + len_start :], False, False + ) + elif index_end != -1: + out += tokenizer.tokenizer.encode(output[:index_end], False, False) + out += [tokenizer.instruct.END_THINK] + out += tokenizer.tokenizer.encode(output[index_end + len_end :], False, False) + else: + out += tokenizer.tokenizer.encode(output, False, False) + return out + + +_TEST_CASES = [ + # v13 special-token encoding: [THINK]…[/THINK] with trailing content. + pytest.param( + False, + "[THINK]r[/THINK]c", + "r", + "c", + id="v13_basic", + ), + pytest.param( + True, + "[THINK]r[/THINK]c", + "r", + "c", + id="v13_basic_streaming", + ), + # No trailing content after reasoning ends. + pytest.param( + False, + "[THINK]r[/THINK]", + "r", + None, + id="v13_no_trailing_content", + ), + pytest.param( + True, + "[THINK]r[/THINK]", + "r", + None, + id="v13_no_trailing_content_streaming", + ), + # Multi-line reasoning and multi-line content. + pytest.param( + False, + "[THINK]a\nb[/THINK]c\nd", + "a\nb", + "c\nd", + id="v13_multiline", + ), + pytest.param( + True, + "[THINK]a\nb[/THINK]c\nd", + "a\nb", + "c\nd", + id="v13_multiline_streaming", + ), + # No reasoning at all: output is pure content. + pytest.param( + False, + "hello world", + None, + "hello world", + id="v13_no_reasoning", + ), + pytest.param( + True, + "hello world", + None, + "hello world", + id="v13_no_reasoning_streaming", + ), + # [THINK] opened but never closed; reasoning accumulates, content is None. + pytest.param( + False, + "[THINK]r", + "r", + None, + id="v13_think_no_end", + ), + pytest.param( + True, + "[THINK]r", + "r", + None, + id="v13_think_no_end_streaming", + ), + # Stray [/THINK] with no matching [THINK]: no reasoning, content follows. + pytest.param( + False, + "[/THINK]This is the rest", + None, + "This is the rest", + id="v13_stray_think_end", + ), + pytest.param( + True, + "[/THINK]This is the rest", + None, + "This is the rest", + id="v13_stray_think_end_streaming", + ), + # Content before [THINK]: leading segment joins with trailing segment. + pytest.param( + False, + "Before\n[THINK]r[/THINK]\nAfter", + "r", + "Before\n\nAfter", + id="v13_leading_content", + ), + pytest.param( + True, + "Before\n[THINK]r[/THINK]\nAfter", + "r", + "Before\n\nAfter", + id="v13_leading_content_streaming", + ), + # Empty output: no reasoning, no content. + pytest.param( + False, + "", + None, + None, + id="v13_empty", + ), + pytest.param( + True, + "", + None, + None, + id="v13_empty_streaming", + ), +] + + +@pytest.mark.parametrize( + "streaming, output, expected_reasoning, expected_content", + _TEST_CASES, +) +def test_mistral_reasoning_v13( + streaming: bool, + output: str, + expected_reasoning: str | None, + expected_content: str | None, + mistral_tokenizer: MistralTokenizer, +) -> None: + output_tokens = _encode_v13(mistral_tokenizer, output) + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(_PARSER_NAME)( + mistral_tokenizer + ) + reasoning, content = run_reasoning_extraction_mistral( + parser, output_tokens, streaming=streaming + ) + assert reasoning == expected_reasoning + assert content == expected_content + + +def test_system_prompt_think_ignored( + mistral_tokenizer: MistralTokenizer, +) -> None: + """Streaming content extraction must not + include `[THINK]`/`[/THINK]` examples from the system prompt. + """ + output = "[THINK]real[/THINK]answer" + output_tokens = _encode_v13(mistral_tokenizer, output) + for streaming in (False, True): + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser( + _PARSER_NAME + )(mistral_tokenizer) + reasoning, content = run_reasoning_extraction_mistral( + parser, output_tokens, streaming=streaming + ) + assert reasoning == "real" + assert content == "answer" + + +def test_reasoning_with_tool_calls( + mistral_tokenizer: MistralTokenizer, +) -> None: + """Reasoning extraction with a tool call in the generated output.""" + tool_calls_token_id = mistral_tokenizer.get_vocab().get("[TOOL_CALLS]") + if tool_calls_token_id is None: + pytest.skip("Tokenizer does not have [TOOL_CALLS] in vocab") + + # Token sequence representing: [THINK]r[/THINK][TOOL_CALLS]f{"a":1} + output_tokens = ( + [mistral_tokenizer.instruct.BEGIN_THINK] + + mistral_tokenizer.tokenizer.encode("r", False, False) + + [mistral_tokenizer.instruct.END_THINK] + + [tool_calls_token_id] + + mistral_tokenizer.tokenizer.encode('f{"a":1}', False, False) + ) + for streaming in (False, True): + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser( + _PARSER_NAME + )(mistral_tokenizer) + reasoning, _content = run_reasoning_extraction_mistral( + parser, output_tokens, streaming=streaming + ) + assert reasoning == "r" + assert "[TOOL_CALLS]" not in (reasoning or "") + + +@pytest.mark.parametrize("streaming", [False, True], ids=["nonstream", "stream"]) +def test_reasoning_v11_plain_text_think( + mistral_v11_tokenizer: MistralTokenizer, streaming: bool +) -> None: + """v11 tokenizers use plain-text ``/`` (no special tokens).""" + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(_PARSER_NAME)( + mistral_v11_tokenizer + ) + + tokens = mistral_v11_tokenizer.tokenizer.encode("rc", False, False) + reasoning, content = run_reasoning_extraction_mistral( + parser, tokens, streaming=streaming + ) + assert reasoning == "r" + assert content == "c" + + +def test_is_reasoning_end( + mistral_tokenizer: MistralTokenizer, +) -> None: + """is_reasoning_end returns True iff reasoning has finished.""" + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(_PARSER_NAME)( + mistral_tokenizer + ) + + # Complete reasoning: END_THINK present → True. + complete_ids = _encode_v13(mistral_tokenizer, "[THINK]r[/THINK]c") + assert parser.is_reasoning_end(complete_ids) + + # Open-only: no END_THINK, BEGIN_THINK present → False. + open_ids = _encode_v13(mistral_tokenizer, "[THINK]r") + assert not parser.is_reasoning_end(open_ids) + + # No reasoning markers at all → False. + no_reasoning_ids = _encode_v13(mistral_tokenizer, "hello") + assert not parser.is_reasoning_end(no_reasoning_ids) + + # Explicit END_THINK closes reasoning → True. + explicit_end_ids = ( + [mistral_tokenizer.instruct.BEGIN_THINK] + + mistral_tokenizer.tokenizer.encode("r", False, False) + + [mistral_tokenizer.instruct.END_THINK] + ) + assert parser.is_reasoning_end(explicit_end_ids) + + # [TOOL_CALLS] acts as implicit reasoning-end marker. + tool_calls_token_id = mistral_tokenizer.get_vocab().get("[TOOL_CALLS]") + if tool_calls_token_id is None: + pytest.skip("Tokenizer does not have [TOOL_CALLS] in vocab") + + implicit_end_ids = ( + [mistral_tokenizer.instruct.BEGIN_THINK] + + mistral_tokenizer.tokenizer.encode("r", False, False) + + [tool_calls_token_id] + ) + assert parser.is_reasoning_end(implicit_end_ids) + + +def _stream_parse_delta(parser, tokenizer, gen_ids, request, prompt_ids): + """Stream `gen_ids` through `parse_delta` and reconstruct reasoning/content.""" + reasoning = "" + content = "" + previous_tokens: list[str] | None = None + prefix_offset = 0 + read_offset = 0 + for i, tok in enumerate(gen_ids): + new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer=tokenizer, + all_input_ids=gen_ids[: i + 1], + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=True, + spaces_between_special_tokens=True, + ) + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + delta_message = parser.parse_delta( + delta_text=delta_text, + delta_token_ids=[tok], + request=request, + prompt_token_ids=prompt_ids, + finished=(i == len(gen_ids) - 1), + ) + if delta_message is not None: + reasoning += delta_message.reasoning or "" + content += delta_message.content or "" + return reasoning, content + + +def test_parse_delta_keeps_reasoning_open_after_closed_think_prompt( + mistral_tokenizer: MistralTokenizer, +) -> None: + """A closed `[THINK]` block in the prompt must not finalize reasoning. + + Regression: reasoning system prompts end the prompt with a closed + `[THINK]...[/THINK]` block. With a parser-injected grammar the generation's + reasoning must still be captured as ``reasoning_content`` rather than + leaking into ``content`` (the prompt-based reasoning-end heuristic must be + skipped). This exercises ``DelegatingParser.parse_delta``'s prompt check, + which the isolated reasoning-parser tests do not cover. + """ + parser_cls = ParserManager.get_parser( + tool_parser_name="mistral", + reasoning_parser_name="mistral", + enable_auto_tools=True, + ) + parser = parser_cls(mistral_tokenizer, None) + request = parser.adjust_request( + ChatCompletionRequest( + messages=[], + model="test", + tools=_SAMPLE_TOOLS, + tool_choice="auto", + ) + ) + assert getattr(request, "_grammar_from_parser", False) + + prompt_ids = _encode_v13(mistral_tokenizer, "[THINK]system reasoning[/THINK]") + gen_ids = _encode_v13( + mistral_tokenizer, + "[THINK]because two plus two is four[/THINK]The answer is 4.", + ) + + reasoning, content = _stream_parse_delta( + parser, mistral_tokenizer, gen_ids, request, prompt_ids + ) + + assert "because two plus two is four" in reasoning + assert "because two plus two is four" not in content + assert "The answer is 4." in content + + +def _encode_gen(tokenizer: MistralTokenizer, text: str) -> list[int]: + """Encode `text` mapping `[THINK]`, `[/THINK]`, `[TOOL_CALLS]`, `[ARGS]` + to their special-token IDs; all other text uses the base tokenizer. + """ + vocab = tokenizer.get_vocab() + markers: dict[str, int] = { + "[THINK]": tokenizer.instruct.BEGIN_THINK, + "[/THINK]": tokenizer.instruct.END_THINK, + "[TOOL_CALLS]": vocab["[TOOL_CALLS]"], + "[ARGS]": vocab["[ARGS]"], + } + out: list[int] = [] + remaining = text + while remaining: + earliest: int | None = None + earliest_marker: str | None = None + for marker in markers: + idx = remaining.find(marker) + if idx != -1 and (earliest is None or idx < earliest): + earliest = idx + earliest_marker = marker + if earliest_marker is None: + out += tokenizer.tokenizer.encode(remaining, False, False) + break + assert earliest is not None # set together with earliest_marker + if earliest > 0: + out += tokenizer.tokenizer.encode(remaining[:earliest], False, False) + out.append(markers[earliest_marker]) + remaining = remaining[earliest + len(earliest_marker) :] + return out + + +def _stream_turn( + parser, + tokenizer: MistralTokenizer, + gen_ids: list[int], + request, + prompt_ids: list[int], +) -> tuple[str, str, list[tuple[str, str]]]: + """Stream `gen_ids` through `parse_delta`, returning `(reasoning, content, + tool_calls)`. + + `tool_calls` is a list of `(name, arguments)` tuples ordered by + `DeltaToolCall.index`. Arguments are concatenated across deltas. + """ + reasoning = "" + content = "" + tool_call_acc: dict[int, dict[str, str]] = {} + previous_tokens: list[str] | None = None + prefix_offset = 0 + read_offset = 0 + for i, tok in enumerate(gen_ids): + new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer=tokenizer, + all_input_ids=gen_ids[: i + 1], + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=True, + spaces_between_special_tokens=True, + ) + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + delta_message = parser.parse_delta( + delta_text=delta_text, + delta_token_ids=[tok], + request=request, + prompt_token_ids=prompt_ids, + finished=(i == len(gen_ids) - 1), + ) + if delta_message is not None: + reasoning += delta_message.reasoning or "" + content += delta_message.content or "" + for tc in delta_message.tool_calls: + if tc.index not in tool_call_acc: + tool_call_acc[tc.index] = {"name": "", "args": ""} + if tc.function: + if tc.function.name: + tool_call_acc[tc.index]["name"] = tc.function.name + if tc.function.arguments: + tool_call_acc[tc.index]["args"] += tc.function.arguments + tool_calls = [ + (tool_call_acc[k]["name"], tool_call_acc[k]["args"]) + for k in sorted(tool_call_acc) + ] + return reasoning, content, tool_calls + + +def test_parse_delta_multi_turn_reason_and_tool( + mistral_tokenizer: MistralTokenizer, +) -> None: + """Two-turn loop: each assistant turn carries reasoning + a tool call. + + For each turn, the mistral-common conversation is encoded into `prompt_ids` + before streaming the next generation. Turn 2's prompt therefore contains + a `[THINK]`/`[/THINK]` block and `[TOOL_CALLS]` from turn 1 — the + accumulated tool-call/result history that single-turn tests never exercise. + + Asserts that across both turns the parser correctly classifies think tokens + as `reasoning_content` (never leaking into `content`), and that the tool + call name/arguments are reconstructed correctly from the streaming deltas. + """ + parser_cls = ParserManager.get_parser( + tool_parser_name="mistral", + reasoning_parser_name="mistral", + enable_auto_tools=True, + ) + request = parser_cls(mistral_tokenizer, None).adjust_request( + ChatCompletionRequest( + messages=[], + model="test", + tools=_SAMPLE_TOOLS, + tool_choice="auto", + ) + ) + assert getattr(request, "_grammar_from_parser", False) + + cities = ["Dallas", "Paris"] + conversation: list = [UserMessage(content="What is the weather in Dallas?")] + + for i, city in enumerate(cities): + # Fresh parser per turn: streaming state must not bleed across turns. + parser = parser_cls(mistral_tokenizer, None) + + prompt_ids = mistral_tokenizer.instruct.encode_instruct( + InstructRequest(messages=conversation) + ).tokens + + gen_str = ( + f"[THINK]I should check {city} weather.[/THINK]" + f'[TOOL_CALLS]get_weather[ARGS]{{"city": "{city}"}}' + ) + gen_ids = _encode_gen(mistral_tokenizer, gen_str) + + reasoning, content, tool_calls = _stream_turn( + parser, mistral_tokenizer, gen_ids, request, prompt_ids + ) + + assert f"check {city} weather" in reasoning, f"turn {i}: reasoning not captured" + assert f"check {city} weather" not in content, ( + f"turn {i}: reasoning leaked into content" + ) + assert "[THINK]" not in content, f"turn {i}: [THINK] leaked" + assert "[TOOL_CALLS]" not in content, f"turn {i}: [TOOL_CALLS] leaked" + assert len(tool_calls) == 1, f"turn {i}: expected 1 tool call" + name, args = tool_calls[0] + assert name == "get_weather", f"turn {i}: wrong tool name {name!r}" + assert city in args, f"turn {i}: city not in args {args!r}" + + # Extend conversation with this turn's history for the next turn. + tc_id = f"tc{i:07d}" # 9-char alphanumeric id: "tc0000000", "tc0000001" + conversation.extend( + [ + AssistantMessage( + content=[ + ThinkChunk( + thinking=f"I should check {city} weather.", + closed=True, + ) + ], + tool_calls=[ + ToolCall( + id=tc_id, + function=FunctionCall( + name="get_weather", + arguments=json.dumps({"city": city}), + ), + ) + ], + ), + ToolMessage( + content=f"{city}: 70F sunny", + tool_call_id=tc_id, + ), + ] + ) + if i < len(cities) - 1: + conversation.append(UserMessage(content=f"And in {cities[i + 1]}?")) diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/parser/mistral/test_tool_calls.py similarity index 53% rename from tests/tool_parsers/test_mistral_tool_parser.py rename to tests/parser/mistral/test_tool_calls.py index 4f75e2f576fe..99f99008d11d 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/parser/mistral/test_tool_calls.py @@ -7,6 +7,7 @@ import partial_json_parser import pytest +import regex as re from mistral_common.protocol.instruct.messages import AssistantMessage from mistral_common.protocol.instruct.request import InstructRequest from mistral_common.protocol.instruct.tool_calls import ( @@ -28,21 +29,26 @@ ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, DeltaMessage, DeltaToolCall, ExtractedToolCallInformation, StructuralTagResponseFormat, ) +from vllm.parser.engine.events import EventType +from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine +from vllm.parser.mistral import _DEFAULT_JSON_SCHEMA, MistralParser, mistral_config +from vllm.parser.parser_manager import ParserManager from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.detokenizer_utils import detokenize_incrementally from vllm.tokenizers.mistral import MistralTokenizer -from vllm.tool_parsers.mistral_tool_parser import ( - _DEFAULT_JSON_SCHEMA, - MistralToolParser, -) +from vllm.tool_parsers.mistral_tool_parser import MistralToolParser _DUMMY_REQUEST = ChatCompletionRequest(messages=[], model="test") +# Minimal request for the engine driver: tool_choice=None falls through to +# auto extraction without triggering "tools must be set" validation. +_ENGINE_REQUEST = ChatCompletionRequest(messages=[], model="test", tool_choice=None) @pytest.fixture(scope="module") @@ -71,6 +77,8 @@ def mistral_tool_parser(mistral_tokenizer): def non_mistral_parser() -> MistralToolParser: mock_tokenizer = MagicMock() mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1} + # Ensure the legacy (non-grammar) path is taken by MistralParser. + mock_tokenizer.supports_grammar = False return MistralToolParser(mock_tokenizer) @@ -96,9 +104,18 @@ def assert_tool_calls( assert actual_tool_call.function.name == expected_tool_call.function.name, ( f"got wrong function name:${actual_tool_call.function.name}" ) - assert ( - actual_tool_call.function.arguments == expected_tool_call.function.arguments - ), f"got wrong function argument:${actual_tool_call.function.arguments}" + actual_args = actual_tool_call.function.arguments + # Streaming deltas (DeltaToolCall) may have partial JSON because the + # engine holds back the final closing brace until finish() is called. + # Apply partial_json_parser so the comparison works for both the + # legacy (complete JSON) and engine (potentially partial) paths. + if isinstance(actual_tool_call, DeltaToolCall) and actual_args is not None: + actual_args = partial_json_parser.ensure_json( + actual_args, Allow.OBJ | Allow.STR + ) + assert actual_args == expected_tool_call.function.arguments, ( + f"got wrong function argument:${actual_tool_call.function.arguments}" + ) def fix_tool_call_tokenization( @@ -106,45 +123,82 @@ def fix_tool_call_tokenization( mistral_tool_parser: MistralToolParser, mistral_tokenizer: TokenizerLike, ): - """ - Replaces the textual token sequence for [TOOL_CALLS] - with its single special token ID. - """ - textual_tool_call_token_ids = mistral_tokenizer.encode( - text=mistral_tool_parser.bot_token, - add_special_tokens=False, - ) - # textual_tool_call_token_ids must not contain special tokens like bos, eos etc - special_tool_call_token_ids = [mistral_tool_parser.bot_token_id] + """Replace textual token sequences for the control markers ([TOOL_CALLS] + and [ARGS]) with their single special token IDs. - # If the input is too short to contain the sequence, no replacement is possible - if not tokens or len(tokens) < len(textual_tool_call_token_ids): - return tokens + Real generation emits these as control tokens; encoding the marker *string* + can instead yield literal text tokens. The engine detects markers by id, so + the test stream must carry the special ids. + """ + vocab = mistral_tokenizer.get_vocab() + markers: list[tuple[str, int]] = [ + ( + mistral_tool_parser._parser_engine.bot_token, + mistral_tool_parser._parser_engine.bot_token_id, + ) + ] + args_id = vocab.get("[ARGS]") + if args_id is not None: + markers.append(("[ARGS]", args_id)) + + # (textual token sequence, special id) pairs to replace. + replacements: list[tuple[list[int], int]] = [] + for text, special_id in markers: + if special_id is None: + continue + textual = mistral_tokenizer.encode(text=text, add_special_tokens=False) + if textual: + replacements.append((textual, special_id)) - result_tokens = [] + result_tokens: list[int] = [] i = 0 - target_len = len(textual_tool_call_token_ids) - while i < len(tokens): - # Check if the slice from the current position matches the target sequence - if tokens[i : i + target_len] == textual_tool_call_token_ids: - # If it matches, add the replacement and jump the index forward - result_tokens.extend(special_tool_call_token_ids) - i += target_len + for textual, special_id in replacements: + if tokens[i : i + len(textual)] == textual: + result_tokens.append(special_id) + i += len(textual) + break else: - # Otherwise, just add the current token and move to the next one result_tokens.append(tokens[i]) i += 1 return result_tokens +def encode_mistral_output(tokenizer: MistralTokenizer, text: str) -> list[int]: + """Encode generated output with control markers as their special-token ids. + + Real generation emits ``[TOOL_CALLS]``/``[ARGS]`` as single control tokens. + Encoding the marker *strings* can instead glue them to neighbours (e.g. + ``}[``) and yield literal text tokens, which the id-based engine never sees. + Split on the markers and insert their special ids so the token stream + matches real generation. + """ + vocab = tokenizer.get_vocab() + marker_ids = { + "[TOOL_CALLS]": vocab.get("[TOOL_CALLS]"), + "[ARGS]": vocab.get("[ARGS]"), + } + ids: list[int] = [] + for part in re.split(r"(\[TOOL_CALLS\]|\[ARGS\])", text): + if part in marker_ids: + # Fall back to text encoding here and the id-based engine path + # silently stops being exercised while the tests still pass. + assert marker_ids[part] is not None, f"{part} missing from vocab" + ids.append(marker_ids[part]) + elif part: + ids.extend(tokenizer.encode(part, add_special_tokens=False)) + return ids + + def stream_delta_message_generator( mistral_tool_parser: MistralToolParser, mistral_tokenizer: TokenizerLike, model_output: str | None, tools: list[tuple[str, str]] | None, chunk_size: int = 1, + driver: str = "legacy", + reasoning_parser: str | None = None, ) -> Generator[DeltaMessage, None, None]: if ( isinstance(mistral_tokenizer, MistralTokenizer) @@ -179,6 +233,33 @@ def stream_delta_message_generator( all_token_ids, mistral_tool_parser, mistral_tokenizer ) + if isinstance(mistral_tokenizer, MistralTokenizer): + # Real serving streams only generated tokens; encode_instruct adds + # framing (BOS/EOS) the parser never receives. Strip them so the + # id-based engine detection sees a realistic stream. + if all_token_ids and all_token_ids[0] == mistral_tokenizer.bos_token_id: + all_token_ids = all_token_ids[1:] + if all_token_ids and all_token_ids[-1] == mistral_tokenizer.eos_token_id: + all_token_ids = all_token_ids[:-1] + + if driver == "engine": + _parser_cls = ParserManager.get_parser( + tool_parser_name="mistral", + reasoning_parser_name=reasoning_parser, + enable_auto_tools=True, + ) + _engine_parser = _parser_cls(mistral_tokenizer, None) + if reasoning_parser is not None: + _req_with_tools = ChatCompletionRequest( + messages=[], + model="test", + tools=SAMPLE_TOOLS_DICTS, + tool_choice="auto", + ) + _engine_req = _engine_parser.adjust_request(_req_with_tools) + else: + _engine_req = _ENGINE_REQUEST + previous_text = "" previous_tokens = None prefix_offset = 0 @@ -215,15 +296,24 @@ def stream_delta_message_generator( 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, - pending_text, - previous_token_ids, - current_token_ids, - pending_token_ids, - request=_DUMMY_REQUEST, - ) + if driver == "legacy": + delta_message = mistral_tool_parser.extract_tool_calls_streaming( + previous_text, + current_text, + pending_text, + previous_token_ids, + current_token_ids, + pending_token_ids, + request=_DUMMY_REQUEST, + ) + else: # engine: per-delta windows, mirroring parse_delta behaviour + delta_message = _engine_parser.parse_delta( + delta_text=pending_text, + delta_token_ids=list(pending_token_ids), + request=_engine_req, + prompt_token_ids=[1], + finished=(i == len(all_token_ids) - 1), + ) if delta_message: yield delta_message @@ -231,6 +321,14 @@ def stream_delta_message_generator( pending_text = "" pending_token_ids = [] + # For the legacy driver the engine holds back the final closing brace + # in _args_buffer until finish_streaming() is called, mirroring how + # the real serving layer flushes at end of stream. + if driver == "legacy": + flush_delta = mistral_tool_parser.finish_streaming() + if flush_delta: + yield flush_delta + @pytest.mark.parametrize( "parser_fixture", @@ -424,8 +522,14 @@ def test_extract_tool_calls_pre_v11_regex_fallback_fails( "single_tool_add", "single_tool_weather", "multiple_tool_calls", + "single_tool_add_args", + "single_tool_weather_args", + "multiple_tool_calls_args", + "content_before_tool_args", "complex", "wrong_json", + "trailing_text_after_json", + "trailing_text_after_args_json", ], argnames=["model_output", "expected_tool_calls", "expected_content"], argvalues=[ @@ -471,6 +575,60 @@ def test_extract_tool_calls_pre_v11_regex_fallback_fails( ], None, ), + ( + # v11+ emits an explicit [ARGS] separator between name and args. + """[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add_this_and_that", + arguments=json.dumps({"a": 3.5, "b": 4}), + ) + ) + ], + None, + ), + ( + """[TOOL_CALLS]get_current_weather[ARGS]{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "San Francisco", "state": "CA", "unit": "celsius"} + ), + ) + ) + ], + None, + ), + ( + """[TOOL_CALLS]add[ARGS]{"a": 3.5, "b": 4}[TOOL_CALLS]multiply[ARGS]{"a": 3, "b": 6}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 3.5, "b": 4}) + ) + ), + ToolCall( + function=FunctionCall( + name="multiply", arguments=json.dumps({"a": 3, "b": 6}) + ) + ), + ], + None, + ), + ( + """hi[TOOL_CALLS]add[ARGS]{"a": 1, "b": 2}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 1, "b": 2}) + ) + ) + ], + "hi", + ), ( # Complex """hi{hi[TOOL_CALLS]bash{"command": "print(\\"hello world!\\")\\nre.compile(r\'{}\')""", # noqa: E501 @@ -501,6 +659,33 @@ def test_extract_tool_calls_pre_v11_regex_fallback_fails( ], "hi{hi", ), + ( + # gh#48975: trailing text after the JSON object must not leak into + # arguments (v11+ name{args} format has no terminator). + """[TOOL_CALLS]get_current_weather{"city": "San Francisco"}Estimating the weather...""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps({"city": "San Francisco"}), + ) + ) + ], + None, + ), + ( + # gh#48975 with the explicit [ARGS] separator. + """[TOOL_CALLS]get_current_weather[ARGS]{"city": "San Francisco"}Estimating the weather...""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps({"city": "San Francisco"}), + ) + ) + ], + None, + ), ], ) def test_extract_tool_calls( @@ -517,17 +702,128 @@ def test_extract_tool_calls( def test_extract_tool_calls_v11_without_args_skipped(mistral_tool_parser): + # Before Stage 2: the legacy path skipped tool calls with no arg separator, + # returning tools_called=True with an empty list. After Stage 2 the engine + # path is used for v11; it recognises the name and emits a tool call with + # empty arguments ("{}"). The semantic result is equivalent — a valid zero- + # argument call — and is strictly more correct than silently dropping it. model_output = "[TOOL_CALLS]toolname_no_args" result = mistral_tool_parser.extract_tool_calls( model_output, request=_DUMMY_REQUEST ) - assert result == ExtractedToolCallInformation( - tools_called=True, tool_calls=[], content=None + assert result.tools_called + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "toolname_no_args" + assert result.tool_calls[0].function.arguments == "{}" + assert result.content is None + + +def test_extract_tool_calls_malformed_name_before_marker_emits_empty_name( + mistral_tool_parser, +): + model_output = 'bash[TOOL_CALLS]{"command": "cat file.py"}' + result = mistral_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "" + assert result.tool_calls[0].function.arguments == json.dumps( + {"command": "cat file.py"} + ) + assert result.content == "bash" + + +def test_extract_tool_calls_well_formed_name_unaffected(mistral_tool_parser): + model_output = '[TOOL_CALLS]bash[ARGS]{"command": "ls -la"}' + result = mistral_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "bash" + assert result.tool_calls[0].function.arguments == json.dumps({"command": "ls -la"}) + assert result.content is None + + +def test_extract_tool_calls_empty_name_unparsable_args_no_crash( + mistral_tool_parser, +): + # Empty name AND args that never form valid JSON: still emitted (raw + # text as arguments) instead of raising. + model_output = "bash[TOOL_CALLS]{not valid json at all" + result = mistral_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "" + assert result.tool_calls[0].function.arguments == "{not valid json at all" + assert result.content == "bash" + + +def test_extract_tool_calls_streaming_malformed_name_before_marker( + mistral_tool_parser, +): + # Streaming counterpart: once the args start the slot can never go back to + # naming, so the (empty) name is already final and streams immediately + # alongside the args. + model_output = 'bash[TOOL_CALLS]{"command": "cat file.py"}' + mid_delta = mistral_tool_parser.extract_tool_calls_streaming( + previous_text="", + current_text=model_output, + delta_text=model_output, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=_DUMMY_REQUEST, + ) + assert mid_delta is not None + assert mid_delta.content == "bash" + assert len(mid_delta.tool_calls) == 1 + assert mid_delta.tool_calls[0].function.name == "" + assert mid_delta.tool_calls[0].function.arguments == '{"command": "cat file.py"' + + final_delta = mistral_tool_parser.finish_streaming() + assert final_delta is not None + assert len(final_delta.tool_calls) == 1 + assert final_delta.tool_calls[0].function.name is None + assert final_delta.tool_calls[0].function.arguments == "}" + + +def test_extract_tool_calls_parallel_malformed_then_well_formed( + mistral_tool_parser, +): + # A malformed empty-name call followed by a well-formed one must not + # regress into brace-splitting: nested braces in the second call's + # arguments must stay intact, and both calls must be correctly separated. + model_output = ( + 'bash[TOOL_CALLS]{"command": "ls"}' + '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}' + ) + result = mistral_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "" + assert result.tool_calls[0].function.arguments == json.dumps({"command": "ls"}) + assert result.tool_calls[1].function.name == "submit" + assert result.tool_calls[1].function.arguments == json.dumps( + {"answer": {"nested": 1}} ) + assert result.content == "bash" def _test_extract_tool_calls_streaming( - tool_parser, tokenizer, model_output, tools, expected_tool_calls, expected_content + tool_parser, + tokenizer, + model_output, + tools, + expected_tool_calls, + expected_content, + driver: str = "legacy", + reasoning_parser: str | None = None, ): other_content: str = "" function_names: list[str] = [] @@ -536,7 +832,12 @@ def _test_extract_tool_calls_streaming( tool_call_ids: list[str | None] = [] for delta_message in stream_delta_message_generator( - tool_parser, tokenizer, model_output, tools + tool_parser, + tokenizer, + model_output, + tools, + driver=driver, + reasoning_parser=reasoning_parser, ): # role should never be streamed from tool parser assert not delta_message.role @@ -547,12 +848,36 @@ def _test_extract_tool_calls_streaming( streamed_tool_calls = delta_message.tool_calls if streamed_tool_calls and len(streamed_tool_calls) > 0: + # On the final finished=True delta the engine's _flush_engine_parsers + # may append the flushed '}' as a separate DeltaToolCall for the + # same index without coalescing. Merge entries sharing an index so + # the "one diff per delta" invariant can be checked after merging. + if len(streamed_tool_calls) > 1: + assert len({tc.index for tc in streamed_tool_calls}) == 1, ( + "only within-chunk finish/flush of one tool is allowed; " + "distinct indices in one delta indicate a regression: " + f"{[tc.index for tc in streamed_tool_calls]}" + ) + by_index: dict[int, DeltaToolCall] = {} + for tc in streamed_tool_calls: + existing = by_index.get(tc.index) + if existing is None: + by_index[tc.index] = tc + else: + if tc.id and not existing.id: + existing.id = tc.id + if tc.type and not existing.type: + existing.type = tc.type + if tc.function and existing.function and tc.function.arguments: + existing.function.arguments = ( + existing.function.arguments or "" + ) + tc.function.arguments + streamed_tool_calls = list(by_index.values()) + # make sure only one diff is present - correct even for parallel assert len(streamed_tool_calls) == 1 tool_call = streamed_tool_calls[0] - assert len(tool_parser.prev_tool_call_arr) > 0 - # if a new tool is being called, set up empty arguments if tool_call.index != tool_call_idx: tool_call_idx = tool_call.index @@ -595,34 +920,8 @@ def _test_extract_tool_calls_streaming( ] assert_tool_calls(actual_tool_calls, expected_tool_calls) - if expected_tool_calls: - assert len(tool_parser.streamed_args_for_tool) == len(expected_tool_calls) - assert len(tool_parser.prev_tool_call_arr) == len(expected_tool_calls) - for i in range(len(expected_tool_calls)): - assert ( - tool_parser.prev_tool_call_arr[i]["arguments"] - == tool_parser.streamed_args_for_tool[i] - ) - assert tool_parser.streamed_args_for_tool[i] == function_args_strs[i] - assert ( - tool_parser.prev_tool_call_arr[i]["name"] - == expected_tool_calls[i].function.name - ) - - # Simulate the serving layer's unstreamed-args check - index = len(tool_parser.prev_tool_call_arr) - 1 - args = tool_parser.prev_tool_call_arr[index].get("arguments", {}) - expected_call = ( - args if isinstance(args, str) else json.dumps(args, ensure_ascii=False) - ) - actual_call = tool_parser.streamed_args_for_tool[index] - remaining_call = expected_call.replace(actual_call, "", 1) - assert remaining_call == "" - else: - assert len(tool_parser.streamed_args_for_tool) == 0 - assert len(tool_parser.prev_tool_call_arr) == 0 - +@pytest.mark.parametrize("driver", ["legacy", "engine"]) @pytest.mark.parametrize( ids=[ "no_tools", @@ -748,6 +1047,7 @@ def test_extract_tool_calls_streaming_pre_v11_tokenizer( model_output, expected_tool_calls, expected_content, + driver, ): _test_extract_tool_calls_streaming( mistral_pre_v11_tool_parser, @@ -756,9 +1056,15 @@ def test_extract_tool_calls_streaming_pre_v11_tokenizer( None, expected_tool_calls, expected_content, + driver=driver, ) +@pytest.mark.parametrize( + "driver,reasoning_parser", + [("legacy", None), ("engine", None), ("engine", "mistral")], + ids=["legacy", "engine_no_reasoning", "engine_with_reasoning"], +) @pytest.mark.parametrize( ids=[ "single_tool_add", @@ -826,6 +1132,8 @@ def test_extract_tool_calls_streaming( tools, expected_tool_calls, expected_content, + driver, + reasoning_parser, ): _test_extract_tool_calls_streaming( mistral_tool_parser, @@ -834,17 +1142,19 @@ def test_extract_tool_calls_streaming( tools, expected_tool_calls, expected_content, + driver=driver, + reasoning_parser=reasoning_parser, ) +# Drives extract_tool_calls_streaming directly (bespoke detokenization loop), +# not via stream_delta_message_generator — driver parametrization not needed. def test_extract_tool_calls_streaming_v11_no_tools( mistral_tool_parser, mistral_tokenizer ): model_output = "This is a test" - if isinstance(mistral_tokenizer, MistralTokenizer): - all_token_ids = mistral_tokenizer.encode(model_output) - else: - all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) + # add_special_tokens=False: real serving streams only generated tokens. + all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) skip_special = isinstance(mistral_tokenizer, MistralTokenizer) collected_content = "" previous_text = "" @@ -887,8 +1197,47 @@ def test_extract_tool_calls_streaming_v11_no_tools( previous_text = current_text assert collected_content == model_output - assert len(mistral_tool_parser.streamed_args_for_tool) == 0 - assert len(mistral_tool_parser.prev_tool_call_arr) == 0 + + +def test_mistral_parser_drops_eos_from_output(mistral_tokenizer): + """EOS token must never surface as content/reasoning; literal EOS string + must be preserved when the EOS token id is absent. + + The engine drops special tokens by id, so: + Case A: the real EOS token id causes the text to be dropped. + Case B: the same EOS string with a non-EOS token id is preserved. + """ + if not isinstance(mistral_tokenizer, MistralTokenizer): + pytest.skip("Requires MistralTokenizer") + + eos_id: int = mistral_tokenizer.eos_token_id + eos_text: str = mistral_tokenizer.decode([eos_id]) + assert eos_text + + # Case A: real EOS token — text must be dropped. + parser_a = MistralParser(mistral_tokenizer) + parser_a.initialize_streaming() + events_a = parser_a._feed(eos_text, [eos_id]) + events_a.extend(parser_a._engine.finish()) + delta_a = parser_a._events_to_delta(events_a, finished=True) + content_a = (delta_a.content if delta_a else None) or "" + reasoning_a = (delta_a.reasoning if delta_a else None) or "" + assert eos_text not in content_a + assert eos_text not in reasoning_a + + # Case B: EOS text with a non-EOS token id — text must be preserved. + # This proves content equal to the EOS string is not silently dropped + # when it did not come from the real EOS token. + non_eos_ids = mistral_tokenizer.encode(text="hello", add_special_tokens=False) + non_eos_id = next(tid for tid in non_eos_ids if tid != eos_id) + parser_b = MistralParser(mistral_tokenizer) + parser_b.initialize_streaming() + events_b = parser_b._feed(eos_text, [non_eos_id]) + events_b.extend(parser_b._engine.finish()) + delta_b = parser_b._events_to_delta(events_b, finished=True) + content_b = (delta_b.content if delta_b else None) or "" + reasoning_b = (delta_b.reasoning if delta_b else None) or "" + assert eos_text in content_b or eos_text in reasoning_b @pytest.mark.parametrize( @@ -961,6 +1310,55 @@ def test_extract_tool_calls_streaming_v11_no_tools( "bla", id="v11-content_before_tool", ), + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", + """[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add_this_and_that", + arguments=json.dumps({"a": 3.5, "b": 4}), + ) + ) + ], + "", + id="v11-single_tool_add_args", + ), + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", + """[TOOL_CALLS]add[ARGS]{"a": 3.5, "b": 4}[TOOL_CALLS]multiply[ARGS]{"a": 3, "b": 6}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 3.5, "b": 4}) + ) + ), + ToolCall( + function=FunctionCall( + name="multiply", arguments=json.dumps({"a": 3, "b": 6}) + ) + ), + ], + "", + id="v11-multiple_tool_calls_args", + ), + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", + """bla[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add_this_and_that", + arguments=json.dumps({"a": 3.5, "b": 4}), + ) + ) + ], + "bla", + id="v11-content_before_tool_args", + ), pytest.param( "mistral_tool_parser", "mistral_tokenizer", @@ -1116,11 +1514,15 @@ def test_extract_tool_calls_streaming_one_chunk( tool_parser = request.getfixturevalue(parser_fixture) tokenizer = request.getfixturevalue(tokenizer_fixture) - if isinstance(tokenizer, MistralTokenizer): - all_token_ids = tokenizer.encode(model_output) + # Real serving streams only generated tokens (no BOS/EOS framing) with the + # control markers as special ids. + if isinstance(tokenizer, MistralTokenizer) and tokenizer.version >= 11: + all_token_ids = encode_mistral_output(tokenizer, model_output) else: all_token_ids = tokenizer.encode(model_output, add_special_tokens=False) - all_token_ids = fix_tool_call_tokenization(all_token_ids, tool_parser, tokenizer) + all_token_ids = fix_tool_call_tokenization( + all_token_ids, tool_parser, tokenizer + ) delta_message = tool_parser.extract_tool_calls_streaming( previous_text="", @@ -1131,6 +1533,26 @@ def test_extract_tool_calls_streaming_one_chunk( delta_token_ids=all_token_ids, request=_DUMMY_REQUEST, ) + # The engine buffers the final closing brace ('}') in _args_buffer until + # finish_streaming() is called, mirroring real serving. Flush it and + # merge the result so the asserted arguments are complete. + flush_delta = tool_parser.finish_streaming() + if flush_delta is not None and flush_delta.tool_calls: + if delta_message is not None and delta_message.tool_calls: + for flush_tc in flush_delta.tool_calls: + for existing_tc in delta_message.tool_calls: + if existing_tc.index == flush_tc.index and flush_tc.function: + existing_tc.function = ( + existing_tc.function or DeltaFunctionCall() + ) + existing_tc.function.arguments = ( + existing_tc.function.arguments or "" + ) + (flush_tc.function.arguments or "") + elif delta_message is not None: + delta_message.tool_calls = flush_delta.tool_calls + else: + delta_message = flush_delta + assert isinstance(delta_message, DeltaMessage) assert len(delta_message.tool_calls) == len(expected_tool_calls) @@ -1145,13 +1567,6 @@ def test_extract_tool_calls_streaming_one_chunk( @pytest.mark.parametrize( "parser_fixture, model_output, fake_count, two_phase", [ - pytest.param( - "mistral_tool_parser", - '[TOOL_CALLS]add{"a": 1, "b": 2}', - 20, - True, - id="v11", - ), pytest.param( "mistral_pre_v11_tool_parser", '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]', @@ -1164,7 +1579,11 @@ def test_extract_tool_calls_streaming_one_chunk( def test_fast_detokenization_text_detection( parser_fixture, model_output, fake_count, two_phase, request ): - """Regression: bot_token in text but not token_ids (PR #37209).""" + """Regression: bot_token in text but not token_ids (PR #37209). + + Only the pre-v11 legacy path detects the marker from text; v11+ routes + through the engine and detects the marker by token id. + """ parser = request.getfixturevalue(parser_fixture) # Token IDs that do NOT contain bot_token_id. fake_token_ids = list(range(99, 99 + fake_count)) @@ -1210,34 +1629,24 @@ def test_fast_detokenization_text_detection( assert delta_message.tool_calls[0].function.name == "add" -@pytest.mark.parametrize( - "parser_fixture, patched_method, current_text", - [ - ( - "mistral_tool_parser", - "_extract_tool_calls_streaming", - "[TOOL_CALLS]add{}", - ), - ( - "mistral_pre_v11_tool_parser", - "_extract_tool_calls_streaming_pre_v11_tokenizer", - '[TOOL_CALLS] [{"name":"a","arguments":{}}]', - ), - ], - ids=["v11", "pre_v11"], -) def test_extract_tool_calls_streaming_exception_returns_none( - parser_fixture, patched_method, current_text, request + mistral_pre_v11_tool_parser, ): - parser = request.getfixturevalue(parser_fixture) - with patch.object(parser, patched_method, side_effect=RuntimeError("boom")): + # v11 routes through the ParserEngine and does not swallow exceptions; + # only the pre-v11 legacy state-machine path has explicit exception handling. + parser = mistral_pre_v11_tool_parser + patched_method = "_extract_tool_calls_streaming_pre_v11_tokenizer" + current_text = '[TOOL_CALLS] [{"name":"a","arguments":{}}]' + with patch.object( + parser._parser_engine, patched_method, side_effect=RuntimeError("boom") + ): result = parser.extract_tool_calls_streaming( previous_text="", current_text=current_text, delta_text=current_text, previous_token_ids=[], - current_token_ids=[parser.bot_token_id], - delta_token_ids=[parser.bot_token_id], + current_token_ids=[parser._parser_engine.bot_token_id], + delta_token_ids=[parser._parser_engine.bot_token_id], request=_DUMMY_REQUEST, ) assert result is None @@ -1368,14 +1777,17 @@ def test_adjust_request_unsupported_grammar_for_tokenizer(mistral_tokenizer) -> @pytest.mark.parametrize( "tool_choice,expected_skip", - [("auto", False), ("none", True)], - ids=["auto_skip_false", "none_skip_true"], + [("auto", False), ("none", False)], + ids=["auto_skip_false", "none_skip_false"], ) def test_adjust_request_non_mistral_tokenizer( non_mistral_parser: MistralToolParser, tool_choice: str, expected_skip: bool, ) -> None: + # MistralParser (ParserEngine) always sets skip_special_tokens=False so + # that special tokens like [TOOL_CALLS] are visible to the parser even + # when tool_choice="none". request = _make_request(tool_choice=tool_choice) result = non_mistral_parser.adjust_request(request) @@ -1572,22 +1984,23 @@ def test_adjust_request_tool_choice_with_json_schema_factory_routing( assert len(result.structured_outputs.grammar) > 0 -def test_grammar_from_tool_parser_default_false() -> None: +def test_grammar_from_parser_default_false() -> None: request = _make_request() - assert request._grammar_from_tool_parser is False + assert request._grammar_from_parser is False -def test_grammar_from_tool_parser_set_by_adjust_request( +def test_grammar_from_parser_set_by_adjust_request( mistral_tool_parser: MistralToolParser, ) -> None: request = _make_request() result = mistral_tool_parser.adjust_request(request) - assert result._grammar_from_tool_parser is True + assert result._grammar_from_parser is True +@pytest.mark.parametrize("driver", ["legacy", "engine"]) @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 + mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size, driver ): """A batched delta spanning the boundary between two parallel calls must keep them on distinct indices (the bug collapsed both onto index 0).""" @@ -1605,6 +2018,7 @@ def test_streaming_pre_v11_parallel_calls_batched_deltas( model_output, tools=None, chunk_size=chunk_size, + driver=driver, ): for tool_call in delta_message.tool_calls or []: if tool_call.index != idx: @@ -1619,3 +2033,623 @@ def test_streaming_pre_v11_parallel_calls_batched_deltas( 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} + + +@pytest.mark.parametrize( + "reasoning_encoding", + ["text", "special_token"], + ids=["text_encoding", "special_token_encoding"], +) +def test_content_tool_calls_transition_emits_reasoning_end( + mistral_tokenizer, + reasoning_encoding, +): + """(CONTENT, TOOL_CALLS) must emit REASONING_END when reasoning is enabled.""" + + cfg = mistral_config(reasoning_encoding=reasoning_encoding) + engine = StreamingParserEngine( + config=cfg, tokenizer=mistral_tokenizer, vocab=mistral_tokenizer.get_vocab() + ) + engine.skip_tool_parsing = True + + bot_token_id = mistral_tokenizer.get_vocab().get("[TOOL_CALLS]") + assert bot_token_id is not None + + events = engine.feed(delta_text="[TOOL_CALLS]", delta_token_ids=[bot_token_id]) + event_types = [e.type for e in events] + + assert EventType.REASONING_END in event_types, ( + f"REASONING_END missing for reasoning_encoding={reasoning_encoding!r}; " + f"got events: {event_types}" + ) + # Tool call start must NOT appear in skip_tool_parsing mode + assert EventType.TOOL_CALL_START not in event_types + + +@pytest.mark.parametrize("chunk_size", [1, 3]) +@pytest.mark.parametrize( + "tools,expected_tool_calls,expected_content", + [ + ( + [("add", '{"a": 3, "b": 4}')], + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 3, "b": 4}) + ) + ) + ], + "", + ), + ( + [ + ("add", '{"a": 3.5, "b": 4}'), + ( + "get_current_weather", + '{"city": "San Francisco", "state": "CA", "unit": "celsius"}', + ), + ], + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 3.5, "b": 4}) + ) + ), + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + { + "city": "San Francisco", + "state": "CA", + "unit": "celsius", + } + ), + ) + ), + ], + "", + ), + ], + ids=["single_tool_add", "parallel_tools"], +) +def test_reasoning_active_no_think_block_no_leak( + mistral_tool_parser, + mistral_tokenizer, + tools, + expected_tool_calls, + expected_content, + chunk_size, +): + """Tool call without a preceding think block must not leak as content + when the reasoning parser is active.""" + accumulated_content = "" + function_names: list[str] = [] + function_args_strs: list[str] = [] + tool_call_idx = -1 + tool_call_ids: list[str | None] = [] + + for delta_message in stream_delta_message_generator( + mistral_tool_parser, + mistral_tokenizer, + model_output=None, + tools=tools, + chunk_size=chunk_size, + driver="engine", + reasoning_parser="mistral", + ): + if delta_message.content: + accumulated_content += delta_message.content + + for tool_call in delta_message.tool_calls or []: + if tool_call.index != tool_call_idx: + tool_call_idx = tool_call.index + function_args_strs.append("") + tool_call_ids.append(None) + if tool_call.id and not tool_call_ids[tool_call.index]: + tool_call_ids[tool_call.index] = tool_call.id + if tool_call.function: + if tool_call.function.name: + function_names.append(tool_call.function.name) + if tool_call.function.arguments: + function_args_strs[tool_call.index] += tool_call.function.arguments + + assert "[TOOL_CALLS]" not in accumulated_content + assert "[ARGS]" not in accumulated_content + assert accumulated_content == expected_content + + actual_tool_calls = [ + ToolCall( + id=tc_id, + function=FunctionCall( + name=name, + arguments=partial_json_parser.ensure_json(args, Allow.OBJ | Allow.STR), + ), + ) + for tc_id, name, args in zip(tool_call_ids, function_names, function_args_strs) + ] + assert_tool_calls(actual_tool_calls, expected_tool_calls) + + +# --------------------------------------------------------------------------- +# Pre-v11 guided schema injection and bare-array extraction tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "tool_choice,expected_items_count", + [ + ("required", 2), # all tools + ( + {"type": "function", "function": {"name": "get_weather"}}, + 1, + ), # named: single tool + ], + ids=["required", "named"], +) +def test_adjust_request_pre_v11_guided_schema_injected( + mistral_pre_v11_tool_parser: MistralToolParser, + tool_choice: object, + expected_items_count: int, +) -> None: + request = _make_request(tool_choice=tool_choice) + result = mistral_pre_v11_tool_parser.adjust_request(request) + + assert result.structured_outputs is not None + schema = result.structured_outputs.json + if isinstance(schema, str): + schema = json.loads(schema) + assert schema["type"] == "array" + assert schema["minItems"] == 1 + items = schema["items"] + assert "anyOf" in items + assert len(items["anyOf"]) == expected_items_count + for entry in items["anyOf"]: + props = entry["properties"] + assert "name" in props + assert "arguments" in props + + +@pytest.mark.parametrize( + "response_format", + [{"type": "json_object"}, {"type": "json_schema", "json_schema": {"name": "s"}}], + ids=["json_object", "json_schema"], +) +def test_adjust_request_pre_v11_required_clears_response_format( + mistral_pre_v11_tool_parser: MistralToolParser, + response_format: dict, +) -> None: + """required + response_format must inject the tool schema and clear + response_format so the tool schema is the sole structured-output + constraint, matching the base ToolParser (otherwise the request either + hits the "multiple constraints" engine error or, with no constraint, + rambles to finish_reason='length').""" + request = _make_request(tool_choice="required", response_format=response_format) + result = mistral_pre_v11_tool_parser.adjust_request(request) + + assert result.response_format is None + assert result.structured_outputs is not None + schema = result.structured_outputs.json + if isinstance(schema, str): + schema = json.loads(schema) + assert schema["type"] == "array" + + +def test_adjust_request_non_mistral_tokenizer_required_injects_schema( + non_mistral_parser: MistralToolParser, +) -> None: + """The guided-schema injection must also fire for non-Mistral (e.g. HF-mode) + tokenizers driving the Mistral tool parser, mirroring the base ToolParser so + required + response_format does not fall back to an unconstrained ramble.""" + request = _make_request( + tool_choice="required", response_format={"type": "json_object"} + ) + result = non_mistral_parser.adjust_request(request) + + assert result.response_format is None + assert result.structured_outputs is not None + assert result.structured_outputs.json is not None + + +@pytest.mark.parametrize( + "tool_choice", + ["auto", "none"], + ids=["auto", "none"], +) +def test_adjust_request_pre_v11_no_injection_for_auto_none( + mistral_pre_v11_tool_parser: MistralToolParser, + tool_choice: str, +) -> None: + request = _make_request(tool_choice=tool_choice) + result = mistral_pre_v11_tool_parser.adjust_request(request) + + assert result.structured_outputs is None + + +def test_legacy_extract_tool_calls_guided_bare_array_required( + mistral_pre_v11_tool_parser: MistralToolParser, +) -> None: + model_output = '[{"name": "get_current_weather", "arguments": {"city": "Dallas"}}]' + request = _make_request(tool_choice="required") + + result = mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=request + ) + + assert result.tools_called + assert len(result.tool_calls) == 1 + tc = result.tool_calls[0] + assert tc.function.name == "get_current_weather" + assert json.loads(tc.function.arguments) == {"city": "Dallas"} + assert isinstance(tc.id, str) + assert len(tc.id) == 9 + assert result.content is None + + +def test_legacy_extract_tool_calls_none_with_tools( + mistral_pre_v11_tool_parser: MistralToolParser, +) -> None: + model_output = ( + '[TOOL_CALLS] [{"name": "get_current_weather",' + ' "arguments": {"city": "Dallas"}}]' + ) + request = _make_request(tool_choice="none") + + result = mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=request + ) + + assert not result.tools_called + assert result.tool_calls == [] + assert result.content == model_output + + +def test_guided_streaming_required_pre_v11( + mistral_pre_v11_tool_parser: MistralToolParser, + mistral_pre_v11_tokenizer, +) -> None: + model_output = '[{"name": "get_current_weather", "arguments": {"city": "Dallas"}}]' + request = _make_request(tool_choice="required") + all_token_ids = mistral_pre_v11_tokenizer.encode( + model_output, add_special_tokens=False + ) + + function_name: str | None = None + function_args = "" + tool_call_id: str | None = None + previous_text = "" + previous_tokens = None + prefix_offset = 0 + read_offset = 0 + + for i, token_id in enumerate(all_token_ids): + (new_tokens, delta_text, prefix_offset, read_offset) = detokenize_incrementally( + tokenizer=mistral_pre_v11_tokenizer, + all_input_ids=all_token_ids[: i + 1], + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=False, + spaces_between_special_tokens=True, + ) + previous_tokens = ( + (previous_tokens + new_tokens) if previous_tokens else new_tokens + ) + current_text = previous_text + delta_text + + delta_message = mistral_pre_v11_tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=all_token_ids[:i], + current_token_ids=all_token_ids[: i + 1], + delta_token_ids=[token_id], + request=request, + ) + previous_text = current_text + + if delta_message and delta_message.tool_calls: + for tc in delta_message.tool_calls: + if tc.id and not tool_call_id: + tool_call_id = tc.id + if tc.function: + if tc.function.name: + function_name = tc.function.name + if tc.function.arguments: + function_args += tc.function.arguments + + assert function_name == "get_current_weather" + assert tool_call_id is not None + assert len(tool_call_id) == 9 + assert json.loads(function_args) == {"city": "Dallas"} + + +# --------------------------------------------------------------------------- +# Malformed-name-before-marker matrix: {malformed, well-formed, mixed +# parallel} x {streaming, non-streaming}. +# +# These tests drive a real MistralParser built from the module's +# `mistral_tokenizer` fixture (grammar-capable v11+ tokenizer, so the +# declarative ParserEngine path is exercised, not the pre-v11 legacy state +# machine) with actual bash/search_replace/submit tools, mirroring the tools +# present in the real captured traces this bug was found in. +# --------------------------------------------------------------------------- + +_BASH_SEARCH_SUBMIT_TOOLS_DICTS = [ + { + "type": "function", + "function": { + "name": "bash", + "description": "Run a bash command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "search_replace", + "description": "Search and replace text in a file", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string"}, + "search": {"type": "string"}, + "replace": {"type": "string"}, + }, + "required": ["path", "search", "replace"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "submit", + "description": "Submit the final answer", + "parameters": { + "type": "object", + "properties": {"answer": {"type": "object"}}, + "required": ["answer"], + }, + }, + }, +] + + +def _malformed_matrix_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + messages=[], + model="test", + tools=_BASH_SEARCH_SUBMIT_TOOLS_DICTS, + tool_choice="auto", + ) + + +def _stream_via_parse_delta( + parser: MistralParser, + tokenizer: MistralTokenizer, + model_output: str, + request: ChatCompletionRequest, +) -> tuple[str, list[str | None], list[str]]: + """Feed `model_output` through `parser.parse_delta` token-by-token. + + This is the exact API production serving drives on every generated + token (see `ParserEngine.parse_delta`), unlike + `extract_tool_calls_streaming` which some legacy call sites drive with + pre-computed previous/current text windows. + """ + parser.initialize_streaming() + all_token_ids = encode_mistral_output(tokenizer, model_output) + + content = "" + tool_call_idx = -1 + names: list[str | None] = [] + args: list[str] = [] + + previous_tokens = None + prefix_offset = 0 + read_offset = 0 + for i, token_id in enumerate(all_token_ids): + new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer=tokenizer, + all_input_ids=all_token_ids[: i + 1], + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=False, + spaces_between_special_tokens=True, + ) + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + finished = i == len(all_token_ids) - 1 + + delta_message = parser.parse_delta( + delta_text=delta_text, + delta_token_ids=[token_id], + request=request, + prompt_token_ids=[1], + finished=finished, + ) + if delta_message is None: + continue + if delta_message.content: + content += delta_message.content + for tool_call in delta_message.tool_calls or []: + if tool_call.index != tool_call_idx: + tool_call_idx = tool_call.index + names.append(None) + args.append("") + if tool_call.function: + if tool_call.function.name is not None: + names[tool_call.index] = tool_call.function.name + if tool_call.function.arguments: + args[tool_call.index] += tool_call.function.arguments + + return content, names, args + + +_MALFORMED_MATRIX_CASES = [ + pytest.param( + # 2x2 corner: name present, `[ARGS]` token present. + '[TOOL_CALLS]bash[ARGS]{"command": "ls -la"}', + ["bash"], + [json.dumps({"command": "ls -la"})], + None, + id="wellformed_single", + ), + pytest.param( + # 2x2 corner: name present, no `[ARGS]` token. + '[TOOL_CALLS]bash{"command": "ls -la"}', + ["bash"], + [json.dumps({"command": "ls -la"})], + None, + id="wellformed_single_no_args_token", + ), + pytest.param( + # 2x2 corner: name skipped, no `[ARGS]` token. "bash" lands in + # content before the marker, then generation jumps straight to + # args (no NAME slot, no [ARGS]). + 'bash[TOOL_CALLS]{"command": "cat file.py"}', + [""], + [json.dumps({"command": "cat file.py"})], + "bash", + id="malformed_single", + ), + pytest.param( + # 2x2 corner: name skipped, `[ARGS]` token present. + 'bash[TOOL_CALLS][ARGS]{"command": "cat file.py"}', + [""], + [json.dumps({"command": "cat file.py"})], + "bash", + id="malformed_single_args_token", + ), + pytest.param( + # parallel well-formed x2; the second call's nested braces in its + # arguments must not confuse call splitting. + '[TOOL_CALLS]bash[ARGS]{"command": "ls"}' + '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', + ["bash", "submit"], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + None, + id="parallel_wellformed", + ), + pytest.param( + # parallel malformed x2: both calls skip the NAME slot entirely + # (marker immediately followed by `{`), each ending with an empty + # name and correctly separated (nested-brace) arguments. + 'bash[TOOL_CALLS]{"command": "ls"}[TOOL_CALLS]{"answer": {"nested": 1}}', + ["", ""], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + "bash", + id="parallel_malformed", + ), + pytest.param( + # parallel malformed x2, `[ARGS]` token variant: both calls skip + # the NAME slot but enter TOOL_ARGS via `[ARGS]` instead of `{`. + 'bash[TOOL_CALLS][ARGS]{"command": "ls"}' + '[TOOL_CALLS][ARGS]{"answer": {"nested": 1}}', + ["", ""], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + "bash", + id="parallel_malformed_args_token", + ), + pytest.param( + # mixed: malformed first call, well-formed second call. The + # important case: proves calls are split on the `[TOOL_CALLS]` + # marker, not on braces — the first call's simple args and the + # second call's nested-brace args must not bleed into each other. + 'bash[TOOL_CALLS]{"command": "ls"}' + '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', + ["", "submit"], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + "bash", + id="mixed_malformed_then_wellformed", + ), + pytest.param( + # mixed, `[ARGS]` token variant: malformed first call enters + # TOOL_ARGS via `[ARGS]` instead of `{`, second call well-formed. + 'bash[TOOL_CALLS][ARGS]{"command": "ls"}' + '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', + ["", "submit"], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + "bash", + id="mixed_malformed_args_token_then_wellformed", + ), + pytest.param( + # mixed: well-formed first call, malformed second call (NAME slot + # skipped entirely for the second call). + '[TOOL_CALLS]bash[ARGS]{"command": "ls"}[TOOL_CALLS]{"answer": {"nested": 1}}', + ["bash", ""], + [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], + None, + id="mixed_wellformed_then_malformed", + ), +] + + +@pytest.mark.parametrize( + "model_output,expected_names,expected_args,expected_content", + _MALFORMED_MATRIX_CASES, +) +def test_malformed_tool_call_matrix_non_streaming( + mistral_tokenizer, + model_output, + expected_names, + expected_args, + expected_content, +): + parser = MistralParser(mistral_tokenizer) + request = _malformed_matrix_request() + + result = parser.extract_tool_calls_from_content(model_output, request) + + assert result.tools_called + assert [tc.function.name for tc in result.tool_calls] == expected_names + assert [tc.function.arguments for tc in result.tool_calls] == expected_args + assert result.content == expected_content + + +@pytest.mark.parametrize( + "model_output,expected_names,expected_args,expected_content", + _MALFORMED_MATRIX_CASES, +) +def test_malformed_tool_call_matrix_streaming( + mistral_tokenizer, + model_output, + expected_names, + expected_args, + expected_content, +): + parser = MistralParser(mistral_tokenizer) + request = _malformed_matrix_request() + + content, names, args = _stream_via_parse_delta( + parser, mistral_tokenizer, model_output, request + ) + + assert content == (expected_content or "") + assert names == expected_names + assert args == expected_args + + +# Real payload captured from a production SWE-bench-agent. +_REAL_LEAK_CONTENT = 'bash[TOOL_CALLS]{"command": "cd /testbed && git diff xarray/core/weighted.py"}ometracediff --git a/xarray/core/weighted.py b/xarray/core/weighted.py\\nindex 5f2d138e..9c38b0b7 100644\\n--- a/xarray/core/weighted.py\\n+++ b/xarray/core/weighted.py\\n@@ -1,4 +1,6 @@\\n from typing import TYPE_CHECKING, Hashable, Iterable, Optional, Union, overload\\n+\\n+import numpy as np\\n \\n from .computation import dot\\n from .options import _get_keep_attrs\\n@@ -105,6 +107,10 @@ class Weighted:\\n )\\n \\n self.obj = obj\\n- self.weights = weights\\n+ # Ensure weights are numeric (float64) to avoid boolean arithmetic issues\\n+ # in dot products and weighted operations\\n+ if not np.issubdtype(weights.dtype, np.number):\\n+ weights = weights.astype(np.float64)\\n+ self.weights = weights\\n "}现在 hareview the changes we made:bash[ARGS]{"command": "cat /testbed/xarray/core/weighted.py | head -120 | tail -20"}' # noqa: E501 + + +def test_extract_tool_calls_real_captured_malformed_output(mistral_tokenizer): + parser = MistralParser(mistral_tokenizer) + request = _malformed_matrix_request() + + result = parser.extract_tool_calls_from_content(_REAL_LEAK_CONTENT, request) + + assert result.tools_called + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "" + assert result.tool_calls[0].function.arguments == json.dumps( + {"command": "cd /testbed && git diff xarray/core/weighted.py"} + ) + assert result.content == "bash" diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index 3e9cff64aa13..028bbed012dc 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -3,6 +3,7 @@ import json from collections.abc import Sequence +from typing import Any, Literal import pytest from openai_harmony import ( @@ -12,14 +13,18 @@ Role, ) from transformers import AutoTokenizer +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import FunctionCall from vllm.entrypoints.openai.parser.harmony_utils import ( get_encoding, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.harmony import HarmonyParser from vllm.parser.parser_manager import ParserManager +from vllm.sampling_params import StructuredOutputsParams REASONING_MODEL_NAME = "openai/gpt-oss-20b" @@ -857,3 +862,335 @@ def test_multi_boundary(self, harmony_parser): ("analysis", "One"), ("final", "Two"), ] + + +class TestAdjustRequest: + REQUEST_TEXT = "Hello" + TOOL_TYPE = "function" + TOOL_1_NAME = "get_user_location" + TOOL_2_NAME = "get_weather" + TOOLS = [ + { + "name": TOOL_1_NAME, + "parameters": { + "type": "object", + "properties": {}, + "required": [], + }, + }, + { + "name": TOOL_2_NAME, + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + ] + OUTPUT_SCHEMA = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + } + + ANALYSIS = "<|channel|>analysis<|message|>analysis message<|end|><|start|>assistant" + COMMENTARY = ( + "<|channel|>commentary<|message|>commentary message<|end|><|start|>assistant" + ) + TOOL_CALL_1 = ( + ANALYSIS + f"<|channel|>commentary to=functions.{TOOL_1_NAME} json<|message|>" + "{}<|call|>" + ) + TOOL_CALL_2 = ( + ANALYSIS + f"<|channel|>commentary to=functions.{TOOL_2_NAME} json<|message|>" + '{"city": "Tokyo"}<|call|>' + ) + FINAL_JSON_SCHEMA = ( + ANALYSIS + + '<|channel|>final <|constrain|>json<|message|>{"answer": "Tokyo"}<|end|>' + ) + FINAL_JSON_OBJECT = ( + ANALYSIS + + '<|channel|>final <|constrain|>json<|message|>{"city": "Tokyo"}<|end|>' + ) + FINAL_TEXT_ONLY = ANALYSIS + "<|channel|>final<|message|>any<|end|>" + FINAL_REGEX = ANALYSIS + "<|channel|>final<|message|>regex<|end|>" + FINAL_CHOICE = ANALYSIS + "<|channel|>final<|message|>choice1<|end|>" + FINAL_GRAMMAR = ANALYSIS + "<|channel|>final<|message|>grammar<|end|>" + FINAL_STRUCTURAL_TAG = ANALYSIS + "<|channel|>final<|message|>tag content<|end|>" + ADMISSION_SAMPLES = ( + "COMMENTARY", + "TOOL_CALL_1", + "TOOL_CALL_2", + "FINAL_JSON_SCHEMA", + "FINAL_JSON_OBJECT", + "FINAL_TEXT_ONLY", + "FINAL_REGEX", + "FINAL_CHOICE", + "FINAL_GRAMMAR", + "FINAL_STRUCTURAL_TAG", + ) + + @staticmethod + def _build_request( + request_kind: Literal["chat", "responses"], + tool_choice: str = "none", + strict_tools: bool = False, + response_format_type: str | None = None, + structured_outputs: StructuredOutputsParams | None = None, + ) -> ChatCompletionRequest | ResponsesRequest: + data: dict[str, Any] = { + "model": REASONING_MODEL_NAME, + } + if request_kind == "chat": + data["messages"] = [ + { + "role": "user", + "content": TestAdjustRequest.REQUEST_TEXT, + } + ] + else: + data["input"] = TestAdjustRequest.REQUEST_TEXT + + if request_kind == "chat": + data["tools"] = [ + { + "type": TestAdjustRequest.TOOL_TYPE, + "function": {"strict": strict_tools, **tool_def}, + } + for tool_def in TestAdjustRequest.TOOLS + ] + data["tool_choice"] = ( + { + "type": TestAdjustRequest.TOOL_TYPE, + "function": {"name": TestAdjustRequest.TOOL_2_NAME}, + } + if tool_choice == "named" + else tool_choice + ) + else: + data["tools"] = [ + { + "type": TestAdjustRequest.TOOL_TYPE, + "strict": strict_tools, + **tool_def, + } + for tool_def in TestAdjustRequest.TOOLS + ] + data["tool_choice"] = ( + { + "type": TestAdjustRequest.TOOL_TYPE, + "name": TestAdjustRequest.TOOL_2_NAME, + } + if tool_choice == "named" + else tool_choice + ) + + if response_format_type == "json_schema": + schema_format = { + "name": "answer_format", + "schema": TestAdjustRequest.OUTPUT_SCHEMA, + "strict": True, + } + if request_kind == "chat": + data["response_format"] = { + "type": "json_schema", + "json_schema": schema_format, + } + else: + data["text"] = { + "format": { + "type": "json_schema", + **schema_format, + } + } + elif response_format_type == "json_object": + if request_kind == "chat": + data["response_format"] = {"type": "json_object"} + else: + data["text"] = {"format": {"type": "json_object"}} + + if structured_outputs is not None: + data["structured_outputs"] = structured_outputs + + if request_kind == "chat": + return ChatCompletionRequest.model_validate(data) + return ResponsesRequest.model_validate(data) + + @staticmethod + def _assert_format_cleared( + adjusted_request: ChatCompletionRequest | ResponsesRequest, + ) -> None: + if isinstance(adjusted_request, ResponsesRequest): + assert adjusted_request.text is None or adjusted_request.text.format is None + else: + assert adjusted_request.response_format is None + + structured_outputs = adjusted_request.structured_outputs + assert structured_outputs is not None + assert structured_outputs.structural_tag is not None + assert structured_outputs.all_non_structural_tag_constraints_none() + + @classmethod + def _assert_structured_outputs_admission( + cls, + adjusted_request: ChatCompletionRequest | ResponsesRequest, + expected_admission: Sequence[str], + ) -> None: + structured_outputs = adjusted_request.structured_outputs + assert structured_outputs is not None + assert structured_outputs.structural_tag is not None + assert structured_outputs.all_non_structural_tag_constraints_none() + + grammar = Grammar.from_structural_tag(structured_outputs.structural_tag) + expected_admission_set = set(expected_admission) + + for sample_name in cls.ADMISSION_SAMPLES: + admitted = _is_grammar_accept_string( + grammar, + getattr(cls, sample_name), + require_termination=False, + ) + should_admit = sample_name in expected_admission_set + assert admitted is should_admit, ( + f"Expected structured_outputs admission for {sample_name} " + f"to be {should_admit}, got {admitted}." + ) + + @pytest.mark.parametrize("request_kind", ["chat", "responses"]) + @pytest.mark.parametrize( + ("request_kwargs", "expected_admission"), + [ + ( + {"tool_choice": "auto", "strict_tools": True}, + [ + "COMMENTARY", + "TOOL_CALL_1", + "TOOL_CALL_2", + "FINAL_JSON_SCHEMA", + "FINAL_JSON_OBJECT", + "FINAL_TEXT_ONLY", + "FINAL_REGEX", + "FINAL_CHOICE", + "FINAL_GRAMMAR", + "FINAL_STRUCTURAL_TAG", + ], + ), + ( + {"tool_choice": "required"}, + ["COMMENTARY", "TOOL_CALL_1", "TOOL_CALL_2"], + ), + ( + {"tool_choice": "named"}, + ["COMMENTARY", "TOOL_CALL_2"], + ), + ( + {"response_format_type": "json_schema"}, + ["FINAL_JSON_SCHEMA"], + ), + ( + {"response_format_type": "json_object"}, + ["FINAL_JSON_SCHEMA", "FINAL_JSON_OBJECT"], + ), + ( + {"structured_outputs": StructuredOutputsParams(json=OUTPUT_SCHEMA)}, + ["FINAL_JSON_SCHEMA"], + ), + ( + {"structured_outputs": StructuredOutputsParams(json_object=True)}, + ["FINAL_JSON_SCHEMA", "FINAL_JSON_OBJECT"], + ), + ( + {"structured_outputs": StructuredOutputsParams(regex=r"regex")}, + ["FINAL_REGEX"], + ), + ( + { + "structured_outputs": StructuredOutputsParams( + choice=["choice1", "choice2"] + ) + }, + ["FINAL_CHOICE"], + ), + ( + { + "structured_outputs": StructuredOutputsParams( + grammar='root ::= "grammar"' + ) + }, + ["FINAL_GRAMMAR"], + ), + ( + { + "structured_outputs": StructuredOutputsParams( + structural_tag=json.dumps( + { + "type": "structural_tag", + "format": { + "type": "json_schema", + "json_schema": OUTPUT_SCHEMA, + }, + } + ) + ) + }, + ["FINAL_JSON_SCHEMA"], + ), + ( + { + "structured_outputs": StructuredOutputsParams( + structural_tag=json.dumps( + { + "type": "structural_tag", + "structures": [ + { + "begin": "", + "schema": {"type": "object"}, + "end": "", + } + ], + "triggers": [""], + } + ) + ) + }, + [ + # Legacy triggered tags allow free text until a trigger, so + # unconstrained final-channel payloads are also admitted. + "FINAL_TEXT_ONLY", + "FINAL_REGEX", + "FINAL_CHOICE", + "FINAL_GRAMMAR", + "FINAL_STRUCTURAL_TAG", + ], + ), + ], + ids=[ + "tool_auto_strict", + "tool_required", + "tool_named", + "response_format_json_schema", + "response_format_json_object", + "structured_outputs_json", + "structured_outputs_json_object", + "structured_outputs_regex", + "structured_outputs_choice", + "structured_outputs_grammar", + "structured_outputs_structural_tag_modern", + "structured_outputs_structural_tag_legacy", + ], + ) + def test_adjust_request( + self, + harmony_parser, + request_kind, + request_kwargs, + expected_admission, + ): + request = self._build_request(request_kind, **request_kwargs) + adjusted_request = harmony_parser.adjust_request(request) + self._assert_format_cleared(adjusted_request) + self._assert_structured_outputs_admission( + adjusted_request, + expected_admission, + ) diff --git a/tests/parser/test_include_reasoning.py b/tests/parser/test_include_reasoning.py index 3d1893577ef0..a75b35173567 100644 --- a/tests/parser/test_include_reasoning.py +++ b/tests/parser/test_include_reasoning.py @@ -7,31 +7,15 @@ """ import json -import os import pytest -_STRICT_TOOL_CALLING_ENV = "VLLM_ENFORCE_STRICT_TOOL_CALLING" -_STRICT_TOOL_CALLING_ENV_VALUE = os.environ.get(_STRICT_TOOL_CALLING_ENV) -os.environ[_STRICT_TOOL_CALLING_ENV] = "0" - -from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage # noqa: E402 -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 -from vllm.parser.abstract_parser import DelegatingParser # noqa: E402 -from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser # noqa: E402 -from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser # noqa: E402 - - -@pytest.fixture(scope="module", autouse=True) -def restore_strict_tool_calling_env(): - yield - if _STRICT_TOOL_CALLING_ENV_VALUE is None: - os.environ.pop(_STRICT_TOOL_CALLING_ENV, None) - else: - os.environ[_STRICT_TOOL_CALLING_ENV] = _STRICT_TOOL_CALLING_ENV_VALUE +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser class ThinkReasoningParser(BaseThinkingReasoningParser): diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py index 2a34ac7eea7b..2379b8a4c2b5 100644 --- a/tests/parser/test_parse.py +++ b/tests/parser/test_parse.py @@ -2,34 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json -import os from types import SimpleNamespace import pytest -_STRICT_TOOL_CALLING_ENV = "VLLM_ENFORCE_STRICT_TOOL_CALLING" -_STRICT_TOOL_CALLING_ENV_VALUE = os.environ.get(_STRICT_TOOL_CALLING_ENV) -os.environ[_STRICT_TOOL_CALLING_ENV] = "0" - -from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 -from vllm.parser.abstract_parser import DelegatingParser # noqa: E402 -from vllm.parser.utils import count_history_tool_calls # noqa: E402 -from vllm.reasoning.basic_parsers import ( # noqa: E402 - BaseThinkingReasoningParser, -) -from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser # noqa: E402 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.utils import count_history_tool_calls +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser -@pytest.fixture(scope="module", autouse=True) -def restore_strict_tool_calling_env(): - yield - if _STRICT_TOOL_CALLING_ENV_VALUE is None: - os.environ.pop(_STRICT_TOOL_CALLING_ENV, None) - else: - os.environ[_STRICT_TOOL_CALLING_ENV] = _STRICT_TOOL_CALLING_ENV_VALUE +@pytest.fixture(autouse=True) +def enable_hermes_required_named_parsing(monkeypatch): + # With VLLM_ENFORCE_STRICT_TOOL_CALLING (default on), Hermes sets + # supports_required_and_named=False and uses structural-tag guided tool + # calling. Force it True to test the non-guided JSON required/named parse path. + monkeypatch.setattr(Hermes2ProToolParser, "supports_required_and_named", True) class ThinkReasoningParser(BaseThinkingReasoningParser): diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index 1e3cca57d6d9..619aa0e491a1 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -14,6 +14,14 @@ from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser +@pytest.fixture(autouse=True) +def enable_hermes_required_named_parsing(monkeypatch): + # With VLLM_ENFORCE_STRICT_TOOL_CALLING (default on), Hermes sets + # supports_required_and_named=False and uses structural-tag guided tool + # calling. Force it True to test the non-guided JSON required/named parse path. + monkeypatch.setattr(Hermes2ProToolParser, "supports_required_and_named", True) + + class ThinkReasoningParser(BaseThinkingReasoningParser): @property def start_token(self) -> str: diff --git a/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py b/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py index a845dae77b48..74cf12250369 100644 --- a/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py +++ b/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py @@ -29,7 +29,7 @@ def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer): self.online_requests: dict[str, SparseEmbeddingCompletionRequestMixin] = {} self.renderer: BaseRenderer = renderer self.default_pooling_params = {} - pooler_config: PoolerConfig = vllm_config.model_config.pooler_config + pooler_config: PoolerConfig | None = vllm_config.model_config.pooler_config if pooler_config is not None: for param in ["use_activation", "dimensions"]: if getattr(pooler_config, param, None) is None: @@ -91,11 +91,11 @@ def _build_sparse_embedding_token_weights( ) -> list[SparseEmbeddingTokenWeight]: token_ids = sparse_embedding.keys() token_weights = sparse_embedding.values() - tokens = [None] * len(token_ids) + tokens: Sequence[str | None] = [None] * len(token_ids) if return_tokens and self.renderer is not None: tokens = convert_ids_list_to_tokens( - self.renderer.get_tokenizer(), token_ids + self.renderer.get_tokenizer(), list(token_ids) ) sparse_embedding_output: list[SparseEmbeddingTokenWeight] = [] for token_id, weight, token in zip(token_ids, token_weights, tokens): diff --git a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py index 89bdab8c39dd..d83399e07aa2 100644 --- a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py +++ b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py @@ -170,7 +170,7 @@ def load_image( list of meta info for each image in *file_paths* """ - imgs = [] + imgs: list[np.ndarray] = [] metas = [] temporal_coords = [] location_coords = [] @@ -209,11 +209,11 @@ def load_image( except Exception: logger.exception("Could not extract timestamp for %s", file) - imgs = np.stack(imgs, axis=0) # num_frames, H, W, C - imgs = np.moveaxis(imgs, -1, 0).astype("float32") # C, num_frames, H, W - imgs = np.expand_dims(imgs, axis=0) # add batch di + stacked = np.stack(imgs, axis=0) # num_frames, H, W, C + stacked = np.moveaxis(stacked, -1, 0).astype("float32") # C, num_frames, H, W + stacked = np.expand_dims(stacked, axis=0) # add batch di - return imgs, temporal_coords, location_coords, metas + return stacked, temporal_coords, location_coords, metas class PrithviMultimodalDataProcessor(IOProcessor[ImagePrompt, ImageRequestOutput]): @@ -327,7 +327,7 @@ def pre_process( } ) - return prompts + return prompts # type: ignore[return-value] def post_process( self, diff --git a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py index 8448003e7053..11e78ccfc15a 100644 --- a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py +++ b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py @@ -6,6 +6,8 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backends.registry import AttentionBackendEnum + from vllm.v1.attention.selector import AttentionSelectorConfig else: VllmConfig = None @@ -20,16 +22,11 @@ class DummyPlatform(Platform): def check_and_update_config(cls, vllm_config: VllmConfig) -> None: vllm_config.compilation_config.custom_ops = ["all"] + @classmethod def get_attn_backend_cls( - self, - backend_name, - head_size, - dtype, - kv_cache_dtype, - block_size, - use_mla, - has_sink, - use_sparse, - use_mm_prefix, - ): + cls, + selected_backend: "AttentionBackendEnum", + attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, + ) -> str: return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501 diff --git a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py index 66ec35c0d5c9..7ba2c09b283c 100644 --- a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py +++ b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py @@ -17,7 +17,9 @@ def __init__(self, vllm_config, engine_idx=0): self.logged = False self.engine_initialized = False - def record(self, scheduler_stats, iteration_stats, mm_cache_stats, engine_idx): + def record( + self, scheduler_stats, iteration_stats, mm_cache_stats=None, engine_idx=0 + ): self.recorded.append( (scheduler_stats, iteration_stats, mm_cache_stats, engine_idx) ) diff --git a/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py b/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py index cc7a021e9812..5ebd7d36e824 100644 --- a/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py +++ b/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py @@ -8,12 +8,12 @@ from typing import Any, NamedTuple import pytest -from huggingface_hub import hf_hub_download from pytest import MarkDecorator from transformers import AutoModelForImageTextToText from vllm.assets.image import ImageAsset from vllm.multimodal.image import rescale_image_size +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import set_default_torch_num_threads from ...conftest import IMAGE_ASSETS, HfRunner, VllmRunner @@ -33,8 +33,8 @@ class GGUFMMTestConfig(NamedTuple): @property def gguf_model(self): - hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj) - return hf_hub_download(self.gguf_repo, filename=self.gguf_backbone) + hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj) + return hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_backbone) # Common prompts aligned with test_common.py "gemma3" entry format @@ -92,7 +92,7 @@ def run_multimodal_gguf_test( num_logprobs: int, ): # Load images at runtime (inside subprocess) to avoid pickle issues - images = [ImageAsset(name).pil_image for name in model.image_names] + images = [ImageAsset(name).pil_image for name in model.image_names] # type: ignore[arg-type] size_factors = [0.25, 0.5, 1.0] inputs_per_image = [ ( diff --git a/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py b/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py index d4adf6f84cf0..4ea8d8605c16 100644 --- a/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py +++ b/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py @@ -4,9 +4,9 @@ import shutil import pytest -from huggingface_hub import snapshot_download from vllm.plugins.lora_resolvers.filesystem_resolver import FilesystemResolver +from vllm.transformers_utils.repo_utils import hf_api MODEL_NAME = "Qwen/Qwen3-0.6B" LORA_NAME = "charent/self_cognition_Alice" @@ -22,12 +22,12 @@ def adapter_cache(request, tmpdir_factory): @pytest.fixture(scope="module") def qwen3_lora_files(): - return snapshot_download(repo_id=LORA_NAME) + return hf_api().snapshot_download(repo_id=LORA_NAME) @pytest.fixture(scope="module") def pa_files(): - return snapshot_download(repo_id=PA_NAME) + return hf_api().snapshot_download(repo_id=PA_NAME) @pytest.mark.asyncio diff --git a/tests/plugins_tests/lora_resolvers/test_hf_hub_resolver.py b/tests/plugins_tests/lora_resolvers/test_hf_hub_resolver.py index 6fa747211aa7..4590ae29c40c 100644 --- a/tests/plugins_tests/lora_resolvers/test_hf_hub_resolver.py +++ b/tests/plugins_tests/lora_resolvers/test_hf_hub_resolver.py @@ -31,6 +31,7 @@ async def test_hf_resolver_with_direct_path(): assert hf_resolver is not None lora_request = await hf_resolver.resolve_lora(LORA_REPO_MODEL_NAME, LORA_REPO) + assert lora_request is not None assert lora_request.lora_name == LORA_REPO assert REPO_DOWNLOAD_DIR in lora_request.lora_path assert "adapter_config.json" in os.listdir(lora_request.lora_path) diff --git a/tests/plugins_tests/test_oot_registration_offline.py b/tests/plugins_tests/test_oot_registration_offline.py index 3f483a570672..ac5361c66f51 100644 --- a/tests/plugins_tests/test_oot_registration_offline.py +++ b/tests/plugins_tests/test_oot_registration_offline.py @@ -50,7 +50,7 @@ def test_oot_registration_embedding( m.setenv("VLLM_PLUGINS", "register_dummy_model") prompts = ["Hello, my name is", "The text does not matter"] llm = LLM( - model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=2048 + model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=512 ) outputs = llm.embed(prompts) diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index b733ee486c18..9eefd88e915e 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -5,13 +5,17 @@ Run `pytest tests/quantization/test_auto_gptq.py -v -s`. """ +from types import SimpleNamespace + import pytest import torch from tests.quantization.utils import is_quant_method_supported +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.auto_gptq import ( AutoGPTQConfig, AutoGPTQLinearMethod, + AutoGPTQMoEMethod, ) PROMPT = "On the surface of Mars, we found" @@ -54,3 +58,78 @@ def check_model(model): def test_auto_gptq_config_get_name(): """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'.""" assert AutoGPTQConfig.get_name() == "auto_gptq" + + +def test_auto_gptq_moe_creates_zero_initialized_expert_biases(): + method = object.__new__(AutoGPTQMoEMethod) + method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {}) + method.input_dtype = None + method.experts_cls = None + layer = torch.nn.Module() + + method.create_weights( + layer=layer, + num_experts=2, + hidden_size=8, + intermediate_size_per_partition=4, + params_dtype=torch.float16, + intermediate_size_full=4, + weight_loader=lambda *args, **kwargs: None, + ) + + assert layer.w13_bias.shape == (2, 8) + assert layer.w2_bias.shape == (2, 8) + assert torch.count_nonzero(layer.w13_bias) == 0 + assert torch.count_nonzero(layer.w2_bias) == 0 + + +def test_routed_experts_loads_per_expert_biases(): + class Loader: + quant_config = None + quant_method = object() + moe_config = SimpleNamespace( + is_act_and_mul=True, + tp_rank=0, + moe_parallel_config=SimpleNamespace(tp_size=1), + ) + _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim) + _narrow_expert_data_for_padding = staticmethod( + RoutedExperts._narrow_expert_data_for_padding + ) + _load_w13 = RoutedExperts._load_w13 + _loaded_expert_biases = set() + + @staticmethod + def _map_global_expert_id_to_local_expert_id(expert_id): + return expert_id + + loader = Loader() + w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False) + w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False) + + for shard_id, loaded in ( + ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])), + ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])), + ): + assert RoutedExperts.weight_loader( + loader, + w13_bias, + loaded, + weight_name="model.layers.0.mlp.experts.w13_bias", + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + + assert RoutedExperts.weight_loader( + loader, + w2_bias, + torch.tensor([9.0, 10.0, 11.0, 12.0]), + weight_name="model.layers.0.mlp.experts.w2_bias", + shard_id="w2", + expert_id=0, + return_success=True, + ) + assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8)) + assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4)) + assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"} diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 6df7211549b8..eef3a91de5e8 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -8,19 +8,27 @@ Run `pytest tests/quantization/test_auto_round.py`. """ +from typing import Any, cast + import pytest import torch from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.inc.config_parser import INCLayerConfig from vllm.model_executor.layers.quantization.inc.inc_linear import INCLinearMethod from vllm.model_executor.layers.quantization.inc.schemes import ( + INCMxfp4Scheme, + INCMxfp8Scheme, INCWna16Scheme, resolve_scheme, ) +from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear import ( + INCMxfp8LinearScheme, +) from vllm.model_executor.layers.quantization.inc.schemes.inc_scheme import ( INCLinearScheme, ) @@ -52,20 +60,61 @@ pytest.param( "Intel/Qwen3-8B-w2g64-for-ut", marks=pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_xpu()) - or current_platform.device_count() < 2, - reason="72B INT2 AutoRound model requires XPU with at least 2 devices.", + not (current_platform.is_xpu()), + reason="INC int2 on XPU requires the ARK backend.", ), id="auto_round:auto_gptq_int2_tp2", ), + pytest.param( + "INC4AI/Qwen3-8B-MXFP8-AR", + marks=pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="MXFP8 AutoRound model only supports CUDA/XPU backend for now.", + ), + id="auto_round:llm_compressor_mxfp8", + ), ] -MODEL_RUNNER_KWARGS = { +QWEN3_AUTOROUND_MODELS = [ + pytest.param( + "INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4", + marks=pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="Qwen3-1.7B MXFP4 AutoRound model requires CUDA/XPU.", + ), + id="auto_round:mxfp4:qwen3-1p7b", + ), + pytest.param( + "INCModel/Qwen3-30B-A3B-12L-W4A16-test", + marks=pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="Qwen3-30B-A3B W4A16 AutoRound model requires CUDA/XPU.", + ), + id="auto_round:w4a16:qwen3-30b-a3b", + ), + pytest.param( + "INCModel/Qwen3-30B-A3B-12L-MXFP4-test", + marks=pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="Qwen3-30B-A3B MXFP4 AutoRound model requires CUDA/XPU.", + ), + id="auto_round:mxfp4:qwen3-30b-a3b", + ), +] + +MODEL_RUNNER_KWARGS: dict[str, dict[str, Any]] = { + "INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4": {"enforce_eager": True}, + "INCModel/Qwen3-30B-A3B-12L-MXFP4-test": {"enforce_eager": True}, "Intel/Qwen3-8B-w2g64-for-ut": { "block_size": 64, "gpu_memory_utilization": 0.8, "max_model_len": 512, }, + "INC4AI/Qwen3-8B-MXFP8-AR": { + "block_size": 64, + "gpu_memory_utilization": 0.8, + "max_model_len": 512, + }, } @@ -77,7 +126,7 @@ ), reason="Only supports CPU/XPU/CUDA backend.", ) -@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS) def test_auto_round_model(vllm_runner, model): with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=8) @@ -128,6 +177,65 @@ def make_layer_config(**overrides) -> INCLayerConfig: return INCLayerConfig(**kwargs) +def make_qwen3_autoround_config(kind: str) -> INCConfig: + configs = { + "qwen3_1p7b_mxfp4": { + "quant_method": "auto-round", + "bits": 4, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:llm_compressor", + "data_type": "mx_fp", + "extra_config": { + "model.layers.0.self_attn.q_proj": { + "bits": 16, + "data_type": "float", + }, + }, + }, + "qwen3_30b_a3b_w4a16": { + "quant_method": "auto-round", + "bits": 4, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:auto_gptq", + "data_type": "int", + "extra_config": { + "model.layers.0.mlp.gate": { + "bits": 16, + "data_type": "float", + }, + }, + }, + "qwen3_30b_a3b_mxfp4": { + "quant_method": "auto-round", + "bits": 4, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:llm_compressor", + "data_type": "mx_fp", + "act_bits": 4, + "act_group_size": 32, + "act_data_type": "mx_fp", + "extra_config": { + "model.layers.0.mlp.gate": { + "bits": 16, + "data_type": "float", + }, + "model.layers.0.self_attn.q_proj": { + "bits": 16, + "data_type": "float", + }, + }, + }, + } + try: + config = configs[kind] + except KeyError as err: + raise AssertionError(f"unknown qwen3 autoround config: {kind}") from err + return INCConfig.from_config(config) + + def test_inc_config_parser_exact_match() -> None: config = make_config( extra_config={ @@ -227,6 +335,27 @@ def test_inc_config_parser_parallel_lm_head_defaults_to_unquantized() -> None: assert layer_config.bits == 16 +def test_inc_config_parser_suffix_match_for_lm_head() -> None: + """Short extra_config key should match fully-qualified lm_head layer name.""" + layer = object.__new__(ParallelLMHead) + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + + layer_config = config.config_parser.resolve(layer, "model.language_model.lm_head") + + assert layer_config.quantized is True + assert layer_config.bits == 4 + assert layer_config.group_size == 128 + assert layer_config.sym is True + + def test_inc_config_parser_fused_moe_requires_consistent_configs() -> None: config = make_config( extra_config={ @@ -273,6 +402,39 @@ def test_inc_config_parser_fused_module_requires_consistent_configs() -> None: config.config_parser.resolve(DummyLayer(), "layers.0.self_attn.qkv_proj") +def test_inc_mxfp8() -> None: + config = make_config( + weight_bits=8, + group_size=32, + sym=True, + packing_format="auto_round:llm_compressor", + data_type="mx_fp", + ) + + assert config.weight_bits == 8 + assert config.group_size == 32 + assert config.data_type == "mx_fp" + assert config.packing_format == "auto_round:llm_compressor" + + +def test_inc_config_rejects_invalid_mxfp8_activation_config() -> None: + with pytest.raises(AssertionError, match="act_dynamic=True"): + INCConfig.from_config( + { + "bits": 8, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:llm_compressor", + "data_type": "mx_fp", + "act_bits": 8, + "act_data_type": "mx_fp", + "act_group_size": 32, + "act_sym": True, + "act_dynamic": False, + } + ) + + def test_inc_layer_config_mx_fp_helpers() -> None: layer_config = INCLayerConfig( bits=4, @@ -304,6 +466,317 @@ def test_inc_resolve_scheme_selects_wna16() -> None: assert isinstance(scheme, INCWna16Scheme) +def test_inc_config_accepts_mxfp_family_llm_compressor() -> None: + config = INCConfig.from_config( + { + "quant_method": "auto-round", + "bits": 4, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:llm_compressor", + "data_type": "mx_fp4e2m1", + } + ) + + layer_config = config.config_parser.resolve( + DummyLayer(), "model.layers.0.mlp.down_proj" + ) + + assert config.sym is True + assert layer_config.is_mxfp4 is True + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme) + + +def test_qwen3_1p7b_mxfp4_autoround_uses_mxfp4_linear_scheme( + monkeypatch, +) -> None: + class DummyKernel: + pass + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes." + "inc_mxfp4_linear.init_mxfp4_linear_kernel", + lambda **kwargs: DummyKernel(), + ) + + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501 + INCMxfp4LinearMethod, + ) + + config = make_qwen3_autoround_config("qwen3_1p7b_mxfp4") + + assert ( + INCConfig.override_quantization_method( + {"quant_method": "auto-round"}, user_quant=None + ) + == "inc" + ) + ignored_method = config.get_quant_method( + object.__new__(LinearBase), "model.layers.0.self_attn.q_proj" + ) + layer_config = config.config_parser.resolve( + DummyLayer(), "model.layers.0.mlp.gate_proj" + ) + method = INCMxfp4Scheme().get_linear_method( + config, + object.__new__(LinearBase), + "model.layers.0.mlp.gate_proj", + layer_config, + ) + + assert isinstance(ignored_method, UnquantizedLinearMethod) + assert layer_config.bits == 4 + assert layer_config.group_size == 32 + assert layer_config.is_mxfp4 is True + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme) + assert isinstance(method, INCLinearMethod) + assert isinstance(method.scheme, INCMxfp4LinearMethod) + assert isinstance(method.scheme.kernel, DummyKernel) + + +def test_qwen3_30b_a3b_w4a16_autoround_routes_to_gptq_moe( + monkeypatch, +) -> None: + captured = {} + expected_method = object() + + class DummyMoeConfig: + pass + + def fake_resolve_gptq_moe(layer, layer_config): + captured["layer"] = layer + captured["layer_config"] = layer_config + return expected_method + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes." + "inc_wna16_scheme._resolve_gptq_moe", + fake_resolve_gptq_moe, + ) + + config = make_qwen3_autoround_config("qwen3_30b_a3b_w4a16") + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + method = config.get_quant_method(layer, "model.layers.0.mlp") + + assert method is expected_method + assert captured["layer"] is layer + assert captured["layer_config"].bits == 4 + assert captured["layer_config"].group_size == 32 + assert captured["layer_config"].is_gptq is True + assert captured["layer_config"].is_wna16_int is True + + +def test_qwen3_30b_a3b_mxfp4_autoround_routes_to_mxfp4_moe( + monkeypatch, +) -> None: + class DummyMoeConfig: + pass + + class DummyMxfp4MoEMethod: + def __init__(self, moe_config) -> None: + self.moe_config = moe_config + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe.INCMxfp4MoEMethod", + DummyMxfp4MoEMethod, + ) + + config = make_qwen3_autoround_config("qwen3_30b_a3b_mxfp4") + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + ignored_method = config.get_quant_method( + object.__new__(LinearBase), "model.layers.0.self_attn.q_proj" + ) + method = config.get_quant_method(layer, "model.layers.0.mlp") + layer_config = config.config_parser.resolve(DummyLayer(), "model.layers.0.mlp") + + assert isinstance(ignored_method, UnquantizedLinearMethod) + assert layer_config.bits == 4 + assert layer_config.group_size == 32 + assert layer_config.is_mxfp4 is True + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme) + assert isinstance(method, DummyMxfp4MoEMethod) + assert method.moe_config is layer.moe_config + + +def test_inc_mxfp4_linear_method_registers_and_processes_weights( + monkeypatch, +) -> None: + captured = {} + + class DummyKernel: + def process_weights_after_loading(self, layer) -> None: + captured["processed_layer"] = layer + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes." + "inc_mxfp4_linear.init_mxfp4_linear_kernel", + lambda **kwargs: DummyKernel(), + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", + lambda: 0, + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", + lambda: 1, + ) + + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501 + INCMxfp4LinearMethod, + ) + + layer = torch.nn.Module() + method = INCMxfp4LinearMethod( + make_layer_config(group_size=32, data_type="mx_fp4e2m1") + ) + + method.create_weights( + layer, + input_size_per_partition=64, + output_partition_sizes=[16, 32], + input_size=64, + output_size=48, + params_dtype=torch.bfloat16, + ) + + assert layer.weight_packed.shape == (48, 32) + assert layer.weight_packed.dtype is torch.uint8 + assert layer.weight_scale.shape == (48, 2) + assert layer.weight_scale.dtype is torch.uint8 + assert layer.logical_widths == [16, 32] + assert layer.input_size_per_partition == 64 + assert layer.output_size_per_partition == 48 + + packed_data = layer.weight_packed.data + method.process_weights_after_loading(layer) + + assert layer.weight.data.data_ptr() == packed_data.data_ptr() + assert not hasattr(layer, "weight_packed") + assert captured["processed_layer"] is layer + + +def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( + monkeypatch, +) -> None: + captured = {} + expected_quant_config = object() + expected_kernel = object() + expected_experts_cls = object() + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." + "CutlassExpertsMxfp4._supports_current_device", + lambda: False, + ) + monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." + "select_mxfp4_moe_backend", + lambda moe: (Mxfp4MoeBackend.XPU, expected_experts_cls), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." + "make_mxfp4_moe_quant_config", + lambda **kwargs: captured.update({"quant_config_kwargs": kwargs}) + or expected_quant_config, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." + "make_mxfp4_moe_kernel", + lambda **kwargs: captured.update({"kernel_kwargs": kwargs}) or expected_kernel, + ) + + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe import ( + INCMxfp4MoEMethod, + ) + + method = INCMxfp4MoEMethod(moe=cast(Any, "moe-config")) + layer = torch.nn.Module() + layer._expert_routing_tables = lambda: "routing-tables" + + method.create_weights( + layer, + num_experts=2, + hidden_size=64, + intermediate_size_per_partition=32, + params_dtype=torch.bfloat16, + ) + + assert method.experts_cls is expected_experts_cls + assert layer.w13_weight_packed.shape == (2, 64, 32) + assert layer.w2_weight_packed.shape == (2, 64, 16) + assert layer.w13_weight_scale.shape == (2, 64, 2) + assert layer.w2_weight_scale.shape == (2, 64, 1) + + w13_packed_data = layer.w13_weight_packed.data + w2_packed_data = layer.w2_weight_packed.data + method.process_weights_after_loading(layer) + + assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr() + assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr() + assert not hasattr(layer, "w13_weight_packed") + assert not hasattr(layer, "w2_weight_packed") + assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale + assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale + assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config + assert captured["kernel_kwargs"]["moe_config"] == "moe-config" + assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls + assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables" + assert method.moe_kernel is expected_kernel + + +def test_wna16_xpu_moe_routes_to_gptq_moe(monkeypatch) -> None: + captured = {} + expected_method = object() + + class DummyMoeConfig: + pass + + 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_scheme._resolve_gptq_moe", + lambda layer, layer_config: captured.update( + {"layer": layer, "layer_config": layer_config} + ) + or expected_method, + ) + + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + method = INCWna16Scheme().get_moe_method( + make_config(), + layer, + "model.layers.0.mlp", + make_layer_config(group_size=32), + ) + + assert method is expected_method + assert captured["layer"] is layer + assert captured["layer_config"].is_gptq is True + + +def test_inc_resolve_scheme_selects_mxfp8() -> None: + layer_config = INCLayerConfig( + bits=8, + group_size=32, + sym=True, + packing_format="auto_round:llm_compressor", + backend="auto", + data_type="mx_fp", + quantized=True, + ) + + scheme = resolve_scheme(layer_config) + + assert isinstance(scheme, INCMxfp8Scheme) + + class DummyLinearScheme(INCLinearScheme): def __init__(self) -> None: self.calls: list[tuple] = [] @@ -323,6 +796,76 @@ def apply_weights(self, layer, x, bias=None): return "applied" +def test_inc_mxfp8_linear_scheme_delegates_to_kernel(monkeypatch) -> None: + class DummyKernel: + def __init__(self) -> None: + self.calls: list[tuple] = [] + + def process_weights_after_loading(self, layer) -> None: + self.calls.append(("process", layer)) + + def apply_weights(self, layer, x, bias=None): + self.calls.append(("apply", layer, x, bias)) + return "applied" + + kernel = DummyKernel() + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", + lambda: kernel, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.ModelWeightParameter", + lambda **kwargs: torch.nn.Parameter(kwargs["data"], requires_grad=False), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.GroupQuantScaleParameter", + lambda **kwargs: torch.nn.Parameter(kwargs["data"], requires_grad=False), + ) + + scheme = INCMxfp8LinearScheme() + layer = torch.nn.Module() + + scheme.create_weights( + layer=layer, + input_size_per_partition=64, + output_partition_sizes=[48, 16], + input_size=64, + output_size=64, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + + assert layer.weight.shape == (64, 64) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale.shape == (64, 2) + assert layer.weight_scale.dtype == torch.uint8 + + scheme.process_weights_after_loading(layer) + result = scheme.apply_weights(layer, torch.randn(1, 64), None) + + assert result == "applied" + assert [call[0] for call in kernel.calls] == ["process", "apply"] + + +def test_inc_mxfp8_linear_scheme_requires_block_32_input(monkeypatch) -> None: + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", + lambda: object(), + ) + scheme = INCMxfp8LinearScheme() + + with pytest.raises(ValueError, match="divisible by 32"): + scheme.create_weights( + layer=torch.nn.Module(), + input_size_per_partition=48, + output_partition_sizes=[32], + input_size=48, + output_size=32, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + + def test_inc_linear_method_delegates() -> None: scheme = DummyLinearScheme() method = INCLinearMethod(scheme) @@ -655,6 +1198,34 @@ def get_linear_method(self, _config, _layer, _prefix, _layer_config): assert method is sentinel +def test_inc_get_quant_method_lm_head_uses_suffix_match(monkeypatch) -> None: + """lm_head extra_config should apply to fully-qualified prefix.""" + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + layer = object.__new__(ParallelLMHead) + sentinel = object() + + class DummyScheme: + def get_linear_method(self, _config, _layer, _prefix, _layer_config): + return sentinel + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.factory.resolve_scheme", + lambda _layer_config: DummyScheme(), + ) + + method = config.get_quant_method(layer, "model.language_model.lm_head") + + assert method is sentinel + + def test_inc_get_quant_method_moe_uses_resolved_scheme(monkeypatch) -> None: config = make_config() layer = object.__new__(RoutedExperts) @@ -694,7 +1265,8 @@ def __init__(self, cfg, moe): captured["moe"] = moe monkeypatch.setattr( - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported", + "vllm.model_executor.layers.quantization.utils.marlin_utils." + "check_moe_marlin_supports_layer", lambda *args, **kwargs: False, ) monkeypatch.setattr( @@ -743,10 +1315,6 @@ def __init__(self, cfg, moe): captured["cfg"] = cfg captured["moe"] = moe - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported", - lambda *args, **kwargs: True, - ) monkeypatch.setattr( "vllm.model_executor.layers.quantization.utils.marlin_utils." "check_moe_marlin_supports_layer", @@ -780,10 +1348,6 @@ def __init__(self, cfg, moe): captured["cfg"] = cfg captured["moe"] = moe - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported", - lambda *args, **kwargs: True, - ) monkeypatch.setattr( "vllm.model_executor.layers.quantization.utils.marlin_utils.check_moe_marlin_supports_layer", lambda *args, **kwargs: True, diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 626717cd4a36..19f63f22862e 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -473,6 +473,11 @@ def check_model(model): "Flat is better than nested.\nSparse is better than dense.", 150.0, ), + ( + "nm-testing/Llama-3.2-1B-Instruct-quipv16-nvfp4", + "Flat is better than nested.\nSparse is better than dense.", + 150.0, + ), ], ) def test_compressed_tensors_transforms_perplexity( @@ -923,12 +928,12 @@ def check_model(model): (None, 32, 64, 128, (False, 64, True)), ], ) -def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected): - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501 - CompressedTensorsWNA16MarlinMoEMethod, +def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected): + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, ) - result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding( + result = CompressedTensorsWNA16MoEMethod._w2_scale_sharding( actorder, group_size, part, full ) assert result == expected diff --git a/tests/quantization/test_experts_int8.py b/tests/quantization/test_experts_int8.py index 7cdb135fa077..6119c1c8468a 100644 --- a/tests/quantization/test_experts_int8.py +++ b/tests/quantization/test_experts_int8.py @@ -12,7 +12,7 @@ from ..models.registry import HF_EXAMPLE_MODELS -MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"] +MODELS = ["ai21labs/Jamba-tiny-random"] @pytest.mark.skipif( diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 0ae51652c627..55ad9c6ddb54 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -21,7 +21,7 @@ Attention, set_default_quant_scales, ) -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.quantization.fp8 import ( Fp8Config, Fp8KVCacheMethod, @@ -93,9 +93,6 @@ def test_online_quantization( use_rocm_aiter: bool, monkeypatch, ) -> None: - if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90): - pytest.skip("FA3 currently rejects FP8 KV cache output dtype on SM90") - if use_rocm_aiter: monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") @@ -105,9 +102,15 @@ def test_online_quantization( if force_marlin: monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1") + model_dtype = "auto" + if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90): + # FA3 requires BF16 output when the query input is FP8. + model_dtype = "bfloat16" + with vllm_runner( "facebook/opt-125m", quantization="fp8", + dtype=model_dtype, enforce_eager=True, kv_cache_dtype=kv_cache_dtype, ) as llm: @@ -390,7 +393,7 @@ def test_fp8_reloading( method.use_marlin = use_marlin else: - layer = FusedMoE( + layer = FusedMoEFactory( num_experts=1, top_k=1, hidden_size=1, diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index 0efcc8a3c62f..c8d34bb0ab50 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -79,21 +79,6 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config): assert experts_cls is not None -@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") -@pytest.mark.skipif( - ROCM_AITER_AVAILABLE, - reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)", -) -def test_w4a4_falls_back_to_triton_unfused_without_aiter(mxfp4_oracle_config): - """Without AITER and no --moe-backend, ROCm falls back to TRITON_UNFUSED.""" - config = _make_w4a4_moe_config() - backend, experts_cls = select_mxfp4_moe_backend( - config, activation_key=kMxfp4Dynamic - ) - assert backend == Mxfp4MoeBackend.TRITON_UNFUSED - assert experts_cls is not None - - @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config): """With --moe-backend emulation, W4A4 selects EMULATION.""" diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py new file mode 100644 index 000000000000..2eb6fa92a547 --- /dev/null +++ b/tests/quantization/test_int8_moe_oracle.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for INT8 (W8A8) fused-MoE oracle backend selection. + +These exercise ``select_int8_moe_backend`` only (no kernels are launched), so +they run on any platform where the Triton INT8 MoE kernel is available — CUDA +(SM >= 7.5) or ROCm — not just gfx950. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + Int8MoeBackend, + select_int8_moe_backend, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt8DynamicTensorSym, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, + kInt8StaticTensorSym, +) +from vllm.platforms import current_platform + +# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on +# ROCm CDNA GPUs. Gate on that rather than on a specific arch. +INT8_MOE_SUPPORTED = ( + current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) +) or current_platform.is_rocm() + +requires_int8_moe = pytest.mark.skipif( + not INT8_MOE_SUPPORTED, + reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)", +) + + +def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + return FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend=moe_backend, + ) + + +@requires_int8_moe +@pytest.mark.parametrize( + "weight_key,activation_key", + [ + # per-channel weight + dynamic per-token activation + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # per-tensor weight + dynamic per-tensor activation + (kInt8StaticTensorSym, kInt8DynamicTensorSym), + ], +) +def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): + """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor + weights) select the Triton backend.""" + config = _make_int8_moe_config() + backend, experts_cls = select_int8_moe_backend( + config, weight_key=weight_key, activation_key=activation_key + ) + assert backend == Int8MoeBackend.TRITON + assert experts_cls is not None + + +@requires_int8_moe +def test_int8_explicit_moe_backend_triton(): + """An explicit --moe-backend triton selects the Triton INT8 backend.""" + config = _make_int8_moe_config(moe_backend="triton") + backend, experts_cls = select_int8_moe_backend( + config, + weight_key=kInt8StaticChannelSym, + activation_key=kInt8DynamicTokenSym, + ) + assert backend == Int8MoeBackend.TRITON + assert experts_cls is not None + + +@requires_int8_moe +def test_int8_unsupported_moe_backend_raises(): + """An unsupported --moe-backend for INT8 MoE raises a clear error.""" + config = _make_int8_moe_config(moe_backend="cutlass") + with pytest.raises(ValueError, match="not supported for Int8 MoE"): + select_int8_moe_backend( + config, + weight_key=kInt8StaticChannelSym, + activation_key=kInt8DynamicTokenSym, + ) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index a08e14c53c8c..4124085c23e3 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -13,14 +13,21 @@ import torch from tests.quantization.utils import is_quant_method_supported +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.model import ModelConfig +from vllm.model_executor.kernels.linear import ( + HummingNvFp4LinearKernel, + MarlinNvFp4LinearKernel, +) from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, + ModelOptFp8LinearMethod, ModelOptMixedPrecisionConfig, ModelOptMxFp8Config, ModelOptNvFp4Config, ModelOptNvFp4LinearMethod, + ModelOptNvFp4W4A16LinearMethod, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -42,12 +49,12 @@ def _skip(msg: str) -> NoReturn: def _snapshot_download_or_skip(model_id: str) -> str: try: - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api except Exception as e: # pragma: no cover _skip(f"huggingface_hub is required to download {model_id}: {e}") try: - return snapshot_download( + return hf_api().snapshot_download( repo_id=model_id, repo_type="model", # These checkpoints are already small; download full repo for simplicity. @@ -108,6 +115,28 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head(): assert isinstance(method, ModelOptNvFp4LinearMethod) +def test_modelopt_fp8_updates_weight_dims_after_transpose(): + layer = torch.nn.Module() + layer.register_parameter( + "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + layer.register_parameter( + "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + + method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod) + method.fp8_linear = Mock() + method.process_weights_after_loading(layer) + + assert layer.weight.shape == (2, 3) + assert layer.weight.input_dim == 0 + assert layer.weight.output_dim == 1 + method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer) + + def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized(): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, @@ -165,6 +194,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling(): assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None +def test_modelopt_mixed_precision_composes_gemma4_mappers(): + from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM + from vllm.model_executor.models.gemma4_mm import ( + Gemma4ForConditionalGeneration, + ) + + config = _mixed_precision_config( + { + "model.language_model.layers.0.experts": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + "model.language_model.layers.1.moe.experts.gate_up_proj": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + } + ) + + config.apply_vllm_mapper( + Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper() + ) + config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper()) + + expected_prefix = "language_model.model.layers.0.moe.experts" + assert set(config.quantized_layers) == { + expected_prefix, + "language_model.model.layers.1.moe.gate_up_proj", + } + assert config._resolve_quant_algo(expected_prefix) == "NVFP4" + + def test_modelopt_mixed_precision_infers_fused_gate_up_projection(): from vllm.model_executor.layers.linear import LinearBase @@ -436,15 +497,13 @@ def test_modelopt_nvfp4_config_dispatches_w4a4_method(): def test_modelopt_nvfp4_config_dispatches_w4a16_method(): - """``quant_method="W4A16_NVFP4"`` routes to the new + """``quant_method="W4A16_NVFP4"`` routes to ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling. Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects one of three FP8 LinearMethods on ``quant_method``); a regression here would mean a W4A16 NVFP4 checkpoint silently loaded under the - W4A4 method, which would try to register an ``input_scale`` runtime - parameter and (more importantly) call the cutlass W4A4 NVFP4 GEMM - instead of FP4 Marlin. + W4A4 activation-quantization path. """ from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, @@ -463,6 +522,21 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method(): assert config.quant_method == "W4A16_NVFP4" +@pytest.mark.parametrize( + ("linear_backend", "kernel_cls"), + [("auto", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel)], +) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): + vllm_config = VllmConfig() + vllm_config.kernel_config.linear_backend = linear_backend + with set_current_vllm_config(vllm_config): + method = ModelOptNvFp4W4A16LinearMethod( + ModelOptNvFp4Config(quant_method="W4A16_NVFP4") + ) + assert isinstance(method.kernel, kernel_cls) + + @pytest.mark.parametrize( "quant_method, expected_use_a16, act_key_is_none", [ diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index c4b0ab5a8464..58faaf4923e3 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -5,45 +5,200 @@ import pytest import torch +from compressed_tensors.quantization import ( + ActivationOrdering, + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method -from vllm.platforms import current_platform +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + _backend_incompatibility_reason, + _convert_moe_wna16_humming_tensors, + convert_to_wna16_moe_kernel_format, + map_wna16_backend, +) +from vllm.model_executor.layers.quantization import moe_wna16 +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.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, +) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") -def test_moe_wna16_apply_passes_layer_activation(monkeypatch): - captured_kwargs = {} +def test_map_wna16_backend_supports_triton(): + assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON - def fake_fused_experts(*args, **kwargs): - captured_kwargs.update(kwargs) - return torch.empty(1, 2) - monkeypatch.setattr( - "vllm.model_executor.layers.fused_moe.fused_experts", - fake_fused_experts, +@pytest.mark.parametrize( + ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"), + [ + ( + WNA16MoEBackend.TRITON, + AutoAWQConfig(4, 128, True, False), + True, + False, + "AutoAWQ weight layout", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, True, True, False, {}, {}), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + actorder=ActivationOrdering.GROUP, + ), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, False, True, False, {}, {}), + False, + True, + "bias", + ), + ( + WNA16MoEBackend.MARLIN, + MoeWNA16Config( + linear_quant_method="gptq", + weight_bits=4, + group_size=128, + has_zp=False, + lm_head_quantized=False, + modules_to_not_convert=None, + full_config={}, + ), + False, + False, + "MoeWNA16 checkpoint layout", + ), + ], +) +def test_wna16_oracle_rejects_incompatible_quant_structures( + backend, quant_config, may_have_zp, may_have_bias, expected +): + from tests.kernels.moe.utils import make_dummy_moe_config + + moe_config = make_dummy_moe_config() + + reason = _backend_incompatibility_reason( + backend=backend, + moe_config=moe_config, + quant_config=quant_config, + may_have_zp=may_have_zp, + may_have_bias=may_have_bias, + allow_tile_padding=True, + ) + + assert reason is not None + assert expected in reason + + +def test_compressed_tensors_weights_are_transposed_for_triton(): + quant_config = QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=32, + ) + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8) + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6) + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8) + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6) + + converted = convert_to_wna16_moe_kernel_format( + backend=WNA16MoEBackend.TRITON, + layer=torch.nn.Module(), + quant_config=quant_config, + input_dtype=None, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, ) + assert converted is not None + assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous()) + assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous()) + + +def test_moe_wna16_setup_forwards_selected_backend(monkeypatch): method = object.__new__(MoeWNA16Method) - method.moe = SimpleNamespace(disable_inplace=False) - method.moe_quant_config = object() - layer = SimpleNamespace( - w13_qweight=torch.empty(1, 2), - w2_qweight=torch.empty(1, 2), - activation=MoEActivation.GELU_TANH, - apply_router_weight_on_input=False, - global_num_experts=1, - expert_map=None, + method.experts_cls = object + method.wna16_backend = WNA16MoEBackend.HUMMING + method.moe = object() + quant_config = object() + method.get_fused_moe_quant_config = lambda layer: quant_config + layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None)) + captured = {} + kernel = object() + + def fake_make_wna16_moe_kernel(**kwargs): + captured.update(kwargs) + return kernel + + monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel) + + method._setup_kernel(layer) + + assert method.moe_kernel is kernel + assert captured["backend"] == WNA16MoEBackend.HUMMING + assert captured["layer"] is layer + + +def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): + qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8) + scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4) + qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2) + + converted = _convert_moe_wna16_humming_tensors( + {"qweight": qweight, "scales": scales, "qzeros": qzeros}, + has_zero_point=True, ) - output = method.apply( - layer, - x=torch.empty(1, 2), - topk_weights=torch.empty(1, 1), - topk_ids=torch.empty(1, 1, dtype=torch.int32), - shared_experts=None, - shared_experts_input=None, + assert torch.equal(converted["weight"], qweight.view(torch.int32)) + assert converted["weight"].shape == (1, 4, 2) + assert torch.equal(converted["weight_scale"], scales) + expected_qzeros = ( + qzeros.transpose(-1, -2) + .contiguous() + .view(torch.int32) + .transpose(-1, -2) + .contiguous() + ) + assert torch.equal(converted["zero_point"], expected_qzeros) + assert converted["zero_point"].shape == (1, 2, 2) + + +def test_moe_wna16_uses_humming_quant_config(monkeypatch): + from vllm.model_executor.layers.quantization.utils import humming_utils + + method = object.__new__(MoeWNA16Method) + method.wna16_backend = WNA16MoEBackend.HUMMING + layer = object() + quant_config = object() + monkeypatch.setattr( + humming_utils, + "get_humming_moe_quant_config", + lambda actual_layer: quant_config if actual_layer is layer else None, ) - assert output.shape == (1, 2) - assert captured_kwargs["activation"] is MoEActivation.GELU_TANH + assert method.get_fused_moe_quant_config(layer) is quant_config diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 3c21441ed65b..d977cff96d02 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -89,6 +89,9 @@ def test_online_quantization( if use_rocm_aiter: monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + if current_platform.is_xpu() and quant_scheme == "fp8_per_block": + pytest.skip("Skip test for online fp8_per_block on XPU platform.") + # `LLM.apply_model` requires pickling a function. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") @@ -121,7 +124,7 @@ def check_model(model): if moe is not None: assert isinstance(moe._quant_method, expected_moe_cls) - if current_platform.is_cuda(): + if current_platform.is_cuda() or current_platform.is_xpu(): assert o_proj.weight.dtype == torch.float8_e4m3fn elif current_platform.is_rocm(): assert o_proj.weight.dtype == current_platform.fp8_dtype() diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 9622944670cf..8440db10719d 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -17,23 +17,32 @@ import torch from packaging import version +from vllm._aiter_ops import is_aiter_found_and_supported from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501 QuarkLinearMethod, QuarkW8A8Fp8, QuarkW8A8Int8, ) from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501 + QuarkW4A8Fp8MoEMethod, QuarkW8A8Int8MoEMethod, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx942, on_gfx950 else: + def on_gfx942() -> bool: + return False + def on_gfx950() -> bool: return False @@ -47,6 +56,8 @@ def on_gfx950() -> bool: importlib.metadata.version("amd-quark") ) >= version.parse(QUARK_MXFP4_MIN_VERSION) +AITER_AVAILABLE = is_aiter_found_and_supported() + DEVICE_TYPE = current_platform.device_type if QUARK_MXFP4_AVAILABLE: @@ -55,7 +66,7 @@ def on_gfx950() -> bool: from quark.torch.quantization.config.config import FP4PerGroupSpec try: - huggingface_hub.list_repo_refs( + hf_api().list_repo_refs( "amd/Llama-3.3-70B-Instruct-WMXFP4-AMXFP4-KVFP8-Scale-UINT8-SQ" ) HF_HUB_AMD_ORG_ACCESS = True @@ -145,7 +156,7 @@ def check_model(model): @pytest.mark.parametrize("tp", [1]) def test_quark_int8_w8a8_moe(vllm_runner, tp): """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model.""" - model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark" + model_path = "amd/tiny-qwen3-moe-w8a8-int8" with vllm_runner( model_path, enforce_eager=True, @@ -170,6 +181,38 @@ def check_model(model): assert output +@pytest.mark.skipif( + not (on_gfx950() or on_gfx942()), + reason="Quark W4A8 (INT4-FP8) MoE requires the AITER kernel on gfx942/gfx950", +) +@pytest.mark.parametrize("tp", [1]) +def test_quark_w4a8_fp8_moe(vllm_runner, monkeypatch, tp): + """Test W4A8 (INT4 weight + FP8 activation) MoE with a tiny Qwen3 MoE model. + + W4A8 dispatches through the AITER fused MoE kernel, so AITER must be on. + """ + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + model_path = "amd/tiny-qwen3-moe-w4a8" + with vllm_runner( + model_path, + enforce_eager=True, + tensor_parallel_size=tp, + gpu_memory_utilization=0.1, + ) as llm: + + def check_model(model): + moe = model.model.layers[0].mlp.experts + assert isinstance(moe._quant_method, QuarkW4A8Fp8MoEMethod), ( + f"Expected QuarkW4A8Fp8MoEMethod, got {type(moe._quant_method)}" + ) + + llm.apply_model(check_model) + + output = llm.generate_greedy("Hello", max_tokens=4) + assert output + + def test_quark_fp8_parity(vllm_runner): quark_model_id = "amd-quark/llama-tiny-fp8-quark-quant-method" fp8_model_id = "amd-quark/llama-tiny-fp8-quant-method" @@ -242,7 +285,7 @@ def get_model_args( excepted_value=10.6, ), AccuracyTestConfig( - model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.4 + model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.45 ), ] @@ -450,6 +493,42 @@ def test_mxfp4_dequant_kernel_match_quark( assert torch.equal(out_hip, out_torch) +@pytest.mark.skipif( + not QUARK_MXFP4_AVAILABLE, + reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available", +) +@pytest.mark.skipif( + not AITER_AVAILABLE, + reason="AITER is not found or not supported on the current platform", +) +@pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("scalings", [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]]) +def test_mxfp4_dynamic_quant_match_quark( + float_dtype: torch.dtype, scalings: list[float] +): + """`AiterMxfp4LinearKernel` quantizes weights dynamically through AITER's + `dynamic_mxfp4_quant`, while the emulation path quantizes/dequantizes + through Quark's `qdq_mxfp4`. Check that both agree on the same input. + """ + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + torch.manual_seed(0) + + hidden_size = 32 * 64 + inp = (torch.rand(48, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2 + for i in range(hidden_size // 32): + inp[:, i * 32 : (i + 1) * 32] = ( + inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)] + ) + + x_q, x_s = dynamic_mxfp4_quant(inp) + out_dynamic_quant = dq_mxfp4_torch(x_q, x_s, float_dtype) + + out_quark_qdq = quant_dequant_mxfp4(inp) + + assert torch.equal(out_dynamic_quant, out_quark_qdq) + + # Unit tests for ``is_layer_skipped`` fused-name handling. FUSED_MAPPING = { diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index 8efc6742a2d9..a724803b9a18 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -1,17 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata import importlib.util import pytest import torch +from packaging import version from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 +else: + + def on_gfx950() -> bool: + return False + + DEVICE_TYPE = current_platform.device_type DTYPE = ["bfloat16"] TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None +TORCHAO_VERSION_0_18_AVAILABLE = TORCHAO_AVAILABLE and version.parse( + importlib.metadata.version("torchao") +) >= version.parse("0.18.0") @pytest.mark.skipif( @@ -90,6 +103,10 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_online_quant_config_dict_json(vllm_runner, enable_pickle): """Testing online quantization, load_weights integration point, with config dict serialized to json string @@ -135,6 +152,10 @@ def load_weights(model): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_online_quant_config_file(vllm_runner): """Testing on the fly quantization, load_weights integration point, with config file @@ -170,6 +191,10 @@ def test_online_quant_config_file(vllm_runner): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_reload_weights(): import json diff --git a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py index 88c9e5f867cd..5a7377439619 100644 --- a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py +++ b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py @@ -1,13 +1,55 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend +from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_nvfp4_moe_layer_for_fi_or_cutlass, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( align_trtllm_fp4_moe_hidden_dim_for_fi, ) +def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch): + monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x) + + num_experts = 3 + layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False)) + w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) + w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) + w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn) + w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn) + weight_scale = torch.ones(num_experts) + + outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass( + backend=NvFp4MoeBackend.FLASHINFER_CUTLASS, + layer=layer, + w13=w13, + w13_scale=w13_scale, + w13_scale_2=weight_scale, + a13_scale=torch.tensor([1.0, 2.0, 3.0]), + w2=w2, + w2_scale=w2_scale, + w2_scale_2=weight_scale, + a2_scale=torch.tensor([4.0, 5.0, 6.0]), + is_act_and_mul=False, + ) + a13_scale, a2_scale = outputs[3], outputs[7] + + torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0)) + torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0)) + distinct_values = torch.arange(num_experts, dtype=torch.float32) + a13_scale.copy_(distinct_values) + a2_scale.copy_(distinct_values) + torch.testing.assert_close(a13_scale, distinct_values) + torch.testing.assert_close(a2_scale, distinct_values) + + def test_align_trtllm_fp4_moe_hidden_dim_noop(): w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256) w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index ccdc69074c7f..f4880abcb5f3 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -277,6 +277,29 @@ def test_no_hybrid_hints_returns_empty(self): assert _get_full_attention_layer_indices(mc) == [] +class TestTurboQuantKVCacheSpec: + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_kv_cache_spec_sets_kv_quant_mode(self, preset): + from vllm.model_executor.layers.attention.attention import Attention + from vllm.v1.kv_cache_interface import KVQuantMode, TQFullAttentionSpec + + layer = SimpleNamespace( + attn_type="decoder", + kv_cache_dtype=preset, + kv_cache_torch_dtype=torch.uint8, + head_size=128, + head_size_v=128, + num_kv_heads=4, + sliding_window=None, + ) + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32)) + + spec = Attention.get_kv_cache_spec(layer, vllm_config) + + assert isinstance(spec, TQFullAttentionSpec) + assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT + + class TestTurboQuantWorkspaceReservation: @staticmethod def _fake_vllm_config( diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index 8ab9c310dcad..5e34199f9421 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -10,15 +10,15 @@ def is_quant_method_supported(quant_method: str) -> bool: - # Currently, all quantization methods require Nvidia or AMD GPUs - if not (current_platform.is_cuda() or current_platform.is_rocm()): + # Currently, quantization tests only run GPUs + if current_platform.is_cpu(): return False - try: current_platform.verify_quantization(quant_method) except ValueError: return False - + if current_platform.is_xpu(): + return True capability = current_platform.get_device_capability() assert capability is not None diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py index a6524ca70728..84bc40afe83e 100644 --- a/tests/reasoning/test_cohere_command_reasoning_parser.py +++ b/tests/reasoning/test_cohere_command_reasoning_parser.py @@ -25,6 +25,8 @@ CohereCommand3ReasoningParser, CohereCommand4ReasoningParser, _has_effective_tools, + _melody_citations_to_vllm, + _melody_sources_to_vllm, _response_format_type, _schema_dict_from_structured_outputs, convert_schema_to_structural_tags, @@ -623,3 +625,68 @@ def test_tools_plus_json_schema_both_kinds(self, parser) -> None: types = _content_types(o.structured_outputs.structural_tag) assert "grammar" in types assert "json_schema" in types + + +class TestMelodySourceResolution: + """Pin the parser-side source resolution (see + ``_melody_sources_to_vllm``). The parser receives melody's numeric + ``(tool_call_index, tool_result_indices)`` addressing and resolves + it against a ``position_to_source`` map handed in by + ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same + map that would otherwise live in the serving layer's resolver. + """ + + @staticmethod + def _fake_melody_source(bucket: int, indices: list[int]) -> Any: + return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) + + def test_multi_index_source_fans_out(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), + (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), + } + raw = [self._fake_melody_source(0, [0, 1])] + out = _melody_sources_to_vllm(raw, position_map) + assert [s.id for s in out] == ["d0", "d1"] + # Verify type / payload were plumbed through, not just the id. + assert out[0].type == "document" + assert out[0].document == {"id": "d0"} + + def test_unresolvable_position_skipped(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [self._fake_melody_source(9, [0])] + assert _melody_sources_to_vllm(raw, position_map) == [] + + def test_missing_position_map_drops_all_sources(self): + # A parser instance without a position map (parser wired + # outside of ``CohereServingChatV2``) can't attribute anything, + # so every source is dropped. Callers downstream will see the + # citation with empty ``sources`` and drop it entirely. + raw = [self._fake_melody_source(0, [0])] + assert _melody_sources_to_vllm(raw, None) == [] + + def test_citations_pass_through_is_thinking_tag(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [ + SimpleNamespace( + start_index=0, + end_index=5, + text="hello", + is_thinking=True, + sources=[self._fake_melody_source(0, [0])], + ) + ] + out = _melody_citations_to_vllm(raw, position_map) + assert out is not None + assert out[0].type == "THINKING_CONTENT" + assert out[0].sources[0].id == "d0" diff --git a/tests/reasoning/test_gptoss_reasoning_parser.py b/tests/reasoning/test_gptoss_reasoning_parser.py index a6f815b6ae5c..e615bc7d73a5 100644 --- a/tests/reasoning/test_gptoss_reasoning_parser.py +++ b/tests/reasoning/test_gptoss_reasoning_parser.py @@ -1,351 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json from unittest.mock import Mock -import pytest -from transformers import AutoTokenizer +from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser -from vllm.entrypoints.mcp.tool_server import ToolServer -from vllm.reasoning import ReasoningParser -from vllm.reasoning.gptoss_reasoning_parser import ( - GptOssReasoningParser, - from_builtin_tool_to_tag, - no_func_reasoning_tag, -) -REASONING_MODEL_NAME = "openai/gpt-oss-120b" - - -@pytest.fixture(scope="module") -def gpt_oss_tokenizer(): - return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) - - -USER_MESSAGE_START = "<|start|>user<|message|>" -REASONING_SECTION_START = "<|end|><|start|>assistant<|channel|>analysis<|message|>" -END = "<|end|>" -ASSISTANT_START = "<|start|>assistant" -ASSISTANT_CONTENT_START_PREFIX = END + ASSISTANT_START + "<|channel|>final" -ASSISTANT_CONTENT_START_SUFFIX = "<|message|>" -ASSISTANT_CONTENT_START = ( - ASSISTANT_CONTENT_START_PREFIX + ASSISTANT_CONTENT_START_SUFFIX -) - -BASIC_CONTENT = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START - + "This is the rest", - "is_reasoning_end": True, -} - -BASIC_REASONING_ONLY = { - "output": REASONING_SECTION_START + "This is reasoning" + "<|end|>", - "is_reasoning_end": False, -} -BASIC_NO_REASONING_NO_ASSISTANT = { - "output": USER_MESSAGE_START + "This is a user message", - "is_reasoning_end": False, -} - -# Edge-case where the model omits the assistant tag entirely. -BASIC_NO_REASONING_ASSISTANT = { - "output": USER_MESSAGE_START + "This is a user message<|end|><|channel|>final", - "is_reasoning_end": True, -} - -COMPLEX_CONTENT_INCOMPLETE_PREFIX_ONLY = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_PREFIX, - "is_reasoning_end": False, -} - -COMPLEX_CONTENT_SUFFIX_ONLY = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_SUFFIX, - "is_reasoning_end": False, -} - -COMPLEX_CONTENT_1_NO_SUFFIX = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_PREFIX - + "<|constrain|> JSON ", - "is_reasoning_end": False, -} - -COMPLEX_CONTENT_1 = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_PREFIX - + "<|constrain|> JSON " - + ASSISTANT_CONTENT_START_SUFFIX, - "is_reasoning_end": True, -} - -COMPLEX_CONTENT_1_WITH_CONTENT = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_PREFIX - + "<|constrain|> JSON " - + ASSISTANT_CONTENT_START_SUFFIX - + "This is the rest", - "is_reasoning_end": True, -} - -COMPLEX_CONTENT_2 = { - "output": REASONING_SECTION_START - + "This is reasoning" - + ASSISTANT_CONTENT_START_PREFIX - + "<|constrain|>ReplyAction " - + ASSISTANT_CONTENT_START_SUFFIX - + "This is the rest", - "is_reasoning_end": True, -} - -MULTI_TURN_CONTENT = { - "output": USER_MESSAGE_START - + "1st turn user message" - + REASONING_SECTION_START - + "1st turn reasoning" - + ASSISTANT_CONTENT_START - + "1st turn response" - + END - + USER_MESSAGE_START - + "2nd turn user message" - + END - + ASSISTANT_START, - "is_reasoning_end": False, -} -TEST_CASES = [ - BASIC_CONTENT, - BASIC_REASONING_ONLY, - COMPLEX_CONTENT_INCOMPLETE_PREFIX_ONLY, - COMPLEX_CONTENT_SUFFIX_ONLY, - COMPLEX_CONTENT_1_NO_SUFFIX, - COMPLEX_CONTENT_1, - COMPLEX_CONTENT_1_WITH_CONTENT, - COMPLEX_CONTENT_2, - MULTI_TURN_CONTENT, -] - - -@pytest.mark.parametrize( - "output, is_reasoning_end", - [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES], -) -def test_gptoss_is_reasoning_end( - output, - is_reasoning_end, - gpt_oss_tokenizer, -): - output = gpt_oss_tokenizer.tokenize(output) - parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer) - - # Test is_reasoning_end - output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output) - actual_is_reasoning_end = parser.is_reasoning_end(output_ids) - assert is_reasoning_end == actual_is_reasoning_end - - -class TestGptOssStructuralTags: - """Test cases for GptOssReasoningParser structural tag functionality.""" - - @pytest.fixture - def mock_tokenizer(self): - """Create a mock tokenizer for testing.""" - tokenizer = Mock() - tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5]) - tokenizer.get_vocab = Mock(return_value={"<|end|>": 6}) - return tokenizer - - @pytest.fixture - def reasoning_parser(self, mock_tokenizer): - """Create a GptOssReasoningParser instance.""" - return GptOssReasoningParser(mock_tokenizer) - - def test_prepare_structured_tag_no_tool_server(self, reasoning_parser): - """Test prepare_structured_tag with no tool server.""" - result = reasoning_parser.prepare_structured_tag(None, None) - expected = json.dumps(no_func_reasoning_tag) - - assert result == expected - - # Verify the structure is correct - parsed = json.loads(result) - assert parsed["type"] == "structural_tag" - assert parsed["format"]["type"] == "triggered_tags" - assert len(parsed["format"]["tags"]) == 1 - assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>" - assert parsed["format"]["triggers"] == ["<|channel|>analysis"] - - def test_prepare_structured_tag_with_original_tag(self, reasoning_parser): - """Test prepare_structured_tag when original_tag is provided.""" - original_tag = '{"custom": "tag"}' - result = reasoning_parser.prepare_structured_tag(original_tag, None) - - # Should return the original tag unchanged - assert result == original_tag - - def test_from_builtin_tool_to_tag(self): - """Test from_builtin_tool_to_tag function.""" - tags = from_builtin_tool_to_tag("python") - - assert len(tags) == 2 - assert tags[0]["begin"] == "<|channel|>commentary to=python" - assert tags[0]["content"]["type"] == "any_text" - assert tags[0]["end"] == "<|end|>" - - assert tags[1]["begin"] == "<|channel|>analysis to=python" - assert tags[1]["content"]["type"] == "any_text" - assert tags[1]["end"] == "<|end|>" - - @pytest.mark.parametrize( - "tools", - [ - [], - ["browser"], - ["python"], - ["container"], - ["browser", "python"], - ["browser", "container"], - ["python", "container"], - ["browser", "python", "container"], - ], - ) - def test_json_validity_comprehensive(self, reasoning_parser, tools): - """Test JSON validity across all possible tool combinations.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools) - - result = reasoning_parser.prepare_structured_tag(None, tool_server) - parsed_result = json.loads(result) - - assert parsed_result["type"] == "structural_tag" - assert "format" in parsed_result - assert "tags" in parsed_result["format"] - assert "triggers" in parsed_result["format"] - - # Tag count should be: 1 (analysis) + 2 * len(tools) - expected_tag_count = 1 + (2 * len(tools)) - assert len(parsed_result["format"]["tags"]) == expected_tag_count - - # Verify triggers are correctly configured - expected_triggers = ["<|channel|>analysis"] - if tools: - expected_triggers.append("<|channel|>commentary to=") - assert set(parsed_result["format"]["triggers"]) == set(expected_triggers) - - def test_no_cross_request_state_pollution(self, reasoning_parser): - """Test that sequential calls with different tool servers produce - independent results, guarding against shared mutable state - (e.g. missing deepcopy in tag_with_builtin_funcs).""" - tool_server_1 = Mock(spec=ToolServer) - tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python") - - tool_server_2 = Mock(spec=ToolServer) - tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser") - - result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1) - result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2) - - tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]] - tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]] - - assert "<|channel|>commentary to=python" in tags_1 - assert "<|channel|>commentary to=browser" not in tags_1 - - assert "<|channel|>commentary to=browser" in tags_2 - assert "<|channel|>commentary to=python" not in tags_2 - - def test_tag_format_consistency(self, reasoning_parser): - """Test that all generated tags follow consistent format, - catching malformed tags from from_builtin_tool_to_tag.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock( - side_effect=lambda tool: tool in ["python", "browser"] - ) - - result = reasoning_parser.prepare_structured_tag(None, tool_server) - parsed_result = json.loads(result) - - for tag in parsed_result["format"]["tags"]: - assert "begin" in tag - assert "content" in tag - assert "end" in tag - assert tag["content"]["type"] == "any_text" - assert tag["end"] == "<|end|>" - assert tag["begin"].startswith("<|channel|>") - - -@pytest.mark.parametrize( - "output, is_reasoning_end", - [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES], -) -def test_gptoss_is_reasoning_end_streaming( - output, - is_reasoning_end, - gpt_oss_tokenizer, -): - """Streaming override must agree with is_reasoning_end for all cases.""" - tokens = gpt_oss_tokenizer.tokenize(output) - parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer) - output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens) - delta_ids = output_ids[-1:] if output_ids else [] - actual = parser.is_reasoning_end_streaming(output_ids, delta_ids) - assert is_reasoning_end == actual - - -@pytest.mark.parametrize( - "output, is_reasoning_end", - [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES], -) -def test_gptoss_is_reasoning_end_streaming_long_prefix( - output, - is_reasoning_end, - gpt_oss_tokenizer, -): - """Windowing must produce correct results even with a long prefix.""" - tokens = gpt_oss_tokenizer.tokenize(output) - parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer) - output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens) - # Prepend 10k dummy reasoning tokens to simulate a long generation - long_prefix = [1] * 10_000 - padded_ids = long_prefix + list(output_ids) - delta_ids = output_ids[-1:] if output_ids else [] - actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids) - assert is_reasoning_end == actual - - -@pytest.mark.parametrize( - "output, is_reasoning_end", - [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES], -) -def test_gptoss_is_reasoning_end_streaming_large_delta( - output, - is_reasoning_end, - gpt_oss_tokenizer, -): - """Simulate speculative decoding where the entire test sequence arrives - as a single large delta appended after a long prefix. The window must - expand to cover delta_ids so the end pattern is never missed.""" - tokens = gpt_oss_tokenizer.tokenize(output) - parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer) - output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens) - long_prefix = [1] * 10_000 - padded_ids = long_prefix + list(output_ids) - # delta_ids = the entire test sequence (as if accepted in one spec step) - delta_ids = list(output_ids) - actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids) - assert is_reasoning_end == actual - - -def test_gptoss_is_reasoning_end_streaming_signature(gpt_oss_tokenizer): - """Verify the method is callable with the expected signature.""" - parser = GptOssReasoningParser(gpt_oss_tokenizer) - result = parser.is_reasoning_end_streaming([], []) - assert result is False +def test_gptoss_reasoning_ended_is_true(): + parser = GptOssReasoningParser(Mock()) + assert parser.is_reasoning_end([]) is True + assert parser.is_reasoning_end_streaming([], []) is True diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py new file mode 100644 index 000000000000..d050c8e3542b --- /dev/null +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.parser.kimi_k3 import KimiK3Parser +from vllm.parser.parser_manager import ParserManager +from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser + +pytestmark = pytest.mark.skip_global_cleanup + +OPEN = "<|open|>" +CLOSE = "<|close|>" +SEP = "<|sep|>" +THINK_OPEN = f"{OPEN}think{SEP}" +THINK_CLOSE = f"{CLOSE}think{SEP}" +RESPONSE_OPEN = f"{OPEN}response{SEP}" + + +class DummyTokenizer: + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + if text == THINK_OPEN: + return [1, 2, 3] + if text == THINK_CLOSE: + return [4, 2, 3] + return [ord(ch) for ch in text] + + +class ReasoningOnlyParser(KimiK3Parser): + reasoning_parser_cls = KimiK3ReasoningParser + + +def test_parser_manager_selects_kimi_k3_parser_for_reasoning_only(): + parser_cls = ParserManager.get_parser(reasoning_parser_name="kimi_k3") + + assert parser_cls is not None + assert issubclass(parser_cls, KimiK3Parser) + assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser + assert parser_cls.tool_parser_cls is None + + +def test_parser_selection_thinking_disabled(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + + assert parser._thinking_enabled is False + + +def test_extract_reasoning_with_xtml_tags(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content = parser.extract_reasoning_content( + f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + + +def test_extract_reasoning_with_generation_prefix_consumed(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content = parser.extract_reasoning_content( + f"step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + + +def test_delegating_parser_strips_response_wrapper_without_tool_parser(): + parser = ReasoningOnlyParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content, tool_calls = parser.parse( + f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + assert tool_calls == [] + + +def test_is_reasoning_end_uses_full_input_ids(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + assert not parser.is_reasoning_end([4, 2]) + assert parser.is_reasoning_end([4, 2, 3]) + + +def test_is_reasoning_end_ignores_stale_close_from_prior_turn(): + # DummyTokenizer: THINK_OPEN -> [1, 2, 3], THINK_CLOSE -> [4, 2, 3]. + # Multi-turn / agent continuation: a prior turn's think channel (its close + # marker) is kept in the prompt, then the current turn opens a new think + # block that has not closed yet. Reasoning must read as NOT ended, otherwise + # the structured-output gate constrains the current turn's reasoning. + parser = KimiK3ReasoningParser(DummyTokenizer()) + + stale_close = [4, 2, 3] + new_open = [1, 2, 3] + # prior close, then current-turn open still unclosed -> not ended + assert not parser.is_reasoning_end([*stale_close, *new_open]) + # ...then the current turn emits its own close -> ended + assert parser.is_reasoning_end([*stale_close, *new_open, *stale_close]) + # open with no close yet -> not ended + assert not parser.is_reasoning_end([*new_open]) + + +def test_streaming_split_open_marker_is_held_back(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + first = parser.extract_reasoning_content_streaming( + previous_text="", + current_text=OPEN, + delta_text=OPEN, + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + ) + second = parser.extract_reasoning_content_streaming( + previous_text=OPEN, + current_text=f"{OPEN}think", + delta_text="think", + previous_token_ids=[1], + current_token_ids=[1, 2], + delta_token_ids=[2], + ) + third = parser.extract_reasoning_content_streaming( + previous_text=f"{OPEN}think", + current_text=THINK_OPEN + "step", + delta_text=f"{SEP}step", + previous_token_ids=[1, 2], + current_token_ids=[1, 2, 3, 9], + delta_token_ids=[3, 9], + ) + + assert first is None + assert second is None + assert isinstance(third, DeltaMessage) + assert third.reasoning == "step" + + +def test_streaming_split_close_marker_hands_content_downstream(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + previous_text = f"{THINK_OPEN}step" + partial_close = parser.extract_reasoning_content_streaming( + previous_text=previous_text, + current_text=previous_text + CLOSE, + delta_text=CLOSE, + previous_token_ids=[1, 2, 3, 9], + current_token_ids=[1, 2, 3, 9, 4], + delta_token_ids=[4], + ) + closed = parser.extract_reasoning_content_streaming( + previous_text=previous_text + CLOSE, + current_text=previous_text + f"{THINK_CLOSE}{RESPONSE_OPEN}answer", + delta_text=f"think{SEP}{RESPONSE_OPEN}answer", + previous_token_ids=[1, 2, 3, 9, 4], + current_token_ids=[1, 2, 3, 9, 4, 2, 3, 10], + delta_token_ids=[2, 3, 10], + ) + + assert partial_close is None + assert isinstance(closed, DeltaMessage) + assert closed.reasoning is None + assert closed.content == f"{RESPONSE_OPEN}answer" + assert parser.extract_content_ids([2, 3, 10]) == [10] + + +def test_thinking_disabled_streams_content(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), chat_template_kwargs={"enable_thinking": False} + ) + + delta = parser.extract_reasoning_content_streaming( + previous_text="", + current_text=f"{RESPONSE_OPEN}answer", + delta_text=f"{RESPONSE_OPEN}answer", + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + ) + + assert isinstance(delta, DeltaMessage) + assert delta.content == f"{RESPONSE_OPEN}answer" + assert delta.reasoning is None + + +def test_delegating_parser_thinking_false_streams_response_content(): + parser = ReasoningOnlyParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + request = ChatCompletionRequest( + model="test-model", + messages=[], + chat_template_kwargs={"thinking": False}, + ) + + first = parser.parse_delta( + delta_text="OK", + delta_token_ids=[10], + request=request, + prompt_token_ids=[1], + finished=False, + ) + partial_close = parser.parse_delta( + delta_text=CLOSE, + delta_token_ids=[2], + request=request, + prompt_token_ids=[1], + finished=False, + ) + closed = parser.parse_delta( + delta_text=f"response{SEP}", + delta_token_ids=[3, 4], + request=request, + prompt_token_ids=[1], + finished=False, + ) + + assert first is not None + assert first.content == "OK" + assert first.reasoning is None + assert partial_close is None + assert closed is None + + +def test_adjust_request_keeps_xtml_markers_contiguous(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + adjusted = parser.adjust_request(request) + + assert adjusted.skip_special_tokens is False + if hasattr(adjusted, "spaces_between_special_tokens"): + assert adjusted.spaces_between_special_tokens is False diff --git a/tests/reasoning/test_mistral_reasoning_parser.py b/tests/reasoning/test_mistral_reasoning_parser.py deleted file mode 100644 index d6da723f80b0..000000000000 --- a/tests/reasoning/test_mistral_reasoning_parser.py +++ /dev/null @@ -1,348 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from tests.reasoning.utils import run_reasoning_extraction_mistral -from vllm.reasoning import ReasoningParser, ReasoningParserManager -from vllm.tokenizers.mistral import MistralTokenizer - -parser_name = "mistral" - - -@pytest.fixture(scope="module") -def mistral_tokenizer(): - mistral_tokenizer = MistralTokenizer.from_pretrained( - "mistralai/Magistral-Small-2509" - ) - return mistral_tokenizer - - -INVALID_SIMPLE_REASONING = { - "output": "This is a reasoning section[/THINK]This is the rest", - "reasoning": None, - "content": "This is a reasoning sectionThis is the rest", - "is_reasoning_end": False, -} -INVALID_COMPLETE_REASONING = { - "output": "This is a reasoning section[/THINK]", - "reasoning": None, - "content": "This is a reasoning section", - "is_reasoning_end": False, -} -NO_CONTENT = { - "output": "[THINK]This is reasoning", - "reasoning": "This is reasoning", - "content": None, - "is_reasoning_end": False, -} -NO_REASONING = { - "output": "This is content", - "reasoning": None, - "content": "This is content", - "is_reasoning_end": False, -} -NO_REASONING_STREAMING = { - "output": "This is a reasoning section", - "reasoning": None, - "content": "This is a reasoning section", - "is_reasoning_end": False, -} -INVALID_MULTIPLE_LINES = { - "output": "This\nThat[/THINK]This is the rest\nThat", - "reasoning": None, - "content": "This\nThatThis is the rest\nThat", - "is_reasoning_end": False, -} -INVALID_SHORTEST_REASONING_NO_STREAMING = { - "output": "[/THINK]This is the rest", - "reasoning": None, - "content": "This is the rest", - "is_reasoning_end": False, -} -INVALID_SHORTEST_REASONING = { - "output": "[/THINK]This is the rest", - "reasoning": None, - "content": "This is the rest", - "is_reasoning_end": False, -} -REASONING_WITH_THINK = { - "output": "[THINK]This is a reasoning section[/THINK]This is the rest", - "reasoning": "This is a reasoning section", - "content": "This is the rest", - "is_reasoning_end": True, -} -COMPLETE_REASONING_WITH_THINK = { - "output": "[THINK]This is a reasoning section[/THINK]", - "reasoning": "This is a reasoning section", - "content": None, - "is_reasoning_end": True, -} -MULTIPLE_LINES_WITH_THINK = { - "output": "[THINK]This\nThat[/THINK]This is the rest\nThat", - "reasoning": "This\nThat", - "content": "This is the rest\nThat", - "is_reasoning_end": True, -} -INVALID_SHORTEST_REASONING_NO_STREAMING_WITH_THINK = { - "output": "[/THINK]This is the rest", - "reasoning": None, - "content": "This is the rest", - "is_reasoning_end": False, -} -INVALID_SHORTEST_REASONING_WITH_THINK = { - "output": "[/THINK]This is the rest", - "reasoning": None, - "content": "This is the rest", - "is_reasoning_end": False, -} -THINK_NO_END = { - "output": "[THINK]This is a reasoning section", - "reasoning": "This is a reasoning section", - "content": None, - "is_reasoning_end": False, -} -EMPTY = { - "output": "", - "reasoning": None, - "content": "", - "is_reasoning_end": False, -} -EMPTY_STREAMING = { - "output": "", - "reasoning": None, - "content": None, - "is_reasoning_end": False, -} -NEW_LINE = { - "output": "Before\n[THINK]This is a reasoning section[/THINK]\nThis is the rest", - "reasoning": "This is a reasoning section", - "content": "Before\n\nThis is the rest", - "is_reasoning_end": True, -} -NEW_LINE_STREAMING = { - "output": "Before\n[THINK]This is a reasoning section[/THINK]\nThis is the rest", - "reasoning": "This is a reasoning section", - "content": "Before\n\nThis is the rest", - "is_reasoning_end": True, -} - -TEST_CASES = [ - pytest.param( - False, - INVALID_SIMPLE_REASONING, - id="invalid_simple_reasoning", - ), - pytest.param( - True, - INVALID_SIMPLE_REASONING, - id="invalid_simple_reasoning_streaming", - ), - pytest.param( - False, - INVALID_COMPLETE_REASONING, - id="invalid_complete_reasoning", - ), - pytest.param( - True, - INVALID_COMPLETE_REASONING, - id="invalid_complete_reasoning_streaming", - ), - pytest.param( - False, - NO_CONTENT, - id="no_content", - ), - pytest.param( - False, - NO_REASONING, - id="no_reasoning", - ), - pytest.param( - True, - NO_REASONING_STREAMING, - id="no_reasoning_token_streaming", - ), - pytest.param( - False, - INVALID_MULTIPLE_LINES, - id="invalid_multiple_lines", - ), - pytest.param( - True, - INVALID_MULTIPLE_LINES, - id="invalid_multiple_lines_streaming", - ), - pytest.param( - True, - INVALID_SHORTEST_REASONING, - id="invalid_shortest", - ), - pytest.param( - False, - INVALID_SHORTEST_REASONING_NO_STREAMING, - id="invalid_shortest_streaming", - ), - pytest.param( - False, - REASONING_WITH_THINK, - id="reasoning_with_think", - ), - pytest.param( - True, - REASONING_WITH_THINK, - id="reasoning_with_think_streaming", - ), - pytest.param( - False, - COMPLETE_REASONING_WITH_THINK, - id="complete_reasoning_with_think", - ), - pytest.param( - True, - COMPLETE_REASONING_WITH_THINK, - id="complete_reasoning_with_think_streaming", - ), - pytest.param( - False, - MULTIPLE_LINES_WITH_THINK, - id="multiple_lines_with_think", - ), - pytest.param( - True, - MULTIPLE_LINES_WITH_THINK, - id="multiple_lines_with_think_streaming", - ), - pytest.param( - False, - INVALID_SHORTEST_REASONING_NO_STREAMING_WITH_THINK, - id="invalid_shortest_with_think", - ), - pytest.param( - True, - INVALID_SHORTEST_REASONING_WITH_THINK, - id="invalid_shortest_with_think_streaming", - ), - pytest.param( - False, - THINK_NO_END, - id="think_no_end", - ), - pytest.param( - True, - THINK_NO_END, - id="think_no_end_streaming", - ), - pytest.param( - False, - EMPTY, - id="empty", - ), - pytest.param( - True, - EMPTY_STREAMING, - id="empty_streaming", - ), - pytest.param( - False, - NEW_LINE, - id="new_line", - ), - pytest.param( - True, - NEW_LINE_STREAMING, - id="new_line_streaming", - ), -] - - -@pytest.mark.parametrize("streaming, param_dict", TEST_CASES) -def test_mistral_reasoning( - streaming: bool, - param_dict: dict, - mistral_tokenizer: MistralTokenizer, -): - output = param_dict["output"] - - index_think = output.find("[THINK]") - len_think = len("[THINK]") - index_end_think = output.find("[/THINK]") - len_end_think = len("[/THINK]") - - # encode everything to tokens ids - output_tokens = [] - if index_think != -1: - output_before_think = output[:index_think] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_before_think, False, False - ) - output_tokens += [mistral_tokenizer.instruct.BEGIN_THINK] - - if index_end_think != -1: - output_middle = output[index_think + len_think : index_end_think] - output_after_think = output[index_end_think + len_end_think :] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_middle, False, False - ) - output_tokens += [mistral_tokenizer.instruct.END_THINK] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_after_think, False, False - ) - else: - output_middle = output[index_think + len_think :] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_middle, False, False - ) - elif index_end_think != -1: - output_before_think = output[:index_end_think] - output_after_think = output[index_end_think + len_end_think :] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_before_think, False, False - ) - output_tokens += [mistral_tokenizer.instruct.END_THINK] - output_tokens += mistral_tokenizer.tokenizer.encode( - output_after_think, False, False - ) - else: - output_tokens += mistral_tokenizer.tokenizer.encode(output, False, False) - - parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( - mistral_tokenizer - ) - - reasoning, content = run_reasoning_extraction_mistral( - parser, output_tokens, streaming=streaming - ) - - assert reasoning == param_dict["reasoning"] - assert content == param_dict["content"] - - # Test is_reasoning_end - is_reasoning_end = parser.is_reasoning_end(output_tokens) - assert is_reasoning_end == param_dict["is_reasoning_end"] - - # Test extract_content - if param_dict["content"] is not None: - # Handle the case where there are tokens outputted before Thinking. - # This should not occur if the model is well trained and prompted. - if "[THINK]" in param_dict["output"] and not param_dict["output"].startswith( - "[THINK]" - ): - before_content = param_dict["output"].split("[THINK]")[0] - before_token_ids = mistral_tokenizer.tokenizer.encode( - before_content, bos=False, eos=False - ) - left_to_encode = param_dict["content"][len(before_content) :] - # Normal situation. - else: - before_token_ids = [] - left_to_encode = param_dict["content"] - - content_tokens = parser.extract_content_ids(output_tokens) - expected_token_ids = before_token_ids + mistral_tokenizer.tokenizer.encode( - left_to_encode, bos=False, eos=False - ) - assert content_tokens == expected_token_ids - else: - content = parser.extract_content_ids(output_tokens) - assert content == [] diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index 2238c41f4989..b537d2a13a02 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -26,6 +26,7 @@ parse_chat_messages, parse_chat_messages_async, ) +from vllm.exceptions import VLLMValidationError from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, _build_mixed_prompt_embeds, @@ -264,7 +265,7 @@ def test_parse_chat_messages_requires_flag(): "content": [{"type": "prompt_embeds", "data": b64}], } ] - with pytest.raises(ValueError, match=_ENABLE_PROMPT_EMBEDS_ERROR): + with pytest.raises(VLLMValidationError, match=_ENABLE_PROMPT_EMBEDS_ERROR): parse_chat_messages( messages, mc, @@ -283,7 +284,7 @@ def test_parse_chat_messages_rejects_missing_data(): "content": [{"type": "prompt_embeds"}], # no `data` } ] - with pytest.raises(ValueError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): + with pytest.raises(VLLMValidationError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): parse_chat_messages( messages, mc, @@ -315,7 +316,7 @@ def test_parse_chat_messages_rejects_missing_data(): def test_parse_chat_messages_rejects_placeholder_in_user_text(content): mc = _make_mock_model_config() # enable_prompt_embeds=True by default messages = [{"role": "user", "content": content}] - with pytest.raises(ValueError, match=_PLACEHOLDER_ERROR_PATTERN): + with pytest.raises(VLLMValidationError, match=_PLACEHOLDER_ERROR_PATTERN): parse_chat_messages(messages, mc, content_format="openai") diff --git a/tests/renderers/test_cohere.py b/tests/renderers/test_cohere.py new file mode 100644 index 000000000000..0ec2156d71d9 --- /dev/null +++ b/tests/renderers/test_cohere.py @@ -0,0 +1,1049 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/renderers/cohere.py``. + +The tests focus on the pure-Python helpers that produce the render-config +dicts passed to ``cohere_melody.render_cmd3`` / ``render_cmd4``. We also +include a class-level instantiation + async-non-blocking test that +mirrors the analogous ``test_mistral.py`` pattern, exercising the +:class:`CohereRenderer` end-to-end with mocked ``model_config`` / +tokenizer / melody bindings. +""" + +import asyncio +import json +import time +from dataclasses import dataclass +from typing import Any +from unittest.mock import Mock + +import pytest + +from vllm.renderers import ChatParams +from vllm.renderers.cohere import ( + CohereRenderer, + MelodyContentType, + _build_render_config, + _content_blocks, + _conversation_to_melody_messages, + _document_to_melody, + _normalize_tool_call, + _role_to_melody, + _tool_to_melody, +) +from vllm.tokenizers.hf import HfTokenizer + +# ====================================================================== +# _role_to_melody +# ====================================================================== + + +class TestRoleToMelody: + def test_assistant_maps_to_chatbot(self): + # melody's templates use the legacy Cohere ``chatbot`` role name. + assert _role_to_melody("assistant") == "chatbot" + + def test_developer_aliases_to_system(self): + # OpenAI's ``developer`` role is documented as high-priority + # instructions; map it onto the ``system`` slot rather than + # letting the templates drop it on the floor. + assert _role_to_melody("developer") == "system" + + @pytest.mark.parametrize("role", ["user", "system", "tool", "chatbot"]) + def test_recognized_roles_passthrough(self, role): + assert _role_to_melody(role) == role + + @pytest.mark.parametrize( + "role,expected", + [ + ("ASSISTANT", "chatbot"), + ("Developer", "system"), + ("User", "user"), + ("SYSTEM", "system"), + ], + ) + def test_role_normalization_is_case_insensitive(self, role, expected): + # cmd3 / cmd4 templates lowercase the role before matching, so + # accept any casing the caller provides. + assert _role_to_melody(role) == expected + + @pytest.mark.parametrize("role", ["function", "moderator", "", "anything"]) + def test_unknown_roles_raise(self, role): + # Silently dropping unknown roles produces malformed prompts + # (the templates' role chain has no else branch). + with pytest.raises(ValueError, match="Unsupported message role"): + _role_to_melody(role) + + def test_non_string_role_rejected(self): + # The function is typed ``role: str`` and the implementation + # relies on Python's attribute lookup (``role.lower()``) to + # reject non-strings — any exception type is acceptable as long + # as we don't silently produce a malformed prompt. + with pytest.raises((AttributeError, TypeError, ValueError)): + _role_to_melody(None) # type: ignore[arg-type] + + +# ====================================================================== +# _normalize_tool_call +# ====================================================================== + + +class TestNormalizeToolCall: + def test_openai_dict_with_dict_arguments_json_encoded(self): + # melody expects ``parameters`` as a JSON-encoded string even when + # OpenAI delivers an already-parsed dict. + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": {"a": 1}}, + } + ) + assert out == {"id": "c1", "name": "f", "parameters": '{"a": 1}'} + + def test_openai_dict_with_string_arguments_preserved(self): + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ) + assert out["parameters"] == '{"a":1}' + + def test_flat_dict_without_function_wrapper(self): + out = _normalize_tool_call({"id": "c1", "name": "f", "arguments": '{"k": 1}'}) + # Falls back to top-level ``name`` / ``arguments``. + assert out == {"id": "c1", "name": "f", "parameters": '{"k": 1}'} + + def test_missing_id_becomes_empty_string(self): + out = _normalize_tool_call({"function": {"name": "f", "arguments": "{}"}}) + assert out["id"] == "" + + def test_pydantic_model_dump_supported(self): + class _Fake: + def model_dump(self): + return { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + + out = _normalize_tool_call(_Fake()) + assert out == {"id": "c1", "name": "f", "parameters": "{}"} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unexpected tool_call value"): + _normalize_tool_call(42) # type: ignore[arg-type] + + +# ====================================================================== +# _content_blocks +# ====================================================================== + + +class TestContentBlocks: + def test_none_returns_empty_list(self): + assert _content_blocks(None) == [] + + def test_string_wrapped_in_text_block(self): + out = _content_blocks("hi") + assert out == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_string_item_in_list_wrapped(self): + out = _content_blocks(["a", "b"]) + assert out == [ + {"type": MelodyContentType.TEXT, "text": "a"}, + {"type": MelodyContentType.TEXT, "text": "b"}, + ] + + @pytest.mark.parametrize( + "part_type", + ["text", "input_text", "output_text", "refusal"], + ) + def test_text_variants_normalized(self, part_type): + out = _content_blocks([{"type": part_type, "text": "hello"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "hello"}] + + def test_thinking_block(self): + out = _content_blocks([{"type": "thinking", "thinking": "thoughts"}]) + assert out == [{"type": MelodyContentType.THINKING, "thinking": "thoughts"}] + + def test_image_block_with_default_placeholder(self): + out = _content_blocks([{"type": "image"}]) + assert out == [ + { + "type": MelodyContentType.IMAGE, + "image": {"template_placeholder": ""}, + } + ] + + def test_image_block_custom_placeholder(self): + out = _content_blocks([{"type": "image", "template_placeholder": "[[IMG]]"}]) + assert out[0]["image"]["template_placeholder"] == "[[IMG]]" + + def test_document_block_dict_passthrough(self): + out = _content_blocks( + [{"type": "document", "document": {"data": {"text": "doc"}}}] + ) + assert out == [ + { + "type": MelodyContentType.DOCUMENT, + "document": {"data": {"text": "doc"}}, + } + ] + + def test_document_block_with_non_dict_falls_back_to_json_text(self): + out = _content_blocks([{"type": "document", "document": "raw string doc"}]) + assert out[0]["type"] == MelodyContentType.TEXT + # JSON-encoded for safety since melody expects a structured doc. + assert out[0]["text"] == json.dumps("raw string doc") + + def test_tool_reference_emitted_as_text(self): + out = _content_blocks([{"type": "tool_reference", "name": "calc"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "calc"}] + + def test_unknown_block_type_fallback_to_text(self): + # Unknown block type with a string value is wrapped in a text block. + out = _content_blocks([{"type": "custom", "custom": "value"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "value"}] + + def test_unknown_block_type_dict_value_json_encoded(self): + out = _content_blocks([{"type": "custom", "custom": {"k": 1}}]) + assert out == [{"type": MelodyContentType.TEXT, "text": json.dumps({"k": 1})}] + + def test_non_string_non_dict_part_rejected(self): + with pytest.raises(TypeError, match="Unexpected content part"): + _content_blocks([42]) # type: ignore[list-item] + + +# ====================================================================== +# _document_to_melody +# ====================================================================== + + +class TestDocumentToMelody: + def test_string_wrapped_in_text_dict(self): + assert _document_to_melody("hello") == {"text": "hello"} + + def test_pure_dict_passthrough(self): + inp = {"text": "x", "id": "d1"} + out = _document_to_melody(inp) + assert out == {"text": "x", "id": "d1"} + # Must be a defensive copy so caller-side mutations of the + # returned dict don't leak back into the input. + out["new_key"] = "value" + assert "new_key" not in inp + + def test_data_wrapper_flattened(self): + # Cohere v2 documents use ``{id, data: {...}}``; melody expects + # the flat shape with ``id`` merged into the payload. + out = _document_to_melody({"id": "d1", "data": {"text": "hello", "title": "t"}}) + assert out == {"id": "d1", "text": "hello", "title": "t"} + + def test_data_wrapper_preserves_inner_id(self): + # If the inner ``data`` already has an ``id``, it wins. + out = _document_to_melody({"id": "outer", "data": {"id": "inner", "text": "x"}}) + assert out["id"] == "inner" + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported document type"): + _document_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _tool_to_melody +# ====================================================================== + + +class TestToolToMelody: + def test_openai_wrapper(self): + out = _tool_to_melody( + { + "type": "function", + "function": { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + }, + } + ) + assert out == { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + } + + def test_flat_dict(self): + out = _tool_to_melody({"name": "calc", "description": "d", "parameters": {}}) + assert out["name"] == "calc" + assert out["parameters"] == {} + + def test_pydantic_like_model_dump(self): + class _Fake: + def model_dump(self): + return { + "type": "function", + "function": { + "name": "calc", + "description": "x", + "parameters": {}, + }, + } + + out = _tool_to_melody(_Fake()) + assert out["name"] == "calc" + + def test_missing_description_becomes_empty(self): + out = _tool_to_melody({"name": "calc"}) + assert out["description"] == "" + assert out["parameters"] == {} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported tool type"): + _tool_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _conversation_to_melody_messages +# ====================================================================== + + +class TestConversationToMelody: + def test_basic_user_assistant_pair(self): + conv = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out == [ + { + "role": "user", + "content": [{"type": MelodyContentType.TEXT, "text": "hi"}], + "tool_calls": [], + }, + { + "role": "chatbot", + "content": [{"type": MelodyContentType.TEXT, "text": "hello"}], + "tool_calls": [], + }, + ] + + def test_assistant_reasoning_prepended_as_thinking_block(self): + # ``reasoning`` (or ``reasoning_content``) is prepended as a + # ``thinking`` block on assistant turns, preserving multi-turn + # chain-of-thought across the rendered prompt. + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [ + {"type": MelodyContentType.THINKING, "thinking": "thoughts"}, + {"type": MelodyContentType.TEXT, "text": "answer"}, + ] + + def test_assistant_reasoning_content_alias_accepted(self): + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning_content": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"][0] == { + "type": MelodyContentType.THINKING, + "thinking": "thoughts", + } + + def test_user_reasoning_ignored(self): + # Only assistant turns get reasoning-as-thinking lifting; user + # turns with a ``reasoning`` key (which shouldn't happen in + # practice) must not produce a phantom thinking block. + conv = [ + { + "role": "user", + "content": "hi", + "reasoning": "should be ignored", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_tool_calls_normalized(self): + conv = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ], + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_calls"] == [ + {"id": "c1", "name": "f", "parameters": '{"a":1}'} + ] + + def test_tool_call_id_preserved_on_tool_role(self): + conv = [ + { + "role": "tool", + "content": "result", + "tool_call_id": "c1", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_call_id"] == "c1" + + def test_messages_citations_attached_by_index(self): + # ``messages_citations`` is a dict keyed by message index. Only + # the message at the matching index should receive the + # citations; other messages must be unaffected. + conv = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + citations = { + 1: [ + { + "start_index": 0, + "end_index": 1, + "text": "a", + "sources": [ + { + "tool_call_index": 0, + "tool_result_indices": [0], + } + ], + "is_thinking": False, + } + ] + } + out = _conversation_to_melody_messages(conv, citations) # type: ignore[arg-type] + assert "citations" not in out[0] + assert out[1]["citations"] == citations[1] + + def test_messages_citations_none_is_a_no_op(self): + conv = [{"role": "assistant", "content": "a"}] + out = _conversation_to_melody_messages(conv, None) # type: ignore[arg-type] + assert "citations" not in out[0] + + def test_messages_citations_missing_index_is_a_no_op(self): + # A ``messages_citations`` dict whose key doesn't hit any + # message must not attach anything (and must not raise). + conv = [{"role": "assistant", "content": "a"}] + out = _conversation_to_melody_messages(conv, {5: [{"anything": 1}]}) # type: ignore[arg-type] + assert "citations" not in out[0] + + +# ====================================================================== +# _build_render_config +# ====================================================================== + + +class TestBuildRenderConfig: + def _conv(self): + return [{"role": "user", "content": "hi"}] + + def test_default_format_is_cmd4(self): + # Bare kwargs -> cmd4 (the current Command A+ prompt format). + # Mirrors ``_DEFAULT_FORMAT`` in ``vllm/renderers/cohere.py`` and + # the ``--cohere-format`` CLI default. + fmt, cfg = _build_render_config(self._conv(), {}) # type: ignore[arg-type] + assert fmt == "cmd4" + assert cfg["use_jinja"] is True + assert isinstance(cfg["messages"], list) + # No additional_template_fields when no extra kwargs are set. + assert "additional_template_fields" not in cfg + + def test_explicit_cmd3(self): + fmt, cfg = _build_render_config(self._conv(), {"cohere_format": "cmd3"}) # type: ignore[arg-type] + assert fmt == "cmd3" + + def test_invalid_format_raises(self): + with pytest.raises(ValueError, match="Invalid cohere_format"): + _build_render_config(self._conv(), {"cohere_format": "cmd5"}) # type: ignore[arg-type] + + def test_documents_converted(self): + _, cfg = _build_render_config( + self._conv(), + { + "documents": [ + "doc text", + {"id": "d1", "data": {"text": "wrapped"}}, + ] + }, + ) # type: ignore[arg-type] + assert cfg["documents"] == [ + {"text": "doc text"}, + {"id": "d1", "text": "wrapped"}, + ] + + def test_available_tools_take_precedence_over_tools(self): + _, cfg = _build_render_config( + self._conv(), + { + "tools": [{"type": "function", "function": {"name": "from_tools"}}], + "available_tools": [ + {"type": "function", "function": {"name": "preferred"}} + ], + }, + ) # type: ignore[arg-type] + names = [t["name"] for t in cfg["available_tools"]] + assert names == ["preferred"] + + def test_tools_used_when_no_available_tools(self): + _, cfg = _build_render_config( + self._conv(), + {"tools": [{"type": "function", "function": {"name": "from_tools"}}]}, + ) # type: ignore[arg-type] + assert [t["name"] for t in cfg["available_tools"]] == ["from_tools"] + + @pytest.mark.parametrize("value", ["enabled", "disabled"]) + def test_reasoning_type_direct(self, value): + _, cfg = _build_render_config(self._conv(), {"reasoning_type": value}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == value + + def test_thinking_dict_shorthand_resolves_reasoning_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "enabled"}}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == "enabled" + + def test_thinking_shorthand_ignores_unknown_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "auto"}}) # type: ignore[arg-type] + assert "reasoning_type" not in cfg + + def test_dev_instruction_forwarded(self): + _, cfg = _build_render_config(self._conv(), {"dev_instruction": "be brief"}) # type: ignore[arg-type] + assert cfg["dev_instruction"] == "be brief" + + def test_response_format_json_object_sets_json_mode(self): + _, cfg = _build_render_config( + self._conv(), {"response_format": {"type": "json_object"}} + ) # type: ignore[arg-type] + assert cfg["json_mode"] is True + assert "json_schema" not in cfg + + def test_response_format_json_schema_sets_json_schema(self): + schema = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + {"response_format": {"type": "json_schema", "schema": schema}}, + ) # type: ignore[arg-type] + # JSON-encoded for melody (string-only schema field). + assert cfg["json_schema"] == json.dumps(schema) + + def test_response_format_nested_json_schema_unwrapped(self): + # When the SDK shape is ``{type: json_schema, schema: {schema: + # {...}}}``, the inner ``schema`` value is used. + inner = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + { + "response_format": { + "type": "json_schema", + "schema": {"schema": inner}, + } + }, + ) # type: ignore[arg-type] + assert cfg["json_schema"] == json.dumps(inner) + + def test_json_schema_kwarg_direct(self): + # Caller can also pass ``json_schema`` directly, both as dict and + # as a pre-stringified value. + _, cfg = _build_render_config(self._conv(), {"json_schema": {"a": 1}}) # type: ignore[arg-type] + assert cfg["json_schema"] == '{"a": 1}' + _, cfg = _build_render_config( + self._conv(), {"json_schema": "raw-string-schema"} + ) # type: ignore[arg-type] + assert cfg["json_schema"] == "raw-string-schema" + + def test_json_mode_kwarg_overrides(self): + _, cfg = _build_render_config(self._conv(), {"json_mode": True}) # type: ignore[arg-type] + assert cfg["json_mode"] is True + + def test_cmd3_safety_mode_lowercased(self): + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "safety_mode": "CONTEXTUAL"}, + ) # type: ignore[arg-type] + assert cfg["safety_mode"] == "contextual" + + def test_cmd3_citation_quality_direct(self): + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "citation_quality": "ACCURATE"}, + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "accurate" + + def test_cmd3_citation_quality_derived_from_citation_options(self): + # When ``citation_quality`` is unset, ``citation_options.mode`` is + # collapsed to on/off so cmd3's binary toggle has a value. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "citation_options": {"mode": "accurate"}}, + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "on" + + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "citation_options": {"mode": "off"}}, + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "off" + + def test_cmd3_skip_preamble_forwarded(self): + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "skip_preamble": True}, + ) # type: ignore[arg-type] + assert cfg["skip_preamble"] is True + + def test_cmd3_no_grounding_field(self): + # cmd3 should never emit a cmd4-only ``grounding`` field. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "grounding": "fast"}, + ) # type: ignore[arg-type] + assert "grounding" not in cfg + + @pytest.mark.parametrize( + "raw,expected", + [ + ("FAST", "enabled"), + ("ACCURATE", "enabled"), + ("OFF", "disabled"), + ("enabled", "enabled"), + ("disabled", "disabled"), + ("unknown", "unknown"), + ], + ) + def test_cmd4_grounding_direct(self, raw, expected): + # melody's cmd4 only accepts ``unknown``/``enabled``/``disabled``, + # so the renderer normalizes any of the v2-facing values into + # that vocab. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": raw}, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + @pytest.mark.parametrize( + "mode,expected", + [ + ("ACCURATE", "enabled"), + ("FAST", "enabled"), + ("OFF", "disabled"), + ], + ) + def test_cmd4_grounding_from_citation_options_mode(self, mode, expected): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "citation_options": {"mode": mode}, + }, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + def test_cmd4_grounding_rejects_unknown_value(self): + with pytest.raises(ValueError, match="Unrecognized cmd4 grounding"): + _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": "foobar"}, + ) # type: ignore[arg-type] + + def test_cmd4_platform_instruction(self): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "platform_instruction": "do this", + }, + ) # type: ignore[arg-type] + assert cfg["platform_instruction"] == "do this" + + def test_cmd4_no_safety_mode_field(self): + # cmd4 should never carry cmd3-only ``safety_mode``/``citation_quality``. + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "safety_mode": "contextual", + "citation_quality": "on", + }, + ) # type: ignore[arg-type] + assert "safety_mode" not in cfg + assert "citation_quality" not in cfg + + def test_extra_kwargs_become_additional_template_fields(self): + # Anything not in the renderer's consumed-keys set is forwarded + # verbatim under ``additional_template_fields`` so jinja templates + # can resolve ``{{ var }}`` directly. + _, cfg = _build_render_config( + self._conv(), + { + "reasoning_effort": "low", + "my_var": "x", + "documents": ["doc"], # consumed, must NOT leak through + }, + ) # type: ignore[arg-type] + extras = cfg["additional_template_fields"] + assert extras == {"reasoning_effort": "low", "my_var": "x"} + # Sanity: the consumed key still produced its dedicated config slot. + assert cfg["documents"] == [{"text": "doc"}] + + def test_template_id_passthrough(self): + # ``template_id`` is a safe selector for one of melody's built-in + # template variants (not raw source) and is still accepted via + # ``chat_template_kwargs``. + _, cfg = _build_render_config( + self._conv(), + {"template_id": "tpl1"}, + ) # type: ignore[arg-type] + assert cfg["template_id"] == "tpl1" + # use_jinja is always True, regardless of caller input. + assert cfg["use_jinja"] is True + + @pytest.mark.parametrize("cohere_only_key", ["template_jinja", "template"]) + def test_cohere_only_template_kwargs_are_rejected(self, cohere_only_key): + # ``chat_template_kwargs.template_jinja`` / ``.template`` are + # accepted at Cohere's own API surface but not at vLLM's -- raw + # template source in vLLM must flow through the standard + # ``chat_template`` request field so the + # ``--trust-request-chat-template`` guard applies uniformly. + # Silently dropping these keys would hide client misconfiguration, + # so we reject them loudly instead. + with pytest.raises( + ValueError, + match=f"chat_template_kwargs.{cohere_only_key!r}", + ): + _build_render_config( + self._conv(), + {cohere_only_key: "raw {{ jinja }}"}, + ) # type: ignore[arg-type] + + @pytest.mark.parametrize("cohere_only_key", ["template_jinja", "template"]) + def test_cohere_only_template_kwargs_none_is_tolerated(self, cohere_only_key): + # An explicit ``None`` (e.g. from ``.model_dump(exclude_none=False)`` + # on an optional field) is treated as absent -- we neither raise + # nor let it fall through as a Jinja variable. + _, cfg = _build_render_config( + self._conv(), + {cohere_only_key: None}, + ) # type: ignore[arg-type] + assert cohere_only_key not in cfg + assert "additional_template_fields" not in cfg + + def test_chat_template_arg_populates_template_jinja(self): + # The standard vLLM ``chat_template`` request field is the sole + # supported channel for raw template source and is forwarded to + # melody under its ``template_jinja`` config key. + _, cfg = _build_render_config( + self._conv(), + {}, + "raw {{ jinja }}", + ) # type: ignore[arg-type] + assert cfg["template_jinja"] == "raw {{ jinja }}" + assert cfg["use_jinja"] is True + + def test_chat_template_arg_none_leaves_template_jinja_unset(self): + _, cfg = _build_render_config( + self._conv(), + {}, + None, + ) # type: ignore[arg-type] + assert "template_jinja" not in cfg + + +# ====================================================================== +# End-to-end async rendering (mirrors ``test_mistral.py``) +# ====================================================================== +# +# Verifies that the synchronous melody bindings run on the renderer's +# thread pool so the asyncio event loop stays responsive under +# concurrent load. Mirrors +# ``test_async_mistral_tokenizer_does_not_block_event_loop`` so future +# regressions in either path are caught uniformly. + + +@dataclass +class _MockHFConfig: + model_type: str = "any" + + +@dataclass +class _MockModelConfig: + runner_type = "generate" + model: str = "cohere-test" + tokenizer: str = "cohere-test" + trust_remote_code: bool = False + max_model_len: int = 100 + tokenizer_revision = None + tokenizer_mode = "cohere" + hf_config = _MockHFConfig() + hf_text_config = _MockHFConfig() + encoder_config: dict[str, Any] | None = None + enable_prompt_embeds: bool = True + skip_tokenizer_init: bool = True + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + + +@dataclass +class _MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class _MockVllmConfig: + model_config: _MockModelConfig + parallel_config: _MockParallelConfig + + +@pytest.mark.asyncio +async def test_async_cohere_renderer_does_not_block_event_loop(): + expected_prompt = "MOCK_RENDERED_PROMPT" + + def slow_render(*_a, **_kw): + time.sleep(2) + return expected_prompt + + mock_tokenizer = Mock(spec=HfTokenizer) + renderer = CohereRenderer( + _MockVllmConfig(_MockModelConfig(), _MockParallelConfig()), + tokenizer=mock_tokenizer, + ) + + # Replace the (already-imported) ``cohere_melody`` bindings with a + # blocking mock. ``_render`` reads ``self._melody`` at call time, so + # this works even though ``_render_async`` was bound at __init__. + fake_melody = Mock() + fake_melody.render_cmd3 = slow_render + fake_melody.render_cmd4 = slow_render + renderer._melody = fake_melody + + task = renderer.render_messages_async([], ChatParams()) + + # Ensure the event loop is not blocked while the (blocking) render + # call is in flight on the thread pool. + blocked_count = 0 + for _ in range(20): # ~2 seconds at 0.1s slices + start = time.perf_counter() + await asyncio.sleep(0) + elapsed = time.perf_counter() - start + if elapsed >= 0.5: + blocked_count += 1 + await asyncio.sleep(0.1) + + _, prompt = await task + assert prompt["prompt"] == expected_prompt, "Mocked blocking render was not called" + assert blocked_count == 0, "Event loop blocked during rendering" + + +# ====================================================================== +# End-to-end: request-level citations -> rendered prompt markup +# ====================================================================== +# +# Ties the whole pipeline together: a Cohere v2 request carrying an +# assistant message with citations must produce a rendered prompt that +# contains melody's inline ``...>`` markup around the +# cited span. Preserving this invariant is the whole point of the +# ``_messages_citations`` chat_template_kwargs entry. + + +class TestRequestCitationsReachRenderedPrompt: + """End-to-end verification that assistant-message citations on the + request survive the OpenAI-shape round-trip and land in the melody- + rendered prompt as inline ``...`` markers. + + The chain covered here: + + CohereChatV2Request + -> CohereServingChatV2._convert_v2_to_chat_completion + -> ChatCompletionRequest.chat_template_kwargs["_messages_citations"] + -> _build_render_config (reads the entry) + -> cohere_melody.render_cmd4 (renders ...) + + A regression that drops the citations or mangles the melody + ``FilterCitation`` payload would remove the citation markers from + the output, so a text-in / text-out assertion is enough. + """ + + @staticmethod + def _melody(): + # Import locally so missing optional deps skip this test class + # rather than failing collection. Two guards: + # * ``cohere_melody`` -- the Rust binding this test drives. + # * ``cohere`` -- required transitively by + # ``vllm.entrypoints.cohere.{protocol,serving}`` (both + # unconditionally ``from cohere.types import ...`` at module + # scope) which every test in this class imports locally. + pytest.importorskip("cohere") + pytest.importorskip("cohere_melody") + import cohere_melody + + return cohere_melody + + @staticmethod + def _openai_msgs_to_conversation( + openai_messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Cheap stand-in for ``parse_chat_messages`` on text-only inputs. + + ``_convert_v2_to_chat_completion`` emits OpenAI-shape assistant + dicts with ``content`` as a plain string. The real chat-utils + pipeline calls ``parse_chat_messages`` (which requires a full + model config / mm tracker), but for text-only content it's + essentially a passthrough -- we normalize ``content`` to the + list-of-parts shape ``_conversation_to_melody_messages`` expects + and preserve every other key the renderer reads. + """ + conv: list[dict[str, Any]] = [] + for m in openai_messages: + entry: dict[str, Any] = dict(m) + content = entry.get("content") + if isinstance(content, str): + entry["content"] = [{"type": "text", "text": content}] + elif content is None: + entry["content"] = [] + conv.append(entry) + return conv + + def test_document_citation_survives_to_prompt(self): + # Prevent this test from being collected/run when the melody + # extension isn't importable. + melody = self._melody() + + from vllm.entrypoints.cohere.protocol import CohereChatV2Request + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + from vllm.renderers.cohere import _build_render_config + + request = CohereChatV2Request( + model="m", + messages=[ + {"role": "user", "content": "Who wrote Hamlet?"}, + { + "role": "assistant", + "content": "Shakespeare wrote it around 1600.", + "citations": [ + { + "start": 0, + "end": 11, # "Shakespeare" + "text": "Shakespeare", + "sources": [{"type": "document", "id": "doc_shakespeare"}], + "type": "TEXT_CONTENT", + } + ], + }, + {"role": "user", "content": "and what year exactly?"}, + ], + documents=[ + { + "id": "doc_shakespeare", + "data": {"text": "Hamlet was written by Shakespeare c. 1600."}, + } + ], + ) + + # Step 1: v2 -> ChatCompletionRequest. Citations must land in + # ``chat_template_kwargs["_messages_citations"]``. + chat_req = CohereServingChatV2._convert_v2_to_chat_completion(request) + assert chat_req.chat_template_kwargs is not None + assert "_messages_citations" in chat_req.chat_template_kwargs + + # Step 2: build the melody render config. The renderer helper + # folds the per-message citations onto the melody message + # dicts. + conversation = self._openai_msgs_to_conversation(chat_req.messages) + fmt, config = _build_render_config(conversation, chat_req.chat_template_kwargs) + + assistant_msg = config["messages"][1] + assert assistant_msg["role"] == "chatbot" + assert "citations" in assistant_msg, ( + "citations were not attached to the melody assistant message dict" + ) + + # Step 3: hand the config to melody and check the rendered + # prompt actually contains inline citation markup around the + # cited span. + # + # The exact id is deterministic for this input. Melody builds + # ``:[]>`` where + # ``tool_call_index=0`` is the reserved bucket for the top-level + # ``documents`` array and ``tool_result_indices`` are positions + # inside it (see ``PromptRenderIds`` in melody/src/templating/ + # util.rs). ``doc_shakespeare`` sits at position 0 in the + # request's ``documents`` list, so we expect ``0:[0]``. Two + # historical regressions this pins: + # * ``0:[]`` -- the source id was never resolved to an index + # (documents didn't flow through) and melody had nothing to + # anchor the marker on. + # * ``1:[0]`` -- the citation was routed through the wrong + # tool-call bucket while documents were present, so it + # pointed at the wrong prompt slot. + # + # Note the same rendered prompt also contains an example + # ``span`` marker baked into melody's + # system-prompt boilerplate (placeholder text ``"span"``); the + # substring below is specific enough to only match the marker + # around the cited text. + if fmt == "cmd4": + rendered = melody.render_cmd4(config) + else: + rendered = melody.render_cmd3(config) + + assert "Shakespeare" in rendered, ( + f"expected inline citation markup around the cited span; " + f"tail of rendered prompt: {rendered[-400:]!r}" + ) + + # And the cited document's text itself must be in the prompt -- + # otherwise the model would have no way to satisfy the citation. + assert "Hamlet" in rendered + + def test_no_markup_when_no_citations(self): + # Control: the same request shape without any citations must + # NOT contain ```` anywhere in the rendered prompt. Guards + # against a false-positive where melody injects citation + # markers regardless of what we passed in. + melody = self._melody() + + from vllm.entrypoints.cohere.protocol import CohereChatV2Request + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + from vllm.renderers.cohere import _build_render_config + + request = CohereChatV2Request( + model="m", + messages=[ + {"role": "user", "content": "Who wrote Hamlet?"}, + { + "role": "assistant", + "content": "Shakespeare wrote it around 1600.", + }, + ], + ) + + chat_req = CohereServingChatV2._convert_v2_to_chat_completion(request) + assert (chat_req.chat_template_kwargs or {}).get("_messages_citations") is None + + conversation = self._openai_msgs_to_conversation(chat_req.messages) + fmt, config = _build_render_config( + conversation, chat_req.chat_template_kwargs or {} + ) + + if fmt == "cmd4": + rendered = melody.render_cmd4(config) + else: + rendered = melody.render_cmd3(config) + + assert "" not in rendered + assert " None: + self.token_ids = token_ids + self.calls: list[dict[str, Any]] = [] + self.conversations: list[list[dict[str, Any]]] = [] + + def apply_chat_template(self, conversation, **kwargs) -> list[int]: + self.conversations.append(conversation) + self.calls.append(kwargs) + return list(self.token_ids) + + +@dataclass +class MockHFConfig: + model_type: str = "kimi_k3" + + +@dataclass +class MockModelConfig: + runner_type: str = "generate" + is_multimodal_model: bool = False + multimodal_config: Any = None + hf_config: MockHFConfig = field(default_factory=MockHFConfig) + allowed_local_media_path: str = "" + allowed_media_domains: Any = None + enable_prompt_embeds: bool = False + renderer_num_workers: int = 1 + + +@dataclass +class MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class MockVllmConfig: + model_config: MockModelConfig + parallel_config: MockParallelConfig + + +def _make_renderer(tokenizer: StubTokenizer) -> KimiK3Renderer: + config = MockVllmConfig(MockModelConfig(), MockParallelConfig()) + return KimiK3Renderer(config, tokenizer) + + +def test_kimi_k3_registered(): + assert RENDERER_REGISTRY.load_renderer_cls("kimi_k3").__name__ == "KimiK3Renderer" + assert ( + TokenizerRegistry.load_tokenizer_cls("kimi_k3").__name__ == "CachedHfTokenizer" + ) + + +def test_k3_media_io_defaults_preserve_original_mode(): + # Default: K3 keeps the original image mode (no background flattening). + assert _merge_k3_media_io_kwargs(None) == {"image": {"image_mode": None}} + + # Server-/request-level values take precedence over the K3 default. + assert _merge_k3_media_io_kwargs({"image": {"image_mode": "RGB"}}) == { + "image": {"image_mode": "RGB"} + } + + # Unrelated image kwargs are merged with the default. + assert _merge_k3_media_io_kwargs( + {"image": {"rgba_background_color": (0, 0, 0)}} + ) == {"image": {"image_mode": None, "rgba_background_color": (0, 0, 0)}} + + +def test_apply_chat_template_forces_tokenize_and_pins_return_dict(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + tools = [{"type": "function", "function": {"name": "search"}}] + params = ChatParams( + chat_template_kwargs={"tools": tools, "tokenize": False, "thinking": True} + ) + + token_ids = renderer._apply_chat_template( + [{"role": "user", "content": "hi"}], params + ) + + assert token_ids == [7, 8, 9] + kwargs = tokenizer.calls[-1] + # tokenize is forced on even though the request asked for False, so K3 keeps + # the special-vs-ordinary token distinction instead of re-tokenizing a string. + assert kwargs["tokenize"] is True + # return_dict is pinned False so we always get a flat list of ids. + assert kwargs["return_dict"] is False + assert kwargs["tools"] == tools + assert kwargs["thinking"] is True + + +def test_apply_chat_template_translates_standard_thinking_kwargs(): + # Standard enable_thinking/reasoning_effort kwargs must be translated + # to K3's native thinking/thinking_effort. + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={"enable_thinking": False, "reasoning_effort": "none"} + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking"] is False + assert "thinking_effort" not in kwargs + assert "enable_thinking" not in kwargs + assert "reasoning_effort" not in kwargs + + +@pytest.mark.parametrize("reasoning_effort", ["low", "high", "max"]) +def test_apply_chat_template_translates_supported_reasoning_effort( + reasoning_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking_effort"] == reasoning_effort + assert "reasoning_effort" not in kwargs + + +@pytest.mark.parametrize("reasoning_effort", ["minimal", "medium", "xhigh"]) +def test_apply_chat_template_rejects_unsupported_reasoning_effort( + reasoning_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) + + with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info: + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert exc_info.value.parameter == "thinking_effort" + assert exc_info.value.value == reasoning_effort + assert tokenizer.calls == [] + + +def test_apply_chat_template_validates_canonical_native_thinking_effort(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={ + "thinking_effort": "low", + "reasoning_effort": "medium", + } + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert tokenizer.calls[-1]["thinking_effort"] == "low" + + +@pytest.mark.parametrize("thinking_effort", ["none", "minimal", "medium", "xhigh"]) +def test_apply_chat_template_rejects_unsupported_native_thinking_effort( + thinking_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"thinking_effort": thinking_effort}) + + with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info: + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert exc_info.value.parameter == "thinking_effort" + assert exc_info.value.value == thinking_effort + assert tokenizer.calls == [] + + +def test_apply_chat_template_native_k3_kwargs_take_precedence(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={ + "thinking": True, + "enable_thinking": False, + "thinking_effort": "low", + "reasoning_effort": "high", + } + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking"] is True + assert kwargs["thinking_effort"] == "low" + + +def test_apply_chat_template_adds_k3_api_metadata(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + response_format = {"type": "json_object"} + params = ChatParams( + tool_choice="required", + response_format=response_format, + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["tool_choice"] == "required" + assert kwargs["response_format"] == response_format + + +def test_apply_chat_template_auto_tool_choice_keeps_template_kwarg(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={"tool_choice": "required"}, + tool_choice="auto", + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert tokenizer.calls[-1]["tool_choice"] == "required" + + +def test_apply_chat_template_omits_tool_choice_without_tools(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + + renderer._apply_chat_template( + [{"role": "user", "content": "hi"}], ChatParams(tool_choice=None) + ) + + assert "tool_choice" not in tokenizer.calls[-1] + + +def test_render_messages_returns_token_prompt(): + renderer = _make_renderer(StubTokenizer([1, 2, 3])) + + conversation, prompt = renderer.render_messages( + [{"role": "user", "content": "hi"}], ChatParams() + ) + + assert prompt == {"prompt_token_ids": [1, 2, 3]} + assert "multi_modal_data" not in prompt + assert conversation[0]["role"] == "user" + + +def test_render_messages_derives_private_xtml_tool_attrs(): + tokenizer = StubTokenizer([1, 2, 3]) + renderer = _make_renderer(tokenizer) + + conversation, _ = renderer.render_messages( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "lookup:0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + { + "id": "lookup:1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "lookup:1", + "tool": "client-supplied-name", + "index": 99, + "content": "second", + }, + { + "role": "tool", + "tool_call_id": "lookup:0", + "content": "first", + }, + ], + ChatParams(), + ) + + assert [message["content"] for message in conversation[1:]] == [ + "first", + "second", + ] + assert conversation[1]["tool"] == "lookup" + assert conversation[1]["index"] == 1 + assert conversation[2]["tool"] == "lookup" + assert conversation[2]["index"] == 2 + assert tokenizer.conversations[-1] == conversation + + +def test_render_messages_ignores_client_supplied_xtml_tool_attrs(): + tokenizer = StubTokenizer([1, 2, 3]) + renderer = _make_renderer(tokenizer) + + conversation, _ = renderer.render_messages( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "lookup:0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "unknown", + "tool": "lookup", + "index": 1, + "content": "result", + }, + ], + ChatParams(), + ) + + assert "tool" not in conversation[1] + assert "index" not in conversation[1] + + +@pytest.mark.asyncio +async def test_render_messages_async_returns_token_prompt(): + renderer = _make_renderer(StubTokenizer([4, 5])) + + conversation, prompt = await renderer.render_messages_async( + [{"role": "user", "content": "hi"}], ChatParams() + ) + + assert prompt == {"prompt_token_ids": [4, 5]} + assert conversation[0]["role"] == "user" diff --git a/tests/samplers/test_non_finite_params.py b/tests/samplers/test_non_finite_params.py index 57fe90f314c3..f982953d6082 100644 --- a/tests/samplers/test_non_finite_params.py +++ b/tests/samplers/test_non_finite_params.py @@ -42,7 +42,7 @@ class TestNonFiniteRepetitionPenalty: ids=["nan", "inf", "-inf", "math.nan", "math.inf"], ) def test_non_finite_repetition_penalty_rejected(self, value: float): - with pytest.raises(ValueError, match="repetition_penalty"): + with pytest.raises(VLLMValidationError, match="repetition_penalty"): SamplingParams(repetition_penalty=value) def test_finite_repetition_penalty_accepted(self): diff --git a/tests/spec_decode/test_custom_proposer.py b/tests/spec_decode/test_custom_proposer.py index acd6089fb9ce..1e42bde8104d 100755 --- a/tests/spec_decode/test_custom_proposer.py +++ b/tests/spec_decode/test_custom_proposer.py @@ -30,6 +30,7 @@ def __init__(self, vllm_config: VllmConfig): Args: vllm_config: vLLM configuration containing model and speculative settings. """ + assert vllm_config.speculative_config is not None self.num_speculative_tokens = ( vllm_config.speculative_config.num_speculative_tokens ) diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py new file mode 100644 index 000000000000..d0673bc462eb --- /dev/null +++ b/tests/test_cmake_utils.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import subprocess +from pathlib import Path + + +def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path): + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_cuda_archs.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +cuda_archs_loose_intersection( + actual "10.0f;10.7f" "10.7") +if(NOT "${{actual}}" STREQUAL "10.7f") + message(FATAL_ERROR "Expected 10.7f, got '${{actual}}'") +endif() +""" + ) + + subprocess.run(["cmake", "-P", script], check=True) + + +def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( + tmp_path: Path, +): + """torch's autodetection can emit a bogus arch=compute_* half (e.g. + capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS + target must win, while PTX-only entries keep the virtual arch.""" + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_extract_archs.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +extract_unique_cuda_archs_ascending(actual + "-gencode arch=compute_20,code=sm_121;\ +-gencode arch=compute_80,code=sm_80;\ +-gencode arch=compute_80,code=compute_80") +if(NOT "${{actual}}" STREQUAL "8.0;12.1") + message(FATAL_ERROR "Expected '8.0;12.1', got '${{actual}}'") +endif() +""" + ) + + subprocess.run(["cmake", "-P", script], check=True) diff --git a/tests/test_config.py b/tests/test_config.py index a785297997b1..21bdde0457d5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -406,26 +406,41 @@ class _TestNestedConfig: a: _TestConfigFields = field(default_factory=lambda: _TestConfigFields(a=0)) +@dataclass +class _TestDerivedConfigFields(_TestConfigFields): + pass + + def test_update_config(): # Simple update config1 = _TestConfigFields(a=0) new_config1 = update_config(config1, {"a": 42}) assert new_config1.a == 42 # Nonexistent field - with pytest.raises(AssertionError): + with pytest.raises(ValueError, match=r"_TestConfigFields\.nonexistent"): new_config1 = update_config(config1, {"nonexistent": 1}) # Nested update with dataclass config2 = _TestNestedConfig() new_inner_config = _TestConfigFields(a=1, c="new_value") new_config2 = update_config(config2, {"a": new_inner_config}) assert new_config2.a == new_inner_config + # Declared field type, not the live value's subtype, defines valid overrides + config_with_derived = _TestNestedConfig(a=_TestDerivedConfigFields(a=0)) + new_config2 = update_config(config_with_derived, {"a": new_inner_config}) + assert new_config2.a is new_inner_config + # Nested update with unrelated dataclass + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a"): + update_config(config2, {"a": _TestNestedConfig()}) # Nested update with dict config3 = _TestNestedConfig() new_config3 = update_config(config3, {"a": {"c": "new_value"}}) assert new_config3.a.c == "new_value" # Nested update with invalid type - with pytest.raises(AssertionError): - new_config3 = update_config(config3, {"a": "new_value"}) + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a"): + update_config(config3, {"a": "new_value"}) + # Invalid nested field preserves its full path + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a\.nonexistent"): + update_config(config3, {"a": {"nonexistent": 1}}) @pytest.mark.parametrize( @@ -1220,6 +1235,26 @@ def test_vllm_config_defaults_are_none(): assert getattr(config.compilation_config, k) is None +def test_validate_mamba_align_subblock_prefill(): + """Align mode permits configured prefill chunks smaller than a block.""" + config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=11392, + mamba_cache_mode="align", + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8192, + long_prefill_token_threshold=4096, + disable_chunked_mm_input=False, + ), + ) + + VllmConfig.validate_block_size(config) + + @pytest.mark.parametrize( ("model_id", "compilation_config", "optimization_level"), [ @@ -1532,21 +1567,42 @@ def test_needs_dp_coordination( assert vllm_config.needs_dp_coordinator == expected_needs_coordinator +def test_fault_tolerance_requires_single_api_server(): + """Fault tolerance assumes one AsyncMPClient manages all engines, so it + is incompatible with API server scale-out (_api_process_count > 1).""" + with pytest.raises(ValueError, match="single API server"): + ParallelConfig(enable_fault_tolerance=True, _api_process_count=2) + + # Single API server (the FT-supported topology) is accepted. + ParallelConfig(enable_fault_tolerance=True, _api_process_count=1) + + def test_renderer_num_workers_with_mm_cache(): - """Disallow renderer_num_workers > 1 when mm processor cache is enabled, - since neither cache type is thread-safe.""" + """Disallow renderer_num_workers > 1 with the mm processor cache only for + pooling models, whose preprocessing runs on the renderer workers.""" mm_model = "Qwen/Qwen2-VL-2B-Instruct" - # Should raise: multi-worker + cache enabled (default cache_gb=4) + # Should raise: pooling + multi-worker + cache enabled (default cache_gb=4) with pytest.raises(ValueError, match="renderer-num-workers"): - ModelConfig(mm_model, renderer_num_workers=4) + ModelConfig(mm_model, runner="pooling", renderer_num_workers=4) - # Should raise: multi-worker + explicit cache size + # Should raise: pooling + multi-worker + explicit cache size with pytest.raises(ValueError, match="renderer-num-workers"): - ModelConfig(mm_model, renderer_num_workers=2, mm_processor_cache_gb=1.0) + ModelConfig( + mm_model, + runner="pooling", + renderer_num_workers=2, + mm_processor_cache_gb=1.0, + ) + + # Should pass: pooling + multi-worker + cache disabled + config = ModelConfig( + mm_model, runner="pooling", renderer_num_workers=4, mm_processor_cache_gb=0 + ) + assert config.renderer_num_workers == 4 - # Should pass: multi-worker + cache disabled - config = ModelConfig(mm_model, renderer_num_workers=4, mm_processor_cache_gb=0) + # Should pass: generate models preprocess on the dedicated mm executor + config = ModelConfig(mm_model, renderer_num_workers=4) assert config.renderer_num_workers == 4 # Should pass: single worker + cache enabled (default) diff --git a/tests/test_envs.py b/tests/test_envs.py index 56c04dd6f2e2..5e0363e33a11 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -15,6 +15,7 @@ env_with_choices, environment_variables, ) +from vllm.exceptions import VLLMValidationError def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch): @@ -145,6 +146,15 @@ def test_precompiled_install_flags_are_orthogonal() -> None: assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is True +def test_rust_bench_auto_path_missing_fails_fast() -> None: + with ( + patch.dict(os.environ, {"VLLM_USE_RUST_BENCH": "1"}, clear=True), + patch("vllm.envs.os.path.isfile", return_value=False), + pytest.raises(FileNotFoundError, match="vllm-rs binary was not found"), + ): + environment_variables["VLLM_RUST_FRONTEND_PATH"]() + + class TestEnvWithChoices: """Test cases for env_with_choices function.""" @@ -538,7 +548,7 @@ def test_sampling_params_respects_limit( max_n = envs.VLLM_MAX_N_SEQUENCES SamplingParams(n=max_n) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): SamplingParams(n=max_n + 1) def test_sampling_params_respects_custom_limit( @@ -554,5 +564,5 @@ def test_sampling_params_respects_custom_limit( SamplingParams(n=128) - with pytest.raises(ValueError, match="n must be at most 128"): + with pytest.raises(VLLMValidationError, match="n must be at most 128"): SamplingParams(n=129) diff --git a/tests/test_jit_monitor.py b/tests/test_jit_monitor.py deleted file mode 100644 index 50261a479d92..000000000000 --- a/tests/test_jit_monitor.py +++ /dev/null @@ -1,534 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import inspect -import os -import sys -from contextlib import contextmanager -from types import ModuleType, SimpleNamespace -from typing import Any, cast -from unittest import mock - -import pytest - -from vllm.utils import jit_monitor - - -@pytest.fixture(autouse=True) -def _reset_monitor(): - """Reset global monitor state between tests.""" - jit_monitor._active = False - jit_monitor._mode = "warn" - jit_monitor._verbose = False - jit_monitor._cutedsl_hook_installed = False - jit_monitor._tilelang_hook_installed = False - jit_monitor._tilelang_jitimpl_compile_depth = 0 - yield - jit_monitor._active = False - jit_monitor._mode = "warn" - jit_monitor._verbose = False - jit_monitor._cutedsl_hook_installed = False - jit_monitor._tilelang_hook_installed = False - jit_monitor._tilelang_jitimpl_compile_depth = 0 - - -# ------------------------------------------------------------------ -# Helpers — lightweight stand-ins for the modules ``activate()`` patches -# ------------------------------------------------------------------ - - -def _make_fake_knobs(*, autotuning_print=False, jit_hook=None): - """Build a minimal fake ``triton.knobs`` namespace.""" - autotuning = SimpleNamespace(print=autotuning_print) - runtime = SimpleNamespace(jit_post_compile_hook=jit_hook) - return SimpleNamespace(autotuning=autotuning, runtime=runtime) - - -def _fake_cute_import_modules(compile_fn): - """Fake Python's parent package + submodule for ``import cutlass.cute``.""" - fake_cute = cast(Any, ModuleType("cutlass.cute")) - fake_cute.compile = compile_fn - fake_parent_package = cast(Any, ModuleType("cutlass")) - fake_parent_package.__path__ = [] - fake_parent_package.cute = fake_cute - return { - "cutlass": fake_parent_package, - "cutlass.cute": fake_cute, - } - - -def _fake_cute_compile(*args, **kwargs): - return "compiled" - - -def _fake_tilelang_import_modules(): - """Fake Python's TileLang modules touched by ``jit_monitor.activate``.""" - - class FakeJITKernel: - def __init__(self, *args, **kwargs): - pass - - class FakeJITImpl: - def __init__(self, func, signature): - self.func = func - self.signature = signature - self.mode = "lazy" - self._kernel_cache = {} - - def __call__(self, *args, **kwargs): - key, _ = self.func.parse_args(*args, **kwargs) - kernel = self._kernel_cache.get(key) - if kernel is None: - kernel = "compiled" - self._kernel_cache[key] = kernel - return kernel - - fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) - fake_kernel.JITKernel = FakeJITKernel - - fake_jit = cast(Any, ModuleType("tilelang.jit")) - fake_jit.JITImpl = FakeJITImpl - fake_jit.kernel = fake_kernel - - fake_tilelang = cast(Any, ModuleType("tilelang")) - fake_tilelang.jit = fake_jit - - return { - "tilelang": fake_tilelang, - "tilelang.jit": fake_jit, - "tilelang.jit.kernel": fake_kernel, - } - - -@contextmanager -def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile): - """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``.""" - fake_triton = cast(Any, ModuleType("triton")) - fake_triton.knobs = fake_knobs - with ( - mock.patch.dict( - sys.modules, - { - "triton": fake_triton, - **_fake_cute_import_modules(cute_compile), - **_fake_tilelang_import_modules(), - }, - ), - mock.patch.object(jit_monitor, "HAS_TRITON", True), - ): - yield - - -# ------------------------------------------------------------------ -# Unit tests (no GPU required, triton is mocked) -# ------------------------------------------------------------------ - - -class TestActivateBasic: - def test_sets_active(self): - assert not jit_monitor.is_active() - with _patch_jit_modules(_make_fake_knobs()): - jit_monitor.activate() - assert jit_monitor.is_active() - - def test_idempotent(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate() - first_hook = fake.runtime.jit_post_compile_hook - jit_monitor.activate() - assert fake.runtime.jit_post_compile_hook is first_hook - - def test_logs_info_on_activation(self): - with ( - mock.patch.object(jit_monitor.logger, "info") as m, - _patch_jit_modules(_make_fake_knobs()), - ): - jit_monitor.activate() - m.assert_called_once() - assert "Kernel JIT monitor activated" in m.call_args[0][0] - - def test_rejects_unknown_mode(self): - with pytest.raises(ValueError, match="Unsupported JIT monitor mode"): - jit_monitor.activate(mode="panic") # type: ignore[arg-type] - - -class TestAutotuningPrint: - def test_enables_autotuning_print(self): - fake = _make_fake_knobs(autotuning_print=False) - with _patch_jit_modules(fake): - jit_monitor.activate() - assert fake.autotuning.print is True - - def test_respects_user_opt_out(self): - fake = _make_fake_knobs(autotuning_print=False) - with ( - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}), - _patch_jit_modules(fake), - ): - jit_monitor.activate() - assert fake.autotuning.print is False - - def test_noop_when_user_already_enabled(self): - fake = _make_fake_knobs(autotuning_print=True) - with ( - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}), - _patch_jit_modules(fake), - ): - jit_monitor.activate() - assert fake.autotuning.print is True - - -class TestTritonJitHook: - def test_hook_registered(self): - fake = _make_fake_knobs() - assert fake.runtime.jit_post_compile_hook is None - with _patch_jit_modules(fake): - jit_monitor.activate() - assert fake.runtime.jit_post_compile_hook is not None - - def test_hook_logs_warning(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="test_kernel") - - with ( - mock.patch.object(jit_monitor.logger, "warning_once") as m, - mock.patch.object(jit_monitor.logger, "warning") as warning, - ): - hook( - key="some_key", - repr="some_repr", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - - m.assert_called_once() - warning.assert_not_called() - msg = m.call_args[0][0] % m.call_args[0][1:] - assert "Triton kernel JIT compilation during inference" in msg - assert "test_kernel" in msg - - def test_hook_chains_existing_hook(self): - existing = mock.MagicMock(return_value="existing_result") - fake = _make_fake_knobs(jit_hook=existing) - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="chained_kernel") - kwargs = dict( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - result = hook(**kwargs) - - existing.assert_called_once() - assert result == "existing_result" - - def test_hook_works_without_existing_hook(self): - fake = _make_fake_knobs(jit_hook=None) - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="solo_kernel") - result = hook( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - assert result is None - - def test_error_mode_raises(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate(mode="error") - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="error_kernel") - with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"): - hook( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - - -class TestNoTritonFallback: - def test_activate_without_triton(self): - with mock.patch.object(jit_monitor, "HAS_TRITON", False): - jit_monitor.activate() - assert jit_monitor.is_active() - - -class TestCuTeDSLHook: - def test_compile_logs_warning(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - result = cute.compile(lambda: None, "arg", option=True) - - assert result == "compiled" - warning_once.assert_called_once() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "CuTeDSL JIT compilation during inference" in msg - - def test_compile_logs_verbose_warning(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate(verbose=True) - with mock.patch.object(jit_monitor.logger, "warning") as warning: - result = cute.compile(lambda: None, "arg", option=True) - - assert result == "compiled" - warning.assert_called_once() - msg = warning.call_args[0][0] % warning.call_args[0][1:] - assert "CuTeDSL JIT compilation during inference" in msg - - def test_error_mode_raises(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate(mode="error") - with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): - cute.compile(lambda: None, "arg", option=True) - - def test_subscripted_compile_is_monitored(self): - """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" - - class FakeCompileCallable: - def __getitem__(self, options): - return self - - def __call__(self, *args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): - import cutlass.cute as cute - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - result = cute.compile[("opt_level", 3)](lambda: None, "arg") - - assert result == "compiled" - warning_once.assert_called_once() - - -class TestTileLangHook: - def test_jit_kernel_logs_warning(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"}) - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi") - - warning_once.assert_called_once() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "TileLang JIT compilation during inference" in msg - assert "tl_kernel" in msg - - def test_jit_impl_logs_warning(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit import JITImpl - - def tilelang_fn( - gemm_out_mul, - hidden_size: int, - n_splits: int = 1, - hc_mult: int = 4, - ): - return None - - class FakeFunc: - orig_func = tilelang_fn - - def parse_args(self, *args, **kwargs): - return ( - ( - "tilelang_key", - kwargs["hidden_size"], - kwargs.get("n_splits", 1), - ), - {}, - ) - - def set_mode(self, mode): - self.mode = mode - - tensor = SimpleNamespace( - shape=(2, 16, 24), - dtype="float32", - device="cuda:0", - ) - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) - - jit_monitor.activate() - with ( - mock.patch.object(jit_monitor.logger, "warning_once") as warning_once, - mock.patch.object(jit_monitor.logger, "warning") as warning, - ): - impl(tensor, hidden_size=7168, n_splits=2) - - warning_once.assert_called_once() - warning.assert_not_called() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "TileLang JIT compilation during inference" in msg - assert "tilelang_fn" in msg - - def test_jit_impl_does_not_log_on_cache_hit(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit import JITImpl - - def tilelang_fn(gemm_out_mul, n_splits: int = 1): - return None - - class FakeFunc: - orig_func = tilelang_fn - - def parse_args(self, *args, **kwargs): - return (("tilelang_key", kwargs.get("n_splits", 1)), {}) - - def set_mode(self, mode): - self.mode = mode - - tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32") - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - impl(tensor, n_splits=2) - impl(tensor, n_splits=2) - - warning_once.assert_called_once() - - def test_from_database_does_not_log(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"}) - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - JITKernel(func=func, from_database=True) - - warning_once.assert_not_called() - - def test_error_mode_raises(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"}) - jit_monitor.activate(mode="error") - with pytest.raises(RuntimeError, match="TileLang JIT compilation"): - JITKernel(func=func) - - -# ------------------------------------------------------------------ -# Integration tests (real Triton + GPU) -# ------------------------------------------------------------------ - -try: - import torch - - _HAS_CUDA = torch.cuda.is_available() -except ImportError: - _HAS_CUDA = False - -try: - import triton - import triton.language as tl - - _HAS_TRITON = True -except ImportError: - _HAS_TRITON = False - -_skip_no_gpu = pytest.mark.skipif( - not (_HAS_CUDA and _HAS_TRITON), - reason="Requires CUDA GPU and Triton", -) - - -if _HAS_TRITON: - - @triton.jit - def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n - x = tl.load(x_ptr + offs, mask=mask) - y = tl.load(y_ptr + offs, mask=mask) - tl.store(out_ptr + offs, x + y, mask=mask) - - -def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None: - """Launch ``_add_kernel`` with vectors of length *n*.""" - x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment - y = torch.randn(n, device="cuda") - out = torch.empty(n, device="cuda") - grid = ((n + block - 1) // block,) - _add_kernel[grid](x, y, out, n, BLOCK=block) - torch.accelerator.synchronize() - - -@_skip_no_gpu -class TestTritonJitHookIntegration: - """End-to-end: real Triton kernel, real GPU, real hook.""" - - def test_no_warning_on_cached_shape(self): - _run_add_kernel(1024) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as w: - _run_add_kernel(1024) - w.assert_not_called() - - def test_warning_on_new_constexpr(self): - _run_add_kernel(1024, block=256) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as w: - # Different BLOCK (a tl.constexpr) forces recompilation. - _run_add_kernel(1024, block=512) - w.assert_called() - msg = w.call_args[0][0] % w.call_args[0][1:] - assert "_add_kernel" in msg - - def test_verbose_warning_on_each_new_pointer_alignment(self): - _run_add_kernel(1024) - - jit_monitor.activate(verbose=True) - with ( - mock.patch.object(jit_monitor.logger, "warning") as w, - mock.patch.object(jit_monitor.logger, "warning_once") as w_once, - ): - _run_add_kernel(1024, offset=1) - assert w.called - w_once.assert_not_called() diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index 17d04078b4e6..f34270d0da52 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -52,13 +52,19 @@ class MockModelConfig: def test_removed_pooling_parameters(parameter: str, value: Any, message: str): data = {"input": "hello", parameter: value} for request_type in (EmbeddingRequest, ClassificationRequest, PoolingRequest): - with pytest.raises(ValidationError, match=message) as exc_info: + with pytest.raises(VLLMValidationError, match=message): TypeAdapter(request_type).validate_python(data) - assert len(exc_info.value.errors()) == 1 - with pytest.raises(ValidationError, match=message) as exc_info: - TypeAdapter(PoolerConfig).validate_python({parameter: value}) - assert len(exc_info.value.errors()) == 1 + # PoolerConfig still raises bare ValueError for `normalize` + # (wrapped to ValidationError by Pydantic), but `check_removed_pooling_task` + # raises VLLMValidationError for removed tasks. + if parameter == "normalize": + with pytest.raises(ValidationError, match=message) as exc_info: + TypeAdapter(PoolerConfig).validate_python({parameter: value}) + assert len(exc_info.value.errors()) == 1 + else: + with pytest.raises(VLLMValidationError, match=message): + TypeAdapter(PoolerConfig).validate_python({parameter: value}) if parameter == "task": with pytest.raises(VLLMValidationError, match=message): @@ -80,7 +86,7 @@ def test_embed(): invalid_parameters = classify_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(embed_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -100,7 +106,7 @@ def test_embed_dimensions(model_info: EmbedModelInfo): pooling_params = PoolingParams(task=task, dimensions=None) pooling_params.verify(model_config) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, dimensions=1) pooling_params.verify(model_config) @@ -131,7 +137,7 @@ def test_embed_dimensions_matryoshka_without_list_upper_bound(): PoolingParams(task=task, dimensions=16).verify(model_config) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): PoolingParams(task=task, dimensions=64).verify(model_config) @@ -150,7 +156,7 @@ def test_classify(task): invalid_parameters = embed_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(classify_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -176,7 +182,7 @@ def test_token_embed(pooling_type: str): invalid_parameters = classify_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(embed_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -202,6 +208,6 @@ def test_token_classify(pooling_type: str): invalid_parameters = embed_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(classify_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index e5d811fbb137..65ab0738c964 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -5,6 +5,7 @@ import pytest from vllm import SamplingParams +from vllm.exceptions import VLLMValidationError @dataclass @@ -32,7 +33,7 @@ def get_vocab_size(self) -> int: ) def test_diffusion_rejects_unsupported_params(kwargs: dict): params = SamplingParams(**kwargs) - with pytest.raises(ValueError, match="not yet supported with diffusion"): + with pytest.raises(VLLMValidationError, match="not yet supported with diffusion"): params.verify(MockModelConfig(is_diffusion=True), None, None, None) diff --git a/tests/tokenizers_/test_mistral.py b/tests/tokenizers_/test_mistral.py index 47abbd812898..fe06f6b32873 100644 --- a/tests/tokenizers_/test_mistral.py +++ b/tests/tokenizers_/test_mistral.py @@ -8,11 +8,13 @@ import pytest from mistral_common.exceptions import InvalidMessageStructureException from mistral_common.guidance.grammar_factory import GrammarFactory -from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy +from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy, SpecialTokens +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tokenizers.mistral import ( MistralTokenizer, _validate_apply_chat_template_args, + validate_request_params, ) @@ -1247,7 +1249,9 @@ def test_convert_tokens_to_string(self, mistral_tokenizer: MistralTokenizer): expected_strings = ( '[{"type": "function", "function": {"name": "get_weather", "description": "Gets the current weather in a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city name"}}, "required": ["city"]}}}] I am an AI\n\nHello world ![TOOL_CALLS][{"name": "get_weather", "arguments": {"city": "Paris"}, "id": "123456789"}] {"content": {"temperature": 20, "unit": "celsius"}, "call_id": "123456789"}', # noqa: E501 - 'I am an AI[{"type": "function", "function": {"name": "get_weather", "description": "Gets the current weather in a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city name"}}, "required": ["city"]}}}]Hello world ![TOOL_CALLS]get_weather{"city": "Paris"}{"temperature": 20, "unit": "celsius"}', # noqa: E501 + # v11+ tool-call decode emits the explicit [ARGS] separator between + # the function name and its JSON arguments (get_weather[ARGS]{...}). + 'I am an AI[{"type": "function", "function": {"name": "get_weather", "description": "Gets the current weather in a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city name"}}, "required": ["city"]}}}]Hello world ![TOOL_CALLS]get_weather[ARGS]{"city": "Paris"}{"temperature": 20, "unit": "celsius"}', # noqa: E501 ) assert ( @@ -1498,6 +1502,7 @@ def test_convert_tokens_to_string(self, mistral_tokenizer: MistralTokenizer): "get", "_", "weather", + "[ARGS]", '{"', "city", '":', @@ -2234,3 +2239,140 @@ def test_apply_chat_template_reasoning_assistant( decoded = mistral_tokenizer.tokenizer.decode(output, SpecialTokenPolicy.KEEP) assert "[THINK]2+2 equals 4[/THINK]" in decoded + + +def test_convert_ids_to_tokens_pre_args_tekken(): + """convert_ids_to_tokens must not raise on pre-[ARGS] Tekken tokenizers. + + Tekken tokenizers before v11 (e.g. Ministral-8B-Instruct-2410) have no + [ARGS] special token. This guards the is_special([ARGS]) gate in + MistralTokenizer.convert_ids_to_tokens, which must skip [ARGS] rather than + call get_special_token(args) unconditionally (which raises on Tekken). + """ + tokenizer = MistralTokenizer.from_pretrained("mistralai/Ministral-8B-Instruct-2410") + assert tokenizer.is_tekken + assert not tokenizer.tokenizer.is_special(SpecialTokens.args) + + ids = tokenizer.encode("Hello world !", add_special_tokens=False) + tokens = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert len(tokens) > 0 + + +# --------------------------------------------------------------------------- +# validate_request_params – reasoning_effort validation (no tokenizer needed) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("reasoning_effort", [None, "none", "high"]) +def test_validate_request_params_valid_effort(reasoning_effort: str | None) -> None: + request = ChatCompletionRequest( + model="test-model", + messages=[], + reasoning_effort=reasoning_effort, + ) + assert validate_request_params(request) is None + + +def test_validate_request_params_rejects_unsupported_effort() -> None: + request = ChatCompletionRequest( + model="test-model", + messages=[], + reasoning_effort="medium", + ) + with pytest.raises(ValueError, match="reasoning_effort"): + validate_request_params(request) + + +# --------------------------------------------------------------------------- +# apply_chat_template reasoning_effort passthrough (v15 tokenizer) +# --------------------------------------------------------------------------- + +_PASSTHROUGH_MESSAGES = [{"role": "user", "content": "Hello"}] + + +@pytest.fixture(scope="module") +def v15_mistral_tokenizer() -> MistralTokenizer: + """Load the v15 Mistral tokenizer; skip if unavailable.""" + try: + return MistralTokenizer.from_pretrained("mistralai/Mistral-Small-4-119B-2603") + except Exception: + pytest.skip("v15 tokenizer unavailable") + + +@pytest.fixture(scope="module") +def v13_mistral_tokenizer() -> MistralTokenizer: + """Load the v13 Mistral tokenizer; skip if unavailable.""" + try: + return MistralTokenizer.from_pretrained("mistralai/Magistral-Small-2509") + except Exception: + pytest.skip("v13 tokenizer unavailable") + + +def test_v15_apply_chat_template_passes_reasoning_effort_high( + monkeypatch: pytest.MonkeyPatch, + v15_mistral_tokenizer: MistralTokenizer, +) -> None: + captured_kwargs: list[dict] = [] + + def fake_apply_chat_template(**kwargs): + captured_kwargs.append(kwargs) + return [0] + + monkeypatch.setattr( + v15_mistral_tokenizer.transformers_tokenizer, + "apply_chat_template", + fake_apply_chat_template, + ) + + v15_mistral_tokenizer.apply_chat_template( + messages=_PASSTHROUGH_MESSAGES, + reasoning_effort="high", + ) + + assert captured_kwargs[-1]["reasoning_effort"] == "high" + + +def test_v15_apply_chat_template_passes_reasoning_effort_none_by_default( + monkeypatch: pytest.MonkeyPatch, + v15_mistral_tokenizer: MistralTokenizer, +) -> None: + captured_kwargs: list[dict] = [] + + def fake_apply_chat_template(**kwargs): + captured_kwargs.append(kwargs) + return [0] + + monkeypatch.setattr( + v15_mistral_tokenizer.transformers_tokenizer, + "apply_chat_template", + fake_apply_chat_template, + ) + + v15_mistral_tokenizer.apply_chat_template(messages=_PASSTHROUGH_MESSAGES) + + assert "reasoning_effort" in captured_kwargs[-1] + assert captured_kwargs[-1]["reasoning_effort"] is None + + +def test_pre_v15_apply_chat_template_omits_reasoning_effort( + monkeypatch: pytest.MonkeyPatch, + v13_mistral_tokenizer: MistralTokenizer, +) -> None: + captured_kwargs: list[dict] = [] + + def fake_apply_chat_template(**kwargs): + captured_kwargs.append(kwargs) + return [0] + + monkeypatch.setattr( + v13_mistral_tokenizer.transformers_tokenizer, + "apply_chat_template", + fake_apply_chat_template, + ) + + v13_mistral_tokenizer.apply_chat_template( + messages=_PASSTHROUGH_MESSAGES, + reasoning_effort="high", + ) + + assert "reasoning_effort" not in captured_kwargs[-1] diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index 9635e9963b5e..0a47426dc974 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -109,6 +109,8 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path): try: def fake_from_pretrained(path_or_repo_id: str, *args, **kwargs): + passed_config = kwargs.pop("config") + assert isinstance(passed_config, Qwen3_5MoeConfig) loaded_config = AutoConfig.from_pretrained( path_or_repo_id, trust_remote_code=False, diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 2e5069dbed94..7fd3860ef719 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from unittest.mock import MagicMock import pytest @@ -10,6 +11,7 @@ ToolParserTests, ) from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser class TestInternLM2ToolParser(ToolParserTests): @@ -120,3 +122,44 @@ def test_config(self) -> ToolParserTestConfig: ), }, ) + + +def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -> None: + """Arguments arriving whole in one delta must not be dropped.""" + tokenizer_vocab = default_tokenizer.get_vocab() + default_tokenizer.get_vocab = MagicMock() + tokenizer_vocab.update( + { + "<|action_start|>": 92540, + "<|plugin|>": 92541, + "<|action_end|>": 92542, + } + ) + default_tokenizer.get_vocab.return_value = tokenizer_vocab + parser = Internlm2ToolParser(default_tokenizer) + + deltas = [ + '<|action_start|><|plugin|>{"name": "get_weather"', + ', "parameters": {"city": "Dallas", "state": "TX"}}<|action_end|>', + ] + + streamed = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed += arguments + + assert json.loads(streamed) == {"city": "Dallas", "state": "TX"} diff --git a/tests/tool_parsers/test_jamba_tool_parser.py b/tests/tool_parsers/test_jamba_tool_parser.py index f0e7899c8aaf..9eb7404209dc 100644 --- a/tests/tool_parsers/test_jamba_tool_parser.py +++ b/tests/tool_parsers/test_jamba_tool_parser.py @@ -306,3 +306,34 @@ def test_extract_tool_calls_streaming( ) ] assert_tool_calls(actual_tool_calls, expected_tool_calls) + + +def test_extract_tool_calls_streaming_arguments_in_single_delta(jamba_tool_parser): + """Arguments delivered whole in one coarse delta must not be dropped.""" + deltas = [ + '[{"name": "get_current_weather"', + ",", + ' "arguments": {"city": "Dallas", "state": "TX"}}]', + "", + ] + + streamed_arguments = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = jamba_tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed_arguments += arguments + + assert json.loads(streamed_arguments) == {"city": "Dallas", "state": "TX"} diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py new file mode 100644 index 000000000000..36b04b67b195 --- /dev/null +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Named tool choice for Kimi K3: allowed when the XTML structural tag is +attached (strict tool calling), rejected otherwise.""" + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.exceptions import VLLMValidationError +from vllm.sampling_params import StructuredOutputsParams + + +class _DummyTokenizer: + def get_vocab(self): + return {} + + def encode(self, text, add_special_tokens=False): + return [ord(c) for c in text] + + +def _request(with_tag: bool) -> ChatCompletionRequest: + req = ChatCompletionRequest( + model="k3", + messages=[{"role": "user", "content": "hi"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + ) + if with_tag: + req.structured_outputs = StructuredOutputsParams( + structural_tag='{"type": "structural_tag", "format": {}}' + ) + return req + + +def _parser(): + from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser + + return KimiK3ToolParser(_DummyTokenizer()) + + +def test_named_choice_allowed_with_structural_tag(): + req = _parser().adjust_request(_request(with_tag=True)) + assert req.skip_special_tokens is False + + +def test_named_choice_rejected_without_structural_tag(): + with pytest.raises(VLLMValidationError): + _parser().adjust_request(_request(with_tag=False)) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 354adab66976..b84dfad0923c 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -5,7 +5,8 @@ from unittest.mock import MagicMock import pytest -from xgrammar import StructuralTag +from xgrammar import Grammar, StructuralTag +from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedFunction, @@ -24,6 +25,7 @@ from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser +from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser @@ -31,7 +33,7 @@ SUPPORTED_STRUCTURAL_TAG_MODELS, VLLM_BUILTIN_STRUCTURAL_TAG_MODELS, XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS, - _get_function_parameters, + get_function_parameters, get_model_structural_tag, ) @@ -165,6 +167,196 @@ def test_hermes_required_tool_calls_use_empty_separator(): assert tag.format.separator == "" +# --------------------------------------------------------------------------- +# Kimi K3 (XTML channel format) structural tag +# --------------------------------------------------------------------------- +_K3_RESPONSE_OPEN = "<|open|>response<|sep|>" +_K3_RESPONSE_CLOSE = "<|close|>response<|sep|>" +_K3_TOOLS_OPEN = "<|open|>tools<|sep|>" +_K3_TOOLS_CLOSE = "<|close|>tools<|sep|>" +_K3_CALL_CLOSE = "<|close|>call<|sep|>" +_K3_ARG_CLOSE = "<|close|>argument<|sep|>" +_K3_MESSAGE_CLOSE = "<|close|>message<|sep|>" + + +def _k3_tools_by_name() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + }, + "required": ["city"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "run_command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, + ), + ] + + +def _k3_arg(key: str, typ: str, val: str) -> str: + return f'<|open|>argument key="{key}" type="{typ}"<|sep|>{val}{_K3_ARG_CLOSE}' + + +def _k3_call(name: str, args: str, idx: int = 1) -> str: + return f'<|open|>call tool="{name}" index="{idx}"<|sep|>{args}{_K3_CALL_CLOSE}' + + +def _k3_response(content: str = "") -> str: + return f"{_K3_RESPONSE_OPEN}{content}{_K3_RESPONSE_CLOSE}" + + +def _k3_tools(*calls: str) -> str: + return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}" + + +def _k3_grammar(tool_choice, tools=None): + tag = get_model_structural_tag( + model="kimi_k3", + tools=tools if tools is not None else _k3_tools_by_name(), + tool_choice=tool_choice, + reasoning=False, + ) + assert isinstance(tag, StructuralTag) + return Grammar.from_structural_tag(tag) + + +def test_kimi_k3_registered_as_vllm_builtin(): + assert "kimi_k3" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS + assert KimiK3ToolParser.structural_tag_model == "kimi_k3" + + +def test_kimi_k3_auto_without_strict_is_unconstrained(): + # auto + no strict tool => no structural tag (matches the strict gate). + tag = get_model_structural_tag( + model="kimi_k3", + tools=_k3_tools_by_name(), + tool_choice="auto", + reasoning=False, + ) + assert tag is None + + +@pytest.mark.parametrize( + "body", + [ + # single required arg + _k3_response() + + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))), + # response content + two args (string + number) + _k3_response("Checking.") + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("city", "string", "Paris") + _k3_arg("days", "number", "3"), + ) + ), + # args in reverse order (parser is order-agnostic) + _k3_response() + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("days", "number", "3") + _k3_arg("city", "string", "Paris"), + ) + ), + # two calls, second tool + _k3_response() + + _k3_tools( + _k3_call("get_weather", _k3_arg("city", "string", "Paris"), 1), + _k3_call("run_command", _k3_arg("command", "string", "ls -la"), 2), + ), + # string value with regex metacharacters / spaces + _k3_response() + + _k3_tools( + _k3_call( + "run_command", _k3_arg("command", "string", "grep -E 'a|b{2,}' x.py") + ) + ), + # trailing message-close marker (model's natural turn terminator) + _k3_response() + + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))) + + _K3_MESSAGE_CLOSE, + # non-thinking mode: response-open is the prompt prefix, so it is absent + _K3_RESPONSE_CLOSE + + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))), + ], +) +def test_kimi_k3_required_accepts_valid_tool_calls(body: str): + assert _is_grammar_accept_string(_k3_grammar("required"), body) + + +@pytest.mark.parametrize( + "body", + [ + # unknown tool name + _k3_response() + + _k3_tools(_k3_call("get_temperature", _k3_arg("city", "string", "x"))), + # number arg given a non-numeric JSON value + _k3_response() + + _k3_tools(_k3_call("get_weather", _k3_arg("days", "number", "abc"))), + # undeclared argument key + _k3_response() + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("city", "string", "Paris") + _k3_arg("zzz", "string", "x"), + ) + ), + # required schema but no argument tags + _k3_response() + _k3_tools(_k3_call("get_weather", "")), + # missing tools close marker + _k3_response() + + _K3_TOOLS_OPEN + + _k3_call("get_weather", _k3_arg("city", "string", "Paris")), + # required but no tool call + _k3_response("hello"), + ], +) +def test_kimi_k3_required_rejects_invalid(body: str): + assert not _is_grammar_accept_string(_k3_grammar("required"), body) + + +def test_kimi_k3_schema_without_required_accepts_empty_call(): + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + ) + ] + grammar = _k3_grammar("required", tools=tools) + body = _k3_response() + _k3_tools(_k3_call("get_weather", "")) + + assert _is_grammar_accept_string(grammar, body) + + +def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict): + # With a strict tool the tag is built; the tools channel is optional so a + # plain response (no tool call) is still valid. + grammar = _k3_grammar("auto", tools=sample_tools_strict) + assert _is_grammar_accept_string(grammar, _k3_response("Just answering.")) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_named_tool_choice( model: str, @@ -338,4 +530,152 @@ def test_get_function_parameters_relaxes_function_strict_false(): strict=False, ) - assert _get_function_parameters(function) is True + assert get_function_parameters(function) is True + + +def _k3_tools_with_root_defs() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "make_config", + "parameters": { + "type": "object", + "properties": { + "config": { + "type": "object", + "properties": { + "build": {"$ref": "#/$defs/build"}, + "index": {"type": "string"}, + }, + "required": ["index"], + "additionalProperties": False, + }, + }, + "required": ["config"], + "$defs": { + "build": { + "type": "object", + "properties": {"outDir": {"type": "string"}}, + "additionalProperties": False, + } + }, + }, + }, + ) + ] + + +def test_kimi_k3_property_ref_to_root_defs_compiles_and_accepts(): + # Root-level $defs referenced from inside a property schema (the walle + # TestReferences shape). Slicing the property out of the parameters + # document orphans "#/$defs/..." unless the builder re-attaches $defs; + # before the fix Grammar.from_structural_tag raised on the dangling ref. + grammar = _k3_grammar("required", tools=_k3_tools_with_root_defs()) + + body = _k3_response() + _k3_tools( + _k3_call( + "make_config", + _k3_arg( + "config", + "object", + '{"build": {"outDir": "dist"}, "index": "a.html"}', + ), + ) + ) + assert _is_grammar_accept_string(grammar, body) + + +def _k3_tools_with_string_enum() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "set_unit", + "parameters": { + "type": "object", + "properties": { + "unit": { + "type": "string", + "enum": ["celsius", " fahrenheit", "\tkelvin"], + }, + }, + "required": ["unit"], + }, + }, + ) + ] + + +@pytest.mark.parametrize("value", ["celsius", " fahrenheit", "\tkelvin"]) +def test_kimi_k3_string_enum_accepts_exact_values(value: str): + # Raw string channel with enum: constrained to the exact enum values, + # including leading-whitespace variants the model otherwise flubs. + grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum()) + body = _k3_response() + _k3_tools( + _k3_call("set_unit", _k3_arg("unit", "string", value)) + ) + assert _is_grammar_accept_string(grammar, body) + + +@pytest.mark.parametrize("value", ["kelvin", "Celsius", "celsius ", ""]) +def test_kimi_k3_string_enum_rejects_non_members(value: str): + grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum()) + body = _k3_response() + _k3_tools( + _k3_call("set_unit", _k3_arg("unit", "string", value)) + ) + assert not _is_grammar_accept_string(grammar, body) + + +def _k3_tools_with_maxlen() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "set_note", + "parameters": { + "type": "object", + "properties": { + "note": {"type": "string", "maxLength": 8, "minLength": 2}, + }, + "required": ["note"], + }, + }, + ) + ] + + +def test_kimi_k3_string_maxlength_bounds_raw_channel(): + # Raw string channel with maxLength/minLength: enforced via a bounded + # regex that keeps the "<|" marker prefix unambiguous but still allows + # a bare '<' inside values. + grammar = _k3_grammar("required", tools=_k3_tools_with_maxlen()) + + def body(val: str) -> str: + return _k3_response() + _k3_tools( + _k3_call("set_note", _k3_arg("note", "string", val)) + ) + + assert _is_grammar_accept_string(grammar, body("ab")) + assert _is_grammar_accept_string(grammar, body("a dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + if text == THINK_OPEN: + return [1, 2, 3] + if text == THINK_CLOSE: + return [4, 2, 3] + return [ord(ch) for ch in text] + + +class KimiK3DelegatingParser(KimiK3Parser): + reasoning_parser_cls = KimiK3ReasoningParser + tool_parser_cls = KimiK3ToolParser + + +def test_parser_manager_selects_kimi_k3_parser(): + parser_cls = ParserManager.get_parser( + tool_parser_name="kimi_k3", + reasoning_parser_name="kimi_k3", + enable_auto_tools=True, + ) + + assert parser_cls is not None + assert issubclass(parser_cls, KimiK3Parser) + assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser + assert parser_cls.tool_parser_cls is KimiK3ToolParser + + +def _request() -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice="auto", + ) + + +def _named_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice={"type": "function", "function": {"name": "calc"}}, + ) + + +def _responses_request(*, tool_choice="auto") -> ResponsesRequest: + return ResponsesRequest.model_validate( + { + "model": "test-model", + "input": "Call the calc tool.", + "tools": [ + { + "type": "function", + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + } + ], + "tool_choice": tool_choice, + } + ) + + +def _arg(key: str, typ: str, value: str) -> str: + return f'{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}' + + +def _call(tool: str, index: int, *args: str) -> str: + body = "".join(args) + return f'{OPEN}call tool="{tool}" index="{index}"{SEP}{body}{CLOSE}call{SEP}' + + +def _response(content: str) -> str: + return f"{OPEN}response{SEP}{content}{RESPONSE_CLOSE}" + + +def _tools(*calls: str) -> str: + return f"{OPEN}tools{SEP}{''.join(calls)}{CLOSE}tools{SEP}" + + +def test_extract_tool_calls_with_response_and_typed_arguments(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _response("answer") + _tools( + _call( + "calc", + 1, + _arg("x", "number", "1"), + _arg("flag", "boolean", "true"), + _arg("text", "string", "raw"), + ) + ) + extracted = parser.extract_tool_calls(output, _request()) + + assert extracted.tools_called is True + assert extracted.content == "answer" + assert len(extracted.tool_calls) == 1 + tool_call = extracted.tool_calls[0] + assert tool_call.function.name == "calc" + assert json.loads(tool_call.function.arguments) == { + "x": 1, + "flag": True, + "text": "raw", + } + + +def test_delegating_parser_preserves_tool_calls_after_reasoning(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + output = ( + f"{THINK_OPEN}step{THINK_CLOSE}" + + _response("answer") + + _tools(_call("calc", 1, _arg("x", "number", "1"))) + ) + + reasoning, content, tool_calls = parser.parse( + output, + _request(), + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content == "answer" + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_required_tool_choice_uses_xtml_parser(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + request = _request().model_copy(update={"tool_choice": "required"}) + output = ( + f"{THINK_OPEN}step{THINK_CLOSE}" + + _response("") + + _tools(_call("calc", 1, _arg("x", "number", "1"))) + ) + + reasoning, content, tool_calls = parser.parse( + output, + request, + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_named_tool_choice_uses_xtml_parser(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + output = ( + f"{THINK_OPEN}step{THINK_CLOSE}" + + _response("") + + _tools(_call("calc", 1, _arg("x", "number", "1"))) + ) + + reasoning, content, tool_calls = parser.parse( + output, + _named_request(), + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_auto_no_call_strips_consumed_response_prefix(): + parser = KimiK3DelegatingParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + request = _request().model_copy( + update={"chat_template_kwargs": {"thinking": False}} + ) + + reasoning, content, tool_calls = parser.parse( + f"answer{RESPONSE_CLOSE}", + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content == "answer" + assert tool_calls is None + + +def test_delegating_parser_required_call_strips_consumed_response_prefix(): + parser = KimiK3DelegatingParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + request = _request().model_copy( + update={ + "tool_choice": "required", + "chat_template_kwargs": {"thinking": False}, + } + ) + output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1"))) + + reasoning, content, tool_calls = parser.parse( + output, + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_truncated_tools_do_not_leak_xtml(): + parser = KimiK3DelegatingParser( + DummyTokenizer(), chat_template_kwargs={"thinking": False} + ) + request = _request().model_copy( + update={ + "tool_choice": "required", + "chat_template_kwargs": {"thinking": False}, + } + ) + + reasoning, content, tool_calls = parser.parse( + (f'{RESPONSE_CLOSE}{OPEN}tools{SEP}{OPEN}call tool="calc" index="1"'), + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content is None + assert tool_calls is None + + +def test_extract_tool_calls_unescapes_attributes(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _tools(_call("a&b"c", 1, _arg("k&q", "string", "v"))) + extracted = parser.extract_tool_calls(output, _request()) + + assert extracted.tools_called is True + assert extracted.tool_calls[0].function.name == 'a&b"c' + assert json.loads(extracted.tool_calls[0].function.arguments) == {"k&q": "v"} + + +def test_extract_tool_calls_allows_less_than_in_attributes(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _tools(_call("calc None: return for repo_id, filename in assets: try: - hf_hub_download(repo_id=repo_id, filename=filename) + hf_api().hf_hub_download(repo_id=repo_id, filename=filename) except Exception as e: logger.warning( "Failed to prefetch %s/%s: %r. Tests depending on this asset may fail.", @@ -1983,7 +1983,10 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator: return pytest.mark.skipif( memory_gb < min_gb, - reason=f"Need at least {min_gb}GB GPU memory to run the test.", + reason=( + f"Need at least {min_gb}GB GPU memory to run the test " + f"(found {memory_gb}GB)." + ), ) diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 421aec3e9b1f..19690d1f34f8 100644 --- a/tests/utils_/test_mem_utils.py +++ b/tests/utils_/test_mem_utils.py @@ -67,6 +67,23 @@ def measure_current_non_torch(): non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa assert abs(non_torch_ratio - 1) <= 0.05 assert result.torch_peak_increase == 1024 * 1024 * 1024 + + expected_total_consumed = (256 + 512) * 1024 * 1024 + total_consumed_ratio = result.total_consumed / expected_total_consumed + assert abs(total_consumed_ratio - 1) <= 0.05, ( + f"total_consumed={result.total_consumed}, " + f"expected={expected_total_consumed}, " + f"ratio={total_consumed_ratio}" + ) + + expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024 + non_kv_ratio = result.non_kv_cache_memory / expected_non_kv + assert abs(non_kv_ratio - 1) <= 0.05, ( + f"non_kv_cache_memory={result.non_kv_cache_memory}, " + f"expected={expected_non_kv}, " + f"ratio={non_kv_ratio}" + ) + del weights lib.cudaFree(handle1) lib.cudaFree(handle2) diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 3110e4b4ee17..87a9c80942ee 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -21,6 +21,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( STR_DTYPE_TO_TORCH_DTYPE, + is_quantized_kv_cache, is_torch_equal_or_newer, set_random_seed, ) @@ -45,6 +46,14 @@ DEVICE_TYPE = current_platform.device_type +# Use the platform's preferred FP8 type so the stored cache matches what the +# backends reinterpret at runtime. On ROCm gfx94x this is e4m3fnuz, not e4m3fn; +# storing e4m3fn bytes there would be re-read as fnuz and produce NaNs. +FP8_KV_CACHE_DTYPES = { + "fp8": current_platform.fp8_dtype(), + "fp8_e4m3": current_platform.fp8_dtype(), +} + # Remove flashinfer from the list if it's not available try: import flashinfer # noqa: F401 @@ -110,6 +119,7 @@ def create_and_prepopulate_kv_cache( num_blocks: int, common_attn_metadata: CommonAttentionMetadata, randomize_blocks: bool = True, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: """Create and prepopulate a KV cache with context data. @@ -140,8 +150,18 @@ def create_and_prepopulate_kv_cache( block_table = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping + # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning + # the higher-precision context tensors quantizes them, mirroring runtime. + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) + storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype + kv_cache = torch.zeros( - num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device + num_blocks, + block_size, + num_kv_heads, + 2 * head_size, + dtype=storage_dtype, + device=device, ) kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size) @@ -195,7 +215,12 @@ def create_and_prepopulate_kv_cache( ] * block_size + token_inter_block_offsets.to(device) # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs) - return kv_cache.transpose(1, 2).contiguous() + kv_cache = kv_cache.transpose(1, 2).contiguous() + + if fp8_kv_cache: + kv_cache = kv_cache.view(torch.uint8) + + return kv_cache class MockAttentionLayer: @@ -224,6 +249,7 @@ def run_attention_backend( kv_cache: torch.Tensor, attn_type: AttentionType = AttentionType.DECODER, sliding_window: int | None = None, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" @@ -291,13 +317,16 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): alibi_slopes=None, sliding_window=sliding_window, attn_type=attn_type, - kv_cache_dtype="auto", + kv_cache_dtype=kv_cache_dtype, ) # Create mock layer and output buffer mock_layer = MockAttentionLayer(device) output = torch.empty_like(query) + if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input: + query = query.to(current_platform.fp8_dtype()) + # Run forward pass # NOTE: The query, key, and value are already shaped correctly # in the calling test function. @@ -324,6 +353,7 @@ def _test_backend_correctness( atol: float = 1e-2, rtol: float = 1e-2, tensor_parallel_size: int = 1, + kv_cache_dtype: str = "auto", ): """ Test that all backends produce similar outputs to a reference implementation @@ -372,6 +402,7 @@ def _test_backend_correctness( num_gpu_blocks=8192, hf_config_override=hf_config_override, ) + vllm_config.cache_config.cache_dtype = kv_cache_dtype device = torch.device(f"{DEVICE_TYPE}:0") kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type) @@ -392,6 +423,13 @@ def _test_backend_correctness( block_size = vllm_config.cache_config.block_size scale = 1.0 / (head_size**0.5) + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) + if fp8_kv_cache: + query_fp8_dtype = current_platform.fp8_dtype() + kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] + atol = max(atol, 6e-2) + rtol = max(rtol, 1e-1) + # 2. Generate data and compute SDPA reference output all_q_vllm, all_k_vllm, all_v_vllm = [], [], [] all_sdpa_outputs = [] @@ -407,10 +445,17 @@ def _test_backend_correctness( k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) + if fp8_kv_cache: + q_ref = q.to(query_fp8_dtype).to(dtype) + k_ref = k_full.to(kv_fp8_dtype).to(dtype) + v_ref = v_full.to(kv_fp8_dtype).to(dtype) + else: + q_ref, k_ref, v_ref = q, k_full, v_full + # SDPA expects (N, H, L, D), so unsqueeze batch and permute - q_sdpa_in = q.unsqueeze(0).transpose(1, 2) - k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2) - v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2) + q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2) + k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2) + v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2) if num_q_heads != num_kv_heads: assert num_q_heads % num_kv_heads == 0, ( @@ -471,6 +516,7 @@ def _test_backend_correctness( num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000, common_attn_metadata=common_attn_metadata, randomize_blocks=True, + kv_cache_dtype=kv_cache_dtype, ) # 4. Run vLLM backends and compare @@ -488,6 +534,12 @@ def _test_backend_correctness( else: backend_cls = None + if is_quantized_kv_cache(kv_cache_dtype) and ( + backend_cls is None + or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype) + ): + continue + if backend_name == AttentionBackendEnum.FLASHINFER: set_kv_cache_layout("HND") reset_kv_cache_layout = True @@ -521,6 +573,7 @@ def _test_backend_correctness( kv_cache_for_backend, sliding_window=sliding_window, attn_type=attn_type, + kv_cache_dtype=kv_cache_dtype, ) finally: if reset_kv_cache_layout: @@ -570,8 +623,13 @@ def error_msg(msg: str, backend_name: str): ) @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"]) def test_causal_backend_correctness( - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int + default_vllm_config, + batch_spec_name: str, + model: str, + tensor_parallel_size: int, + kv_cache_dtype: str, ): """Test backend's correctness with causal attention.""" @@ -612,6 +670,7 @@ def causal_mask_mod( SMALL_BLOCK_BACKENDS, causal_mask_mod, tensor_parallel_size=tensor_parallel_size, + kv_cache_dtype=kv_cache_dtype, ) # Fast FlexAttention needs to run with block_size=128 @@ -623,6 +682,7 @@ def causal_mask_mod( causal_mask_mod, block_size=128, tensor_parallel_size=tensor_parallel_size, + kv_cache_dtype=kv_cache_dtype, ) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 159bb8af3fb9..03cbd0dc364a 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -1,13 +1,36 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from tests.v1.attention.utils import create_vllm_config from vllm.v1.attention.backend import CommonAttentionMetadata -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.indexer import ( + BuildPrefillChunkMetadataKernel, + DeepseekV32IndexerMetadataBuilder, +) from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.worker.block_table import get_block_table_width + + +def test_indexer_warmup_normalizes_zero_compress_ratios(): + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0]) + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + ) + + keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config) + + assert {key.COMPRESS_RATIO for key in keys} == {1, 4, 128} @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -26,11 +49,14 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_ compress_ratio=4, ) vllm_config = create_vllm_config(max_model_len=1024) + max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024) + block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size) builder = DeepseekV32IndexerMetadataBuilder( kv_cache_spec=kv_cache_spec, layer_names=["dummy"], vllm_config=vllm_config, device=device, + block_table_width=block_table_width, ) # Construct a single request where: diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 99dcb09ab154..4ec138270203 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -42,6 +42,8 @@ def _make_vllm_config( cache_config=SimpleNamespace( block_size=block_size, mamba_cache_mode="all", + use_replayssm=False, + replayssm_buffer_len=16, ), compilation_config=SimpleNamespace( cudagraph_mode=CUDAGraphMode.FULL, diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index aca09c3db3d6..b90600ce40aa 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -40,12 +40,17 @@ MLAPrefillBackendEnum, get_mla_prefill_backend, ) +from vllm.v1.attention.backends.mla.prefill.base import MLADimensions +from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported from vllm.v1.kv_cache_interface import ( KVQuantMode, MLAAttentionSpec, ) +from vllm.v1.worker.block_table import get_block_table_width BACKENDS_TO_TEST = [ AttentionBackendEnum.CUTLASS_MLA, @@ -69,7 +74,11 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( cache_dtype: str, expected_quant_mode: KVQuantMode ): - layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576) + layer = SimpleNamespace( + kv_cache_dtype=cache_dtype, + head_size=576, + non_causal_multi_token_decode=False, + ) vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=64), model_config=None ) @@ -119,6 +128,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): layer.kv_b_proj.quant_method = None layer.is_aiter_triton_fp4_bmm_enabled = False layer.is_aiter_triton_fp8_bmm_enabled = False + layer.dcp_q_replicate = False layer.quant_config = None layer.layer_name = "test" @@ -144,14 +154,67 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100) -# Filtered per-test via validate_configuration (capability/deps/dims). +# Validate parameter combinations during collection, before GPU fixtures run. PREFILL_BACKENDS_TO_TEST = [ + MLAPrefillBackendEnum.ROCM_AITER_FA, MLAPrefillBackendEnum.FLASH_ATTN, MLAPrefillBackendEnum.FLASHINFER, MLAPrefillBackendEnum.TRTLLM_RAGGED, MLAPrefillBackendEnum.TOKENSPEED_MLA, ] +MLA_DIMENSIONS_TO_TEST = [ + ("deepseek", 128, 128), + ("glm", 192, 256), +] + + +def _prefill_backend_dimension_params(): + device_capability = current_platform.get_device_capability() + params = [] + for prefill_backend in PREFILL_BACKENDS_TO_TEST: + for dimensions_id, qk_nope_head_dim, v_head_dim in MLA_DIMENSIONS_TO_TEST: + if device_capability is None: + invalid_reasons = ["device capability unavailable"] + else: + try: + invalid_reasons = ( + prefill_backend.get_class().validate_configuration( + device_capability, + MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=64, + v_head_dim=v_head_dim, + ), + ), + ) + ) + except ImportError: + invalid_reasons = ["ImportError"] + + marks = [] + if invalid_reasons: + marks.append( + pytest.mark.skip( + reason=( + f"Prefill backend {prefill_backend.name} unavailable: " + f"{invalid_reasons}" + ) + ) + ) + params.append( + pytest.param( + prefill_backend, + qk_nope_head_dim, + v_head_dim, + id=f"{dimensions_id}-{prefill_backend}", + marks=marks, + ) + ) + return params + SPEC_DECODE_BACKENDS = [] for backend in BACKENDS_TO_TEST: @@ -1097,11 +1160,9 @@ def run_attention_backend( @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"]) @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)]) -@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST) @pytest.mark.parametrize( - ("qk_nope_head_dim", "v_head_dim"), - [(128, 128), (192, 256)], - ids=["deepseek", "glm"], + ("prefill_backend", "qk_nope_head_dim", "v_head_dim"), + _prefill_backend_dimension_params(), ) def test_backend_correctness( default_vllm_config, @@ -1152,32 +1213,6 @@ def test_backend_correctness( if not backends_to_test: pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}") - # Skip prefill backends that can't satisfy capability/deps/dimension constraints. - from vllm.v1.attention.backends.mla.prefill.base import MLADimensions - from vllm.v1.attention.backends.mla.prefill.selector import ( - MLAPrefillSelectorConfig, - ) - - try: - prefill_invalid_reasons = prefill_backend.get_class().validate_configuration( - current_platform.get_device_capability(), - MLAPrefillSelectorConfig( - dtype=torch.bfloat16, - mla_dimensions=MLADimensions( - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=64, - v_head_dim=v_head_dim, - ), - ), - ) - except ImportError: - prefill_invalid_reasons = ["ImportError"] - if prefill_invalid_reasons: - pytest.skip( - f"Prefill backend {prefill_backend.name} unavailable: " - f"{prefill_invalid_reasons}" - ) - batch_spec = BATCH_SPECS[batch_spec_name] is_spec_decode_test = batch_spec_name.startswith("spec_decode") unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test)) @@ -1461,16 +1496,9 @@ def test_backend_correctness( batch_spec, block_size, device ) - # Pad block table to meet requirement: - # block_num % (128 / block_size) == 0 - required_divisor = int(128 / block_size) current_block_num = common_attn_metadata.block_table_tensor.shape[1] - if current_block_num % required_divisor != 0: - # Pad to next multiple of required_divisor - padded_block_num = ( - (current_block_num + required_divisor - 1) // required_divisor - ) * required_divisor - padding_cols = padded_block_num - current_block_num + padded_block_num = get_block_table_width(current_block_num, block_size) + if padding_cols := padded_block_num - current_block_num: padding = torch.zeros( (common_attn_metadata.block_table_tensor.shape[0], padding_cols), dtype=torch.int32, diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py new file mode 100644 index 000000000000..7aca8400f993 --- /dev/null +++ b/tests/v1/attention/test_mla_noncausal.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonMetadata, + MLACommonMetadataBuilder, + QueryLenSupport, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import MLAAttentionSpec + + +class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + supports_non_causal_multi_token_decode = True + + +def _metadata( + query_start_loc: list[int], num_tokens: int | None = None +) -> CommonAttentionMetadata: + num_reqs = len(query_start_loc) - 1 + num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens + return CommonAttentionMetadata( + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), + query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), + seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8, + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=max( + end - start for start, end in zip(query_start_loc, query_start_loc[1:]) + ), + max_seq_len=num_reqs * 100 + 8, + block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view( + num_reqs, 3 + ), + slot_mapping=torch.arange(num_tokens), + causal=False, + seq_lens_cpu_upper_bound=None, + ) + + +def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: + builder = object.__new__(_NonCausalMLAMetadataBuilder) + builder.device = torch.device("cpu") + builder.reorder_batch_threshold = 1 + builder.query_len_support = QueryLenSupport.SINGLE_ONLY + builder.non_causal_multi_token_decode = marked + builder.dcp_world_size = 1 + builder.metadata_cls = MLACommonMetadata + builder.model_config = SimpleNamespace( + dtype=torch.bfloat16, get_head_size=lambda: 576 + ) + return builder + + +def test_noncausal_block_uses_decode_without_cpu_lengths(): + common_metadata = _metadata([0, 8, 16]) + metadata = _builder().build(0, common_metadata) + + assert metadata.num_decodes == 2 + assert metadata.num_decode_tokens == 16 + assert metadata.num_prefills == 0 + assert metadata.prefill is None + assert metadata.decode is not None + assert metadata.decode.block_table.shape == (2, 3) + assert metadata.decode.seq_lens.shape == (2,) + assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor) + assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens) + assert not metadata.causal + + +def test_noncausal_support_is_explicit_and_uniform(): + with pytest.raises(ValueError, match="explicitly supported"): + _builder(marked=False).build(0, _metadata([0, 8, 16])) + with pytest.raises(ValueError, match="uniform query block"): + _builder().build(0, _metadata([0, 3, 8])) + + +def test_noncausal_block_allows_trailing_cudagraph_padding(): + common_metadata = _metadata([0, 8, 16, 16], num_tokens=24) + common_metadata.seq_lens[-1] = 0 + + metadata = _builder().build(0, common_metadata) + + assert metadata.num_decodes == 3 + assert metadata.num_decode_tokens == 24 + assert metadata.decode is not None + assert metadata.decode.seq_lens.tolist() == [108, 208, 0] + + +def test_noncausal_block_rejects_non_trailing_padding(): + with pytest.raises(ValueError, match="uniform query block"): + _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24)) + + +def test_noncausal_decode_metadata_keeps_live_request_buffers(): + common_metadata = _metadata([0, 8, 16]) + metadata = _builder().build(0, common_metadata) + + assert metadata.decode is not None + assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr() + assert ( + metadata.decode.block_table.data_ptr() + == common_metadata.block_table_tensor.data_ptr() + ) + + +def test_mla_cache_marker_is_promoted_to_group_capability(): + kwargs = { + "block_size": 64, + "num_kv_heads": 1, + "head_size": 576, + "dtype": torch.bfloat16, + } + marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True) + unmarked = MLAAttentionSpec(**kwargs) + + assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode + assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode + assert not MLAAttentionSpec.merge( + [unmarked, unmarked] + ).non_causal_multi_token_decode diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d82397591f77..e6d9f939ea58 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -8,6 +8,7 @@ import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -163,7 +164,7 @@ def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int): class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" - def test_blackwell_glm_dimensions_fall_back_to_trtllm(self): + def test_blackwell_glm_dimensions_use_trtllm(self): capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, @@ -184,11 +185,6 @@ def test_blackwell_glm_dimensions_fall_back_to_trtllm(self): with ( patch("vllm.platforms.current_platform") as mock_platform, patch.object(flash_attn_cls, "is_available", return_value=True), - patch( - "vllm.v1.attention.backends.mla.prefill.flash_attn." - "get_flash_attn_version", - return_value=4, - ), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): # Force the non-ROCm priority on the Blackwell. @@ -292,6 +288,11 @@ def test_backend_supported_dimension_validation(self): assert invalid_reasons == [] +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Imports vllm.platforms.rocm, whose module init requires a CUDA or " + "ROCm torch build; not importable on XPU/CPU/TPU.", +) class TestROCmAiterFAPrefillSelection: """Tests for the ROCm AITER FlashAttention MLA prefill backend.""" @@ -300,7 +301,12 @@ def test_rocm_priorities_prefer_aiter_fa(self): with patch("vllm.platforms.current_platform") as mock_platform: mock_platform.is_rocm.return_value = True priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5) + DeviceCapability(major=9, minor=5), + MLADimensions( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), ) assert priorities == [ diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py new file mode 100644 index 000000000000..cbbfedee964a --- /dev/null +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba2 ReplaySSM decode write-position derivation in +BaseMambaAttentionMetadataBuilder: write_pos and is_flush computed from the +per-request ring origin (replayssm_decode_base) and num_computed. +""" + +from dataclasses import dataclass + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + MockMambaBuilder, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") + + +@dataclass +class ReplaySSMBuildCase: + """A decode batch and its expected per-row write_pos / is_flush. + + num_computed = seq_len - query_len; write_pos = + (num_computed - decode_base) % buffer_len; is_flush = write_pos == + buffer_len - 1 (or a forced one-token flush when num_computed < decode_base). + """ + + seq_lens: list[int] + query_lens: list[int] + is_prefilling: list[bool] + decode_base: list[int] + buffer_len: int + expected_write_pos: list[int] + expected_is_flush: list[int] + mamba_cache_mode: str = "none" + + +REPLAYSSM_BUILD_CASES = { + # decode_base == num_prompt (fresh request). + "fresh_decode": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + ), + # decode_base > num_prompt anchors write_pos at the resume point. + "resumed_reanchors_to_zero": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + ), + # write_pos == buffer_len - 1 flushes. + "flush_boundary": ReplaySSMBuildCase( + seq_lens=[116], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Resumed request landing on a flush boundary. + "resumed_flush_boundary": ReplaySSMBuildCase( + seq_lens=[121], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Per-row write_pos / is_flush are independent. + "mixed_rows": ReplaySSMBuildCase( + seq_lens=[104, 106, 216], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 105, 200], + buffer_len=16, + expected_write_pos=[3, 0, 15], + expected_is_flush=[0, 0, 1], + ), + # write_pos wraps within the buffer (6 % 4 == 2). + "small_buffer_wrap": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=4, + expected_write_pos=[2], + expected_is_flush=[0], + ), + # Single-token prefill-as-decode still in the prompt (num_computed < + # decode_base): forced one-token flush. + "leftover_prompt_one_token_flush": ReplaySSMBuildCase( + seq_lens=[100], + query_lens=[1], + is_prefilling=[True], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[1], + ), + # Align mode (block_size 16). Past the first boundary the ring re-anchors at + # the block start: num_computed 117 -> block_start 112, write_pos 5 (vs 1 in + # none mode). + "align_reanchor_past_boundary": ReplaySSMBuildCase( + seq_lens=[118], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # First-block boundary (num_computed+1 == 112) forces a flush even though + # write_pos (11) != buffer_len - 1. + "align_first_block_boundary_flush": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[11], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # First step of a new block re-anchors write_pos to 0. + "align_new_block_start_zero": ReplaySSMBuildCase( + seq_lens=[113], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # block_size % buffer_len == 0: a later boundary lands on write_pos == + # buffer_len - 1, so the boundary flush coincides with the natural flush. + "align_boundary_coincides_natural_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # block_size % buffer_len != 0 (buffer_len 6): the boundary step still flushes + # although write_pos (3) != buffer_len - 1. + "align_unaligned_buffer_forces_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=6, + expected_write_pos=[3], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # Per-row independence in align mode: partial-block / new-block / boundary. + "align_mixed_rows": ReplaySSMBuildCase( + seq_lens=[105, 113, 112], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 100, 100], + buffer_len=16, + expected_write_pos=[4, 0, 11], + expected_is_flush=[0, 0, 1], + mamba_cache_mode="align", + ), +} + + +def _make_mamba_spec(buffer_len: int) -> MambaSpec: + # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups). + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=( + (1, 1), + (1, 1, 1), + (1, buffer_len, 1), + (1, buffer_len), + (1, buffer_len, 1), + ), + dtypes=(torch.float32,), + ) + + +def _create_replayssm_builder( + buffer_len: int, mamba_cache_mode: str = "none" +) -> MockMambaBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE + ) + # Set the flags after construction to skip validate_mamba_cached_kernel + # (it requires a Triton backend) on the mock model. + vllm_config.cache_config.use_replayssm = True + vllm_config.cache_config.replayssm_buffer_len = buffer_len + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + return MockMambaBuilder( + _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE + ) + + +def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase): + batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool), + replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32), + ) + return builder.build(0, common) + + +@pytest.mark.parametrize( + "case", REPLAYSSM_BUILD_CASES.values(), ids=REPLAYSSM_BUILD_CASES.keys() +) +def test_replayssm_write_pos(case: ReplaySSMBuildCase): + builder = _create_replayssm_builder(case.buffer_len, case.mamba_cache_mode) + meta = _build(builder, case) + + assert meta.write_pos_d is not None + assert meta.is_flush_d is not None + n = len(case.expected_write_pos) + assert meta.write_pos_d[:n].tolist() == case.expected_write_pos + assert meta.is_flush_d[:n].tolist() == case.expected_is_flush + + +def test_resumed_request_differs_from_fresh(): + """Same token count, different decode_base: fresh (base 100) -> write_pos 5, + resumed (base 105) -> write_pos 0.""" + builder = _create_replayssm_builder(16) + batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1]) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([False, False]), + replayssm_decode_base_cpu=torch.tensor([100, 105], dtype=torch.int32), + ) + meta = builder.build(0, common) + + assert meta.write_pos_d.tolist()[:2] == [5, 0] + assert meta.is_flush_d.tolist()[:2] == [0, 0] diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py new file mode 100644 index 000000000000..844bee0f71c0 --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm AITER MLA tests", allow_module_level=True) + +from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 +from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402 + AiterMLAMetadataBuilder, +) + + +class _NoOpTritonKernel: + def __getitem__(self, grid): + self.grid = grid + return self + + def __call__(self, *args, **kwargs): + pass + + +class _ExpandPageIndicesKernel: + def __getitem__(self, grid): + self.grid = grid + return self + + def __call__( + self, + page_indices, + block_table_tensor, + stride, + paged_kv_indptr, + seq_lens_for_kernel, + *, + KERNEL_BLOCK_SIZE, + BLOCK_SIZE, + ): + self.kernel_block_size = KERNEL_BLOCK_SIZE + for req_idx in range(self.grid[0]): + out_start = int(paged_kv_indptr[req_idx].item()) + seq_len = int(seq_lens_for_kernel[req_idx].item()) + for token_idx in range(seq_len): + block_id = int( + block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item() + ) + page_indices[out_start + token_idx] = ( + block_id * KERNEL_BLOCK_SIZE + token_idx % KERNEL_BLOCK_SIZE + ) + + +def _builder( + *, + mtp_decode_qlen: int, + has_full_cudagraphs: bool = False, + kernel_block_size: int = 1, + max_decode_rows: int = 32, + num_heads: int = 16, +): + return SimpleNamespace( + device=torch.device("cpu"), + num_heads=num_heads, + paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32), + paged_kv_indices=torch.empty(1024, dtype=torch.int32), + paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), + qo_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace( + has_full_cudagraphs=lambda: has_full_cudagraphs + ) + ), + _mtp_decode_qlen=mtp_decode_qlen, + _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), + _use_persistent_metadata=False, + kernel_block_size=kernel_block_size, + _num_attention_heads=16, + _mla_work_meta_data=torch.empty(1, dtype=torch.int32), + _mla_work_info_set=torch.empty(1, dtype=torch.int32), + _mla_work_indptr=torch.empty(1, dtype=torch.int32), + _mla_reduce_indptr=torch.empty(1, dtype=torch.int32), + _mla_reduce_final_map=torch.empty(1, dtype=torch.int32), + _mla_reduce_partial_map=torch.empty(1, dtype=torch.int32), + _mla_q_dtype=torch.bfloat16, + _mla_kv_dtype=torch.bfloat16, + decode_attn_out_dtype=torch.bfloat16, + ) + + +def test_backend_declares_uniform_batch_support(): + # UNIFORM/UNIFORM_BATCH is unconditional: MTP yields uniform qlen>1 and + # non-MTP yields qlen==1, both uniform batches. + assert ( + AiterMLAMetadataBuilder.query_len_support + == rocm_aiter_mla.QueryLenSupport.UNIFORM + ) + assert ( + AiterMLAMetadataBuilder._cudagraph_support + == rocm_aiter_mla.AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128]) +def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): + """Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355. + + Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the + metadata, covering the fp8 nhead=32 (TP4) fold path. + """ + + dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") + info_calls = [] + + def get_mla_metadata_info_v1( + max_batch_size, + max_qo_len, + num_attention_heads, + q_dtype, + kv_dtype, + *, + is_sparse, + fast_mode, + ): + info_calls.append( + { + "max_batch_size": max_batch_size, + "max_qo_len": max_qo_len, + "num_attention_heads": num_attention_heads, + "q_dtype": q_dtype, + "kv_dtype": kv_dtype, + "is_sparse": is_sparse, + "fast_mode": fast_mode, + } + ) + return tuple((1, torch.int32) for _ in range(6)) + + def init_common_builder(self, *args, **kwargs): + self.num_heads = num_heads + + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace( + dtypes=dtypes, + get_mla_metadata_info_v1=get_mla_metadata_info_v1, + ), + ) + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, + "__init__", + init_common_builder, + ) + monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) + + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method="deepseek_mtp", + num_speculative_tokens=3, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=8), + model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16), + scheduler_config=SimpleNamespace(max_num_seqs=2), + cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False) + ), + ) + builder = AiterMLAMetadataBuilder( + kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), + layer_names=["layer.0"], + vllm_config=config, + device=torch.device("cpu"), + ) + + assert info_calls == [ + { + "max_batch_size": config.scheduler_config.max_num_seqs, + "max_qo_len": config.speculative_config.num_speculative_tokens + 1, + "num_attention_heads": max(16, num_heads), + "q_dtype": dtypes.fp8, + "kv_dtype": dtypes.fp8, + "is_sparse": False, + "fast_mode": True, + } + ] + assert builder._mla_q_dtype == dtypes.fp8 + assert builder._mla_kv_dtype == dtypes.fp8 + + +def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch): + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + metadata = AiterMLAMetadataBuilder._build_decode( + _builder(mtp_decode_qlen=4), + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), + max_seq_len=7, + query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32), + num_decode_tokens=8, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == 4 + assert torch.equal(metadata.seq_lens, torch.tensor([7, 5], dtype=torch.int32)) + assert torch.equal(metadata.qo_indptr, torch.tensor([0, 4, 8], dtype=torch.int32)) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == 4 + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4 + + +def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr( + monkeypatch, +): + """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759.""" + + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + mtp_qlen = 4 + seq_lens = torch.tensor([7, 0], dtype=torch.int32) + qo_lens = torch.tensor([mtp_qlen, 0], dtype=torch.int32) + expected_seq_lens = torch.where(qo_lens > 0, seq_lens, mtp_qlen) + expected_paged_kv_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + expected_seq_lens.cumsum(dim=0, dtype=torch.int32), + ] + ) + expected_qo_indptr = torch.arange( + 0, + (seq_lens.numel() + 1) * mtp_qlen, + step=mtp_qlen, + dtype=torch.int32, + ) + + builder = _builder( + mtp_decode_qlen=mtp_qlen, + has_full_cudagraphs=True, + max_decode_rows=4, + ) + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=seq_lens, + max_seq_len=int(seq_lens.max().item()), + query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), + num_decode_tokens=seq_lens.numel() * mtp_qlen, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == mtp_qlen + assert torch.equal(metadata.seq_lens, expected_seq_lens) + assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) + assert torch.equal(metadata.qo_indptr, expected_qo_indptr) + assert torch.all( + builder.paged_kv_indptr[expected_paged_kv_indptr.numel() :] + == expected_paged_kv_indptr[-1] + ) + assert torch.all( + builder.qo_indptr[expected_qo_indptr.numel() :] == expected_qo_indptr[-1] + ) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == mtp_qlen + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == mtp_qlen + + +def test_decode_expands_kernel_block_page_indices(monkeypatch): + """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704.""" + + expand_kernel = _ExpandPageIndicesKernel() + monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel) + # qlen==1 now takes the persistent-metadata path, which imports + # get_mla_metadata_v1 from aiter; mock it so this test needs no real kernel. + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()), + ) + + kernel_block_size = 2 + block_table = torch.tensor( + [ + [10, 11, 99], + [20, 21, 22], + ], + dtype=torch.int32, + ) + seq_lens = torch.tensor([3, 5], dtype=torch.int32) + expected_paged_kv_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + seq_lens.cumsum(dim=0, dtype=torch.int32), + ] + ) + expected_indices = torch.tensor( + [ + int(block_table[req_idx, token_idx // kernel_block_size].item()) + * kernel_block_size + + token_idx % kernel_block_size + for req_idx, seq_len in enumerate(seq_lens.tolist()) + for token_idx in range(seq_len) + ], + dtype=torch.int32, + ) + + metadata = AiterMLAMetadataBuilder._build_decode( + _builder( + mtp_decode_qlen=1, + kernel_block_size=kernel_block_size, + ), + block_table_tensor=block_table, + seq_lens_device=seq_lens, + max_seq_len=int(seq_lens.max().item()), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32), + num_decode_tokens=seq_lens.numel(), + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == 1 + assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) + assert torch.equal( + metadata.paged_kv_indices[: expected_indices.numel()], + expected_indices, + ) + assert expand_kernel.grid == (seq_lens.numel(),) + assert expand_kernel.kernel_block_size == kernel_block_size + + +@pytest.mark.parametrize( + "mtp_decode_qlen, qo_len, num_heads, expect_persistent", + [ + (1, 1, 16, True), # non-MTP decode + (4, 2, 16, True), # MTP deployment, in-range step + (4, 4, 16, True), # MTP deployment, full-qlen verification step + (2, 4, 16, False), # step demand exceeds provisioned K -> fallback + (1, 1, 8, False), # small head count -> Gluon decode owns qlen==1 + (4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1 + ], +) +def test_persistent_metadata_gate( + monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent +): + """Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K. + + K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets + the pre-built schedule only when its qlen fits those buffers, otherwise it + falls back to the kernel computing its own. qlen==1 (non-MTP) must stay + in-range -- dropping it is the regression this guards. Fewer than 16 heads + is served by the Gluon decode paths, which never read this schedule. + """ + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + + # Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so + # num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in. + num_reqs = 2 + query_start_loc = torch.arange( + 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 + ) + metadata = AiterMLAMetadataBuilder._build_decode( + _builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads), + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), + max_seq_len=8, + query_start_loc_cpu=query_start_loc, + query_start_loc_device=query_start_loc, + num_decode_tokens=num_reqs * qo_len, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.max_qo_len == qo_len + assert metadata.has_persistent_metadata is expect_persistent + assert get_mla_metadata_v1.called is expect_persistent + if expect_persistent: + assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len + assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8f9e8acac60e..36d5cca9a4e9 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -28,16 +28,9 @@ def mock_vllm_config(): @pytest.fixture -def mock_on_gfx9(): - """Mock gfx9 arch detection to return True.""" - with patch("vllm.platforms.rocm.on_gfx9", return_value=True): - yield - - -@pytest.fixture -def mock_on_mi3xx(): - """Mock mi3xx arch detection to return True.""" - with patch("vllm.platforms.rocm.on_mi3xx", return_value=True): +def mock_get_cdna_version(): + """Mock cdna version arch detection to return True.""" + with patch("vllm.platforms.rocm.get_cdna_version", return_value=3): yield @@ -111,8 +104,7 @@ def test_standard_attention_backend_selection( selected_backend, expected_backend_path, mock_vllm_config, - mock_on_gfx9, - mock_on_mi3xx, + mock_get_cdna_version, monkeypatch, ): """Test standard attention backend selection with various configurations.""" @@ -305,12 +297,12 @@ def test_mla_backend_selection( def test_aiter_fa_requires_mi3xx(mock_vllm_config): - """Test that ROCM_AITER_FA requires mi3xx architecture.""" + """Test that ROCM_AITER_FA requires CDNA3+ architecture.""" from vllm.platforms.rocm import RocmPlatform - # Mock on_mi3xx to return False (used by supports_compute_capability) + # Mock cdna version to return 1 (used by supports_compute_capability) with ( - patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + patch("vllm.platforms.rocm.get_cdna_version", return_value=1), pytest.raises( ValueError, match="compute capability not supported", diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 3fe5df918ab2..0c7b645c1484 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -862,6 +862,8 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected): PREFILL_BATCH_SPECS = { "short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]), "short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]), + "masked_mha": BatchSpec(seq_lens=[256], query_lens=[256]), + "masked_mha_chunked_context": BatchSpec(seq_lens=[448, 384], query_lens=[256, 256]), } @@ -878,7 +880,7 @@ def test_sparse_backend_prefill_correctness( kv_cache_dtype, workspace_init, ): - """Test single-pass FA4 dense forward_mha for sparse MLA prefill.""" + """Test single-pass dense and masked MHA for sparse MLA prefill.""" backend_cls = FlashMLASparseBackend batch_spec = PREFILL_BATCH_SPECS[batch_name] @@ -892,7 +894,8 @@ def test_sparse_backend_prefill_correctness( qk_rope_head_dim = 64 v_head_dim = 128 head_size = kv_lora_rank + qk_rope_head_dim - topk_tokens = 512 + masked_mha = batch_name.startswith("masked_mha") + topk_tokens = 200 if masked_mha else 512 max_seqlen = max(batch_spec.seq_lens) total_cache_tokens = sum(batch_spec.seq_lens) @@ -918,6 +921,7 @@ def test_sparse_backend_prefill_correctness( model_type="deepseek_v2", ) model_config.dtype = dtype + model_config.model_arch_config.total_num_attention_heads = num_heads model_config.get_num_attention_heads = MethodType( lambda self, parallel_config: num_heads, model_config ) @@ -942,13 +946,14 @@ def test_sparse_backend_prefill_correctness( # Compute dense reference outputs. total_query_tokens = sum(query_lens) - sparse_indices = torch.zeros( - total_query_tokens, topk_tokens, dtype=torch.int32, device=device + sparse_indices = torch.full( + (total_query_tokens, topk_tokens), -1, dtype=torch.int32, device=device ) all_q, all_kv_c_new, all_k_pe_new = [], [], [] kv_c_contexts, k_pe_contexts = [], [] reference_outputs = [] + global_token_idx = 0 for i in range(batch_spec.batch_size): s_len = seq_lens[i] @@ -981,8 +986,15 @@ def test_sparse_backend_prefill_correctness( for j in range(q_len): attend_end = ctx_len + j + 1 q_tok = q_mha[j : j + 1] # (1, H, D_qk) - k_attend = k_all[:attend_end] # (N, H, D_qk) - v_attend = v_all[:attend_end] # (N, H, D_v) + if masked_mha: + actual_topk = min(topk_tokens, attend_end) + attend_indices = torch.randperm(attend_end, device=device)[:actual_topk] + sparse_indices[global_token_idx, :actual_topk] = attend_indices + k_attend = k_all[attend_indices] + v_attend = v_all[attend_indices] + else: + k_attend = k_all[:attend_end] # (N, H, D_qk) + v_attend = v_all[:attend_end] # (N, H, D_v) q_sdpa = q_tok.unsqueeze(0).transpose(1, 2).float() k_sdpa = k_attend.unsqueeze(0).transpose(1, 2).float() @@ -993,6 +1005,7 @@ def test_sparse_backend_prefill_correctness( ) out = out.transpose(1, 2).squeeze(0) # (1, H, D_v) reference_outputs.append(out.to(dtype).flatten(start_dim=-2)) + global_token_idx += 1 all_q.append(q_mha) all_kv_c_new.append(kv_c_full[ctx_len:]) @@ -1047,6 +1060,13 @@ def test_sparse_backend_prefill_correctness( builder_cls = backend_cls.get_builder_cls() builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + if batch_name == "masked_mha_chunked_context": + builder.chunked_prefill_workspace_size = block_size * batch_spec.batch_size + builder.chunked_prefill_workspace = torch.empty( + (builder.chunked_prefill_workspace_size, head_size), + dtype=dtype, + device=device, + ) # Drive the queries through the dense-MHA prefill path directly (the routing # threshold would otherwise classify these short queries as MQA decodes). builder.reorder_batch_threshold = 1 diff --git a/tests/v1/attention/test_sparse_mla_mask.py b/tests/v1/attention/test_sparse_mla_mask.py new file mode 100644 index 000000000000..3ecc179be4b4 --- /dev/null +++ b/tests/v1/attention/test_sparse_mla_mask.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytest.importorskip("cutlass") + +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + _build_topk_mask, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_build_topk_mask_single_request_matches_generic_path() -> None: + topk = torch.tensor( + [[0, 31, 32, 63, -1], [1, 64, 127, -1, -1]], + dtype=torch.int32, + device="cuda", + ) + + num_words = (128 + 31) // 32 + single_out = torch.zeros(1, 2, num_words, dtype=torch.int32, device="cuda") + generic_out = torch.zeros(2, 1, num_words, dtype=torch.int32, device="cuda") + single_req = _build_topk_mask([topk], [2], 2, 128, single_out) + generic = _build_topk_mask([topk[:1], topk[1:]], [1, 1], 1, 128, generic_out) + + torch.testing.assert_close(single_req[0, 0], generic[0, 0]) + torch.testing.assert_close(single_req[0, 1], generic[1, 0]) diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 3daafcdc6ff8..20562cfe85f5 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -33,6 +33,7 @@ EncoderOnlyAttentionSpec, FullAttentionSpec, MambaSpec, + get_kv_quant_mode, ) @@ -178,6 +179,7 @@ def create_standard_kv_cache_spec( head_size=vllm_config.model_config.get_head_size(), dtype=vllm_config.model_config.dtype, sliding_window=vllm_config.model_config.get_sliding_window(), + kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype), ) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index a42fc4150110..ec3f293c0a29 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -40,6 +40,8 @@ def test_mamba_align_split_partial_tail_schedule(): hash_block_size = 32 mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=8192, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=hash_block_size, mamba_partial_cache_hit=True, @@ -75,6 +77,78 @@ def test_mamba_align_split_partial_tail_schedule(): assert split(self=mock, request=req2, num_new_tokens=1000) == 512 +def test_mamba_align_split_when_block_exceeds_scheduling_budget(): + """Sub-block chunks make progress only when no step can fit a full block.""" + block_size = 11392 + token_budget = 8192 + prompt_length = 30000 + mock = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=token_budget, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + use_eagle=False, + hash_block_size=32, + mamba_partial_cache_hit=False, + ) + req = make_request("0", [0] * prompt_length, 32, sha256) + split = Scheduler._mamba_block_aligned_split + + mock.max_num_scheduled_tokens = block_size + assert split(self=mock, request=req, num_new_tokens=token_budget) == 0 + mock.max_num_scheduled_tokens = token_budget + + scheduled_chunks = [] + while req.num_computed_tokens < prompt_length: + num_new_tokens = min(token_budget, prompt_length - req.num_computed_tokens) + num_scheduled_tokens = split( + self=mock, + request=req, + num_new_tokens=num_new_tokens, + ) + assert 0 < num_scheduled_tokens <= token_budget + scheduled_chunks.append(num_scheduled_tokens) + req.num_computed_tokens += num_scheduled_tokens + + assert scheduled_chunks == [8192, 3200, 8192, 3200, 7216] + + +def test_mamba_align_split_when_block_exceeds_long_prefill_threshold(): + """A long-prefill cap below the block size permits sub-block progress.""" + block_size = 512 + token_budget = 8192 + long_prefill_threshold = 384 + prompt_length = 1300 + mock = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=token_budget, + scheduler_config=SimpleNamespace( + long_prefill_token_threshold=long_prefill_threshold + ), + use_eagle=False, + hash_block_size=32, + mamba_partial_cache_hit=False, + ) + req = make_request("0", [0] * prompt_length, 32, sha256) + split = Scheduler._mamba_block_aligned_split + + scheduled_chunks = [] + while req.num_computed_tokens < prompt_length: + num_new_tokens = min( + long_prefill_threshold, + prompt_length - req.num_computed_tokens, + ) + num_scheduled_tokens = split( + self=mock, + request=req, + num_new_tokens=num_new_tokens, + ) + assert 0 < num_scheduled_tokens <= long_prefill_threshold + scheduled_chunks.append(num_scheduled_tokens) + req.num_computed_tokens += num_scheduled_tokens + + assert scheduled_chunks == [384, 128, 384, 128, 276] + + def test_hybrid_mamba_align_partial_hash_hit(): hash_block_size = 2 mamba_block_size = 2 * hash_block_size @@ -223,6 +297,299 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): assert moved[0].block_hash_num_tokens == 6 +def test_take_partial_tail_offloads_returns_cow_target(): + """The connector offload hand-off exposes the mamba CoW *target* block Y + (the durable boundary state), not the overwritten source X, and only at + the CoW step.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + + # Step A registered the partial tail but has not CoW'd yet: no offload. + assert manager.take_partial_tail_offloads() == {} + + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1] + source_block = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert source_block is not None + source_block_id = source_block[0].block_id + + # Step B: the producer continues, triggering the CoW X->Y. + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + offloads = manager.take_partial_tail_offloads() + assert list(offloads.keys()) == ["0"] + assert len(offloads["0"]) == 1 + group_id, block_id, boundary_tokens = offloads["0"][0] + assert group_id == 1 # the mamba group + assert boundary_tokens == 6 + copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in copies if c.src_block_id == source_block_id) + # The offload points at the durable CoW target Y, not the overwritten X. + assert block_id == cow_copy.dst_block_id + assert block_id != source_block_id + # Draining clears it. + assert manager.take_partial_tail_offloads() == {} + + # The hand-off pinned Y (its CoW retention is released after this step, + # and Y is off the request block table); freeing the request unpins it. + cow_block = manager.block_pool.blocks[block_id] + pinned_ref = cow_block.ref_cnt + assert pinned_ref >= 1 + manager.free(req0) + assert cow_block.ref_cnt == pinned_ref - 1 + + +def test_partial_tail_pin_survives_released_cow_retention(): + """If the CoW retention is released before the hand-off is drained + (immediate-free mode), the drain must rescue the cow block from the free + queue: a raw ref increment would leave a ref>0 block allocatable, and the + next allocation would pop it and assert.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + # Retention released before the drain (defer_block_free=False ordering). + _copies, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + + offloads = manager.take_partial_tail_offloads() + ((_group_id, block_id, boundary_tokens),) = offloads["0"] + assert boundary_tokens == 6 + cow_block = manager.block_pool.blocks[block_id] + assert cow_block.ref_cnt == 1 + + # The pinned block is out of the free queue: draining every free block + # neither trips the allocator's ref_cnt assert nor hands it out. + new_blocks = manager.block_pool.get_new_blocks( + manager.block_pool.get_num_free_blocks() + ) + assert block_id not in {b.block_id for b in new_blocks} + + +def test_partial_tail_offload_dropped_when_request_freed_before_drain(): + """A hand-off recorded in the same scheduling pass as the request's death + must not be drained: its release hook has already run, so draining would + leak a pinned block.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + # The request dies (preempt/abort) before the scheduler drains. + manager.block_pool.free_blocks(manager.pop_blocks_for_free(req0)) + assert manager.take_partial_tail_offloads() == {} + + +def test_take_partial_tail_offloads_empty_without_partial_tail(): + """A prompt ending on a block boundary registers no partial tail, so there + is nothing to offload.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + # 4-token prompt ends exactly on the mamba block boundary (block_size=4). + req0 = make_request("0", [0, 0, 1, 1], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + assert manager.take_partial_tail_offloads() == {} + + req0.num_computed_tokens = 4 + req0.append_output_token_ids([2]) + assert manager.allocate_slots(req0, 1) is not None + assert manager.take_partial_tail_offloads() == {} + + +def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): + """truncate_computed_blocks slices each group by its own block size, + keeps null placeholders in the retained prefix, and leaves the original + lookup result untouched (pure view, no refcount changes).""" + hash_block_size = 2 + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=2 * hash_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + blocks, num_computed, _ = manager.get_computed_blocks(producer) + assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None + manager.free(producer) + manager.new_step_starts() + + consumer = make_request( + "consumer", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256 + ) + blocks, num_computed, _ = manager.get_computed_blocks(consumer) + assert num_computed == 6 + assert [len(group) for group in blocks.blocks] == [3, 2] + assert blocks.blocks[1][0].is_null + + truncated = manager.truncate_computed_blocks(blocks, 4) + + assert [len(group) for group in truncated.blocks] == [2, 1] + assert truncated.blocks[1][0].is_null + assert [len(group) for group in blocks.blocks] == [3, 2] + + def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): hash_block_size = 2 block_size = 2 * hash_block_size diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index e34a0da54d8d..1b01d184dc8a 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections import deque +from collections import defaultdict, deque +from collections.abc import Callable from unittest.mock import Mock import pytest @@ -329,46 +330,384 @@ def free_request(req, delay_free_blocks=False): assert not scheduler.running -def test_no_placeholder_underflow_on_discarded_spec_frame(): - num_spec = 5 +class PipelinedEngine: + """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue: + schedule until the batch queue is full, then process the oldest step's + output. Async PP runs pp_size+1 concurrent batches, so up to pp_size + steps are in flight at each schedule() call -- the window in which + preemption must handle output that has not yet returned. (Single-GPU e2e + tests can never create this window: at PP=1, exactly one step is in + flight and it is processed before a preempted request can resume.) + + The model runner is emulated with the V2 runner's own bookkeeping, from + only what the scheduler serializes to it: slots flushed on + preempted_req_ids, resumed requests re-added from the NewRequestData + snapshot, sampling when a step reaches the end of the runner's own view + of the sequence. This makes preemption races observable: a stale token + delivered after a resume is scheduled extends the scheduler's sequence + but not the runner's. + + Every sample emits a globally unique token tagged with its sampled + position, so tests can assert exact delivery. + """ + + def __init__( + self, + scheduler: AsyncScheduler, + queue_size: int, + accept_drafts: Callable[[int, str, int], int] | None = None, + ): + self.scheduler = scheduler + self.queue_size = queue_size + self.accept_drafts = accept_drafts + # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order. + self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque() + # Runner-side request state: req_id -> [seq_len, num_computed] as the + # runner sees them (its own sampled tokens, not the scheduler's). + self.runner_view: dict[str, list[int]] = {} + # All tokens the fake runner ever sampled, per request, in order. + self.emitted: dict[str, list[int]] = defaultdict(list) + # Sequence position each (globally unique) token was sampled for. + self.emitted_position: dict[int, int] = {} + self.step_idx = 0 + self._next_token = 1000 + + def _schedule(self) -> bool: + scheduler_output = self.scheduler.schedule() + self.step_idx += 1 + # Snapshot what NewRequestData serializes at schedule time (both new + # and resumed requests for the V2 runner). + new_reqs = [ + (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens) + for r in scheduler_output.scheduled_new_reqs + ] + # Enqueue empty steps too (the engine executes them), so the runner + # still observes their preempted/finished request ids in step order. + self.queue.appendleft((scheduler_output, new_reqs)) + return True + + def _process_oldest_step(self) -> None: + scheduler_output, new_reqs = self.queue.pop() + # Worker-side state updates, in step order: flush preempted/finished + # slots, then (re-)add new/resumed requests. + for req_id in scheduler_output.preempted_req_ids or (): + self.runner_view.pop(req_id, None) + for req_id in scheduler_output.finished_req_ids or (): + self.runner_view.pop(req_id, None) + for req_id, seq_len, num_computed in new_reqs: + self.runner_view[req_id] = [seq_len, num_computed] + + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + sampled_token_ids: list[list[int]] = [] + for req_id in req_ids: + num_scheduled = scheduler_output.num_scheduled_tokens[req_id] + view = self.runner_view.get(req_id) + if view is None: + # Slot already flushed (request finished/aborted mid-flight). + sampled_token_ids.append([]) + continue + seq_len, num_computed = view + end = num_computed + num_scheduled + if end < seq_len: + # Partial prefill by the runner's own bookkeeping: no sample. + view[1] = end + sampled_token_ids.append([]) + continue + drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()) + num_accepted = ( + min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts)) + if drafts and self.accept_drafts + else 0 + ) + num_rejected = len(drafts) - num_accepted + tokens = list(range(self._next_token, self._next_token + 1 + num_accepted)) + self._next_token += 1 + num_accepted + self.emitted[req_id].extend(tokens) + sampled_token_ids.append(tokens) + # Rejected drafts roll back computed; the sampled tokens extend + # the runner's sequence. + view[1] = end - num_rejected + view[0] = view[1] + 1 + for offset, token in enumerate(tokens): + self.emitted_position[token] = view[0] - len(tokens) + offset + model_runner_output = ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=sampled_token_ids, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + self.scheduler.update_from_output(scheduler_output, model_runner_output) + + def run( + self, + max_steps: int = 2000, + before_step: Callable[[int, "PipelinedEngine"], None] | None = None, + ) -> None: + for i in range(max_steps): + if not self.scheduler.has_requests() and not self.queue: + return + if before_step is not None: + before_step(i, self) + scheduled = ( + self.scheduler.has_requests() + and len(self.queue) < self.queue_size + and self._schedule() + ) + if scheduled and len(self.queue) < self.queue_size: + # Queue not yet full: the engine returns without blocking. + continue + if self.queue: + self._process_oldest_step() + raise AssertionError("engine loop did not converge") + + +def _create_async_pp_scheduler( + num_spec: int, pp_size: int = 3, num_blocks: int = 5 +) -> AsyncScheduler: scheduler = create_scheduler( async_scheduling=True, - num_speculative_tokens=num_spec, - speculative_method="ngram_gpu", + num_speculative_tokens=num_spec or None, + speculative_method="ngram_gpu" if num_spec else None, + use_v2_model_runner=True, + num_blocks=num_blocks, + block_size=16, + max_num_batched_tokens=512, ) - 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 + # Emulate PP at the scheduler level; constructing with + # pipeline_parallel_size>1 requires that many visible GPUs. Drive with + # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches). + scheduler.pp_size = pp_size + scheduler.use_pp = pp_size > 1 + return scheduler + + +def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None: + """Delivered tokens must be an order-preserving subset of the emitted + tokens with no duplicates (tokens are globally unique).""" + it = iter(emitted) + for token in delivered: + assert token in it, f"token {token} delivered out of order or twice" + + +def _assert_positions_consistent(req, engine: PipelinedEngine) -> None: + """The i-th delivered output token must be one the runner sampled for + exactly sequence position prompt_len + i: catches a preempted request's + stale output landing on a position the resumed request resampled (or + vice versa), which token-stream equality alone cannot see.""" + for i, token in enumerate(req.output_token_ids): + expected = req.num_prompt_tokens + i + actual = engine.emitted_position[token] + assert actual == expected, ( + f"output {i} of {req.request_id}: token sampled for position " + f"{actual}, delivered as position {expected}" + ) - 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=[], +@pytest.mark.parametrize("num_spec", [0, 3]) +def test_kv_pressure_preemption_with_inflight_output(num_spec: int): + """KV-pressure preemption of requests with in-flight async output. + + PP=3 + async scheduling (batch queue of 4), a block pool small enough + that decodes contend and preempt mid-flight, and staggered arrivals so + the batch queue actually pipelines. A preempted request's in-flight steps + still return: their tokens must be delivered exactly once, their stale + spec-rejection counts must not corrupt the rolled-back counters, and the + resume must not resample a position that output later delivers. + + Regression for the num_output_placeholders underflow EngineCore crash: + with the fix reverted, the num_spec=3 variant fails with exactly + ``assert request.num_output_placeholders >= 0`` when a stale spec output + returns after the preempted request was resumed and sampled. + """ + max_tokens = 24 + scheduler = _create_async_pp_scheduler(num_spec) + requests = create_requests( + num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True ) - 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=[], + pending = list(requests) + for _ in range(2): + scheduler.add_request(pending.pop(0)) + + # Observe that the scenario under test actually occurs. + preempts_with_inflight_output = 0 + orig_preempt = scheduler._preempt_request + + def counting_preempt(request, timestamp, **kwargs): + nonlocal preempts_with_inflight_output + if request.num_in_flight_tokens > 0: + preempts_with_inflight_output += 1 + return orig_preempt(request, timestamp, **kwargs) + + scheduler._preempt_request = counting_preempt + + def add_requests(step: int, engine: PipelinedEngine): + if pending: + scheduler.add_request(pending.pop(0)) + + engine = PipelinedEngine( + scheduler, + queue_size=4, + # Deterministically vary spec acceptance so stale outputs carry + # nonzero rejection counts. + accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1), ) + engine.run(before_step=add_requests) + + assert preempts_with_inflight_output > 0, ( + "test did not exercise preemption with in-flight output" + ) + for req in requests: + assert req.is_finished() + assert req.num_output_tokens == max_tokens + # Lossless: delivered tokens are exactly the sampled tokens, in order + # (the excluded tail was emitted after the request finished). + emitted = engine.emitted[req.request_id] + assert list(req.output_token_ids) == emitted[:max_tokens] + _assert_positions_consistent(req, engine) + + +@pytest.mark.parametrize("pp_size", [1, 3]) +def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int): + """reset_prefix_cache(reset_running_requests=True) resumes requests in + the same step it preempts them, so in-flight output must be dropped (the + resume resamples those positions). + + pp_size=1: regression for the frame-based discard this fix replaces, + which with spec decode drained one *token* count per output frame and + over-discarded, corrupting the fresh frames after the resume. + pp_size=3: back-to-back resets, so the second re-preempts requests whose + dropped stale share is still in flight -- it must be recorded once (not + accumulated) and stay dropped. + """ + max_tokens = 24 + scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size) + requests = create_requests( + num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True + ) + pending = list(requests) + for _ in range(2): + scheduler.add_request(pending.pop(0)) + + # Observe re-preemptions with an undrained stale share (the + # double-count hazard). + repreempts_with_stale = 0 + orig_preempt = scheduler._preempt_request + + def counting_preempt(request, timestamp, **kwargs): + nonlocal repreempts_with_stale + if getattr(request, "num_stale_output_tokens", 0) > 0: + repreempts_with_stale += 1 + return orig_preempt(request, timestamp, **kwargs) + + scheduler._preempt_request = counting_preempt + + resets = 0 + reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19} + + def before_step(step: int, engine: PipelinedEngine): + nonlocal resets + if pending: + scheduler.add_request(pending.pop(0)) + if step in reset_steps and (engine.queue or scheduler.running): + scheduler.reset_prefix_cache(reset_running_requests=True) + resets += 1 + + engine = PipelinedEngine( + scheduler, + queue_size=pp_size + 1, + accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1), + ) + engine.run(before_step=before_step) + + assert resets > 0, "test did not exercise reset_prefix_cache" + if pp_size > 1: + # The re-preempt-while-stale-pending window needs pipeline depth. + assert repreempts_with_stale > 0, ( + "test did not exercise re-preemption with an undrained stale share" + ) + for req in requests: + assert req.is_finished() + assert req.num_output_tokens == max_tokens + # Dropped tokens are never delivered; order must be preserved with + # no duplicates. + _assert_ordered_subset( + list(req.output_token_ids), engine.emitted[req.request_id] + ) + _assert_positions_consistent(req, engine) + # All stale shares fully drained by the end. + assert getattr(req, "num_stale_output_tokens", 0) == 0 + + +def test_requires_kv_delivery_defaults_to_producer_role(): + # No connector: nothing is handed off, so keep the lossless deliver-stale + # path on preemption. + assert create_scheduler(async_scheduling=True).requires_kv_delivery is False + # Only a producer hands KV off when a request completes. + for role, expected in ( + ("kv_producer", True), + ("kv_both", True), + ("kv_consumer", False), + ): + scheduler = create_scheduler( + async_scheduling=True, use_kv_connector=True, kv_role=role + ) + assert scheduler.requires_kv_delivery is expected, role - 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 +@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"]) +def test_kv_pressure_preempt_mid_handoff(kv_role: str): + """P/D race: a request is KV-pressure preempted while the output of its + final prefill chunk -- the hand-off token that would finish it -- is in + flight. + + On a producer, that output must be dropped so the request recomputes; + delivering it would finish the request and hand off blocks the preemption + already freed, so the consumer pulls garbage. A consumer hands nothing off, + so it keeps the lossless deliver-stale path. + """ + is_producer = kv_role == "kv_producer" + scheduler = create_scheduler( + async_scheduling=True, + use_kv_connector=True, + kv_role=kv_role, + num_blocks=5, + block_size=16, + max_num_batched_tokens=512, + ) + assert scheduler.requires_kv_delivery is is_producer + + # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the + # next decode allocation preempts the tail of the running queue (the handoff + # request) while its prefill output is still in flight. + decoder = create_requests( + num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"] + )[0] + handoff = create_requests( + num_requests=1, num_tokens=32, max_tokens=1, req_ids=["handoff"] + )[0] + scheduler.add_request(decoder) + scheduler.add_request(handoff) + sched_output = scheduler.schedule() + assert handoff.status == RequestStatus.RUNNING + assert handoff.num_output_placeholders == 1 + + scheduler.schedule() + assert handoff.status == RequestStatus.PREEMPTED + assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens + assert handoff.drop_stale_output is is_producer + + scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) + + assert handoff.num_stale_output_tokens == 0 + if is_producer: + # Dropped: recomputed from the waiting queue, so the hand-off happens + # against real KV. + assert not handoff.is_finished() + assert handoff.status == RequestStatus.PREEMPTED + assert handoff.num_output_tokens == 0 + assert handoff.request_id in scheduler.requests + else: + assert handoff.is_finished() + assert handoff.num_output_tokens == 1 diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 647241ce73cb..88d17e9acdc7 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -9,6 +9,7 @@ from vllm.v1.core.kv_cache_utils import ( _get_kv_cache_config_packed, + _get_kv_cache_groups_uniform_groups, get_kv_cache_config_from_groups, ) from vllm.v1.kv_cache_interface import ( @@ -16,6 +17,7 @@ KVCacheGroupSpec, KVCacheTensor, MLAAttentionSpec, + SlidingWindowMLASpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, ) @@ -109,7 +111,132 @@ def _page_sizes_by_layer( return page_sizes +def _packing_by_layer( + tensors: list[KVCacheTensor], +) -> dict[str, tuple[int, int]]: + return { + layer_name: (tensor.offset, tensor.block_stride) + for tensor in tensors + for layer_name in tensor.shared_by + } + + +def _make_views( + groups: list[KVCacheGroupSpec], + num_blocks: int, + tensors: list[KVCacheTensor], +) -> dict[str, torch.Tensor]: + page_sizes = _page_sizes_by_layer(groups) + packing = _packing_by_layer(tensors) + backing = torch.zeros(tensors[0].size, dtype=torch.uint8) + return { + layer_name: torch.as_strided( + backing, + size=(num_blocks, page_size), + stride=(packing[layer_name][1], 1), + storage_offset=packing[layer_name][0], + ) + for layer_name, page_size in page_sizes.items() + } + + +def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec: + specs = { + f"{prefix}.{i}": MagicMock(page_size_bytes=page_size) + for i, page_size in enumerate(page_sizes) + } + return KVCacheGroupSpec( + layer_names=list(specs), + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs), + ) + + class TestInterleavedPacking: + def test_compact_cache_overlays_fp32_state_group(self): + full_specs = {} + state_specs = {} + for i in range(2): + full_specs[f"mla.{i}"] = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + page_size_padded=32768, + indexes_kv_by_block_stride=True, + compress_ratio=4, + ) + full_specs[f"indexer.{i}"] = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=68, + dtype=torch.uint8, + page_size_padded=4608, + compress_ratio=4, + ) + state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec( + block_size=4, + num_kv_heads=1, + head_size=2048, + dtype=torch.float32, + sliding_window=8, + indexes_kv_by_block_stride=True, + ) + state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec( + block_size=4, + num_kv_heads=1, + head_size=512, + dtype=torch.float32, + sliding_window=8, + indexes_kv_by_block_stride=True, + ) + + grouped_specs = [ + UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs), + UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs), + ] + groups = _get_kv_cache_groups_uniform_groups(grouped_specs) + + assert len(groups) == 2 + assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608} + assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == { + 4352 + } + assert { + state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2) + } == {8192} + + full_group_bytes = 2 * (32768 + 4608) + state_group_bytes = 2 * (32768 + 8192) + bytes_per_block = max(full_group_bytes, state_group_bytes) + num_blocks, tensors = _get_kv_cache_config_packed( + _mock_vllm_config(), groups, bytes_per_block * 32 + ) + assert num_blocks == 32 + assert {tensor.block_stride for tensor in tensors} == {bytes_per_block} + + packing = _packing_by_layer(tensors) + assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0 + assert packing["indexer.0"][0] == 32768 + assert packing["indexer_state.0"][0] == 32768 + + def test_deepseek_v4_pro_stride(self): + groups = [ + _make_page_group("full", [32768, 4608] * 30 + [1024] * 31), + _make_page_group("c4_state", [32768, 8192] * 30), + _make_page_group("c128_state", [32768] * 31), + _make_page_group("swa.0", [32768] * 31), + _make_page_group("swa.1", [32768] * 30), + ] + expected_stride = 1_228_800 + + num_blocks, tensors = _get_kv_cache_config_packed( + _mock_vllm_config(), groups, expected_stride * 32 + ) + + assert num_blocks == 32 + assert {tensor.block_stride for tensor in tensors} == {expected_stride} + assert {tensor.size for tensor in tensors} == {expected_stride * 32} + def test_all_tensors_have_block_stride(self): _, tensors = _run() for t in tensors: @@ -122,9 +249,30 @@ def test_all_tensors_share_same_size(self): assert sizes.pop() > 0 def test_offsets_within_one_block(self): - _, tensors = _run() - for t in tensors: - assert t.offset < t.block_stride + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) + _, tensors = _get_kv_cache_config_packed( + _mock_vllm_config(), groups, 100 * 1024 * 1024 + ) + page_sizes = _page_sizes_by_layer(groups) + packing = _packing_by_layer(tensors) + for layer_name, page_size in page_sizes.items(): + offset, block_stride = packing[layer_name] + assert offset + page_size <= block_stride + + def test_layouts_are_disjoint_within_each_group(self): + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) + _, tensors = _get_kv_cache_config_packed( + _mock_vllm_config(), groups, 100 * 1024 * 1024 + ) + page_sizes = _page_sizes_by_layer(groups) + packing = _packing_by_layer(tensors) + + for group in groups: + ranges = sorted( + (packing[name][0], packing[name][0] + page_sizes[name]) + for name in group.layer_names + ) + assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:])) def test_all_layers_accounted_for(self): n_c4, n_c128, n_swa = 5, 4, 7 @@ -135,29 +283,29 @@ def test_all_layers_accounted_for(self): expected = n_c4 * 2 + n_c128 + n_swa assert len(all_names) == expected - def test_strided_views_are_independent(self): + def test_group_owned_blocks_do_not_alias(self): groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) - page_sizes = _page_sizes_by_layer(groups) num_blocks, tensors = _get_kv_cache_config_packed( _mock_vllm_config(), groups, 100 * 1024 * 1024 ) - backing = torch.zeros(tensors[0].size, dtype=torch.uint8) - views = [] - for t in tensors: - page_size = page_sizes[t.shared_by[0]] - v = torch.as_strided( - backing, - size=(num_blocks, page_size), - stride=(t.block_stride, 1), - storage_offset=t.offset, - ) - views.append(v) - - for i, v in enumerate(views): - v.fill_(i + 1) - - for i, v in enumerate(views): - assert (v == i + 1).all(), f"View {i} was corrupted" + views = _make_views(groups, num_blocks, tensors) + + expected = {} + value = 1 + for block_id, group in enumerate(groups): + for layer_name in group.layer_names: + views[layer_name][block_id].fill_(value) + expected[layer_name] = (block_id, value) + value += 1 + + for layer_name, (block_id, value) in expected.items(): + assert (views[layer_name][block_id] == value).all() + + # Once the first group releases its block, another group may reuse it. + for layer_name in groups[1].layer_names: + views[layer_name][0].fill_(255) + for layer_name in groups[1].layer_names: + assert (views[layer_name][0] == 255).all() def test_hma_attention_groups_keep_default_backing(self): full = _make_full_spec() diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 6c4779ddda36..4ee2ea9ef4bd 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy import hashlib import importlib from collections.abc import Callable +from types import SimpleNamespace from typing import Any import pytest @@ -30,6 +32,7 @@ generate_scheduler_kv_cache_config, get_kv_cache_capacity, get_kv_cache_configs, + get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, group_and_unify_kv_cache_specs, @@ -1477,6 +1480,128 @@ def test_get_max_concurrency_for_kv_cache_config(): assert num_tokens == max_concurrency_hybrid_model * max_model_len assert max_concurrency == max_concurrency_hybrid_model + # Unequal group sizes in the standard layout: each group's pages cost + # whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the + # same as the equal-hybrid case above, regardless of the second group + # holding only 2 layers. + kv_cache_config_unequal_groups = KVCacheConfig( + num_blocks=1153 * 3, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec), + KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec), + ], + ) + assert ( + get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config_unequal_groups + ) + == 3 + ) + + # UniformTypeKVCacheSpecs group (worker config shape): the aggregated + # spec's memory/page ratio equals a single layer's page count, so the + # group needs 1024 blocks and the request 1153 in total. The previous + # formula normalized both groups' memory by the first group's page size, + # reporting 3459/1057 = 3.27 here instead of 3 — and a different value + # again for the scheduler-config shape below. + uniform_full_spec = UniformTypeKVCacheSpecs( + block_size=full_attention_spec.block_size, + kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)}, + ) + kv_cache_config_uniform_group = KVCacheConfig( + num_blocks=1153 * 3, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec), + KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec), + ], + ) + assert ( + get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config_uniform_group + ) + == 3 + ) + + # Scheduler-config shape: generate_scheduler_kv_cache_config replaces the + # uniform-type group's spec with a representative per-layer spec. + # Capacity must not change between the two shapes (the engine computes + # on the scheduler config, the worker loop on the worker config). + kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config( + [copy.deepcopy(kv_cache_config_uniform_group)] + ) + assert get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config_scheduler_shape + ) == get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config_uniform_group + ) + + +def test_get_max_concurrency_packed_kv_cache_config(): + from vllm.v1.core.kv_cache_utils import ( + _get_kv_cache_config_packed, + _use_packed_kv_cache_config, + ) + + model_config = ModelConfig( + "Qwen/Qwen1.5-7B", + runner="generate", + dtype="float16", + max_model_len=16384, + ) + scheduler_config = SchedulerConfig( + max_num_batched_tokens=1024, + enable_chunked_prefill=True, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + async_scheduling=False, + ) + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + ) + + # All-UniformTypeKVCacheSpecs groups select the packed layout. + mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)} + swa_specs = { + f"layer_{i}": SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.float32, + sliding_window=128, + ) + for i in range(4, 6) + } + kv_cache_groups = [ + KVCacheGroupSpec( + list(mla_specs), + UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs), + ), + KVCacheGroupSpec( + list(swa_specs), + UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs), + ), + ] + assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups) + num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( + vllm_config, kv_cache_groups, 2 * GiB_bytes + ) + assert num_blocks > 0 + kv_cache_config_packed = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + ) + # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages; + # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The + # previous formula normalized by the first group's page size and gave + # 1061 blocks per request instead of 1097. + assert get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config_packed + ) == num_blocks / (1024 + 73) + def test_allocate_with_lookahead(): """Verify that lookahead tokens correctly affect block allocation""" @@ -1944,10 +2069,37 @@ def test_generate_scheduler_kv_cache_config(): ) -def new_mla_spec(cache_dtype_str=None): +def test_mixed_precision_kv_cache_with_uniform_type_specs(): + fp8_spec = new_kv_cache_spec(dtype=torch.float8_e4m3fn) + bf16_spec = new_kv_cache_spec(dtype=torch.bfloat16) + worker_config = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["fp8_layer"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs={"fp8_layer": fp8_spec} + ), + ), + KVCacheGroupSpec( + ["bf16_layer"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs={"bf16_layer": bf16_spec} + ), + ), + ], + ) + scheduler_config = generate_scheduler_kv_cache_config([worker_config]) + + assert worker_config.needs_kv_cache_zeroing + assert scheduler_config.needs_kv_cache_zeroing + + +def new_mla_spec(cache_dtype_str=None, block_size=16): # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 return MLAAttentionSpec( - block_size=16, + block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.float32, @@ -1997,6 +2149,66 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups(): assert layer_names == {"mla.0", "mla.1", "swa.0"} +def new_indexer_mla_spec(block_size=16): + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller + # page size than the main MLA attention (uint8, small head), so their pages + # cannot be unified. + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + ) + + +def _grouping_config(): + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + ) + + +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): + specs = { + "target.0.attn": new_mla_spec(), + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups(_grouping_config(), specs) + + assert len(groups) == 2 + assert all( + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + + +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention + # compute sliding-window while promoting only its allocation semantics lets + # every layer share the target's block table and remain contiguous. + draft = new_sliding_window_spec(block_size=16) + specs = { + "target.0.attn": new_mla_spec(block_size=64), + "target.0.indexer": new_indexer_mla_spec(block_size=64), + "draft.0": draft, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + group_spec = groups[0].kv_cache_spec + assert isinstance(group_spec, UniformTypeKVCacheSpecs) + assert group_spec.block_size == 64 + promoted_draft = group_spec.kv_cache_specs["draft.0"] + assert isinstance(promoted_draft, FullAttentionSpec) + assert not isinstance(promoted_draft, SlidingWindowSpec) + assert promoted_draft.block_size == 64 + assert promoted_draft.sliding_window == draft.sliding_window + assert specs["draft.0"] is draft + assert "attention compute is unchanged" in caplog_vllm.text + + def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION @@ -2496,19 +2708,22 @@ def test_page_size_padded_wins(): def test_unify_hybrid_kv_cache_specs(): # 1. has_full_attention and has_sliding_window - before_spec_1 = new_kv_cache_spec() + before_spec_1 = new_kv_cache_spec(block_size=64) before_spec_2 = new_sliding_window_spec( - page_size_padded=32 * 1024, sliding_window=1024 + block_size=16, page_size_padded=32 * 1024, sliding_window=1024 ) kv_cache_spec = { "layer_1": before_spec_1, "layer_2": before_spec_2, } kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) + expected_spec_1 = new_kv_cache_spec(block_size=64) + expected_spec_2 = new_kv_cache_spec( + block_size=64, page_size_padded=64 * 1024, sliding_window=1024 + ) assert kv_cache_spec["layer_1"] == expected_spec_1 assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 # 2. has_full_attention and has_chunked_local_attention before_spec_1 = new_kv_cache_spec() diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py new file mode 100644 index 000000000000..31671caf766b --- /dev/null +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba "align" prefill chunk splitting (`_mamba_block_aligned_split`). + +Invariant: slot `p` holds the state after exactly `(p + 1) * block_size` tokens. +State is written at chunk ends, so a chunk ending mid-block leaves its slot at +the wrong offset, and a later chunk crossing that boundary publishes it anyway. +Requests resuming from it then restore a truncated state (#43559). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.utils.math_utils import cdiv +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.request import Request + +from .utils import create_requests + +pytestmark = pytest.mark.cpu_test + +# Mirrors the deployment where the poisoning was observed (Qwen3.6-27B): mamba +# block 1600, MTP with 3 draft tokens, prompts shorter than 2 mamba blocks. +ATTN_BLOCK_SIZE = 16 +MAMBA_BLOCK_SIZE = 1600 +NUM_SPEC = 3 +PROMPT_LEN = 2002 +MAMBA_GROUP_ID = 1 + + +def _make_hybrid_kv_cache_manager() -> KVCacheManager: + config = KVCacheConfig( + num_blocks=10000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_layer"], + FullAttentionSpec( + block_size=ATTN_BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba_layer"], + MambaSpec( + block_size=MAMBA_BLOCK_SIZE, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=NUM_SPEC, + ), + ), + ], + ) + return KVCacheManager( + config, + max_model_len=262144, + scheduler_block_size=MAMBA_BLOCK_SIZE, + hash_block_size=ATTN_BLOCK_SIZE, + enable_caching=True, + use_eagle=True, + ) + + +def _split( + request: Request, + num_new_tokens: int, + use_eagle: bool = True, + partial_hit: bool = False, +) -> int: + """Call the real `Scheduler._mamba_block_aligned_split` on a stub self.""" + stub = SimpleNamespace( + cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE), + use_eagle=use_eagle, + max_num_scheduled_tokens=16384, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + # `prefix_match_unit` finer than the block size (#46384). + mamba_partial_cache_hit=partial_hit, + hash_block_size=ATTN_BLOCK_SIZE, + ) + return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens) + + +def _run_chunked_prefill( + manager: KVCacheManager, request: Request, budgets: list[int] +) -> dict[int, int]: + """Prefill `request`, one step per entry in `budgets`. + + A zero-token split means the budget cannot fund an aligned chunk; the + scheduler defers the request to a later step, so this does too. + + Returns physical mamba block id -> the token offset of the state it holds, + mirroring the GDN kernel: the running slot ends up at the chunk end. + """ + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID] + state_at: dict[int, int] = {} + # `budgets` fragments the first steps; afterwards the request is alone and + # gets as much as it can use, so the prefill always finishes. + for step in range(len(budgets) + 64): + computed = request.num_computed_tokens + if computed >= request.num_tokens: + break + budget = budgets[step] if step < len(budgets) else request.num_tokens + num_new = _split(request, min(request.num_tokens - computed, budget)) + if num_new == 0: + continue + assert ( + manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC) + is not None + ) + request.num_computed_tokens = computed + num_new + blocks = mamba_manager.req_to_blocks[request.request_id] + running = cdiv(request.num_computed_tokens, MAMBA_BLOCK_SIZE) - 1 + state_at[blocks[running].block_id] = request.num_computed_tokens + return state_at + + +def _count_cached_boundary_states( + manager: KVCacheManager, request: Request, state_at: dict[int, int] +) -> int: + """Assert every hash-cached mamba slot holds the state its hash claims. + + Covers both full-block snapshots (`(p + 1) * block_size`) and the + partial-tail entries align mode registers at an exact token count. + + Returns the number of cached slots checked. + """ + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID] + checked = 0 + for pos, block in enumerate(mamba_manager.req_to_blocks[request.request_id]): + if block.is_null or block.block_hash is None: + continue + claimed = block.block_hash_num_tokens + assert state_at.get(block.block_id) == claimed, ( + f"mamba slot {pos} is hashed as state@{claimed} but holds " + f"state@{state_at.get(block.block_id)}" + ) + checked += 1 + return checked + + +def _prefill(prompt_len: int, budgets: list[int]) -> int: + manager = _make_hybrid_kv_cache_manager() + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE) + state_at = _run_chunked_prefill(manager, request, budgets) + assert request.num_computed_tokens == prompt_len, "prefill did not complete" + return _count_cached_boundary_states(manager, request, state_at) + + +def test_fragmented_first_chunk_does_not_poison_mamba_prefix_cache() -> None: + """EAGLE zeroes `last_cache_position` for prompts under two blocks. + + Past that point any chunk end used to be accepted, so a short first chunk + (concurrent prefills sharing the budget) left slot 0 at state@364 while the + next chunk crossed 1600 and published slot 0 as state@1600. + """ + _prefill(PROMPT_LEN, budgets=[364, PROMPT_LEN]) + + +def test_fragmented_tail_chunk_does_not_poison_mamba_prefix_cache() -> None: + """Same poisoning one block in, where a hit is still cacheable. + + `last_cache_position` is 1600, so the chunk ending there is cached. The next + chunk used to be free to stop mid-block (slot 1 at state@2600) and the one + after it crossed 3200, publishing slot 1 as state@3200. + """ + assert _prefill(3602, budgets=[1600, 1000, 3602]) > 0 + + +@pytest.mark.parametrize("first_chunk", [800, 900, 1599, 1601, 2000]) +def test_intermediate_chunk_ends_stay_block_aligned(first_chunk: int) -> None: + """Every non-final prefill chunk must end on a mamba block boundary.""" + _prefill(PROMPT_LEN, budgets=[first_chunk, PROMPT_LEN, PROMPT_LEN]) + + +@pytest.mark.parametrize( + ("block_size", "prompt_len", "budgets"), + [ + # Kimi-K3-scale mamba blocks: TP8 shards the recurrent state 8 ways + # (~1.5k block), DEP16 keeps it whole (~12k block). The first budget + # walks the request up to `last_cache_position`; the second is the + # sub-block fragment that lands in the unguarded tail region. + (1536, 30000, [27648, 1024, 30000]), + (12288, 30000, [12288, 4000, 30000]), + (12288, 41000, [24576, 8000, 41000]), + ], +) +def test_poisoning_is_block_size_independent( + monkeypatch: pytest.MonkeyPatch, + block_size: int, + prompt_len: int, + budgets: list[int], +) -> None: + """The invariant is per-block, so large mamba blocks are not safer.""" + import sys + + monkeypatch.setattr(sys.modules[__name__], "MAMBA_BLOCK_SIZE", block_size) + assert _prefill(prompt_len, budgets=budgets) > 0 + + +@pytest.mark.parametrize("partial_hit", [False, True]) +@pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011]) +def test_unaligned_resume_never_runs_past_its_block( + partial_hit: bool, resume_at: int +) -> None: + """A prefill resuming mid-block must re-align before crossing a boundary. + + Reachable with a finer `prefix_match_unit` (its partial-tail stop ends a + chunk off-grid by design) and with unaligned external tokens from a KV + connector. + """ + prompt_len = 3602 + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE) + tail_boundary = prompt_len // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE + + pos, ends = resume_at, [] + while pos < prompt_len: + request.num_computed_tokens = pos + num_new = _split(request, prompt_len - pos, partial_hit=partial_hit) + assert num_new > 0, f"no progress at {pos}" + if pos % MAMBA_BLOCK_SIZE != 0: + block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE + assert pos + num_new <= block_end, ( + f"chunk [{pos}, {pos + num_new}) starts mid-block and runs past " + f"{block_end}; the slot holding state@{pos} gets hashed as " + f"state@{block_end}" + ) + pos += num_new + ends.append(pos) + + for end in ends[:-1]: + aligned = end % MAMBA_BLOCK_SIZE == 0 + assert aligned or (partial_hit and end == tail_boundary), ( + f"intermediate chunk end {end} is neither block-aligned nor the " + f"partial-tail boundary" + ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index dd89769315b9..13ed7c7b9d8b 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1078,6 +1078,8 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): # Create minimal mock with just the needed attributes mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=3 * block_size, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=block_size, mamba_partial_cache_hit=False, diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index b8430a997783..ac966ad11b1a 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -26,14 +26,17 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.utils.hashing import sha256 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.engine import FinishReason from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + MambaSpec, ) from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus @@ -1069,6 +1072,156 @@ def test_preempt_during_execution(): assert requests[1].output_token_ids[0] == 42 +def test_prefix_cache_query_not_inflated_by_connector_defer(): + """The GPU prefix-cache query is recorded at admission, so a request the + connector defers several times is counted once, not once per retry.""" + num_defers_before_matching = 3 + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv( + matched_tokens=0, + is_async=False, + num_defers_before_matching=num_defers_before_matching, + ), + ) + request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0] + scheduler.add_request(request) + + # Each deferred step re-runs the lookup but records nothing. + for _ in range(num_defers_before_matching): + assert not scheduler.schedule().scheduled_new_reqs + + output = scheduler.schedule() + assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs) + + stats = scheduler.kv_cache_manager.prefix_cache_stats + assert stats is not None + assert stats.requests == 1 + assert stats.queries == request.num_tokens + + +def test_preemption_re_records_prefix_cache_query(): + """A preempted request re-enters the lookup on resume, so its recomputation + is counted again into the preempted stats.""" + scheduler = create_scheduler(enable_prefix_caching=True) + request = create_requests(num_requests=1)[0] + scheduler.add_request(request) + + scheduler_output = scheduler.schedule() + stats = scheduler.kv_cache_manager.prefix_cache_stats + assert stats is not None + assert (stats.requests, stats.preempted_requests) == (1, 0) + + scheduler.running.remove(request) + scheduler._preempt_request(request, 0.0) + assert request.status == RequestStatus.PREEMPTED + + scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[1000]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + assert request.num_stale_output_tokens == 0 + stats = scheduler.kv_cache_manager.prefix_cache_stats + assert stats is not None + + scheduler.schedule() + assert request.status == RequestStatus.RUNNING + assert stats.preempted_requests == 1 + + +def test_prefix_cache_stats_not_recorded_when_caching_disabled(): + """With prefix caching off there is no local lookup, so admitting a request + records no phantom miss.""" + scheduler = create_scheduler(enable_prefix_caching=False) + for request in create_requests(num_requests=2): + scheduler.add_request(request) + + scheduler.schedule() + + stats = scheduler.kv_cache_manager.prefix_cache_stats + assert stats is not None + assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0) + + +def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request(): + """A real cache hit rejected once by allocate_slots and admitted on the next + step is counted exactly once, hits included.""" + block_size = 16 + scheduler = create_scheduler( + enable_prefix_caching=True, + enable_chunked_prefill=False, + block_size=block_size, + ) + + # Seed the cache so the next request with the same prompt hits it. + seed = create_requests( + num_requests=1, + num_tokens=block_size * 2, + max_tokens=2, + same_prompt=True, + block_size=block_size, + req_ids=["seed"], + )[0] + scheduler.add_request(seed) + _step_until_done( + scheduler, + scheduler.schedule(), + ModelRunnerOutput( + req_ids=["seed"], + req_id_to_index={"seed": 0}, + sampled_token_ids=[[1000]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + # The seeding step swapped in a fresh accumulator, so re-read it; the + # retried request must be the only thing recorded from here on. + stats = scheduler.kv_cache_manager.prefix_cache_stats + assert stats is not None + assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0) + + retried = create_requests( + num_requests=1, + num_tokens=block_size * 3, + max_tokens=1, + same_prompt=True, + block_size=block_size, + req_ids=["retried"], + )[0] + scheduler.add_request(retried) + + # Reject the first allocation attempt, then delegate to the real one. + orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots + allocate_results: list = [] + + def spy_allocate_slots(*args, **kwargs): + result = None if not allocate_results else orig_allocate_slots(*args, **kwargs) + allocate_results.append(result) + return result + + scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots + + assert not scheduler.schedule().scheduled_new_reqs + assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0) + + assert "retried" in scheduler.schedule().num_scheduled_tokens + assert allocate_results[0] is None and allocate_results[1] is not None + assert (stats.requests, stats.queries, stats.hits) == ( + 1, + retried.num_tokens, + block_size * 2, + ) + + def test_scheduler_reset_prefix_cache(): scheduler = create_scheduler(enable_prefix_caching=True) requests = create_requests(num_requests=10) @@ -5383,3 +5536,197 @@ def test_async_load_reservation_prevents_wedge_e2e(): assert b.status == RequestStatus.WAITING assert b.num_preemptions == 0 assert b.request_id not in req_to_blocks + + +def _create_hybrid_mamba_connector_scheduler( + matched_tokens: int, + block_size: int = 16, + num_blocks: int = 100, +) -> Scheduler: + """FA + Mamba ("all" cache mode) scheduler with a MockKVConnector.""" + model_config = ModelConfig( + model="facebook/opt-125m", + trust_remote_code=True, + dtype="float16", + seed=42, + skip_tokenizer_init=True, + ) + vllm_config = VllmConfig( + scheduler_config=SchedulerConfig( + max_num_seqs=4, + max_num_batched_tokens=8192, + max_model_len=8192, + enable_chunked_prefill=True, + is_encoder_decoder=False, + watermark=0.0, + ), + model_config=model_config, + cache_config=CacheConfig( + block_size=block_size, + enable_prefix_caching=True, + mamba_cache_mode="all", + ), + kv_transfer_config=KVTransferConfig( + kv_connector="MockKVConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "matched_tokens": matched_tokens, + "is_async": False, + }, + ), + ) + vllm_config.cache_config.num_gpu_blocks = num_blocks + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["fa"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="all", + ), + ), + ], + ) + register_all_kvcache_specs(vllm_config) + return Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + structured_output_manager=StructuredOutputManager(vllm_config), + block_size=block_size, + hash_block_size=block_size, + log_stats=True, + ) + + +@pytest.mark.parametrize( + "matched_tokens,expected_num_computed", + [ + # No external hit: resume on the deepest locally-consistent boundary + # (block 0's state survives for both groups). + (0, 16), + # One external block on top of the reconciled local boundary. + (16, 32), + ], +) +def test_hybrid_per_group_hit_divergence_with_connector( + matched_tokens: int, expected_num_computed: int +): + """Per-group prefix hits can diverge for hybrid models with a connector + (#46453): under block pressure the FA prefix tail is evicted while a + deeper Mamba state block survives. The scheduler must not report the + deeper hit as locally computed (evicted FA blocks are not resident -> + engine crash / dirty KV); it falls back to the reconciled boundary that + every group is consistent at. + """ + block_size = 16 + scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens) + manager = scheduler.kv_cache_manager + assert isinstance(manager.coordinator, HybridKVCacheCoordinator) + + # Seed a 4-block prefix so both groups cache all four boundaries + # (mamba cache mode "all" caches every block's state densely). + [fill] = create_requests( + num_requests=1, + num_tokens=4 * block_size, + max_tokens=1, + same_prompt=True, + block_size=block_size, + req_ids=["fill"], + ) + computed_blocks, num_computed, _ = manager.get_computed_blocks(fill) + blocks = manager.allocate_slots( + fill, fill.num_tokens, num_computed, computed_blocks + ) + fa_ids = [b.block_id for b in blocks.blocks[0]] + mamba_ids = [b.block_id for b in blocks.blocks[1]] + manager.free(fill) + + # Evict the FA tail and the middle mamba states; block 0 (both groups) + # and the deep mamba state at block 3 survive. + manager.block_pool.evict_blocks({fa_ids[2], fa_ids[3], mamba_ids[1], mamba_ids[2]}) + + # A replay of the prefix plus one extra block now sees diverged + # per-group hits: FA stops at the evicted tail, while the mamba lookup + # finds the deeper surviving state. + [replay] = create_requests( + num_requests=1, + num_tokens=5 * block_size, + max_tokens=1, + same_prompt=True, + block_size=block_size, + req_ids=["replay"], + ) + _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group( + replay.block_hashes, replay.num_tokens - 1 + ) + assert per_group_hits == (2 * block_size, 4 * block_size) # diverged + + scheduler.add_request(replay) + output = scheduler.schedule() + num_scheduled = output.num_scheduled_tokens[replay.request_id] + assert replay.num_tokens - num_scheduled == expected_num_computed + + +def test_hybrid_per_group_hit_divergence_fa_deeper_no_external(): + """The opposite divergence: the FA prefix survives deeper than the Mamba + state and the connector supplies nothing (ext == 0). Reporting the deep FA + hit as locally computed would resume with no valid Mamba state at that + boundary (silent bad output). The scheduler must fall back to the + convergent boundary that every group agrees on (block 0's surviving state). + """ + block_size = 16 + scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0) + manager = scheduler.kv_cache_manager + assert isinstance(manager.coordinator, HybridKVCacheCoordinator) + + # Seed a 4-block prefix in both groups. + [fill] = create_requests( + num_requests=1, + num_tokens=4 * block_size, + max_tokens=1, + same_prompt=True, + block_size=block_size, + req_ids=["fill"], + ) + computed_blocks, num_computed, _ = manager.get_computed_blocks(fill) + blocks = manager.allocate_slots( + fill, fill.num_tokens, num_computed, computed_blocks + ) + mamba_ids = [b.block_id for b in blocks.blocks[1]] + manager.free(fill) + + # Keep all FA blocks; evict every mamba state but block 0. FA reaches 4 + # blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba). + manager.block_pool.evict_blocks({mamba_ids[1], mamba_ids[2], mamba_ids[3]}) + + [replay] = create_requests( + num_requests=1, + num_tokens=5 * block_size, + max_tokens=1, + same_prompt=True, + block_size=block_size, + req_ids=["replay"], + ) + _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group( + replay.block_hashes, replay.num_tokens - 1 + ) + assert per_group_hits == (4 * block_size, 1 * block_size) # FA deeper + + scheduler.add_request(replay) + output = scheduler.schedule() + num_scheduled = output.num_scheduled_tokens[replay.request_id] + # Must resume at the convergent boundary (block 0), not the deep FA hit. + assert replay.num_tokens - num_scheduled == block_size diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 19beba1a53dd..5807574c6444 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -38,8 +38,12 @@ EOS_TOKEN_ID = 50256 -def mock_kv(matched_tokens: int, is_async: bool): - return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async) +def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0): + return MockKVConfig( + matched_tokens=matched_tokens, + is_async=is_async, + num_defers_before_matching=num_defers_before_matching, + ) def create_scheduler( @@ -51,6 +55,7 @@ def create_scheduler( long_prefill_token_threshold: int = 0, disable_chunked_mm_input: bool = False, use_kv_connector: None | bool | str | MockKVConfig = None, + kv_role: str = "kv_both", num_blocks: int = 10000, block_size: int = 16, max_model_len: int | None = None, @@ -111,21 +116,24 @@ def create_scheduler( if isinstance(use_kv_connector, MockKVConfig): kv_transfer_config = KVTransferConfig( kv_connector="MockKVConnector", - kv_role="kv_both", + kv_role=kv_role, kv_connector_extra_config={ "matched_tokens": use_kv_connector.matched_tokens, "is_async": use_kv_connector.is_async, + "num_defers_before_matching": ( + use_kv_connector.num_defers_before_matching + ), }, ) elif isinstance(use_kv_connector, str): kv_transfer_config = KVTransferConfig( kv_connector=use_kv_connector, - kv_role="kv_both", + kv_role=kv_role, ) elif use_kv_connector: kv_transfer_config = KVTransferConfig( kv_connector="ExampleConnector", - kv_role="kv_both", + kv_role=kv_role, kv_connector_extra_config={"shared_storage_path": "local_storage"}, ) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 742aafd3890e..dc9c1ffb3067 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -6,7 +6,6 @@ from __future__ import annotations -import os import threading from contextlib import nullcontext from unittest.mock import patch @@ -14,7 +13,21 @@ import pytest import torch -os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" + +@pytest.fixture(autouse=True) +def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch): + """Enable breakable cudagraphs for this module's tests only. + + eager_break_during_capture reads the env at decoration time, which + happens inside the test bodies, so a per-test fixture suffices. + monkeypatch restores the env so other test files running in the same + pytest process are unaffected (a module-level os.environ assignment + used to leak into test_cudagraph_dispatch.py and break it). + """ + import vllm.envs as envs + + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") + envs.disable_envs_cache() def test_piecewise_capture_builds_fresh_metadata_for_both_passes(): @@ -86,10 +99,19 @@ def cuda_capture_stream(): """ if not torch.cuda.is_available(): pytest.skip("CUDA required") + from vllm.utils.torch_utils import _current_stream_tls + + prev_stream = getattr(_current_stream_tls, "value", None) stream = torch.cuda.Stream() with torch.cuda.stream(stream): yield stream torch.cuda.current_stream().wait_stream(stream) + # Exiting torch.cuda.stream() records the default stream in vllm's + # patched set_stream cache. A later CUDAGraphWrapper capture in the same + # process would then run on the default stream, which cannot capture + # (this broke test_cudagraph_dispatch.py when run in-process after this + # file). Restore the pre-fixture value so this module leaves no trace. + _current_stream_tls.value = prev_stream # --------------------------------------------------------------------------- @@ -311,6 +333,54 @@ def attention_like(t: torch.Tensor) -> None: assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) +def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): + """Handle an eager attention break inside a multi-stream overlap region.""" + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) + from vllm.utils.multi_stream_utils import maybe_execute_in_parallel + + x = torch.zeros(1024, device="cuda") + output = torch.empty_like(x) + aux_stream = torch.cuda.Stream() + main_event = torch.cuda.Event() + aux_event = torch.cuda.Event() + + @eager_break_during_capture + def attention(query: torch.Tensor, out: torch.Tensor) -> None: + torch.add(query, 3.0, out=out) + + def attention_frontend() -> torch.Tensor: + query = x * 2.0 + attention_output = torch.empty_like(x) + attention(query, attention_output) + return attention_output + + cap = BreakableCUDAGraphCapture() + with cap: + attention_output, gate = maybe_execute_in_parallel( + attention_frontend, + lambda: x * 5.0, + main_event, + aux_event, + aux_stream, + ) + torch.add(attention_output, gate, out=output) + + assert cap.num_graphs == 2 + assert cap.num_eager_breaks == 1 + + for value in (1.0, 7.0, 19.0): + x.fill_(value) + cap.replay() + cuda_capture_stream.synchronize() + torch.testing.assert_close( + output, + torch.full_like(output, value * 7.0 + 3.0), + ) + + # --------------------------------------------------------------------------- # Replay ordering # --------------------------------------------------------------------------- diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py new file mode 100644 index 000000000000..655fcf0d2e92 --- /dev/null +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.config import ( + CompilationConfig, + CUDAGraphMode, + ParallelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.distributed.device_communicators import pynccl_allocator +from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture(autouse=True) +def _reset_graph_pool_id(): + pynccl_allocator._graph_pool_id = None + yield + pynccl_allocator._graph_pool_id = None + + +def _create_vllm_config() -> MagicMock: + compilation_config = CompilationConfig( + cudagraph_mode="FULL", + cudagraph_capture_sizes=[4], + ) + compilation_config.max_cudagraph_capture_size = 4 + compilation_config.post_init_cudagraph_sizes() + + vllm_config = MagicMock(spec=VllmConfig) + vllm_config.compilation_config = compilation_config + vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=4) + vllm_config.parallel_config = ParallelConfig() + vllm_config.speculative_config = None + vllm_config.num_speculative_tokens = 0 + return vllm_config + + +def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch): + """FULL capture must set graph_pool_id before entering torch.cuda.graph(). + + NCCL symmetric memory checks this global during graph capture; without + it, capture fails with: + AssertionError: graph_pool_id is not set under graph capture + """ + graph_pool = object() + monkeypatch.setattr( + gpu_cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr( + gpu_cudagraph_utils.current_platform, + "get_global_graph_pool", + lambda: graph_pool, + ) + + manager = gpu_cudagraph_utils.CudaGraphManager( + vllm_config=_create_vllm_config(), + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL, + decode_query_len=1, + ) + + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=4, + num_reqs=4, + uniform_token_count=1, + ) + manager._capture_descs[CUDAGraphMode.FULL] = [desc] + + def create_forward_fn(desc, warmup): + return lambda _mode: None + + @contextmanager + def fake_graph_capture(*args, **kwargs): + yield SimpleNamespace(stream=MagicMock()) + + fake_offloader = MagicMock() + + def cuda_graph_enter(*args, **kwargs): + assert pynccl_allocator._graph_pool_id is graph_pool + + mock_cuda_graph_ctx = MagicMock() + mock_cuda_graph_ctx.__enter__ = cuda_graph_enter + mock_cuda_graph_ctx.__exit__ = MagicMock(return_value=False) + + with ( + patch.object(gpu_cudagraph_utils, "graph_capture", fake_graph_capture), + patch.object(gpu_cudagraph_utils, "get_offloader", lambda: fake_offloader), + patch.object(gpu_cudagraph_utils.torch.cuda, "CUDAGraph"), + patch.object( + gpu_cudagraph_utils.torch.cuda, + "graph", + return_value=mock_cuda_graph_ctx, + ) as mock_cuda_graph, + ): + manager.capture(create_forward_fn) + + mock_cuda_graph.assert_called_once() diff --git a/tests/v1/cudagraph/test_cudagraph_mode.py b/tests/v1/cudagraph/test_cudagraph_mode.py index f4f74d16c701..e413fd91e35a 100644 --- a/tests/v1/cudagraph/test_cudagraph_mode.py +++ b/tests/v1/cudagraph/test_cudagraph_mode.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import weakref from contextlib import ExitStack import pytest -from tests.utils import wait_for_gpu_memory_to_clear +from tests.utils import create_new_process_for_each_test from tests.v1.attention.utils import full_cg_backend_configs as backend_configs from vllm import LLM from vllm.config import CompilationConfig, CompilationMode @@ -32,6 +31,7 @@ @pytest.mark.parametrize("backend_name, cudagraph_mode, supported", combo_cases_1) +@create_new_process_for_each_test("spawn") def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supported): if backend_name == "FlashInfer": try: @@ -64,17 +64,6 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte ), ) llm.generate(["Hello, my name is"] * 10) - # when above code raises, `llm` may be undefined, so we need to catch that - try: - llm = weakref.proxy(llm) - del llm - except UnboundLocalError: - pass - - wait_for_gpu_memory_to_clear( - devices=[0], - threshold_ratio=0.1, - ) # test cudagraph_mode with different compilation mode. @@ -98,6 +87,7 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte @pytest.mark.parametrize( "backend_name,cudagraph_mode,compilation_mode,supported", combo_cases_2 ) +@create_new_process_for_each_test("spawn") def test_cudagraph_compilation_combo( backend_name, cudagraph_mode, compilation_mode, supported ): @@ -120,14 +110,3 @@ def test_cudagraph_compilation_combo( ), ) llm.generate(["Hello, my name is"] * 10) - # when above code raises, `llm` may be undefined, so we need to catch that - try: - llm = weakref.proxy(llm) - del llm - except UnboundLocalError: - pass - finally: - wait_for_gpu_memory_to_clear( - devices=[0], - threshold_ratio=0.1, - ) diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index ed816d817c15..f2a54c55c809 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -106,8 +106,14 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager: """ mgr = object.__new__(EncoderCudaGraphManager) mgr.token_budgets = sorted(budgets) + mgr.path_token_budgets = {"default": mgr.token_budgets} mgr.max_batch_size = 16 mgr.use_dp = False + mgr.config = EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[], + out_hidden_size=32, + ) mgr.budget_graphs = {"default": {}} mgr.graph_pool = None mgr.graph_hits = 0 @@ -412,6 +418,7 @@ def _make_manager_for_gpu( """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests.""" mgr = object.__new__(EncoderCudaGraphManager) mgr.token_budgets = sorted(token_budgets) + mgr.path_token_budgets = {"default": mgr.token_budgets} mgr.max_batch_size = max_batch_size mgr.max_frames_per_batch = ( max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2 diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b2706ed89b7b..37fd5cba6a56 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -27,8 +27,10 @@ "backend", BACKENDS, ) +@pytest.mark.parametrize("rms_norm_impl", ["default", "vllm_c"]) def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( backend, + rms_norm_impl, ): """ Ensures that the same request (the 'needle' prompt) yields identical output @@ -60,6 +62,16 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( random.seed(seed) attention_config = {"backend": backend} + # Force the C++ RMSNorm implementation so we actually exercise the + # num_tokens-dependent block-size branches. + kernel_config = None + if rms_norm_impl == "vllm_c": + kernel_config = { + "ir_op_priority": { + "rms_norm": ["vllm_c"], + "fused_add_rms_norm": ["vllm_c"], + } + } # Allow overrides from environment (useful for CI tuning) # "facebook/opt-125m" is too small, doesn't reliably test determinism model = TEST_MODEL @@ -96,6 +108,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( gpu_memory_utilization=gpu_mem_util, max_model_len=max_model_len, attention_config=attention_config, + kernel_config=kernel_config, ) # Baseline generation for the needle prompt alone. @@ -923,11 +936,15 @@ def LLM_with_max_seqs( gpu_memory_utilization: float, max_model_len: int, attention_config: dict | None = None, + kernel_config: dict | None = None, ) -> LLM: """ Helper to construct an LLM with a specific max_num_seqs (batch-size limit) using the high-level v1 LLM API, while constraining memory usage. """ + extra_kwargs: dict = {} + if kernel_config is not None: + extra_kwargs["kernel_config"] = kernel_config return LLM( model=model, max_num_seqs=max_num_seqs, @@ -939,4 +956,5 @@ def LLM_with_max_seqs( attention_config=attention_config, # Enable for MOE models # enable_expert_parallel=True, + **extra_kwargs, ) diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 5b3b7a8758b1..232a43b1f98d 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -1,11 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test batch-invariant RMS normalization against standard implementations. - -This test compares the Triton-based batch-invariant RMS norm implementation -with the standard CUDA-based implementation to ensure numerical accuracy. -""" +"""Test batch-invariant RMS normalization against a PyTorch reference.""" import pytest import torch @@ -14,43 +9,52 @@ from vllm.model_executor.layers.batch_invariant import ( rms_norm_batch_invariant, ) -from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform DEVICE_TYPE = current_platform.device_type +def _rms_norm_reference( + input_tensor: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Compute RMSNorm independently using PyTorch operations.""" + input_fp32 = input_tensor.float() + output = input_fp32 * torch.rsqrt( + input_fp32.square().mean(dim=-1, keepdim=True) + eps + ) + return (output * weight.float()).to(input_tensor.dtype) + + @skip_if_not_cuda -@pytest.mark.parametrize("batch_size", [1, 4, 16, 64]) +@pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6, 1e-5]) -def test_rms_norm_batch_invariant_vs_standard( +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_batch_invariant_vs_reference( default_vllm_config, batch_size: int, hidden_size: int, dtype: torch.dtype, eps: float, + seed: int, ): """ - Compare batch-invariant Triton RMS norm against standard CUDA implementation. + Compare batch-invariant Triton RMS norm against a PyTorch reference. Tests that the Triton-based batch-invariant RMS norm produces numerically - equivalent results to the standard CUDA implementation across various - configurations. + equivalent results to an independent implementation across configurations. """ device = torch.device(DEVICE_TYPE) # Create test input and weight - torch.manual_seed(42) + torch.manual_seed(seed) input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) - # Standard implementation (CUDA ops) - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device) - rms_norm_layer.weight.data = weight.clone() - - standard_output = rms_norm_layer.forward_cuda(input_tensor) + reference_output = _rms_norm_reference(input_tensor, weight, eps) # Batch-invariant implementation (Triton) triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) @@ -64,12 +68,12 @@ def test_rms_norm_batch_invariant_vs_standard( torch.testing.assert_close( triton_output, - standard_output, + reference_output, rtol=rtol, atol=atol, msg=f"RMS norm mismatch for batch_size={batch_size}, " f"hidden_size={hidden_size}, " - f"dtype={dtype}, eps={eps}", + f"dtype={dtype}, eps={eps}, seed={seed}", ) @@ -77,17 +81,21 @@ def test_rms_norm_batch_invariant_vs_standard( @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("n_extra", [3, 299]) +@pytest.mark.parametrize("seed", list(range(16))) def test_fused_add_rms_norm_batch_invariant_residual_path( hidden_size: int, dtype: torch.dtype, eps: float, + n_extra: int, + seed: int, ): """ Test the batch-invariant fused residual-add + RMSNorm helper directly. """ device = torch.device(DEVICE_TYPE) - torch.manual_seed(42) + torch.manual_seed(seed) x_single = torch.randn(1, hidden_size, dtype=dtype, device=device) residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) @@ -95,14 +103,14 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( x_batch = torch.cat( [ x_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) residual_batch = torch.cat( [ residual_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) @@ -127,7 +135,7 @@ def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]: ) merged_single = x_single + residual_single - ref_out = rms_norm_batch_invariant(merged_single, weight, eps=eps) + ref_out = _rms_norm_reference(merged_single, weight, eps) torch.testing.assert_close( residual_out_single, @@ -162,7 +170,139 @@ def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]: rtol=rtol, atol=atol, msg="Fused add RMSNorm output should stay numerically close to the " - "batch-invariant RMSNorm reference", + "PyTorch RMSNorm reference", + ) + + +FP8_DTYPE = current_platform.fp8_dtype() + +# The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block +# 256, while the small launch (255 rows) stays under the threshold and keeps the +# larger block (1024, or 512 for per-block quant). Under the pin the two launches +# use the same block, so the shared first 255 rows must match bit-for-bit; 255 is +# the most rows a single small launch can hold (< 256, and <= 256 for per-block). +_LARGE_TOKENS = 300 +_SMALL_TOKENS = 255 + + +def _assert_rows_bit_identical(small, large, msg): + if small.dtype == FP8_DTYPE: + assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg + else: + torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_batch_invariant_nonresidual_kernel( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm`` (no residual) must be batch invariant across the block + threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]`` + (default priority is ``native``/inductor codegen when compiling). + """ + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + + def rms_norm(x): + out = torch.empty_like(x) + ops.rms_norm(out, x, weight, 1e-6) + return out + + large = rms_norm(rows.clone()) + small = rms_norm(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "rms_norm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +@pytest.mark.parametrize("add_residual", [False, True]) +def test_rms_norm_static_fp8_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool +): + """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across + the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with + ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8 + path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them). + """ + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + residual = ( + torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + if add_residual + else None + ) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + + def quant(x, res): + out = torch.empty_like(x, dtype=FP8_DTYPE) + if add_residual: + torch.ops._C.fused_add_rms_norm_static_fp8_quant( + out, x, res, weight, quant_scale, 1e-6 + ) + else: + torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6) + return out + + large = quant(rows.clone(), residual.clone() if residual is not None else None) + small = quant( + rows[:_SMALL_TOKENS].clone(), + residual[:_SMALL_TOKENS].clone() if residual is not None else None, + ) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "static-fp8-quant RMSNorm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_per_block_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm_per_block_quant`` must be batch invariant across the + block threshold (compiled fp8 block-quant path; block pinned to 512).""" + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + group_size = [1, 128] + + def per_block_quant(x): + return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size) + + out_large, scale_large = per_block_quant(rows.clone()) + out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + out_small, + out_large[:_SMALL_TOKENS], + "rms_norm_per_block_quant output depends on num_tokens (block size)", + ) + torch.testing.assert_close( + scale_small, + scale_large[:_SMALL_TOKENS], + rtol=0.0, + atol=0.0, + msg="rms_norm_per_block_quant scales depend on num_tokens (block size)", ) @@ -189,10 +329,7 @@ def test_rms_norm_3d_input( ) weight = torch.randn(hidden_size, dtype=dtype, device=device) - # Standard implementation - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device) - rms_norm_layer.weight.data = weight.clone() - standard_output = rms_norm_layer.forward_cuda(input_tensor) + reference_output = _rms_norm_reference(input_tensor, weight, eps) # Batch-invariant implementation triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) @@ -202,7 +339,7 @@ def test_rms_norm_3d_input( torch.testing.assert_close( triton_output, - standard_output, + reference_output, rtol=rtol, atol=atol, msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, " @@ -238,20 +375,17 @@ def test_rms_norm_numerical_stability(default_vllm_config): weight = torch.ones(hidden_size, dtype=dtype, device=device) for idx, input_tensor in enumerate(test_cases): - # Standard implementation - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device) - rms_norm_layer.weight.data = weight.clone() - standard_output = rms_norm_layer.forward_cuda(input_tensor) + reference_output = _rms_norm_reference(input_tensor, weight, eps) # Batch-invariant implementation triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) # Check for NaN or Inf - assert not torch.isnan(standard_output).any(), ( - f"Standard RMS norm produced NaN for test case {idx}" + assert not torch.isnan(reference_output).any(), ( + f"Reference RMS norm produced NaN for test case {idx}" ) - assert not torch.isinf(standard_output).any(), ( - f"Standard RMS norm produced Inf for test case {idx}" + assert not torch.isinf(reference_output).any(), ( + f"Reference RMS norm produced Inf for test case {idx}" ) assert not torch.isnan(triton_output).any(), ( f"Triton RMS norm produced NaN for test case {idx}" @@ -263,7 +397,7 @@ def test_rms_norm_numerical_stability(default_vllm_config): # Compare outputs - very lenient for extreme values with float16 torch.testing.assert_close( triton_output, - standard_output, + reference_output, rtol=2e-1, # 20% tolerance for extreme values atol=2e-1, msg=f"RMS norm mismatch for extreme value test case {idx}", @@ -321,10 +455,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) - # Standard implementation - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device) - rms_norm_layer.weight.data = weight.clone() - standard_output = rms_norm_layer.forward_cuda(input_tensor) + reference_output = _rms_norm_reference(input_tensor, weight, eps) # Batch-invariant implementation triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) @@ -334,7 +465,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): torch.testing.assert_close( triton_output, - standard_output, + reference_output, rtol=rtol, atol=atol, msg=f"RMS norm mismatch for hidden_size={hidden_size}", @@ -420,21 +551,18 @@ def test_rms_norm_batch_invariance(dtype): input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) - # Standard implementation - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device) - rms_norm_layer.weight.data = weight.clone() - standard_output = rms_norm_layer.forward_cuda(input_tensor) + reference_output = _rms_norm_reference(input_tensor, weight, eps) # Batch-invariant implementation triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) # Compare - max_diff = (triton_output - standard_output).abs().max().item() - mean_diff = (triton_output - standard_output).abs().mean().item() + max_diff = (triton_output - reference_output).abs().max().item() + mean_diff = (triton_output - reference_output).abs().mean().item() print(f"Max difference: {max_diff:.6e}") print(f"Mean difference: {mean_diff:.6e}") - print(f"Standard output sample: {standard_output[0, :5].tolist()}") + print(f"Reference output sample: {reference_output[0, :5].tolist()}") print(f"Triton output sample: {triton_output[0, :5].tolist()}") if max_diff < 1e-3: diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 9269f294b8bd..711daaef7400 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -12,13 +12,14 @@ from vllm import SamplingParams from vllm.config import VllmConfig +from vllm.config.parallel import DataParallelBackend from vllm.engine.arg_utils import AsyncEngineArgs from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind from vllm.v1.engine.async_llm import AsyncLLM -from vllm.v1.engine.core_client import DPAsyncMPClient +from vllm.v1.engine.core_client import DPLBAsyncMPClient from vllm.v1.metrics.loggers import StatLoggerBase from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats @@ -82,7 +83,7 @@ async def generate( async def test_load( model: str, output_kind: RequestOutputKind, - data_parallel_backend: str, + data_parallel_backend: DataParallelBackend, async_scheduling: bool, ): if async_scheduling and data_parallel_backend == "ray": @@ -162,7 +163,8 @@ def log_engine_initialized(self): assert not engine.output_processor.has_unfinished_requests() # testing internals here which may break - core_client: DPAsyncMPClient = engine.engine_core + core_client = engine.engine_core + assert isinstance(core_client, DPLBAsyncMPClient) # the engines only synchronize stopping every N steps so # allow a small amount of time here. for _ in range(10): diff --git a/tests/v1/distributed/test_eagle_dp.py b/tests/v1/distributed/test_eagle_dp.py index 019f9c6f213d..bb2b9f9fae16 100644 --- a/tests/v1/distributed/test_eagle_dp.py +++ b/tests/v1/distributed/test_eagle_dp.py @@ -61,7 +61,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str): data_parallel_backend="mp", # ray takes more time trust_remote_code=True, max_model_len=16384, - attention_config={"backend": attn_backend}, + attention_config={"backend": attn_backend}, # type: ignore[arg-type] ) eagle_engine_args = replace( diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index cd0aff79de83..955e5c5bf025 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -59,7 +59,7 @@ def test_decoder_max_context_length_validation( "Make sure that `max_model_len` is no smaller than the number of " "text tokens (prompt + requested output tokens)." ) - with pytest.raises(ValueError) as excinfo: + with pytest.raises(VLLMValidationError) as excinfo: vllm_model.generate_greedy(prompt_ids, max_tokens) assert expected_msg in str(excinfo.value) 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 11f77492d2e9..3155da82ea85 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -18,7 +18,7 @@ @pytest.fixture def test_prompts(): """ - Adapted from tests/v1/e2e/spec_decode/test_spec_decode.py + Adapted from tests/v1/e2e/spec_decode/utils.py """ prompt_types = ["repeat", "sentence"] # Setting higher num prompts increases the chance of numerics mismatch diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 5a7af6f22c5d..db0cab471b4b 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -338,6 +338,30 @@ def check_copy_info( assert copy_info[0][-1] == expected_temporal_src assert copy_info[1][-1] == expected_temporal_dest + def check_fused_copy_info( + action: tuple[int, int], + align_ctx: mamba_utils.MambaSpecDecodeGPUContext, + ): + # Align + spec-decode on a hybrid model routes the pre-copy through the + # fused kernel (preprocess_mamba -> run_fused_precopy) instead of + # do_mamba_copy_block, so copy_info is never populated. Verify from the + # fused buffers (req-0 scope, mirroring check_copy_info): + # - the copy DECISION: src_col is -1 iff no pre-copy is scheduled; + # - the DESTINATION column: state_idx == action[1] (curr_state_idx; + # maps directly through block_ids, exactly as check_copy_info's dst). + # The source column is NOT asserted here: on the scalar path the source + # address is produced by the per-state copy func with an accept-token + # bias offset (collect_mamba_copy_meta), so prev_state_idx does not map + # to action[0] by plain equality. Source block-level exactness (incl. + # the accept-bias) is covered by test_precopy_mamba_align.py. + src_col = int(align_ctx.precopy_src_col_buf.np[0]) + state_idx = int(align_ctx.mamba_state_idx_buf.np[0]) + if action == (-1, -1): + assert src_col == -1 + else: + assert src_col != -1 + assert state_idx == action[1] + def fake_preprocess_mamba_fn( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, @@ -348,6 +372,7 @@ def fake_preprocess_mamba_fn( forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], copy_bufs: mamba_utils.MambaCopyBuffers, + align_ctx: mamba_utils.MambaSpecDecodeGPUContext | None = None, ): nonlocal copy_info copy_info = None @@ -361,14 +386,21 @@ def fake_preprocess_mamba_fn( forward_context, mamba_state_copy_funcs, copy_bufs, + align_ctx, ) if cur_step_action is not None: - check_copy_info( - cur_step_action.preprocess_copy_idx, - kv_cache_config, - forward_context, - input_batch, - ) + if align_ctx is not None: + check_fused_copy_info( + cur_step_action.preprocess_copy_idx, + align_ctx, + ) + else: + check_copy_info( + cur_step_action.preprocess_copy_idx, + kv_cache_config, + forward_context, + input_batch, + ) return ret def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers): diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py index bb041cd38627..c5b6341fb98b 100644 --- a/tests/v1/e2e/general/test_min_tokens.py +++ b/tests/v1/e2e/general/test_min_tokens.py @@ -16,6 +16,7 @@ import pytest from vllm import LLM, SamplingParams +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput # Test configuration @@ -479,13 +480,13 @@ def test_min_tokens_validation(): # Invalid cases with pytest.raises( - ValueError, + VLLMValidationError, match="min_tokens must be greater than or equal to 0", ): SamplingParams(min_tokens=-1, max_tokens=10) with pytest.raises( - ValueError, + VLLMValidationError, match="min_tokens must be less than or equal to max_tokens", ): SamplingParams(min_tokens=15, max_tokens=10) diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py index 01c5fe6f8eb0..1954ce6a7dcf 100644 --- a/tests/v1/e2e/general/test_streaming_input.py +++ b/tests/v1/e2e/general/test_streaming_input.py @@ -20,6 +20,7 @@ from vllm import SamplingParams from vllm.engine.protocol import StreamingInput +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind @@ -571,13 +572,17 @@ async def dummy_generator() -> AsyncGenerator[StreamingInput, None]: yield StreamingInput(prompt="test") # Test n > 1 is rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_n2 = SamplingParams(max_tokens=10, n=2) async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"): pass # Test FINAL_ONLY is rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_final = SamplingParams( max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY ) @@ -585,7 +590,9 @@ async def dummy_generator() -> AsyncGenerator[StreamingInput, None]: pass # Test stop strings are rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_stop = SamplingParams(max_tokens=10, stop=["stop"]) async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"): pass diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py new file mode 100644 index 000000000000..a4c9c67a0523 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +import pytest +import torch + +from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline +from tests.utils import single_gpu_only +from vllm import LLM +from vllm.distributed import cleanup_dist_env_and_memory + +from ...utils import compute_acceptance_len +from ..utils import run_acceptance_length_eval + + +@dataclass(frozen=True) +class DFlashCorrectnessConfig: + model: str + draft_model: str + expected_accuracy: float + expected_acceptance_len: float + num_speculative_tokens: int = 16 + max_model_len: int = 4096 + max_num_seqs: int = 128 + num_questions: int = 1319 + use_chat_completions: bool = False + enforce_eager: bool = False + disable_flashinfer_sampler: bool = False + + +QWEN3_DFLASH = DFlashCorrectnessConfig( + model="Qwen/Qwen3-8B", + draft_model="z-lab/Qwen3-8B-DFlash-b16", + expected_accuracy=0.8, + expected_acceptance_len=3.5, +) + +LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig( + model="poolside/Laguna-XS-2.1-NVFP4", + draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", + expected_accuracy=0.7, # Standard GSM8K sanity floor. + expected_acceptance_len=3.55 * 0.9, + num_speculative_tokens=15, + max_model_len=8192, + max_num_seqs=32, + num_questions=200, + use_chat_completions=True, + enforce_eager=True, + disable_flashinfer_sampler=True, +) + + +@pytest.mark.parametrize("use_mrv2", [False, True]) +def test_dflash_reference_acceptance_lengths( + monkeypatch: pytest.MonkeyPatch, + use_mrv2: bool, +): + run_acceptance_length_eval( + monkeypatch, + spec_config={ + "model": QWEN3_DFLASH.model, + "trust_remote_code": True, + "speculative_config": { + "method": "dflash", + "model": QWEN3_DFLASH.draft_model, + "num_speculative_tokens": QWEN3_DFLASH.num_speculative_tokens, + "max_model_len": 32768, + }, + "max_model_len": 32768, + "max_num_seqs": 128, + "gpu_memory_utilization": 0.85, + "enforce_eager": False, + "disable_log_stats": False, + }, + # Table 1 in https://arxiv.org/pdf/2602.06036. + expected_acceptance_lengths={ + "mt-bench": 4.24, + "humaneval": 6.50, + "gsm8k": 6.54 * 0.975, + }, + chat_template_kwargs={"enable_thinking": False}, + use_mrv2=use_mrv2, + ) + + +@single_gpu_only +@pytest.mark.parametrize( + ("config", "use_mrv2"), + [ + pytest.param( + QWEN3_DFLASH, + False, + id="qwen3-mrv1", + ), + pytest.param( + QWEN3_DFLASH, + True, + id="qwen3-mrv2", + ), + pytest.param( + LAGUNA_DFLASH_NVFP4, + True, + id="laguna-nvfp4-mrv2", + ), + ], +) +def test_dflash_correctness( + monkeypatch: pytest.MonkeyPatch, + config: DFlashCorrectnessConfig, + use_mrv2: bool, +): + """Guard GSM8K accuracy and batched acceptance length for DFlash models.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") + if config.disable_flashinfer_sampler: + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + + spec_llm = LLM( + model=config.model, + trust_remote_code=True, + speculative_config={ + "method": "dflash", + "model": config.draft_model, + "num_speculative_tokens": config.num_speculative_tokens, + "max_model_len": config.max_model_len, + }, + max_model_len=config.max_model_len, + max_num_seqs=config.max_num_seqs, + gpu_memory_utilization=0.85, + enforce_eager=config.enforce_eager, + disable_log_stats=False, + ) + + results = evaluate_gsm8k_offline( + spec_llm, + num_questions=config.num_questions, + use_chat_completions=config.use_chat_completions, + ) + accuracy = results["accuracy"] + acceptance_len = compute_acceptance_len(spec_llm.get_metrics()) + print( + f"{config.model}: GSM8K accuracy={accuracy:.3f}, " + f"acceptance_len={acceptance_len:.2f}" + ) + + assert accuracy >= config.expected_accuracy + assert acceptance_len >= config.expected_acceptance_len + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py new file mode 100644 index 000000000000..0827d80f669d --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline +from vllm import LLM +from vllm.distributed import cleanup_dist_env_and_memory + +from ...utils import compute_acceptance_len, compute_acceptance_rate + + +def test_gemma4_dspark_correctness_and_acceptance_rate( + monkeypatch: pytest.MonkeyPatch, +): + """ + E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length + regression coverage plus GSM8K correctness, at temperature=1.0 to exercise + the probabilistic draft-sampling/rejection-sampling path (not just greedy). + + gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template + (use_chat_completions=True; raw few-shot completion collapses to a few + percent). Reference: measured over 5 runs of 200 GSM8K questions at + temperature=1.0 (prefix caching disabled): + accuracy: min=0.900 max=0.955 mean=0.937 + acceptance_rate: min=0.578 max=0.595 mean=0.588 + acceptance_len: min=5.044 max=5.167 mean=5.116 + Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling + """ + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + + spec_llm = LLM( + model="RedHatAI/gemma-4-12B-it-NVFP4", + trust_remote_code=True, + speculative_config={ + "method": "dspark", + "model": "deepseek-ai/dspark_gemma4_12b_block7", + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + }, + max_model_len=4096, + max_num_seqs=32, + gpu_memory_utilization=0.85, + enforce_eager=True, + enable_prefix_caching=False, + disable_log_stats=False, + ) + try: + results = evaluate_gsm8k_offline( + spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True + ) + gsm8k_accuracy = results["accuracy"] + + metrics = spec_llm.get_metrics() + acceptance_rate = compute_acceptance_rate(metrics) + acceptance_len = compute_acceptance_len(metrics) + print( + f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, " + f"acceptance_len={acceptance_len:.2f}, " + f"gsm8k_accuracy={gsm8k_accuracy:.3f}" + ) + + assert acceptance_rate >= 0.58 * 0.9 + assert acceptance_len >= 5.0 * 0.9 + assert gsm8k_accuracy >= 0.92 * 0.9 + finally: + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + +@pytest.fixture +def dspark_config(): + target_model = "Qwen/Qwen3-4B-FP8" + draft_model = "deepseek-ai/dspark_qwen3_4b_block7" + + return dict( + model=target_model, + trust_remote_code=True, + speculative_config={ + "method": "dspark", + "model": draft_model, + "num_speculative_tokens": 7, + "attention_backend": "FLASH_ATTN", + "draft_sample_method": "probabilistic", + }, + max_model_len=4096, + disable_log_stats=False, + ) + + +def test_dspark_correctness_and_acceptance_rate(dspark_config): + """ + E2E test for DSpark speculative decoding: acceptance rate/length + regression coverage plus GSM8K correctness, at temperature=1.0 to + exercise the probabilistic draft-sampling/rejection-sampling path + (not just greedy). + + Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft + model. Reference: measured over 12 runs of the full GSM8K set at + temperature=1.0 (prefix caching disabled to avoid cross-run reuse): + accuracy: min=0.782 max=0.814 mean=0.801 + acceptance_rate: min=0.418 max=0.434 mean=0.428 + acceptance_len: min=3.928 max=4.037 mean=3.994 + Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling + """ + spec_llm = LLM(**dspark_config) + + results = evaluate_gsm8k_offline(spec_llm, temperature=1.0) + gsm8k_accuracy = results["accuracy"] + + metrics = spec_llm.get_metrics() + acceptance_rate = compute_acceptance_rate(metrics) + acceptance_len = compute_acceptance_len(metrics) + + print( + f"DSpark acceptance_rate={acceptance_rate:.2f}, " + f"acceptance_len={acceptance_len:.2f}, " + f"gsm8k_accuracy={gsm8k_accuracy:.3f}" + ) + + assert acceptance_rate >= 0.428 * 0.9 + assert acceptance_len >= 3.994 * 0.9 + assert gsm8k_accuracy >= 0.801 * 0.9 + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py new file mode 100644 index 000000000000..cd9c115a9dc3 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory + +from ...utils import compute_acceptance_rate + + +def test_medusa_acceptance_rate( + sampling_config: SamplingParams, +): + """Verify a trained Medusa checkpoint achieves nonzero acceptance rate. + + Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the + speculation path actually accepts tokens — unlike test_medusa_correctness, + which uses a random head and only validates output correctness. + """ + target_model = "lmsys/vicuna-7b-v1.3" + medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3" + prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0] + + spec_llm = LLM( + model=target_model, + speculative_config={ + "method": "medusa", + "model": medusa_model, + "num_speculative_tokens": 3, + }, + max_model_len=1024, + enforce_eager=True, + disable_log_stats=False, + ) + spec_llm.generate(prompts, sampling_config) + metrics = spec_llm.get_metrics() + acceptance_rate = compute_acceptance_rate(metrics) + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + min_acceptance_rate = 0.198 + print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})") + + # Regression guard at 90% of the measured baseline. + assert acceptance_rate >= min_acceptance_rate, ( + f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}" + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py new file mode 100644 index 000000000000..831501361ac7 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from ..utils import run_acceptance_length_eval + + +@pytest.mark.parametrize("use_mrv2", [False, True]) +def test_gemma4_mtp_acceptance_lengths( + monkeypatch: pytest.MonkeyPatch, + use_mrv2: bool, +): + run_acceptance_length_eval( + monkeypatch, + spec_config={ + "model": "google/gemma-4-E4B-it", + "trust_remote_code": True, + "speculative_config": { + "method": "mtp", + "model": "google/gemma-4-E4B-it-assistant", + "num_speculative_tokens": 2, + "max_model_len": 32768, + }, + "max_model_len": 32768, + "limit_mm_per_prompt": {"image": 0, "audio": 0}, + "disable_log_stats": False, + }, + expected_acceptance_lengths={ + "mt-bench": 2.28, + "humaneval": 2.68, + "gsm8k": 2.67 * 0.975, + }, + chat_template_kwargs={}, + use_mrv2=use_mrv2, + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py new file mode 100644 index 000000000000..aef0805513ad --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.utils import single_gpu_only +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory + +from ...utils import compute_acceptance_len, get_test_prompts + + +@single_gpu_only +def test_synthetic_acceptance_rate(): + """Verify that synthetic rejection sampling produces an acceptance + length close to the requested mean acceptance length.""" + num_spec_tokens = 3 + expected_acceptance_len = 1.875 + tolerance = 0.15 + + spec_llm = LLM( + model="meta-llama/Llama-3.2-1B-Instruct", + trust_remote_code=True, + speculative_config={ + "method": "eagle3", + "model": "nm-testing/Llama3_2_1B_speculator.eagle3", + "num_speculative_tokens": num_spec_tokens, + "max_model_len": 2048, + "rejection_sample_method": "synthetic", + "synthetic_acceptance_length": expected_acceptance_len, + }, + max_model_len=2048, + enforce_eager=True, + disable_log_stats=False, + ) + + test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50) + spec_llm.chat( + test_prompts, + SamplingParams(temperature=0, max_tokens=64, ignore_eos=True), + ) + + metrics = spec_llm.get_metrics() + acceptance_len = compute_acceptance_len(metrics) + + print( + f"Synthetic acceptance length: {acceptance_len:.3f}" + f" (expected={expected_acceptance_len:.3f}," + f" tolerance=±{tolerance})" + ) + assert abs(acceptance_len - expected_acceptance_len) <= tolerance, ( + f"Synthetic acceptance length {acceptance_len:.3f} is not within" + f" ±{tolerance} of expected {expected_acceptance_len:.3f}" + ) + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py new file mode 100644 index 000000000000..631dc95b8186 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +import torch +from tqdm import tqdm + +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory + +from ..utils import compute_acceptance_len + + +def load_and_process_dataset(data_name: str): + from datasets import load_dataset + + if data_name == "gsm8k": + dataset = load_dataset("openai/gsm8k", "main", split="test") + prompt_fmt = ( + "{question}\nPlease reason step by step," + " and put your final answer within \\boxed{{}}." + ) + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + elif data_name == "mt-bench": + dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") + dataset = dataset.map(lambda x: {"turns": x["prompt"]}) + elif data_name == "humaneval": + dataset = load_dataset("openai/openai_humaneval", split="test") + prompt_fmt = ( + "Write a solution to the following problem and make sure" + " that it passes the tests:\n```python\n{prompt}\n```" + ) + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + return dataset + + +def run_acceptance_length_eval( + monkeypatch: pytest.MonkeyPatch, + spec_config: dict[str, Any], + expected_acceptance_lengths: dict[str, float], + chat_template_kwargs: dict[str, Any], + use_mrv2: bool, +): + """ + E2E acceptance-rate validation for speculative decoding. + + Drives one or more datasets (keyed in ``expected_acceptance_lengths``) + through the spec decode engine and asserts the mean acceptance length + stays within tolerance of the reference figure for each dataset. + """ + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") + spec_llm = LLM(**spec_config) + + max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k + + tokenizer = spec_llm.get_tokenizer() + for dataset_name, expected_len in expected_acceptance_lengths.items(): + dataset = load_and_process_dataset(dataset_name) + prev_metrics = None + acceptance_lengths = [] + for i in tqdm( + range(min(max_prompts_per_dataset, len(dataset))), + desc=f"Processing {dataset_name}", + ): + user_content = dataset[i]["turns"][0] + prompt_text = tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + **chat_template_kwargs, + ) + + # Greedy (temp=0) so acceptance length is deterministic and comparable + # across runs. + spec_llm.generate( + [prompt_text], + SamplingParams(temperature=0, max_tokens=2048), + use_tqdm=False, + ) + current_metrics = spec_llm.get_metrics() + acceptance_len = compute_acceptance_len(current_metrics, prev_metrics) + prev_metrics = current_metrics + acceptance_lengths.append(acceptance_len) + + mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) + # Fairly tight tolerance of 95% against the reference figures, + # watching for regressions. Can be relaxed if test is flaky but be sure to + # check for genuine issues such as #40727. + expected_len = expected_len * 0.95 + print( + f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" + f" (expected at least {expected_len:.2f})" + ) + + assert mean_acceptance_length >= expected_len, ( + f"acceptance_len for {dataset_name} is below expected threshold: " + f"{mean_acceptance_length:.2f} < {expected_len:.2f}" + ) + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/conftest.py b/tests/v1/e2e/spec_decode/conftest.py new file mode 100644 index 000000000000..969450ad96d1 --- /dev/null +++ b/tests/v1/e2e/spec_decode/conftest.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm import SamplingParams + +from .utils import greedy_sampling + + +@pytest.fixture +def sampling_config() -> SamplingParams: + return greedy_sampling() + + +@pytest.fixture +def model_name() -> str: + return "meta-llama/Llama-3.1-8B-Instruct" + + +@pytest.fixture(autouse=True) +def reset_torch_dynamo(): + yield + torch._dynamo.reset() diff --git a/tests/v1/e2e/spec_decode/draft_model/__init__.py b/tests/v1/e2e/spec_decode/draft_model/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/e2e/spec_decode/test_async_spec_decode.py b/tests/v1/e2e/spec_decode/draft_model/test_async.py similarity index 100% rename from tests/v1/e2e/spec_decode/test_async_spec_decode.py rename to tests/v1/e2e/spec_decode/draft_model/test_async.py diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py new file mode 100644 index 000000000000..94e32b2375a6 --- /dev/null +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +import pytest +import torch + +from tests.utils import multi_gpu_only, single_gpu_only +from vllm import LLM, SamplingParams +from vllm.config import VllmConfig, replace +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.engine.arg_utils import EngineArgs + +from ..utils import ( + Messages, + _skip_if_insufficient_gpus_for_tp, + compute_acceptance_len, + compute_acceptance_rate, + evaluate_llm_for_gsm8k, + get_instruct_coder_messages, + get_test_prompts, + greedy_sampling, + stochastic_sampling, +) + + +class AsyncSchedulingNotEnabledError(AssertionError): + """Raised when draft-model spec decode does not enable async scheduling.""" + + +@dataclass +class ArgsTest: + target_model: str + draft_model: str + sampling_config: SamplingParams + num_speculative_tokens: int + expected_acceptance_rate: float + expected_acceptance_len: float + expected_gsm8k_accuracy: float = 0.0 # skip by default + # Defaults + enforce_eager: bool = True + parallel_drafting: bool = False + target_tensor_parallel_size: int = 1 + draft_tensor_parallel_size: int = 1 + max_model_len: int = 2048 + gpu_memory_utilization: float = 0.5 + dataset: str = "test_prompts" + num_prompts: int = 100 + + +def get_messages(dataset: str, n: int) -> list[Messages]: + if dataset == "test_prompts": + return get_test_prompts(mm_enabled=False, num_prompts=n) + if dataset == "likaixin/InstructCoder": + return get_instruct_coder_messages(n=n) + raise NotImplementedError(f"Dataset '{dataset}' not implemented") + + +def some_high_acceptance_metrics() -> dict: + return { + "sampling_config": greedy_sampling(), + "num_speculative_tokens": 3, + "expected_acceptance_len": 3.4, # ref: 3.75 + "expected_acceptance_rate": 0.8, # ref: 0.9 + } + + +cases = [ + # Same model for draft and target, greedy sampling. + ArgsTest( + target_model="Qwen/Qwen3-0.6B", + draft_model="Qwen/Qwen3-0.6B", + sampling_config=greedy_sampling(), + num_speculative_tokens=3, # K + expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1 + expected_acceptance_rate=0.98, # slight epsilon + expected_gsm8k_accuracy=0.25, # ref: 35-40% + ), + # Smaller draft model, stochastic sampling. + ArgsTest( + target_model="Qwen/Qwen3-1.7B", + draft_model="Qwen/Qwen3-0.6B", + sampling_config=stochastic_sampling(), + num_speculative_tokens=3, + expected_acceptance_len=3.4, # ref: 3.7 + expected_acceptance_rate=0.80, # ref: 0.90 + expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling + ), +] + + +@pytest.mark.parametrize("args", cases) +@pytest.mark.parametrize("enforce_eager", [True, False]) +@single_gpu_only +# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 +@pytest.mark.xfail( + raises=AsyncSchedulingNotEnabledError, + reason="draft_model does not yet enable async_scheduling: issue #38929", +) +def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool): + args.enforce_eager = enforce_eager + assert_draft_model_correctness(args) + + +@single_gpu_only +# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 +@pytest.mark.xfail( + raises=AsyncSchedulingNotEnabledError, + reason="draft_model does not yet enable async_scheduling: issue #38929", +) +def test_draft_model_realistic_example(): + args = ArgsTest( + target_model="Qwen/Qwen3-1.7B", + draft_model="Qwen/Qwen3-0.6B", + dataset="likaixin/InstructCoder", + num_speculative_tokens=3, + sampling_config=greedy_sampling(), + enforce_eager=False, + expected_acceptance_len=2.6, # ref: 2.86 + expected_acceptance_rate=0.5, # ref: 0.62 + ) + assert_draft_model_correctness(args) + + +@single_gpu_only +# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 +@pytest.mark.xfail( + raises=AsyncSchedulingNotEnabledError, + reason="draft_model does not yet enable async_scheduling: issue #38929", +) +def test_draft_model_parallel_drafting(): + args = ArgsTest( + target_model="Qwen/Qwen3-1.7B", + draft_model="amd/PARD-Qwen3-0.6B", + dataset="likaixin/InstructCoder", + num_speculative_tokens=3, + sampling_config=greedy_sampling(), + parallel_drafting=True, + enforce_eager=False, + expected_acceptance_len=2.3, # ref: 2.52 + expected_acceptance_rate=0.4, # ref: 0.51 + ) + assert_draft_model_correctness(args) + + +@pytest.mark.parametrize( + "models", + [ + # target_model, draft_model + ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized + ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized + ], + ids=["target_quantized", "draft_quantized"], +) +@pytest.mark.parametrize("enforce_eager", [True, False]) +@single_gpu_only +# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 +@pytest.mark.xfail( + raises=AsyncSchedulingNotEnabledError, + reason="draft_model does not yet enable async_scheduling: issue #38929", +) +def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): + tgt_model, draft_model = models + sd_case = ArgsTest( + target_model=tgt_model, + draft_model=draft_model, + **some_high_acceptance_metrics(), + enforce_eager=enforce_eager, + ) + assert_draft_model_correctness(sd_case) + + +@multi_gpu_only(num_gpus=2) +# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 +@pytest.mark.xfail( + raises=AsyncSchedulingNotEnabledError, + reason="draft_model does not yet enable async_scheduling: issue #38929", +) +def test_draft_model_tensor_parallelism(): + """Ensure spec decode works when running with TP > 1.""" + _skip_if_insufficient_gpus_for_tp(2) + sd_case = ArgsTest( + target_model="Qwen/Qwen3-1.7B", + target_tensor_parallel_size=2, + draft_model="Qwen/Qwen3-0.6B", + draft_tensor_parallel_size=2, + **some_high_acceptance_metrics(), + enforce_eager=False, + expected_gsm8k_accuracy=0.5, + ) + assert_draft_model_correctness(sd_case) + + +@multi_gpu_only(num_gpus=2) +def test_draft_model_engine_args_tensor_parallelism(): + """Ensure the vllm_config for the draft model is created correctly, + and independently of the target model (quantization, TP, etc.)""" + _skip_if_insufficient_gpus_for_tp(2) + + engine_args = EngineArgs( + model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized + tensor_parallel_size=2, + speculative_config={ + "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized + "method": "draft_model", + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 1, # <<< valid arg name + }, + ) + target_config: VllmConfig = engine_args.create_engine_config() + assert target_config.parallel_config.tensor_parallel_size == 2 + assert target_config.quant_config.get_name() == "fp8" + + speculative_config = target_config.speculative_config + draft_config: VllmConfig = replace( + target_config, + quant_config=None, + parallel_config=replace( + speculative_config.draft_parallel_config, + rank=target_config.parallel_config.rank, + ), + model_config=speculative_config.draft_model_config, + ) + assert draft_config.parallel_config.tensor_parallel_size == 1 + assert draft_config.quant_config is None + + +def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: + """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic + so we can test it without instantiating a full proposer.""" + spec_cfg = vllm_config.speculative_config + if spec_cfg.moe_backend is not None: + return replace( + vllm_config, + kernel_config=replace( + vllm_config.kernel_config, + moe_backend=spec_cfg.moe_backend, + ), + ) + return vllm_config + + +def test_draft_model_moe_backend_override(): + """When moe_backend is set in speculative_config, the draft VllmConfig + should use it while the target keeps its own setting.""" + engine_args = EngineArgs( + model="Qwen/Qwen3-1.7B", + tensor_parallel_size=1, + moe_backend="flashinfer_trtllm", + speculative_config={ + "model": "Qwen/Qwen3-0.6B", + "method": "draft_model", + "num_speculative_tokens": 3, + "moe_backend": "triton", + }, + ) + tgt_config: VllmConfig = engine_args.create_engine_config() + assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" + assert tgt_config.speculative_config.moe_backend == "triton" + + draft_config = _apply_draft_moe_backend(tgt_config) + assert draft_config.kernel_config.moe_backend == "triton" + # Target config must be unaffected. + assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" + + +def test_draft_model_moe_backend_inherits_target(): + """When moe_backend is not set in speculative_config, the draft should + inherit the target's moe_backend.""" + engine_args = EngineArgs( + model="Qwen/Qwen3-1.7B", + tensor_parallel_size=1, + moe_backend="flashinfer_cutlass", + speculative_config={ + "model": "Qwen/Qwen3-0.6B", + "method": "draft_model", + "num_speculative_tokens": 3, + }, + ) + tgt_config: VllmConfig = engine_args.create_engine_config() + assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass" + assert tgt_config.speculative_config.moe_backend is None + + draft_config = _apply_draft_moe_backend(tgt_config) + assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass" + assert draft_config is tgt_config + + +def test_draft_model_moe_backend_default_auto(): + """When neither target nor draft set moe_backend explicitly, both should + default to 'auto'.""" + engine_args = EngineArgs( + model="Qwen/Qwen3-1.7B", + tensor_parallel_size=1, + speculative_config={ + "model": "Qwen/Qwen3-0.6B", + "method": "draft_model", + "num_speculative_tokens": 3, + }, + ) + tgt_config: VllmConfig = engine_args.create_engine_config() + assert tgt_config.kernel_config.moe_backend == "auto" + assert tgt_config.speculative_config.moe_backend is None + + draft_config = _apply_draft_moe_backend(tgt_config) + assert draft_config.kernel_config.moe_backend == "auto" + assert draft_config is tgt_config + + +def test_draft_model_engine_args_rejects_invalid_tp_argname(): + """The user should pass "draft_tensor_parallel_size" rather than + "tensor_parallel_size". We enforce this with validation.""" + + engine_args = EngineArgs( + model="Qwen/Qwen3-1.7B", + tensor_parallel_size=1, + speculative_config={ + "model": "Qwen/Qwen3-0.6B", + "method": "draft_model", + "num_speculative_tokens": 3, + "tensor_parallel_size": 1, # <<< invalid arg name + }, + ) + with pytest.raises(ValueError): + engine_args.create_engine_config() + + +def assert_draft_model_correctness(args: ArgsTest): + """Compare the outputs using and not using speculative decoding. + In the greedy decoding case, the outputs must match EXACTLY.""" + test_prompts: list[Messages] = get_messages( + dataset=args.dataset, n=args.num_prompts + ) + + spec_llm = LLM( + model=args.target_model, + speculative_config={ + "model": args.draft_model, + "method": "draft_model", + "num_speculative_tokens": args.num_speculative_tokens, + "max_model_len": args.max_model_len, + "enforce_eager": args.enforce_eager, + "draft_tensor_parallel_size": args.draft_tensor_parallel_size, + "parallel_drafting": args.parallel_drafting, + }, + max_num_seqs=100, # limit cudagraph capture runtime + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + tensor_parallel_size=args.target_tensor_parallel_size, + enforce_eager=args.enforce_eager, + disable_log_stats=False, # enables get_metrics() + ) + + # we don't check the outputs, only check the metrics + spec_llm.chat(test_prompts, args.sampling_config) + metrics = spec_llm.get_metrics() + acceptance_rate: float = compute_acceptance_rate(metrics) + acceptance_len: float = compute_acceptance_len(metrics) + + # Need to evaluate after getting metrics to avoid polluting the AR + evaluate_llm_for_gsm8k( + spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy + ) + + print( + f"spec-decode: target={args.target_model}, draft={args.draft_model}, " + f"temperature={args.sampling_config.temperature:.2f}, " + f"acceptance_rate={acceptance_rate:.2f}, " + f"acceptance_len={acceptance_len:.2f}, " + ) + + assert acceptance_rate >= args.expected_acceptance_rate + assert acceptance_len >= args.expected_acceptance_len + # draft_model supports async scheduling; assert it is active by default. + # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that + # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this + # specific failure — leaving all other assertion failures (e.g. correctness or + # acceptance-rate checks above) visible as real test failures. + has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling + del spec_llm # CLEANUP + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + if not has_async: + raise AsyncSchedulingNotEnabledError( + "Expected async_scheduling=True for draft_model spec decode, got False." + " See https://github.com/vllm-project/vllm/issues/38929" + ) diff --git a/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py b/tests/v1/e2e/spec_decode/draft_model/test_lora.py similarity index 100% rename from tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py rename to tests/v1/e2e/spec_decode/draft_model/test_lora.py diff --git a/tests/v1/e2e/spec_decode/eagle/__init__.py b/tests/v1/e2e/spec_decode/eagle/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py new file mode 100644 index 000000000000..c5c32d80e594 --- /dev/null +++ b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.utils import ( + get_attn_backend_list_based_on_platform, + single_gpu_only, +) +from vllm import SamplingParams +from vllm.platforms import current_platform + +from .utils import _run_eagle_correctness + + +@single_gpu_only +@pytest.mark.skipif( + current_platform.is_device_capability_family(100), + reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)", +) +@pytest.mark.parametrize( + [ + "model_setup", + "mm_enabled", + "enable_chunked_prefill", + "model_impl", + "expected_accuracy_threshold", + ], + [ + ( + ( + "eagle", + "eagle618/deepseek-v3-random", + "eagle618/eagle-deepseek-v3-random", + 1, + ), + False, + False, + "auto", + 0.0, + ), + ], + ids=["deepseek_eagle"], +) +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) +def test_eagle_correctness_light( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + _run_eagle_correctness( + monkeypatch, + sampling_config, + model_setup, + mm_enabled, + expected_accuracy_threshold, + enable_chunked_prefill, + model_impl, + attn_backend, + ) + + +@single_gpu_only +@pytest.mark.parametrize( + [ + "model_setup", + "mm_enabled", + "enable_chunked_prefill", + "model_impl", + "expected_accuracy_threshold", + ], + [ + ( + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), + False, + False, + "auto", + 0.8, + ), + pytest.param( + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), + False, + False, + "transformers", + 0.8, + # TODO(hmellor): figure out why memory usage is so high + marks=pytest.mark.skip( + reason="Feature is experimental and uses too much memory in CI", + ), + ), + pytest.param( + ( + "eagle3", + "Qwen/Qwen3-VL-8B-Instruct", + "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3", + 1, + ), + False, + False, + "auto", + 0.8, + marks=pytest.mark.skip( + reason="architecture of its eagle3 is LlamaForCausalLMEagle3" + ), + ), + pytest.param( + ( + "eagle3", + "Qwen/Qwen2.5-VL-7B-Instruct", + "Rayzl/qwen2.5-vl-7b-eagle3-sgl", + 1, + ), + False, + False, + "auto", + 0.7, + marks=pytest.mark.skip( + reason="Skipping due to its head_dim not being a multiple of 32" + ), + ), + ( + ( + "eagle3", + "meta-llama/Llama-3.1-8B-Instruct", + "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + 1, + ), + False, + False, + "auto", + 0.7, + ), + ], + ids=[ + "qwen3_eagle3", + "qwen3_eagle3-transformers", + "qwen3_vl_eagle3", + "qwen2_5_vl_eagle3", + "llama3_eagle3", + ], +) +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) +def test_eagle_correctness_medium( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + _run_eagle_correctness( + monkeypatch, + sampling_config, + model_setup, + mm_enabled, + expected_accuracy_threshold, + enable_chunked_prefill, + model_impl, + attn_backend, + ) diff --git a/tests/v1/e2e/spec_decode/eagle/utils.py b/tests/v1/e2e/spec_decode/eagle/utils.py new file mode 100644 index 000000000000..998be71af685 --- /dev/null +++ b/tests/v1/e2e/spec_decode/eagle/utils.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from tests.utils import wait_for_rocm_memory_to_settle +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform + +from ..utils import ( + _skip_if_insufficient_gpus_for_tp, + evaluate_llm_for_gsm8k, + get_test_prompts, +) + + +def _run_eagle_correctness( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + """ + Compare the outputs of an original LLM and a speculative LLM + which should be the same when using eagle speculative decoding. + """ + if model_impl == "transformers": + import transformers + from packaging.version import Version + + installed = Version(transformers.__version__) + required = Version("5.0.0") + if installed < required: + pytest.skip( + "Eagle3 with the Transformers modeling backend requires " + f"transformers>={required}, but got {installed}" + ) + + test_prompts = get_test_prompts(mm_enabled) + + if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN": + if current_platform.is_rocm(): + print( + "FLASH_ATTN for spec_decode not supported on " + "ROCm currently. Changing to FLEX_ATTENTION backend." + ) + attention_config = {"backend": "FLEX_ATTENTION"} + else: + attention_config = None + else: + attention_config = {"backend": attn_backend} + + if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm(): + pytest.skip( + "TRITON_ATTN does not support " + "multi-token eagle spec decode on current platform" + ) + + with monkeypatch.context() as m: + m.setenv("VLLM_MLA_DISABLE", "1") + + if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): + if "deepseek" in model_setup[1].lower(): + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.delenv("VLLM_MLA_DISABLE", raising=False) + attention_config = {"backend": "ROCM_AITER_MLA"} + else: + m.setenv("VLLM_ROCM_USE_AITER", "1") + + method, model_name, spec_model_name, tp_size = model_setup + _skip_if_insufficient_gpus_for_tp(tp_size) + + max_model_len = 2048 + max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len + + ref_llm = LLM( + model=model_name, + max_model_len=max_model_len, + tensor_parallel_size=tp_size, + attention_config=attention_config, + ) + evaluate_llm_for_gsm8k( + ref_llm, expected_accuracy_threshold=expected_accuracy_threshold + ) + ref_outputs = ref_llm.chat(test_prompts, sampling_config) + del ref_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + # ROCm frees VRAM lazily; wait so the spec engine started right after + # does not OOM on its startup memory guard. + wait_for_rocm_memory_to_settle() + + spec_llm = LLM( + model=model_name, + trust_remote_code=True, + tensor_parallel_size=tp_size, + speculative_config={ + "method": method, + "model": spec_model_name, + "num_speculative_tokens": 3, + "max_model_len": max_model_len, + }, + max_model_len=max_model_len, + max_num_batched_tokens=max_num_batched_tokens, + enable_chunked_prefill=enable_chunked_prefill, + model_impl=model_impl, + attention_config=attention_config, + ) + # EAGLE/EAGLE3 supports async scheduling; assert it is active by default. + assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling + evaluate_llm_for_gsm8k( + spec_llm, expected_accuracy_threshold=expected_accuracy_threshold + ) + spec_outputs = spec_llm.chat(test_prompts, sampling_config) + matches = 0 + misses = 0 + for ref_output, spec_output in zip(ref_outputs, spec_outputs): + if ref_output.outputs[0].text == spec_output.outputs[0].text: + matches += 1 + else: + misses += 1 + print(f"ref_output: {ref_output.outputs[0].text}") + print(f"spec_output: {spec_output.outputs[0].text}") + + assert matches > int(0.6 * len(ref_outputs)) + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + # ROCm frees VRAM lazily; wait so the next parametrization's engine does + # not OOM on its startup memory guard. + wait_for_rocm_memory_to_settle() diff --git a/tests/v1/e2e/spec_decode/mtp/__init__.py b/tests/v1/e2e/spec_decode/mtp/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/e2e/spec_decode/mtp/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/test_mtp.py new file mode 100644 index 000000000000..45a82367949c --- /dev/null +++ b/tests/v1/e2e/spec_decode/mtp/test_mtp.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +from typing import Any + +import pytest +import torch + +from tests.utils import single_gpu_only +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform + +from ..utils import ( + _skip_if_insufficient_gpus_for_tp, + evaluate_llm_for_gsm8k, + get_test_prompts, +) + + +@pytest.mark.parametrize( + ["model_setup", "mm_enabled", "expected_accuracy_threshold"], + [ + (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70% + pytest.param( + ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), + False, + 0.0, + marks=pytest.mark.skipif( + current_platform.is_device_capability_family(100), + reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell", + ), + ), # dummy model + ( + ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1), + False, + 0.20, + ), # hybrid + MTP, ref: ~34%-35% + ( + ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"), + False, + 0.50, + ), # gemma4 MTP with assistant model, ref: ~62% + ], + ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"], +) +@single_gpu_only +def test_mtp_correctness( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, int] | tuple[str, str, int, str], + mm_enabled: bool, + expected_accuracy_threshold: float, +): + """ + Compare the outputs of a original LLM and a speculative LLM + which should be the same when using MTP speculative decoding. Due to some variance + in the engine, it is possible for some outputs to differ, so we expect that at least + 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed + reference threshold for each model. + """ + # Generate test prompts inside the function instead of using fixture + test_prompts = get_test_prompts(mm_enabled) + with monkeypatch.context() as m: + m.setenv("VLLM_MLA_DISABLE", "1") + + if len(model_setup) == 4: + method, model_name, tp_size, draft_model = model_setup + else: + method, model_name, tp_size = model_setup + draft_model = None + _skip_if_insufficient_gpus_for_tp(tp_size) + + if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"): + pytest.skip( + "Model Runner V2 does not yet support hybrid models " + "(Qwen3.5 mixes Mamba-style GDN with attention layers)." + ) + + attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto" + + # Skip multimodal profiling for models that don't need it in this test. + extra_kwargs: dict[str, Any] = {} + if "Qwen3.5" in model_name: + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0} + elif "gemma-4" in model_name: + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0} + + if draft_model is not None and "gemma-4" in draft_model: + import transformers + from packaging.version import Version + + if Version(transformers.__version__) < Version("5.8.0"): + pytest.skip( + "Gemma4 MTP assistant requires transformers>=5.8.0, " + f"got {transformers.__version__}" + ) + + ref_llm = LLM( + model=model_name, + max_model_len=2048, + tensor_parallel_size=tp_size, + trust_remote_code=True, + attention_backend=attn_backend, + **extra_kwargs, + ) + ref_outputs = ref_llm.chat(test_prompts, sampling_config) + evaluate_llm_for_gsm8k( + ref_llm, expected_accuracy_threshold=expected_accuracy_threshold + ) + del ref_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + speculative_config: dict[str, Any] = { + "method": method, + "num_speculative_tokens": 1, + "max_model_len": 2048, + } + if draft_model is not None: + speculative_config["model"] = draft_model + speculative_config["num_speculative_tokens"] = 2 + + spec_llm = LLM( + model=model_name, + trust_remote_code=True, + tensor_parallel_size=tp_size, + speculative_config=speculative_config, + max_model_len=2048, + attention_backend=attn_backend, + **extra_kwargs, + ) + # MTP supports async scheduling; assert it is active by default. + assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling + evaluate_llm_for_gsm8k( + spec_llm, expected_accuracy_threshold=expected_accuracy_threshold + ) + spec_outputs = spec_llm.chat(test_prompts, sampling_config) + matches = 0 + misses = 0 + for ref_output, spec_output in zip(ref_outputs, spec_outputs): + if ref_output.outputs[0].text == spec_output.outputs[0].text: + matches += 1 + else: + misses += 1 + print(f"ref_output: {ref_output.outputs[0].text}") + print(f"spec_output: {spec_output.outputs[0].text}") + + # Heuristic: expect at least 80% of the prompts to match exactly + # Upon failure, inspect the outputs to check for inaccuracy. + assert matches > int(0.8 * len(ref_outputs)) + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/__init__.py b/tests/v1/e2e/spec_decode/ngram_suffix/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py new file mode 100644 index 000000000000..b6dc946a98d7 --- /dev/null +++ b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from tests.utils import single_gpu_only +from vllm import LLM, SamplingParams +from vllm.config import CompilationConfig +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform + +from ..utils import evaluate_llm_for_gsm8k, get_test_prompts + + +@pytest.fixture +def disable_vllm_compile_cache_on_rocm(request: pytest.FixtureRequest) -> None: + if current_platform.is_rocm(): + request.getfixturevalue("disable_vllm_compile_cache") + + +@pytest.mark.parametrize( + "speculative_config", + [ + { + "method": "ngram", + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": 3, + }, + { + "method": "suffix", + "suffix_decoding_max_spec_factor": 2.0, + }, + ], +) +@pytest.mark.usefixtures("disable_vllm_compile_cache_on_rocm") +@single_gpu_only +def test_ngram_and_suffix_correctness( + speculative_config: dict, + model_name: str, + vllm_runner, +): + with vllm_runner( + model_name, + # Keep LLM defaults; VllmRunner only provides lifecycle cleanup here. + trust_remote_code=False, + enable_chunked_prefill=None, + speculative_config=speculative_config, + max_model_len=4096, + # Preserve LLM's default compilation/cudagraph configuration. Without + # this, VllmRunner injects its reduced test-only capture sizes. + compilation_config=CompilationConfig(), + ) as runner: + evaluate_llm_for_gsm8k(runner.llm) + + +@pytest.mark.parametrize("async_scheduling", [True], ids=["async"]) +@single_gpu_only +def test_ngram_gpu_default_with_async_scheduling( + async_scheduling: bool, +): + """ + Test ngram_gpu speculative decoding (k=3) correctness with and without + async scheduling, validated via GSM8K accuracy. + Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%). + """ + qwen3_model = "Qwen/Qwen3-8B" + spec_llm = LLM( + model=qwen3_model, + speculative_config={ + "method": "ngram_gpu", + "prompt_lookup_max": 3, + "prompt_lookup_min": 2, + "num_speculative_tokens": 2, + }, + max_model_len=4096, + async_scheduling=async_scheduling, + ) + # Assert the resolved async_scheduling config matches what was requested. + assert ( + spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling + == async_scheduling + ) + evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8) + del spec_llm + cleanup_dist_env_and_memory() + + +@single_gpu_only +def test_suffix_decoding_acceptance( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_name: str, +): + """ + Check that suffix decoding caching takes effect and improves acceptance + lengths and acceptance rates over multiple runs of the same prompts. + """ + test_prompts = get_test_prompts(mm_enabled=False) + + spec_llm = LLM( + model=model_name, + speculative_config={ + "method": "suffix", + "suffix_decoding_max_spec_factor": 2.0, + "suffix_decoding_max_cached_requests": 1000, + }, + max_model_len=1024, + disable_log_stats=False, + ) + + # Run several times and check that the accepted tokens increase. + num_draft = [] + num_accept = [] + for i in range(10): # Run multiple times to warm up the cache. + spec_llm.chat(test_prompts, sampling_config) + # Collect draft and acceptance stats. + metrics = spec_llm.get_metrics() + for metric in metrics: + if metric.name == "vllm:spec_decode_num_draft_tokens": + num_draft.append(metric.value) + if metric.name == "vllm:spec_decode_num_accepted_tokens": + num_accept.append(metric.value) + + # Calculate the acceptance rates for the first and last runs. + first_accept_tokens = num_accept[0] + first_draft_tokens = num_draft[0] + first_accept_rate = first_accept_tokens / first_draft_tokens + + # Take the diff since the stats are cumulative. + last_accept_tokens = num_accept[-1] - num_accept[-2] + last_draft_tokens = num_draft[-1] - num_draft[-2] + last_accept_rate = last_accept_tokens / last_draft_tokens + + # Expect the acceptance length to improve. + assert first_accept_tokens < last_accept_tokens + + # Expect the acceptance rate to improve. + assert first_accept_rate < last_accept_rate + + # Heuristic: expect at least 80.0% acceptance rate at the end. + assert last_accept_rate > 0.80 + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/speculators/__init__.py b/tests/v1/e2e/spec_decode/speculators/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/e2e/spec_decode/speculators/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/speculators/test_speculators.py b/tests/v1/e2e/spec_decode/speculators/test_speculators.py new file mode 100644 index 000000000000..27fa21fd7669 --- /dev/null +++ b/tests/v1/e2e/spec_decode/speculators/test_speculators.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.utils import single_gpu_only +from vllm import SamplingParams +from vllm.config import CompilationConfig + +from ..utils import evaluate_llm_for_gsm8k, get_test_prompts + + +@pytest.mark.parametrize( + ["model_path", "expected_accuracy_threshold"], + [ + ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80% + ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92% + ], + ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"], +) +@single_gpu_only +def test_speculators_model_integration( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_path: str, + expected_accuracy_threshold: float, + vllm_runner, +): + """ + Test that speculators models work with the simplified integration. + + This verifies the `vllm serve ` use case where + speculative config is automatically detected from the model config + without requiring explicit --speculative-config argument. + + Tests: + 1. Speculator model is correctly detected + 2. Verifier model is extracted from speculator config + 3. Speculative decoding is automatically enabled + 4. Text generation works correctly + 5. GSM8k accuracy of the model passes a sanity check when speculative decoding on + 6. Output matches reference (non-speculative) generation + """ + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + # Generate test prompts + test_prompts = get_test_prompts(mm_enabled=False) + + # First run: Direct speculator model (simplified integration) + with vllm_runner( + model_path, + trust_remote_code=False, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + max_model_len=4096, + gpu_memory_utilization=0.92, + ) as spec_runner: + evaluate_llm_for_gsm8k( + spec_runner.llm, expected_accuracy_threshold=expected_accuracy_threshold + ) + spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config) + + # Verify speculative config was auto-detected + assert spec_runner.llm.llm_engine.vllm_config.speculative_config is not None, ( + f"Speculative config should be auto-detected for {model_path}" + ) + + spec_config = spec_runner.llm.llm_engine.vllm_config.speculative_config + assert spec_config.num_speculative_tokens > 0, ( + f"Expected positive speculative tokens, " + f"got {spec_config.num_speculative_tokens}" + ) + + # Verify draft model is set to the speculator model + assert spec_config.model == model_path, ( + f"Draft model should be {model_path}, got {spec_config.model}" + ) + + # Extract verifier model for reference run + verifier_model = spec_runner.llm.llm_engine.vllm_config.model_config.model + + # Second run: Reference without speculative decoding + with vllm_runner( + verifier_model, + trust_remote_code=False, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + max_model_len=4096, + gpu_memory_utilization=0.92, + ) as ref_runner: + ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config) + + # Compare outputs + matches = sum( + 1 + for ref, spec in zip(ref_outputs, spec_outputs) + if ref.outputs[0].text == spec.outputs[0].text + ) + + # Heuristic: expect at least 66% of prompts to match exactly + assert matches >= int(0.66 * len(ref_outputs)), ( + f"Only {matches}/{len(ref_outputs)} outputs matched. " + f"Expected at least {int(0.66 * len(ref_outputs))} matches." + ) diff --git a/tests/v1/e2e/spec_decode/test_laguna_dflash.py b/tests/v1/e2e/spec_decode/test_laguna_dflash.py deleted file mode 100644 index 1ba9b9749f15..000000000000 --- a/tests/v1/e2e/spec_decode/test_laguna_dflash.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from tests.utils import large_gpu_mark -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory - - -def _get_counter(metrics, name: str) -> float: - metric = next((m for m in metrics if m.name == name), None) - assert metric is not None, f"Missing metric: {name}" - return metric.value - - -@pytest.mark.slow_test -@large_gpu_mark(min_gb=80) -def test_laguna_dflash_hf_pair_smoke(monkeypatch): - """Smoke-test the public Laguna XS-2.1 base/DFlash checkpoint pair.""" - monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") - - llm = LLM( - model="poolside/Laguna-XS-2.1", - trust_remote_code=True, - speculative_config={ - "method": "dflash", - "model": "poolside/Laguna-XS-2.1-DFlash", - "num_speculative_tokens": 15, - }, - max_model_len=8192, - max_num_batched_tokens=65536, - max_num_seqs=32, - enforce_eager=True, - disable_log_stats=False, - ) - - try: - outputs = llm.generate( - [ - "What is the capital of the United Kingdom?", - "Write a Python function that returns the square of a number.", - ], - SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True), - ) - assert len(outputs) == 2 - assert all(output.outputs[0].text for output in outputs) - - metrics = llm.get_metrics() - num_drafts = _get_counter(metrics, "vllm:spec_decode_num_drafts") - num_accepted = _get_counter(metrics, "vllm:spec_decode_num_accepted_tokens") - - assert num_drafts > 0 - acceptance_len = 1 + (num_accepted / num_drafts) - assert acceptance_len > 1.0 - finally: - del llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py new file mode 100644 index 000000000000..186faa6529b1 --- /dev/null +++ b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +from contextlib import AsyncExitStack +from dataclasses import dataclass, replace +from typing import Any + +import pytest +import torch + +from tests.utils import multi_gpu_marks +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.reader import Metric + +DEEPSEEK_MTP_MAIN_RANDOM = "luccafong/deepseek_mtp_main_random" +DEEPSEEK_MTP_DRAFT_RANDOM = "luccafong/deepseek_mtp_draft_random" + +PROMPT = "The capital of France is" +MAX_TOKENS = 8 +MAX_MODEL_LEN = 2048 +GPU_MEM_UTIL = 0.85 +# Spec / no-spec greedy output should match, but exact match is not guaranteed +# across parallelism layouts (reduction order can flip near-tie argmaxes), so +# gate on a similarity ratio like the spec-decode E2E tests. +MIN_MATCH_RATIO = 0.8 + + +def _token_match_ratio(a: tuple[int, ...], b: tuple[int, ...]) -> float: + n = min(len(a), len(b)) + return sum(x == y for x, y in zip(a, b)) / n if n else 0.0 + + +@dataclass(frozen=True) +class InlineConfig: + id: str + tp_size: int = 1 + pp_size: int = 1 + enable_expert_parallel: bool = False + enable_eplb: bool = False + skip_reason: str | None = None + + +INLINE_CONFIGS = [ + InlineConfig(id="tp2", tp_size=2), + InlineConfig(id="ep2", tp_size=2, enable_expert_parallel=True), + InlineConfig( + id="ep2_eplb", + tp_size=2, + enable_expert_parallel=True, + enable_eplb=True, + ), + InlineConfig( + id="pp2", + pp_size=2, + skip_reason=( + "DeepSeek MTP pipeline-parallel support is in flight upstream; " + "see https://github.com/vllm-project/vllm/pull/38104" + ), + ), +] + + +def _inline_kwargs(config: InlineConfig, *, with_spec: bool) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "model": DEEPSEEK_MTP_MAIN_RANDOM, + "tensor_parallel_size": config.tp_size, + "pipeline_parallel_size": config.pp_size, + "enable_expert_parallel": config.enable_expert_parallel, + "max_model_len": MAX_MODEL_LEN, + "gpu_memory_utilization": GPU_MEM_UTIL, + "enforce_eager": True, + "trust_remote_code": True, + "disable_log_stats": False, + # MLA + batch invariance currently disables prefix caching at runtime + # (see vllm/v1/attention/backends/mla/common.py); set explicitly so the + # configuration is unambiguous. + "enable_prefix_caching": False, + } + if config.enable_eplb: + kwargs["enable_eplb"] = True + # Rearrangement first fires after step_interval // 4 forward steps + # (the step counter starts at 3/4 of step_interval, see + # vllm/distributed/eplb/eplb_state.py). Keep these small so the EPLB + # routine actually runs within the short MAX_TOKENS generation instead + # of never triggering. + kwargs["eplb_config"] = { + "num_redundant_experts": config.tp_size, + "window_size": 2, + "step_interval": 4, + "log_balancedness": False, + } + if with_spec: + kwargs["speculative_config"] = { + "method": "mtp", + "model": DEEPSEEK_MTP_DRAFT_RANDOM, + "num_speculative_tokens": 1, + } + return kwargs + + +def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]: + outputs = llm.generate( + [PROMPT], + SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True), + ) + assert outputs and outputs[0].outputs, "expected one completion" + return tuple(outputs[0].outputs[0].token_ids) + + +def _spec_decode_num_drafts(metrics: list[Metric]) -> int: + name2metric = {m.name: m for m in metrics} + counter = name2metric.get("vllm:spec_decode_num_drafts") + assert counter is not None, ( + "spec_decode_num_drafts metric missing; check disable_log_stats=False" + ) + return int(counter.value) + + +@pytest.mark.parametrize( + "config", + [ + pytest.param(c, id=c.id, marks=multi_gpu_marks(num_gpus=2)) + for c in INLINE_CONFIGS + ], +) +def test_deepseek_mtp_load_inline( + monkeypatch: pytest.MonkeyPatch, + config: InlineConfig, +): + """MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy.""" + if config.skip_reason is not None: + pytest.skip(config.skip_reason) + + # Reduces run-to-run nondeterminism so spec / no-spec stay close; + # see tests/v1/distributed/test_eagle_dp.py for the same pattern. + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + + spec_llm = LLM(**_inline_kwargs(config, with_spec=True)) + try: + spec_tokens = _generate_token_ids_inline(spec_llm) + n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics()) + finally: + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False)) + try: + no_spec_tokens = _generate_token_ids_inline(no_spec_llm) + finally: + del no_spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + # Non-vacuity: a silently-broken MTP drafter that falls back to + # verifier-only decoding would still produce matching output below. + assert n_drafts > 0, ( + f"MTP drafter never fired under {config.id}: vllm:spec_decode_num_drafts == 0" + ) + match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens) + print(f"\n{config.id}: spec/no-spec match_ratio={match_ratio:.3f}") + assert match_ratio >= MIN_MATCH_RATIO, ( + f"Spec / no-spec output divergence under {config.id}: " + f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n" + f" spec_tokens = {spec_tokens}\n" + f" no_spec_tokens= {no_spec_tokens}" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "dp_size", + [ + pytest.param(2, marks=pytest.mark.distributed(num_gpus=2), id="dp2"), + pytest.param(1, id="dp1"), + ], +) +async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: int): + """MTP loads under DP (on and off) via AsyncLLM; spec matches no-spec greedy.""" + if torch.accelerator.device_count() < dp_size: + pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs") + + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + + base_args = AsyncEngineArgs( + model=DEEPSEEK_MTP_MAIN_RANDOM, + tensor_parallel_size=1, + data_parallel_size=dp_size, + data_parallel_backend="mp", + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=GPU_MEM_UTIL, + enforce_eager=True, + trust_remote_code=True, + enable_prefix_caching=False, + ) + spec_args = replace( + base_args, + speculative_config={ + "method": "mtp", + "model": DEEPSEEK_MTP_DRAFT_RANDOM, + "num_speculative_tokens": 1, + }, + ) + + sampling_params = SamplingParams( + max_tokens=MAX_TOKENS, + ignore_eos=True, + output_kind=RequestOutputKind.FINAL_ONLY, + temperature=0.0, + ) + + async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]: + try: + async with AsyncExitStack() as stack: + engine = AsyncLLM.from_engine_args(args) + stack.callback(engine.shutdown) + async for out in engine.generate( + request_id=request_id, + prompt=PROMPT, + sampling_params=sampling_params, + ): + token_ids = tuple(out.outputs[0].token_ids) + assert len(token_ids) == MAX_TOKENS + return token_ids + raise AssertionError("AsyncLLM produced no output") + finally: + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + spec_tokens = await asyncio.wait_for( + _generate(spec_args, f"deepseek-mtp-dp{dp_size}-spec"), timeout=600 + ) + no_spec_tokens = await asyncio.wait_for( + _generate(base_args, f"deepseek-mtp-dp{dp_size}-no-spec"), timeout=600 + ) + + match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens) + print(f"\ndp{dp_size}: spec/no-spec match_ratio={match_ratio:.3f}") + assert match_ratio >= MIN_MATCH_RATIO, ( + f"Spec / no-spec output divergence under dp{dp_size}: " + f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n" + f" spec_tokens = {spec_tokens}\n" + f" no_spec_tokens= {no_spec_tokens}" + ) diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py deleted file mode 100644 index 01bec69640b8..000000000000 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ /dev/null @@ -1,1629 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os -import random -from collections.abc import Iterable -from dataclasses import dataclass -from typing import Any - -import pytest -import torch -from tqdm import tqdm - -from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline -from tests.utils import ( - get_attn_backend_list_based_on_platform, - large_gpu_mark, - multi_gpu_marks, - multi_gpu_only, - single_gpu_only, - wait_for_rocm_memory_to_settle, -) -from vllm import LLM, SamplingParams -from vllm.assets.base import VLLM_S3_BUCKET_URL -from vllm.assets.image import VLM_IMAGES_DIR -from vllm.benchmarks.datasets import InstructCoderDataset -from vllm.config import VllmConfig, replace -from vllm.distributed import cleanup_dist_env_and_memory -from vllm.engine.arg_utils import EngineArgs -from vllm.platforms import current_platform -from vllm.v1.metrics.reader import Metric - -MTP_SIMILARITY_RATE = 0.8 - - -class AsyncSchedulingNotEnabledError(AssertionError): - """Raised when async_scheduling is expected to be True for draft_model - spec decode but is False. Tracked in: - https://github.com/vllm-project/vllm/issues/38929 - """ - - -def _skip_if_insufficient_gpus_for_tp(tp_size: int): - """Skip test if available GPUs < tp_size on ROCm.""" - available_gpus = torch.accelerator.device_count() - if available_gpus < tp_size: - pytest.skip( - f"Test requires {tp_size} GPUs, but only {available_gpus} available" - ) - - -Messages = list[dict[str, Any]] - - -def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]: - prompt_types = ["repeat", "gsm8k"] - if mm_enabled: - prompt_types.append("mm") - prompts: list[Messages] = [] - - num_repeat_prompts = num_prompts // len(prompt_types) - if mm_enabled: - num_gsm8k_prompts = num_prompts // len(prompt_types) - num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts - else: - num_mm_prompts = 0 - num_gsm8k_prompts = num_prompts - num_repeat_prompts - - # Generate a mixed batch of prompts, some of which can be easily - # predicted by n-gram matching and some which likely cannot. - random.seed(0) - for _ in range(num_repeat_prompts): - word_choices = ["test", "temp", "hello", "where"] - word = random.choice(word_choices) - prompts.append( - [ - { - "role": "user", - "content": f""" - please repeat the word '{word}' 10 times. - give no other output than the word at least ten times in a row, - in lowercase with spaces between each word and without quotes. - """, - } - ] - ) - prompts.extend( - [{"role": "user", "content": prompt}] - for prompt in _build_gsm8k_prompts( - num_questions=num_gsm8k_prompts, num_shots=5 - )[0] - ) - for _ in range(num_mm_prompts): - placeholders = [ - { - "type": "image_url", - "image_url": { - "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg" - }, - } - ] - prompt = [ - *placeholders, - {"type": "text", "text": "The meaning of the image is"}, - ] - prompts.append([{"role": "user", "content": prompt}]) - - return prompts - - -def get_instruct_coder_messages(n: int) -> list[Messages]: - dataset = InstructCoderDataset( - dataset_path="likaixin/InstructCoder", dataset_split="train" - ) - prompts: Iterable[str] = dataset.sample_prompts(n=n) - return [[{"role": "user", "content": prompt}] for prompt in prompts] - - -@pytest.fixture -def sampling_config(): - return greedy_sampling() - - -def greedy_sampling() -> SamplingParams: - return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False) - - -def stochastic_sampling() -> SamplingParams: - return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False) - - -@pytest.fixture -def model_name(): - return "meta-llama/Llama-3.1-8B-Instruct" - - -def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None: - """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold. - - The default threshold assumes the LLM uses the same target model as the "model_name" - fixture, with max model len == 4096. Precomputed reference value is 75% to 80% - on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70% - to verify that the model is correct. - """ - if expected_accuracy_threshold <= 0.0: - print("Skipping GSM8K evaluation") - return - results = evaluate_gsm8k_offline(llm) - accuracy = results["accuracy"] - print(f"GSM8K accuracy: {accuracy:.3f}") - assert accuracy >= expected_accuracy_threshold, ( - f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}" - ) - - -@pytest.fixture(autouse=True) -def reset_torch_dynamo(): - """Reset torch dynamo cache before each test""" - yield - # Cleanup after test - torch._dynamo.reset() - - -@pytest.mark.parametrize( - "speculative_config", - [ - { - "method": "ngram", - "prompt_lookup_max": 5, - "prompt_lookup_min": 3, - "num_speculative_tokens": 3, - }, - { - "method": "suffix", - "suffix_decoding_max_spec_factor": 2.0, - }, - ], -) -@single_gpu_only -@large_gpu_mark(min_gb=20) -def test_ngram_and_suffix_correctness( - speculative_config: dict, - model_name: str, -): - spec_llm = LLM( - model=model_name, - speculative_config=speculative_config, - max_model_len=4096, - ) - evaluate_llm_for_gsm8k(spec_llm) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@pytest.mark.parametrize("async_scheduling", [True], ids=["async"]) -@single_gpu_only -@large_gpu_mark(min_gb=20) -def test_ngram_gpu_default_with_async_scheduling( - async_scheduling: bool, -): - """ - Test ngram_gpu speculative decoding (k=3) correctness with and without - async scheduling, validated via GSM8K accuracy. - Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%). - """ - qwen3_model = "Qwen/Qwen3-8B" - spec_llm = LLM( - model=qwen3_model, - speculative_config={ - "method": "ngram_gpu", - "prompt_lookup_max": 3, - "prompt_lookup_min": 2, - "num_speculative_tokens": 2, - }, - max_model_len=4096, - async_scheduling=async_scheduling, - ) - # Assert the resolved async_scheduling config matches what was requested. - assert ( - spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - == async_scheduling - ) - evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8) - del spec_llm - cleanup_dist_env_and_memory() - - -@single_gpu_only -@large_gpu_mark(min_gb=20) -def test_suffix_decoding_acceptance( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_name: str, -): - """ - Check that suffix decoding caching takes effect and improves acceptance - lengths and acceptance rates over multiple runs of the same prompts. - """ - test_prompts = get_test_prompts(mm_enabled=False) - - spec_llm = LLM( - model=model_name, - speculative_config={ - "method": "suffix", - "suffix_decoding_max_spec_factor": 2.0, - "suffix_decoding_max_cached_requests": 1000, - }, - max_model_len=1024, - disable_log_stats=False, - ) - - # Run several times and check that the accepted tokens increase. - num_draft = [] - num_accept = [] - for i in range(10): # Run multiple times to warm up the cache. - spec_llm.chat(test_prompts, sampling_config) - # Collect draft and acceptance stats. - metrics = spec_llm.get_metrics() - for metric in metrics: - if metric.name == "vllm:spec_decode_num_draft_tokens": - num_draft.append(metric.value) - if metric.name == "vllm:spec_decode_num_accepted_tokens": - num_accept.append(metric.value) - - # Calculate the acceptance rates for the first and last runs. - first_accept_tokens = num_accept[0] - first_draft_tokens = num_draft[0] - first_accept_rate = first_accept_tokens / first_draft_tokens - - # Take the diff since the stats are cumulative. - last_accept_tokens = num_accept[-1] - num_accept[-2] - last_draft_tokens = num_draft[-1] - num_draft[-2] - last_accept_rate = last_accept_tokens / last_draft_tokens - - # Expect the acceptance length to improve. - assert first_accept_tokens < last_accept_tokens - - # Expect the acceptance rate to improve. - assert first_accept_rate < last_accept_rate - - # Heuristic: expect at least 80.0% acceptance rate at the end. - assert last_accept_rate > 0.80 - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@pytest.mark.slow_test -@large_gpu_mark(min_gb=80) -def test_gemma4_dspark_correctness_and_acceptance_rate( - monkeypatch: pytest.MonkeyPatch, -): - """ - E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length - regression coverage plus GSM8K correctness, at temperature=1.0 to exercise - the probabilistic draft-sampling/rejection-sampling path (not just greedy). - - Uses google/gemma-4-12B-it as target with the dspark_gemma4_12b_block7 draft - model. Exercises the full Gemma4 DSpark path (draft build over the reused - base classes DFlashQwen3Model / Qwen3DSparkForCausalLM / Gemma4MTP* / - DSparkMarkovHead, the fused context-KV precompute, the Markov head, and - rejection sampling), so it fails if any base class drifts out of sync. - - gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template - (use_chat_completions=True; raw few-shot completion collapses to a few - percent). Reference: measured over 5 runs of 200 GSM8K questions at - temperature=1.0 (prefix caching disabled): - accuracy: min=0.900 max=0.955 mean=0.937 - acceptance_rate: min=0.578 max=0.595 mean=0.588 - acceptance_len: min=5.044 max=5.167 mean=5.116 - Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling - """ - monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") - - spec_llm = LLM( - model="google/gemma-4-12B-it", - trust_remote_code=True, - speculative_config={ - "method": "dspark", - "model": "deepseek-ai/dspark_gemma4_12b_block7", - "num_speculative_tokens": 7, - "draft_sample_method": "probabilistic", - }, - max_model_len=8192, - max_num_seqs=32, - gpu_memory_utilization=0.8, - enforce_eager=True, - enable_prefix_caching=False, - disable_log_stats=False, - ) - try: - results = evaluate_gsm8k_offline( - spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True - ) - gsm8k_accuracy = results["accuracy"] - - metrics = spec_llm.get_metrics() - acceptance_rate = compute_acceptance_rate(metrics) - acceptance_len = compute_acceptance_len(metrics) - print( - f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, " - f"acceptance_len={acceptance_len:.2f}, " - f"gsm8k_accuracy={gsm8k_accuracy:.3f}" - ) - - assert acceptance_rate >= 0.588 * 0.9 - assert acceptance_len >= 5.116 * 0.9 - assert gsm8k_accuracy >= 0.937 * 0.9 - finally: - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@pytest.mark.parametrize( - ["model_path", "expected_accuracy_threshold"], - [ - ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80% - ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92% - ], - ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"], -) -@single_gpu_only -@large_gpu_mark(min_gb=24) -def test_speculators_model_integration( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_path: str, - expected_accuracy_threshold: float, -): - """ - Test that speculators models work with the simplified integration. - - This verifies the `vllm serve ` use case where - speculative config is automatically detected from the model config - without requiring explicit --speculative-config argument. - - Tests: - 1. Speculator model is correctly detected - 2. Verifier model is extracted from speculator config - 3. Speculative decoding is automatically enabled - 4. Text generation works correctly - 5. GSM8k accuracy of the model passes a sanity check when speculative decoding on - 6. Output matches reference (non-speculative) generation - """ - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - - # Generate test prompts - test_prompts = get_test_prompts(mm_enabled=False) - - # First run: Direct speculator model (simplified integration) - spec_llm = LLM(model=model_path, max_model_len=4096, gpu_memory_utilization=0.92) - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - spec_outputs = spec_llm.chat(test_prompts, sampling_config) - - # Verify speculative config was auto-detected - assert spec_llm.llm_engine.vllm_config.speculative_config is not None, ( - f"Speculative config should be auto-detected for {model_path}" - ) - - spec_config = spec_llm.llm_engine.vllm_config.speculative_config - assert spec_config.num_speculative_tokens > 0, ( - f"Expected positive speculative tokens, " - f"got {spec_config.num_speculative_tokens}" - ) - - # Verify draft model is set to the speculator model - assert spec_config.model == model_path, ( - f"Draft model should be {model_path}, got {spec_config.model}" - ) - - # Extract verifier model for reference run - verifier_model = spec_llm.llm_engine.vllm_config.model_config.model - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - # Second run: Reference without speculative decoding - ref_llm = LLM(model=verifier_model, max_model_len=4096, gpu_memory_utilization=0.92) - ref_outputs = ref_llm.chat(test_prompts, sampling_config) - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - # Compare outputs - matches = sum( - 1 - for ref, spec in zip(ref_outputs, spec_outputs) - if ref.outputs[0].text == spec.outputs[0].text - ) - - # Heuristic: expect at least 66% of prompts to match exactly - assert matches >= int(0.66 * len(ref_outputs)), ( - f"Only {matches}/{len(ref_outputs)} outputs matched. " - f"Expected at least {int(0.66 * len(ref_outputs))} matches." - ) - - -def _run_eagle_correctness( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_setup: tuple[str, str, str, int], - mm_enabled: bool, - expected_accuracy_threshold: float, - enable_chunked_prefill: bool, - model_impl: str, - attn_backend: str, -): - """ - Compare the outputs of an original LLM and a speculative LLM - which should be the same when using eagle speculative decoding. - """ - if model_impl == "transformers": - import transformers - from packaging.version import Version - - installed = Version(transformers.__version__) - required = Version("5.0.0") - if installed < required: - pytest.skip( - "Eagle3 with the Transformers modeling backend requires " - f"transformers>={required}, but got {installed}" - ) - - test_prompts = get_test_prompts(mm_enabled) - - if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN": - if current_platform.is_rocm(): - print( - "FLASH_ATTN for spec_decode not supported on " - "ROCm currently. Changing to FLEX_ATTENTION backend." - ) - attention_config = {"backend": "FLEX_ATTENTION"} - else: - attention_config = None - else: - attention_config = {"backend": attn_backend} - - if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm(): - pytest.skip( - "TRITON_ATTN does not support " - "multi-token eagle spec decode on current platform" - ) - - with monkeypatch.context() as m: - m.setenv("VLLM_MLA_DISABLE", "1") - - if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): - if "deepseek" in model_setup[1].lower(): - m.setenv("VLLM_ROCM_USE_AITER", "1") - m.delenv("VLLM_MLA_DISABLE", raising=False) - attention_config = {"backend": "ROCM_AITER_MLA"} - else: - m.setenv("VLLM_ROCM_USE_AITER", "1") - - method, model_name, spec_model_name, tp_size = model_setup - _skip_if_insufficient_gpus_for_tp(tp_size) - - max_model_len = 2048 - max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len - - ref_llm = LLM( - model=model_name, - max_model_len=max_model_len, - tensor_parallel_size=tp_size, - attention_config=attention_config, - ) - evaluate_llm_for_gsm8k( - ref_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - ref_outputs = ref_llm.chat(test_prompts, sampling_config) - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - # ROCm frees VRAM lazily; wait so the spec engine started right after - # does not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle() - - spec_llm = LLM( - model=model_name, - trust_remote_code=True, - tensor_parallel_size=tp_size, - speculative_config={ - "method": method, - "model": spec_model_name, - "num_speculative_tokens": 3, - "max_model_len": max_model_len, - }, - max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, - enable_chunked_prefill=enable_chunked_prefill, - model_impl=model_impl, - attention_config=attention_config, - ) - # EAGLE/EAGLE3 supports async scheduling; assert it is active by default. - assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - spec_outputs = spec_llm.chat(test_prompts, sampling_config) - matches = 0 - misses = 0 - for ref_output, spec_output in zip(ref_outputs, spec_outputs): - if ref_output.outputs[0].text == spec_output.outputs[0].text: - matches += 1 - else: - misses += 1 - print(f"ref_output: {ref_output.outputs[0].text}") - print(f"spec_output: {spec_output.outputs[0].text}") - - assert matches > int(0.6 * len(ref_outputs)) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - # ROCm frees VRAM lazily; wait so the next parametrization's engine does - # not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle() - - -@single_gpu_only -@pytest.mark.skipif( - current_platform.is_device_capability_family(100), - reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)", -) -@pytest.mark.parametrize( - [ - "model_setup", - "mm_enabled", - "enable_chunked_prefill", - "model_impl", - "expected_accuracy_threshold", - ], - [ - ( - ( - "eagle", - "eagle618/deepseek-v3-random", - "eagle618/eagle-deepseek-v3-random", - 1, - ), - False, - False, - "auto", - 0.0, - ), - ], - ids=["deepseek_eagle"], -) -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) -def test_eagle_correctness_light( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_setup: tuple[str, str, str, int], - mm_enabled: bool, - expected_accuracy_threshold: float, - enable_chunked_prefill: bool, - model_impl: str, - attn_backend: str, -): - _run_eagle_correctness( - monkeypatch, - sampling_config, - model_setup, - mm_enabled, - expected_accuracy_threshold, - enable_chunked_prefill, - model_impl, - attn_backend, - ) - - -@single_gpu_only -@large_gpu_mark(min_gb=24) -@pytest.mark.parametrize( - [ - "model_setup", - "mm_enabled", - "enable_chunked_prefill", - "model_impl", - "expected_accuracy_threshold", - ], - [ - ( - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), - False, - False, - "auto", - 0.8, - ), - pytest.param( - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), - False, - False, - "transformers", - 0.8, - # TODO(hmellor): figure out why memory usage is so high - marks=pytest.mark.skip( - reason="Feature is experimental and uses too much memory in CI", - ), - ), - pytest.param( - ( - "eagle3", - "Qwen/Qwen3-VL-8B-Instruct", - "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3", - 1, - ), - False, - False, - "auto", - 0.8, - marks=pytest.mark.skip( - reason="architecture of its eagle3 is LlamaForCausalLMEagle3" - ), - ), - pytest.param( - ( - "eagle3", - "Qwen/Qwen2.5-VL-7B-Instruct", - "Rayzl/qwen2.5-vl-7b-eagle3-sgl", - 1, - ), - False, - False, - "auto", - 0.7, - marks=pytest.mark.skip( - reason="Skipping due to its head_dim not being a multiple of 32" - ), - ), - ( - ( - "eagle3", - "meta-llama/Llama-3.1-8B-Instruct", - "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", - 1, - ), - False, - False, - "auto", - 0.7, - ), - ], - ids=[ - "qwen3_eagle3", - "qwen3_eagle3-transformers", - "qwen3_vl_eagle3", - "qwen2_5_vl_eagle3", - "llama3_eagle3", - ], -) -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) -def test_eagle_correctness_medium( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_setup: tuple[str, str, str, int], - mm_enabled: bool, - expected_accuracy_threshold: float, - enable_chunked_prefill: bool, - model_impl: str, - attn_backend: str, -): - _run_eagle_correctness( - monkeypatch, - sampling_config, - model_setup, - mm_enabled, - expected_accuracy_threshold, - enable_chunked_prefill, - model_impl, - attn_backend, - ) - - -@pytest.mark.parametrize( - [ - "model_setup", - "mm_enabled", - "enable_chunked_prefill", - "model_impl", - "expected_accuracy_threshold", - ], - [ - pytest.param( - ( - "eagle", - "meta-llama/Llama-3.1-8B-Instruct", - "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", - 1, - ), - False, - True, - "auto", - 0.7, - marks=large_gpu_mark(min_gb=40), - id="llama3_eagle", - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - False, - False, - "auto", - 0.8, - marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=40)], - id="llama4_eagle", - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - True, - True, - "auto", - 0.8, - marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)], - id="llama4_eagle_mm", - ), - ], -) -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) -def test_eagle_correctness_heavy( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_setup: tuple[str, str, str, int], - mm_enabled: bool, - expected_accuracy_threshold: float, - enable_chunked_prefill: bool, - model_impl: str, - attn_backend: str, -): - _run_eagle_correctness( - monkeypatch, - sampling_config, - model_setup, - mm_enabled, - expected_accuracy_threshold, - enable_chunked_prefill, - model_impl, - attn_backend, - ) - - -@large_gpu_mark(min_gb=24) -def test_medusa_acceptance_rate( - sampling_config: SamplingParams, -): - """Verify a trained Medusa checkpoint achieves nonzero acceptance rate. - - Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the - speculation path actually accepts tokens — unlike test_medusa_correctness, - which uses a random head and only validates output correctness. - """ - target_model = "lmsys/vicuna-7b-v1.3" - medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3" - prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0] - - spec_llm = LLM( - model=target_model, - speculative_config={ - "method": "medusa", - "model": medusa_model, - "num_speculative_tokens": 3, - }, - max_model_len=1024, - enforce_eager=True, - disable_log_stats=False, - ) - spec_llm.generate(prompts, sampling_config) - metrics = spec_llm.get_metrics() - acceptance_rate = compute_acceptance_rate(metrics) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - min_acceptance_rate = 0.198 - print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})") - - # Regression guard at 90% of the measured baseline. - assert acceptance_rate >= min_acceptance_rate, ( - f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}" - ) - - -@pytest.mark.parametrize( - ["model_setup", "mm_enabled", "expected_accuracy_threshold"], - [ - (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70% - pytest.param( - ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), - False, - 0.0, - marks=pytest.mark.skipif( - current_platform.is_device_capability_family(100), - reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell", - ), - ), # dummy model - ( - ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1), - False, - 0.20, - ), # hybrid + MTP, ref: ~34%-35% - ( - ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"), - False, - 0.50, - ), # gemma4 MTP with assistant model, ref: ~62% - ], - ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"], -) -@single_gpu_only -@large_gpu_mark(min_gb=20) -def test_mtp_correctness( - monkeypatch: pytest.MonkeyPatch, - sampling_config: SamplingParams, - model_setup: tuple[str, str, int] | tuple[str, str, int, str], - mm_enabled: bool, - expected_accuracy_threshold: float, -): - """ - Compare the outputs of a original LLM and a speculative LLM - which should be the same when using MTP speculative decoding. Due to some variance - in the engine, it is possible for some outputs to differ, so we expect that at least - 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed - reference threshold for each model. - """ - # Generate test prompts inside the function instead of using fixture - test_prompts = get_test_prompts(mm_enabled) - with monkeypatch.context() as m: - m.setenv("VLLM_MLA_DISABLE", "1") - - if len(model_setup) == 4: - method, model_name, tp_size, draft_model = model_setup - else: - method, model_name, tp_size = model_setup - draft_model = None - _skip_if_insufficient_gpus_for_tp(tp_size) - - if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"): - pytest.skip( - "Model Runner V2 does not yet support hybrid models " - "(Qwen3.5 mixes Mamba-style GDN with attention layers)." - ) - - attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto" - - # Skip multimodal profiling for models that don't need it in this test. - extra_kwargs: dict[str, Any] = {} - if "Qwen3.5" in model_name: - extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0} - elif "gemma-4" in model_name: - extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0} - - if draft_model is not None and "gemma-4" in draft_model: - import transformers - from packaging.version import Version - - if Version(transformers.__version__) < Version("5.8.0"): - pytest.skip( - "Gemma4 MTP assistant requires transformers>=5.8.0, " - f"got {transformers.__version__}" - ) - - ref_llm = LLM( - model=model_name, - max_model_len=2048, - tensor_parallel_size=tp_size, - trust_remote_code=True, - attention_backend=attn_backend, - **extra_kwargs, - ) - ref_outputs = ref_llm.chat(test_prompts, sampling_config) - evaluate_llm_for_gsm8k( - ref_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - speculative_config: dict[str, Any] = { - "method": method, - "num_speculative_tokens": 1, - "max_model_len": 2048, - } - if draft_model is not None: - speculative_config["model"] = draft_model - speculative_config["num_speculative_tokens"] = 2 - - spec_llm = LLM( - model=model_name, - trust_remote_code=True, - tensor_parallel_size=tp_size, - speculative_config=speculative_config, - max_model_len=2048, - attention_backend=attn_backend, - **extra_kwargs, - ) - # MTP supports async scheduling; assert it is active by default. - assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - spec_outputs = spec_llm.chat(test_prompts, sampling_config) - matches = 0 - misses = 0 - for ref_output, spec_output in zip(ref_outputs, spec_outputs): - if ref_output.outputs[0].text == spec_output.outputs[0].text: - matches += 1 - else: - misses += 1 - print(f"ref_output: {ref_output.outputs[0].text}") - print(f"spec_output: {spec_output.outputs[0].text}") - - # Heuristic: expect at least 80% of the prompts to match exactly - # Upon failure, inspect the outputs to check for inaccuracy. - assert matches > int(MTP_SIMILARITY_RATE * len(ref_outputs)) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@dataclass -class ArgsTest: - target_model: str - draft_model: str - sampling_config: SamplingParams - num_speculative_tokens: int - expected_acceptance_rate: float - expected_acceptance_len: float - expected_gsm8k_accuracy: float = 0.0 # skip by default - # Defaults - enforce_eager: bool = True - parallel_drafting: bool = False - target_tensor_parallel_size: int = 1 - draft_tensor_parallel_size: int = 1 - max_model_len: int = 2048 - gpu_memory_utilization: float = 0.5 - dataset: str = "test_prompts" - num_prompts: int = 100 - - -cases = [ - # Same model for draft and target, greedy sampling. - ArgsTest( - target_model="Qwen/Qwen3-0.6B", - draft_model="Qwen/Qwen3-0.6B", - sampling_config=greedy_sampling(), - num_speculative_tokens=3, # K - expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1 - expected_acceptance_rate=0.98, # slight epsilon - expected_gsm8k_accuracy=0.25, # ref: 35-40% - ), - # Smaller draft model, stochastic sampling. - ArgsTest( - target_model="Qwen/Qwen3-1.7B", - draft_model="Qwen/Qwen3-0.6B", - sampling_config=stochastic_sampling(), - num_speculative_tokens=3, - expected_acceptance_len=3.4, # ref: 3.7 - expected_acceptance_rate=0.80, # ref: 0.90 - expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling - ), -] - - -@pytest.mark.parametrize("args", cases) -@pytest.mark.parametrize("enforce_eager", [True, False]) -@single_gpu_only -# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 -@pytest.mark.xfail( - raises=AsyncSchedulingNotEnabledError, - reason="draft_model does not yet enable async_scheduling: issue #38929", -) -def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool): - args.enforce_eager = enforce_eager - assert_draft_model_correctness(args) - - -@single_gpu_only -# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 -@pytest.mark.xfail( - raises=AsyncSchedulingNotEnabledError, - reason="draft_model does not yet enable async_scheduling: issue #38929", -) -def test_draft_model_realistic_example(): - args = ArgsTest( - target_model="Qwen/Qwen3-1.7B", - draft_model="Qwen/Qwen3-0.6B", - dataset="likaixin/InstructCoder", - num_speculative_tokens=3, - sampling_config=greedy_sampling(), - enforce_eager=False, - expected_acceptance_len=2.6, # ref: 2.86 - expected_acceptance_rate=0.5, # ref: 0.62 - ) - assert_draft_model_correctness(args) - - -@single_gpu_only -# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 -@pytest.mark.xfail( - raises=AsyncSchedulingNotEnabledError, - reason="draft_model does not yet enable async_scheduling: issue #38929", -) -def test_draft_model_parallel_drafting(): - args = ArgsTest( - target_model="Qwen/Qwen3-1.7B", - draft_model="amd/PARD-Qwen3-0.6B", - dataset="likaixin/InstructCoder", - num_speculative_tokens=3, - sampling_config=greedy_sampling(), - parallel_drafting=True, - enforce_eager=False, - expected_acceptance_len=2.3, # ref: 2.52 - expected_acceptance_rate=0.4, # ref: 0.51 - ) - assert_draft_model_correctness(args) - - -@pytest.mark.parametrize( - "models", - [ - # target_model, draft_model - ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized - ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized - ], - ids=["target_quantized", "draft_quantized"], -) -@pytest.mark.parametrize("enforce_eager", [True, False]) -@single_gpu_only -# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 -@pytest.mark.xfail( - raises=AsyncSchedulingNotEnabledError, - reason="draft_model does not yet enable async_scheduling: issue #38929", -) -def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): - tgt_model, draft_model = models - sd_case = ArgsTest( - target_model=tgt_model, - draft_model=draft_model, - **some_high_acceptance_metrics(), - enforce_eager=enforce_eager, - ) - assert_draft_model_correctness(sd_case) - - -@multi_gpu_only(num_gpus=2) -# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929 -@pytest.mark.xfail( - raises=AsyncSchedulingNotEnabledError, - reason="draft_model does not yet enable async_scheduling: issue #38929", -) -def test_draft_model_tensor_parallelism(): - """Ensure spec decode works when running with TP > 1.""" - _skip_if_insufficient_gpus_for_tp(2) - sd_case = ArgsTest( - target_model="Qwen/Qwen3-1.7B", - target_tensor_parallel_size=2, - draft_model="Qwen/Qwen3-0.6B", - draft_tensor_parallel_size=2, - **some_high_acceptance_metrics(), - enforce_eager=False, - expected_gsm8k_accuracy=0.5, - ) - assert_draft_model_correctness(sd_case) - - -@multi_gpu_only(num_gpus=2) -def test_draft_model_engine_args_tensor_parallelism(): - """Ensure the vllm_config for the draft model is created correctly, - and independently of the target model (quantization, TP, etc.)""" - _skip_if_insufficient_gpus_for_tp(2) - - engine_args = EngineArgs( - model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized - tensor_parallel_size=2, - speculative_config={ - "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized - "method": "draft_model", - "num_speculative_tokens": 3, - "draft_tensor_parallel_size": 1, # <<< valid arg name - }, - ) - target_config: VllmConfig = engine_args.create_engine_config() - assert target_config.parallel_config.tensor_parallel_size == 2 - assert target_config.quant_config.get_name() == "fp8" - - speculative_config = target_config.speculative_config - draft_config: VllmConfig = replace( - target_config, - quant_config=None, - parallel_config=replace( - speculative_config.draft_parallel_config, - rank=target_config.parallel_config.rank, - ), - model_config=speculative_config.draft_model_config, - ) - assert draft_config.parallel_config.tensor_parallel_size == 1 - assert draft_config.quant_config is None - - -def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: - """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic - so we can test it without instantiating a full proposer.""" - spec_cfg = vllm_config.speculative_config - if spec_cfg.moe_backend is not None: - return replace( - vllm_config, - kernel_config=replace( - vllm_config.kernel_config, - moe_backend=spec_cfg.moe_backend, - ), - ) - return vllm_config - - -def test_draft_model_moe_backend_override(): - """When moe_backend is set in speculative_config, the draft VllmConfig - should use it while the target keeps its own setting.""" - engine_args = EngineArgs( - model="Qwen/Qwen3-1.7B", - tensor_parallel_size=1, - moe_backend="flashinfer_trtllm", - speculative_config={ - "model": "Qwen/Qwen3-0.6B", - "method": "draft_model", - "num_speculative_tokens": 3, - "moe_backend": "triton", - }, - ) - tgt_config: VllmConfig = engine_args.create_engine_config() - assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" - assert tgt_config.speculative_config.moe_backend == "triton" - - draft_config = _apply_draft_moe_backend(tgt_config) - assert draft_config.kernel_config.moe_backend == "triton" - # Target config must be unaffected. - assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" - - -def test_draft_model_moe_backend_inherits_target(): - """When moe_backend is not set in speculative_config, the draft should - inherit the target's moe_backend.""" - engine_args = EngineArgs( - model="Qwen/Qwen3-1.7B", - tensor_parallel_size=1, - moe_backend="flashinfer_cutlass", - speculative_config={ - "model": "Qwen/Qwen3-0.6B", - "method": "draft_model", - "num_speculative_tokens": 3, - }, - ) - tgt_config: VllmConfig = engine_args.create_engine_config() - assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass" - assert tgt_config.speculative_config.moe_backend is None - - draft_config = _apply_draft_moe_backend(tgt_config) - assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass" - assert draft_config is tgt_config - - -def test_draft_model_moe_backend_default_auto(): - """When neither target nor draft set moe_backend explicitly, both should - default to 'auto'.""" - engine_args = EngineArgs( - model="Qwen/Qwen3-1.7B", - tensor_parallel_size=1, - speculative_config={ - "model": "Qwen/Qwen3-0.6B", - "method": "draft_model", - "num_speculative_tokens": 3, - }, - ) - tgt_config: VllmConfig = engine_args.create_engine_config() - assert tgt_config.kernel_config.moe_backend == "auto" - assert tgt_config.speculative_config.moe_backend is None - - draft_config = _apply_draft_moe_backend(tgt_config) - assert draft_config.kernel_config.moe_backend == "auto" - assert draft_config is tgt_config - - -def test_draft_model_engine_args_rejects_invalid_tp_argname(): - """The user should pass "draft_tensor_parallel_size" rather than - "tensor_parallel_size". We enforce this with validation.""" - - engine_args = EngineArgs( - model="Qwen/Qwen3-1.7B", - tensor_parallel_size=1, - speculative_config={ - "model": "Qwen/Qwen3-0.6B", - "method": "draft_model", - "num_speculative_tokens": 3, - "tensor_parallel_size": 1, # <<< invalid arg name - }, - ) - with pytest.raises(ValueError): - engine_args.create_engine_config() - - -def assert_draft_model_correctness(args: ArgsTest): - """Compare the outputs using and not using speculative decoding. - In the greedy decoding case, the outputs must match EXACTLY.""" - test_prompts: list[Messages] = get_messages( - dataset=args.dataset, n=args.num_prompts - ) - - spec_llm = LLM( - model=args.target_model, - speculative_config={ - "model": args.draft_model, - "method": "draft_model", - "num_speculative_tokens": args.num_speculative_tokens, - "max_model_len": args.max_model_len, - "enforce_eager": args.enforce_eager, - "draft_tensor_parallel_size": args.draft_tensor_parallel_size, - "parallel_drafting": args.parallel_drafting, - }, - max_num_seqs=100, # limit cudagraph capture runtime - max_model_len=args.max_model_len, - gpu_memory_utilization=args.gpu_memory_utilization, - tensor_parallel_size=args.target_tensor_parallel_size, - enforce_eager=args.enforce_eager, - disable_log_stats=False, # enables get_metrics() - ) - - # we don't check the outputs, only check the metrics - spec_llm.chat(test_prompts, args.sampling_config) - metrics = spec_llm.get_metrics() - acceptance_rate: float = compute_acceptance_rate(metrics) - acceptance_len: float = compute_acceptance_len(metrics) - - # Need to evaluate after getting metrics to avoid polluting the AR - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy - ) - - print( - f"spec-decode: target={args.target_model}, draft={args.draft_model}, " - f"temperature={args.sampling_config.temperature:.2f}, " - f"acceptance_rate={acceptance_rate:.2f}, " - f"acceptance_len={acceptance_len:.2f}, " - ) - - assert acceptance_rate >= args.expected_acceptance_rate - assert acceptance_len >= args.expected_acceptance_len - # draft_model supports async scheduling; assert it is active by default. - # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that - # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this - # specific failure — leaving all other assertion failures (e.g. correctness or - # acceptance-rate checks above) visible as real test failures. - has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - del spec_llm # CLEANUP - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - if not has_async: - raise AsyncSchedulingNotEnabledError( - "Expected async_scheduling=True for draft_model spec decode, got False." - " See https://github.com/vllm-project/vllm/issues/38929" - ) - - -def get_messages(dataset: str, n: int) -> list[Messages]: - if dataset == "test_prompts": - return get_test_prompts(mm_enabled=False, num_prompts=n) - elif dataset == "likaixin/InstructCoder": - return get_instruct_coder_messages(n=n) - else: - raise NotImplementedError(f"Dataset '{dataset}' not implemented") - - -def some_high_acceptance_metrics() -> dict: - return { - "sampling_config": greedy_sampling(), - "num_speculative_tokens": 3, - "expected_acceptance_len": 3.4, # ref: 3.75 - "expected_acceptance_rate": 0.8, # ref: 0.9 - } - - -def compute_acceptance_rate( - metrics: list[Metric], prev_metrics: list[Metric] | None = None -) -> float: - name2metric = {metric.name: metric for metric in metrics} - n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value - if n_draft_toks == 0: - return float("nan") - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value - if prev_metrics is not None: - prev_name2metric = {metric.name: metric for metric in prev_metrics} - n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value - n_accepted_toks -= prev_name2metric[ - "vllm:spec_decode_num_accepted_tokens" - ].value - if n_draft_toks <= 0: - return float("nan") - return n_accepted_toks / n_draft_toks - - -def compute_acceptance_len( - metrics: list[Metric], prev_metrics: list[Metric] | None = None -) -> float: - name2metric = {metric.name: metric for metric in metrics} - n_drafts = name2metric["vllm:spec_decode_num_drafts"].value - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value - if n_drafts == 0: - return 1 - if prev_metrics is not None: - prev_name2metric = {metric.name: metric for metric in prev_metrics} - n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value - n_accepted_toks -= prev_name2metric[ - "vllm:spec_decode_num_accepted_tokens" - ].value - if n_drafts <= 0: - return 1 - return 1 + (n_accepted_toks / n_drafts) - - -# Datasets in the format used in DFlash validations -def load_and_process_dataset(data_name: str): - from datasets import load_dataset - - if data_name == "gsm8k": - dataset = load_dataset("openai/gsm8k", "main", split="test") - prompt_fmt = ( - "{question}\nPlease reason step by step," - " and put your final answer within \\boxed{{}}." - ) - dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) - elif data_name == "mt-bench": - dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") - dataset = dataset.map(lambda x: {"turns": x["prompt"]}) - elif data_name == "humaneval": - dataset = load_dataset("openai/openai_humaneval", split="test") - prompt_fmt = ( - "Write a solution to the following problem and make sure" - " that it passes the tests:\n```python\n{prompt}\n```" - ) - dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) - - return dataset - - -@pytest.fixture -def dflash_config(): - target_model = "Qwen/Qwen3-8B" - draft_model = "z-lab/Qwen3-8B-DFlash-b16" - - return dict( - model=target_model, - trust_remote_code=True, - speculative_config={ - "method": "dflash", - "model": draft_model, - "num_speculative_tokens": 16, - "max_model_len": 32768, - }, - max_model_len=32768, - max_num_seqs=128, - gpu_memory_utilization=0.85, - enforce_eager=False, - disable_log_stats=False, - ) - - -@pytest.mark.parametrize("use_mrv2", [False, True]) -def test_dflash_acceptance_rates( - monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config -): - """ - E2E test for DFlash (block diffusion) speculative decoding. - Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval - comparing against baseline results from the paper (Table 1). - See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology. - """ - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") - - spec_llm = LLM(**dflash_config) - - max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k - - # All scores from Table 1 in https://arxiv.org/pdf/2602.06036 - expected_acceptance_lengths = { - "mt-bench": 4.24, - "humaneval": 6.50, - "gsm8k": 6.54 * 0.975, # runs with a subset of prompts so extra wide tol here - } - - tokenizer = spec_llm.get_tokenizer() - for dataset_name, expected_len in expected_acceptance_lengths.items(): - dataset = load_and_process_dataset(dataset_name) - prev_metrics = None - acceptance_lengths = [] - for i in tqdm( - range(min(max_prompts_per_dataset, len(dataset))), - desc=f"Processing {dataset_name}", - ): - user_content = dataset[i]["turns"][0] - prompt_text = tokenizer.apply_chat_template( - [{"role": "user", "content": user_content}], - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - - # Temp=0, MaxTokens=2048 from the paper - spec_llm.generate( - [prompt_text], - SamplingParams(temperature=0, max_tokens=2048), - use_tqdm=False, - ) - current_metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(current_metrics, prev_metrics) - prev_metrics = current_metrics - acceptance_lengths.append(acceptance_len) - - mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) - # Fairly tight tolerance of 95% against the paper's figures, - # watching for regressions. Can be relaxed if test is flaky but be sure to - # check for genuine issues such as #40727. - expected_len = expected_len * 0.95 - print( - f"DFlash acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" - f" (expected at least {expected_len:.2f})" - ) - - assert mean_acceptance_length >= expected_len, ( - f"DFlash acceptance_len for {dataset_name} is below expected threshold:" - f"{mean_acceptance_length:.2f} < {expected_len:.2f}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@pytest.fixture -def dspark_config(): - target_model = "Qwen/Qwen3-4B-FP8" - draft_model = "deepseek-ai/dspark_qwen3_4b_block7" - - return dict( - model=target_model, - trust_remote_code=True, - speculative_config={ - "method": "dspark", - "model": draft_model, - "num_speculative_tokens": 7, - "attention_backend": "FLASH_ATTN", - "draft_sample_method": "probabilistic", - }, - max_model_len=4096, - disable_log_stats=False, - ) - - -@single_gpu_only -@large_gpu_mark(min_gb=24) -def test_dspark_correctness_and_acceptance_rate(dspark_config): - """ - E2E test for DSpark speculative decoding: acceptance rate/length - regression coverage plus GSM8K correctness, at temperature=1.0 to - exercise the probabilistic draft-sampling/rejection-sampling path - (not just greedy). - - Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft - model. Reference: measured over 12 runs of the full GSM8K set at - temperature=1.0 (prefix caching disabled to avoid cross-run reuse): - accuracy: min=0.782 max=0.814 mean=0.801 - acceptance_rate: min=0.418 max=0.434 mean=0.428 - acceptance_len: min=3.928 max=4.037 mean=3.994 - Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling - """ - spec_llm = LLM(**dspark_config) - - results = evaluate_gsm8k_offline(spec_llm, temperature=1.0) - gsm8k_accuracy = results["accuracy"] - - metrics = spec_llm.get_metrics() - acceptance_rate = compute_acceptance_rate(metrics) - acceptance_len = compute_acceptance_len(metrics) - - print( - f"DSpark acceptance_rate={acceptance_rate:.2f}, " - f"acceptance_len={acceptance_len:.2f}, " - f"gsm8k_accuracy={gsm8k_accuracy:.3f}" - ) - - assert acceptance_rate >= 0.428 * 0.9 - assert acceptance_len >= 3.994 * 0.9 - assert gsm8k_accuracy >= 0.801 * 0.9 - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@single_gpu_only -def test_synthetic_acceptance_rate(): - """Verify that synthetic rejection sampling produces an acceptance - length close to the requested mean acceptance length.""" - num_spec_tokens = 3 - expected_acceptance_len = 1.875 - tolerance = 0.15 - - spec_llm = LLM( - model="meta-llama/Llama-3.2-1B-Instruct", - trust_remote_code=True, - speculative_config={ - "method": "eagle3", - "model": "nm-testing/Llama3_2_1B_speculator.eagle3", - "num_speculative_tokens": num_spec_tokens, - "max_model_len": 2048, - "rejection_sample_method": "synthetic", - "synthetic_acceptance_length": expected_acceptance_len, - }, - max_model_len=2048, - enforce_eager=True, - disable_log_stats=False, - ) - - test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50) - spec_llm.chat( - test_prompts, - SamplingParams(temperature=0, max_tokens=64, ignore_eos=True), - ) - - metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(metrics) - - print( - f"Synthetic acceptance length: {acceptance_len:.3f}" - f" (expected={expected_acceptance_len:.3f}," - f" tolerance=±{tolerance})" - ) - assert abs(acceptance_len - expected_acceptance_len) <= tolerance, ( - f"Synthetic acceptance length {acceptance_len:.3f} is not within" - f" ±{tolerance} of expected {expected_acceptance_len:.3f}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@pytest.mark.parametrize("use_mrv2", [False, True]) -def test_dflash_correctness( - monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config -): - """ - E2E test for DFlash (block diffusion) speculative decoding. - Ensures output correctness on GSM8k, with cudagraphs and batching on. - """ - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") - - spec_llm = LLM(**dflash_config) - - # Evaluate GSM8k accuracy (Qwen3-8B ref: ~87-92% on GSM8k) - evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8) - - current_metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(current_metrics) - - # AR is thoroughly validated in test_dflash_acceptance_rates, in a manner consistent - # with the DFlash paper. However, that test measures AL per-request and thus runs - # with a batch size of 1. To ensure that AL does not collapse with large batch sizes - # we enforce a baseline on the AL over the full lm-eval-style GSM8k test. - expected_len = 3.5 # Measured is 3.9 to 4.0 - print(f"DFlash GSM8k correctness test got AL {acceptance_len}") - assert acceptance_len >= expected_len, ( - "DFlash correctness check failed with" - f" {acceptance_len=}, expected at least {expected_len}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/utils.py b/tests/v1/e2e/spec_decode/utils.py new file mode 100644 index 000000000000..ccd5232eaaea --- /dev/null +++ b/tests/v1/e2e/spec_decode/utils.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import random +from collections.abc import Iterable +from typing import Any + +import pytest +import torch + +from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline +from vllm import LLM, SamplingParams +from vllm.assets.base import VLLM_S3_BUCKET_URL +from vllm.assets.image import VLM_IMAGES_DIR +from vllm.v1.metrics.reader import Metric + + +def _skip_if_insufficient_gpus_for_tp(tp_size: int): + """Skip test if available GPUs < tp_size on ROCm.""" + available_gpus = torch.accelerator.device_count() + if available_gpus < tp_size: + pytest.skip( + f"Test requires {tp_size} GPUs, but only {available_gpus} available" + ) + + +Messages = list[dict[str, Any]] + + +def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]: + prompt_types = ["repeat", "gsm8k"] + if mm_enabled: + prompt_types.append("mm") + prompts: list[Messages] = [] + + num_repeat_prompts = num_prompts // len(prompt_types) + if mm_enabled: + num_gsm8k_prompts = num_prompts // len(prompt_types) + num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts + else: + num_mm_prompts = 0 + num_gsm8k_prompts = num_prompts - num_repeat_prompts + + # Generate a mixed batch of prompts, some of which can be easily + # predicted by n-gram matching and some which likely cannot. + random.seed(0) + for _ in range(num_repeat_prompts): + word_choices = ["test", "temp", "hello", "where"] + word = random.choice(word_choices) + prompts.append( + [ + { + "role": "user", + "content": f""" + please repeat the word '{word}' 10 times. + give no other output than the word at least ten times in a row, + in lowercase with spaces between each word and without quotes. + """, + } + ] + ) + prompts.extend( + [{"role": "user", "content": prompt}] + for prompt in _build_gsm8k_prompts( + num_questions=num_gsm8k_prompts, num_shots=5 + )[0] + ) + for _ in range(num_mm_prompts): + placeholders = [ + { + "type": "image_url", + "image_url": { + "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg" + }, + } + ] + prompt = [ + *placeholders, + {"type": "text", "text": "The meaning of the image is"}, + ] + prompts.append([{"role": "user", "content": prompt}]) + + return prompts + + +def get_instruct_coder_messages(n: int) -> list[Messages]: + from vllm.benchmarks.datasets import InstructCoderDataset + + dataset = InstructCoderDataset( + dataset_path="likaixin/InstructCoder", dataset_split="train" + ) + prompts: Iterable[str] = dataset.sample_prompts(n=n) + return [[{"role": "user", "content": prompt}] for prompt in prompts] + + +def greedy_sampling() -> SamplingParams: + return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False) + + +def stochastic_sampling() -> SamplingParams: + return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False) + + +def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None: + """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold. + + The default threshold assumes the LLM uses the same target model as the "model_name" + fixture, with max model len == 4096. Precomputed reference value is 75% to 80% + on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70% + to verify that the model is correct. + """ + if expected_accuracy_threshold <= 0.0: + print("Skipping GSM8K evaluation") + return + results = evaluate_gsm8k_offline(llm) + accuracy = results["accuracy"] + print(f"GSM8K accuracy: {accuracy:.3f}") + assert accuracy >= expected_accuracy_threshold, ( + f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}" + ) + + +def compute_acceptance_rate( + metrics: list[Metric], prev_metrics: list[Metric] | None = None +) -> float: + name2metric = {metric.name: metric for metric in metrics} + n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value + if n_draft_toks == 0: + return float("nan") + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value + if prev_metrics is not None: + prev_name2metric = {metric.name: metric for metric in prev_metrics} + n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value + n_accepted_toks -= prev_name2metric[ + "vllm:spec_decode_num_accepted_tokens" + ].value + if n_draft_toks <= 0: + return float("nan") + return n_accepted_toks / n_draft_toks + + +def compute_acceptance_len( + metrics: list[Metric], prev_metrics: list[Metric] | None = None +) -> float: + name2metric = {metric.name: metric for metric in metrics} + n_drafts = name2metric["vllm:spec_decode_num_drafts"].value + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value + if n_drafts == 0: + return 1 + if prev_metrics is not None: + prev_name2metric = {metric.name: metric for metric in prev_metrics} + n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value + n_accepted_toks -= prev_name2metric[ + "vllm:spec_decode_num_accepted_tokens" + ].value + if n_drafts <= 0: + return 1 + return 1 + (n_accepted_toks / n_drafts) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py new file mode 100644 index 000000000000..4fc6768e7446 --- /dev/null +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel.""" + +import pytest + +from vllm.v1.metrics.reader import Counter + +from ...models.utils import check_logprobs_close +from ...utils import large_gpu_mark, multi_gpu_test + +# Mamba2 (Nemotron-3) hybrid. +MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +MODELS = [ + pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), +] + +PROMPTS = [ + "The capital of France is", + "Once upon a time, in a small village,", +] + + +def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1): + # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a + # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are + # common-mode and only ReplaySSM varies. + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + tensor_parallel_size=tensor_parallel_size, + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline", + name_1="replayssm", + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_decode_matches_baseline(vllm_runner, model_name): + _check_replayssm_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): + # Tensor-parallel correctness: ReplaySSM's caches and checkpoint state are + # sharded per rank, so TP2 decode must still match the baseline at TP2. + _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) + + +# Prefix spans several mamba blocks; prefix caching only reuses full blocks. +_PC_SENTENCE = ( + "In a detailed survey of state space models, the authors compared many " + "architectures across a wide range of long-context language tasks and " + "measured their throughput, memory use, and accuracy in careful detail. " +) +_PC_PREFIX = _PC_SENTENCE * 120 +PREFIX_CACHING_PROMPTS = [ + _PC_PREFIX + "The most important conclusion was that", + _PC_PREFIX + "Surprisingly, the experiments showed that", + _PC_PREFIX + "The most important conclusion was that", +] + + +def _prefix_cache_hits(llm) -> int: + return sum( + m.value + for m in llm.llm.get_metrics() + if isinstance(m, Counter) and m.name == "vllm:prefix_cache_hits" + ) + + +def _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, *, tensor_parallel_size=1 +): + # align mode materializes the exact SSM state at each block boundary, so + # ReplaySSM's cached prefixes must match the always-materialized baseline. + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode="align", + disable_log_stats=False, # required for llm.get_metrics() + tensor_parallel_size=tensor_parallel_size, + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + # Prime the cache, then measure, so cache hits are deterministic. + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay_hits = _prefix_cache_hits(llm) + + # Without real cache hits the cached path is never exercised. + assert replay_hits > 0, ( + "ReplaySSM align-mode run produced no prefix-cache hits; the shared " + "prefix may be shorter than one mamba block, so prefix caching is inert" + ) + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline_align_pc", + name_1="replayssm_align_pc", + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, tensor_parallel_size=2 + ) diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index 65716444a57c..c58df0c076a5 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh @@ -21,12 +21,19 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # Model to test MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}" # Set 1 to use multimodal prompts; else to use text-only USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}" -MM_FLAG="" -if [ "$USE_MM_PROMPTS" = "1" ]; then - MM_FLAG="--use_mm_prompts" +USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}" +TEST_FLAGS=() +if [[ "$USE_MM_PROMPTS" == "1" ]]; then + TEST_FLAGS+=(--use_mm_prompts) +fi +if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then + TEST_FLAGS+=(--skip_two_image_prompt) fi # GPU configuration @@ -36,6 +43,16 @@ GPU_D="${GPU_D:-2}" GPU_SINGLE="${GPU_SINGLE:-$GPU_P}" GPU_PD="${GPU_PD:-$GPU_P}" +# Device platform and affinity environment variable +DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}" +if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then + if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then + DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK" + else + DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES" + fi +fi + # Port ENCODE_PORT="${ENCODE_PORT:-19534}" PREFILL_PORT="${PREFILL_PORT:-19535}" @@ -87,11 +104,12 @@ run_baseline() { # Start baseline instance echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT" - CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" vllm serve "$MODEL" \ --port "$PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ - --max-num-seqs 128 \ + --gpu-memory-utilization 0.9 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ > "$LOG_PATH"/baseline.log 2>&1 & @@ -112,7 +130,7 @@ run_baseline() { --model_name "$MODEL" \ --mode baseline \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup baseline echo "Stopping baseline instance..." @@ -139,14 +157,15 @@ run_epd_1e_1pd() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -160,12 +179,13 @@ run_epd_1e_1pd() { # Start prefill+decode instance echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_PD" vllm serve "$MODEL" \ --port "$PREFILL_DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -212,7 +232,7 @@ run_epd_1e_1pd() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1PD Correctness Test finished" @@ -242,14 +262,15 @@ run_baseline_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ vllm serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -260,14 +281,15 @@ run_baseline_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ vllm serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -309,7 +331,7 @@ run_baseline_1p_1d() { --model_name "$MODEL" \ --mode baseline_pd \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "Stopping PD (1P+1D) instances..." @@ -339,14 +361,15 @@ run_epd_1e_1p_1d() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -360,14 +383,15 @@ run_epd_1e_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ vllm serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -385,14 +409,15 @@ run_epd_1e_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ vllm serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -438,7 +463,7 @@ run_epd_1e_1p_1d() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1P+1D Correctness Test finished" @@ -465,7 +490,7 @@ run_epd_1e_1pd # Step 3: Test baseline 1P + 1D run_baseline_1p_1d -# Step 4: Test 1E + 1P + 1D +# # Step 4: Test 1E + 1P + 1D run_epd_1e_1p_1d # Cleanup output file diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index eae4b7427240..ece73efa42d3 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/tests/v1/ec_connector/integration/test_epd_correctness.py @@ -192,6 +192,12 @@ def main(): help="Use multimodal prompts (default: use text-only for quick testing)", ) + parser.add_argument( + "--skip_two_image_prompt", + action="store_true", + help="Skip the two-image multimodal prompt", + ) + args = parser.parse_args() print(f"Service URL: {args.service_url}") @@ -221,7 +227,9 @@ def main(): # Select prompts to use if args.use_mm_prompts: - test_prompts = SAMPLE_PROMPTS_MM + test_prompts = ( + SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM + ) print("Using multimodal prompts") else: test_prompts = SAMPLE_PROMPTS_TEXT diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index afb6e4c98b78..3845d6297a7c 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -19,6 +19,7 @@ from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.exceptions import VLLMValidationError from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform @@ -485,7 +486,7 @@ async def test_dp_rank_argument(): pass # Test with out-of-range DP rank. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): async for _ in engine.generate( request_id="request-35", prompt=TEXT_PROMPT, @@ -554,8 +555,8 @@ async def test_header_dp_rank_argument(): # Test 2: Out-of-range DP rank (1) mock_raw_request.headers = {"X-data-parallel-rank": "1"} - # should raise ValueError for out-of-range rank - with pytest.raises(ValueError): + # should raise VLLMValidationError for out-of-range rank + with pytest.raises(VLLMValidationError): await serving_chat.create_chat_completion(req, mock_raw_request) diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index a986bc07a3e8..a7a25ef0d6a7 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -308,7 +308,6 @@ def create_dp_placement_groups(vllm_config: Any): executor_class=_DummyExecutor, log_stats=False, addresses=addresses, - num_api_servers=2, ) as ( engine_manager, _coordinator, diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 0b44b205cd4d..64adf7a8b3c8 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -8,6 +8,7 @@ import signal import time import uuid +from collections import Counter from concurrent.futures import Future from dataclasses import dataclass from threading import Thread @@ -27,7 +28,11 @@ from vllm.pooling_params import LateInteractionParams, PoolingParams from vllm.usage.usage_lib import UsageContext from vllm.utils.torch_utils import set_default_torch_num_threads -from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest +from vllm.v1.engine import ( + EngineCoreOutputs, + EngineCoreReadyResponse, + EngineCoreRequest, +) from vllm.v1.engine.core import EngineCore from vllm.v1.engine.core_client import ( AsyncMPClient, @@ -198,13 +203,19 @@ def _make_pooling_request( ) -def test_dplb_late_interaction_sticky_routing(): +def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncMPClient: client = object.__new__(DPLBAsyncMPClient) - client.client_count = 1 + client.client_count = client_count client.reqs_in_flight = {} - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"] - client.lb_engines = [[0, 0], [0, 0], [0, 0]] + client.engine_inflight = Counter() + client.core_engines = [bytes([i, 0]) for i in range(num_engines)] + client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)] client.eng_start_index = 0 + return client + + +def test_dplb_late_interaction_sticky_routing(): + client = _make_dplb_client() query_key = "rerank-abc-query-0" query_request = _make_pooling_request( @@ -223,12 +234,8 @@ def test_dplb_late_interaction_sticky_routing(): def test_dplb_non_late_interaction_still_uses_lb(): - client = object.__new__(DPLBAsyncMPClient) - client.client_count = 1 - client.reqs_in_flight = {} - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"] - client.lb_engines = [[2, 1], [0, 0], [1, 0]] - client.eng_start_index = 0 + client = _make_dplb_client() + client.lb_engines = [[2, 1, 0.0], [0, 0, 0.0], [1, 0, 0.0]] request = make_request(SamplingParams(max_tokens=1)) chosen_engine = client.get_core_engine_for_request(request) @@ -237,6 +244,70 @@ def test_dplb_non_late_interaction_still_uses_lb(): assert client.lb_engines[1][0] == 1 +def test_dplb_burst_round_robins_despite_snapshot_rebinds(): + """A stats snapshot rebind wipes the optimistic lb_engines increments; + the exact in-flight floor must keep a burst spreading round-robin.""" + client = _make_dplb_client(num_engines=4) + + for _ in range(4): + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1))) + # Coordinator snapshot arrives, not yet reflecting the 4 routed requests. + client.lb_engines = [[0, 0, 0.0] for _ in range(4)] + for _ in range(4): + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1))) + + assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2] + + +def test_dplb_snapshot_backpressure_overrides_inflight(): + """An engine reported heavily loaded by the coordinator is avoided even + when this client has routed nothing to it.""" + client = _make_dplb_client(num_engines=2) + client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + + assert chosen == client.core_engines[1] + + +def test_dplb_kv_pressure_amplifies_waiting_penalty(): + """A waiting queue on a KV-bound engine (slow drain) is penalized, while + the same queue with low KV usage is not (e.g. transient burst).""" + client = _make_dplb_client(num_engines=2) + # Engine 0 has a smaller total but is KV-bound with a queue. + client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + assert chosen == client.core_engines[1] + + # Same counts without KV pressure: the smaller total wins. + client = _make_dplb_client(num_engines=2) + client.lb_engines = [[5, 10, 0.2], [0, 20, 0.2]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + assert chosen == client.core_engines[0] + + +def test_dplb_finished_requests_release_inflight(): + client = _make_dplb_client(num_engines=2) + + req = make_request(SamplingParams(max_tokens=1)) + engine = client.get_core_engine_for_request(req) + assert client.engine_inflight[engine] == 1 + + outputs = EngineCoreOutputs(finished_requests={req.request_id}) + asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs)) + + assert client.engine_inflight[engine] == 0 + assert req.request_id not in client.reqs_in_flight + + def test_apply_ready_response_syncs_block_size(): import msgspec @@ -257,6 +328,13 @@ def test_apply_ready_response_syncs_block_size(): vllm_version="test", world_size=1, data_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", ) ) client._apply_ready_response(payload) diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 1919349790fa..51dbb9c98952 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -141,6 +141,82 @@ def test_incremental_detokenization( assert not output_processor.has_unfinished_requests() +def test_request_stream_interval_raises_but_not_below_engine_default( + dummy_test_vectors, +): + """A per-request stream_interval can raise the interval above the engine + default but not below it (values under the default clamp up), without + altering the generated text.""" + engine_stream_interval = 5 + # Request 0 (below the default) clamps up to 5; request 1 raises it to 10. + request_stream_intervals = [1, 10] + output_processor = OutputProcessor( + dummy_test_vectors.tokenizer, + log_stats=False, + stream_interval=engine_stream_interval, + ) + + requests = [ + EngineCoreRequest( + request_id=f"request-{idx}-int", + external_req_id=f"request-{idx}", + prompt_token_ids=prompt_tokens, + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=SamplingParams( + skip_special_tokens=False, + spaces_between_special_tokens=False, + output_kind=RequestOutputKind.DELTA, + stop=[], + include_stop_str_in_output=False, + stream_interval=request_stream_intervals[idx], + ), + pooling_params=None, + ) + for idx, prompt_tokens in enumerate( + dummy_test_vectors.prompt_tokens[: len(request_stream_intervals)] + ) + ] + + num_requests = len(requests) + engine_core = MockEngineCore( + tokens_list=dummy_test_vectors.generation_tokens[:num_requests], + prompts_list=dummy_test_vectors.prompt_tokens[:num_requests], + request_ids=[req.request_id for req in requests], + ) + + for request, prompt in zip(requests, dummy_test_vectors.prompt_strings): + output_processor.add_request(request, prompt) + + gen_strings: dict[str, str] = {} + gen_tokens: dict[str, list[int]] = {} + while outputs := engine_core.get_outputs(): + for request_output in output_processor.process_outputs(outputs).request_outputs: + request_id = request_output.request_id + new_tokens = request_output.outputs[0].token_ids + if request_id not in gen_strings: + gen_strings[request_id] = request_output.outputs[0].text + gen_tokens[request_id] = list(new_tokens) + assert len(new_tokens) == 1, f"{len(new_tokens)=}" + continue + gen_strings[request_id] += request_output.outputs[0].text + gen_tokens[request_id].extend(new_tokens) + if not request_output.finished: + requested = request_stream_intervals[int(request_id.split("-")[1])] + interval = max(requested, engine_stream_interval) + assert len(new_tokens) == interval, f"{len(new_tokens)=}, {interval=}" + + for idx in range(num_requests): + request_id = f"request-{idx}" + assert gen_strings[request_id] == dummy_test_vectors.generation_strings[idx] + assert gen_tokens[request_id] == dummy_test_vectors.generation_tokens[idx] + + assert not output_processor.has_unfinished_requests() + + def _validate_logprobs( gen_tokens: dict[str, list[int]], gen_logprobs: dict[str, SampleLogprobs | None], diff --git a/tests/v1/executor/test_multiproc_executor_timeout.py b/tests/v1/executor/test_multiproc_executor_timeout.py new file mode 100644 index 000000000000..8d7998920cf6 --- /dev/null +++ b/tests/v1/executor/test_multiproc_executor_timeout.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for stale multiproc RPC deadlines.""" + +import time +from collections import deque +from concurrent.futures import Future, InvalidStateError +from contextlib import suppress +from unittest.mock import patch + + +def _dequeue_timeout(deadline: float | None) -> float | None: + return None if deadline is None else max(0.0, deadline - time.monotonic()) + + +class FutureWrapper(Future): + def __init__( + self, + futures_queue: deque["FutureWrapper"], + get_response, + aggregate=lambda x: x, + ): + self.futures_queue = futures_queue + self.get_response = get_response + self.aggregate = aggregate + super().__init__() + self.futures_queue.appendleft(self) + + def result(self, timeout=None): + if timeout is not None: + raise RuntimeError("timeout not implemented") + + while not self.done(): + future = self.futures_queue.pop() + future._wait_for_response() + return super().result() + + def _wait_for_response(self): + try: + response = self.aggregate(self.get_response()) + with suppress(InvalidStateError): + self.set_result(response) + except Exception as e: + with suppress(InvalidStateError): + self.set_exception(e) + + +class FakeClock: + def __init__(self, start: float = 0.0) -> None: + self._now = start + + def monotonic(self) -> float: + return self._now + + def advance(self, seconds: float) -> None: + self._now += seconds + + +class FakeResponseMQ: + def __init__(self, response: object = ("SUCCESS", "dummy_response")) -> None: + self.response = response + self.timeouts: list[float | None] = [] + + def dequeue(self, timeout: float | None = None) -> object: + assert timeout is None or timeout >= 0.0, ( + f"dequeue received negative timeout: {timeout}" + ) + self.timeouts.append(timeout) + return self.response + + +def test_future_wrapper_stale_deadline_never_passes_negative_timeout() -> None: + clock = FakeClock(100.0) + futures_queue: deque[FutureWrapper] = deque() + response_mq = FakeResponseMQ() + + deadline = clock.monotonic() + 1.0 + + def get_response() -> object: + return response_mq.dequeue(timeout=_dequeue_timeout(deadline)) + + future = FutureWrapper( + futures_queue, + get_response=get_response, + aggregate=lambda x: x, + ) + + clock.advance(5.0) + + with patch("time.monotonic", clock.monotonic): + future.result() + + assert response_mq.timeouts == [0.0] + + +def test_future_wrapper_non_expired_deadline_passes_positive_timeout() -> None: + clock = FakeClock(100.0) + futures_queue: deque[FutureWrapper] = deque() + response_mq = FakeResponseMQ() + + deadline = clock.monotonic() + 5.0 + + def get_response() -> object: + return response_mq.dequeue(timeout=_dequeue_timeout(deadline)) + + future = FutureWrapper( + futures_queue, + get_response=get_response, + aggregate=lambda x: x, + ) + + clock.advance(2.0) + + with patch("time.monotonic", clock.monotonic): + future.result() + + assert len(response_mq.timeouts) == 1 + timeout = response_mq.timeouts[0] + assert timeout is not None + assert timeout > 0.0 + assert timeout <= 5.0 + + +def test_future_wrapper_deadline_none_passes_none() -> None: + futures_queue: deque[FutureWrapper] = deque() + response_mq = FakeResponseMQ() + + def get_response() -> object: + return response_mq.dequeue(timeout=_dequeue_timeout(None)) + + future = FutureWrapper( + futures_queue, + get_response=get_response, + aggregate=lambda x: x, + ) + + future.result() + + assert response_mq.timeouts == [None] + + +def test_future_wrapper_drains_pending_before_own_get_response() -> None: + futures_queue: deque[FutureWrapper] = deque() + call_order: list[str] = [] + + def make_get_response(label: str): + def _get() -> object: + call_order.append(label) + return FakeResponseMQ().dequeue(timeout=None) + + return _get + + first = FutureWrapper(futures_queue, make_get_response("first")) + second = FutureWrapper(futures_queue, make_get_response("second")) + + second.result() + + assert first.done() + assert second.done() + assert call_order == ["first", "second"] + + +def test_recv_timeout_ms_clamps_negative_timeout() -> None: + def recv_timeout_ms(timeout: float | None) -> int | None: + return None if timeout is None else max(0, int(timeout * 1000)) + + assert recv_timeout_ms(None) is None + assert recv_timeout_ms(-1.0) == 0 + assert recv_timeout_ms(-0.001) == 0 + assert recv_timeout_ms(0.0) == 0 + assert recv_timeout_ms(0.001) == 1 + assert recv_timeout_ms(2.5) == 2500 diff --git a/tests/v1/fault_tolerance/__init__.py b/tests/v1/fault_tolerance/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/fault_tolerance/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py new file mode 100644 index 000000000000..f6d15343b56f --- /dev/null +++ b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the elastic fault-tolerance framework. + +Requires nixl_ep FT hardware; gated behind ``has_nixl_ep()``. +""" + +import contextlib +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import psutil +import pytest +import requests + +from tests.utils import RemoteOpenAIServer, multi_gpu_test +from vllm.utils.import_utils import has_nixl_ep + +MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") +DP_SIZE = 2 + +# Fault-detection timeout budget: +# - CPU: Gloo DP allreduce timeout (30s) detects the dead peer. +# - nixl_ep: kernel masks the dead rank after Buffer's default timeout_ms=30000 (30s). +# - Deadline (45s): slowest fallback (30s) + margin. +CPU_DISTRIBUTED_TIMEOUT_S = 30 +FAULT_DETECTION_DEADLINE_S = 45 + + +# Patches ``gpu.dp_utils.sync_cudagraph_and_dp_padding`` to raise on ``rank`` at +# a chosen step. Gated on VLLM_FT_TEST_INJECT_FAULT. +_FAULT_INJECT_SITECUSTOMIZE = """\ +import builtins +import os +import sys + +_SPEC = os.environ.get("VLLM_FT_TEST_INJECT_FAULT") +_MODULE = "vllm.v1.worker.gpu.dp_utils" +_ATTR = "sync_cudagraph_and_dp_padding" + +if _SPEC: + _f = dict(kv.split("=", 1) for kv in _SPEC.split(",")) + _RANK, _STEP = int(_f["rank"]), int(_f["step"]) + _steps = [0] + + def _patch(m): + import inspect + _orig = getattr(m, _ATTR) + _sig = inspect.signature(_orig) + def _wrapped(*args, **kwargs): + result = _orig(*args, **kwargs) + bound = _sig.bind(*args, **kwargs) + bound.apply_defaults() + dp_rank = bound.arguments.get("dp_rank") + if dp_rank == _RANK: + _steps[0] += 1 + if _steps[0] == _STEP: + raise RuntimeError( + "FT test fault injection (rank=%d step=%d)" % (_RANK, _STEP) + ) + return result + + setattr(m, _ATTR, _wrapped) + + _real_import = builtins.__import__ + + def _hook(name, *a, **k): + module = _real_import(name, *a, **k) + m = sys.modules.get(_MODULE) + # During vLLM's circular import the module lands in sys.modules before + # its functions are defined; hasattr guards against patching too early. + if ( + m is not None + and hasattr(m, _ATTR) + and not getattr(m, "_ft_patched", False) + ): + m._ft_patched = True + _patch(m) + return module + + builtins.__import__ = _hook +""" + + +def _install_fault_injection(monkeypatch, tmp_path, rank: int, step: int) -> None: + """Arrange for the DP-sync fn to raise on ``rank`` at serving ``step``. + + Writes a ``sitecustomize.py`` and prepends its dir to PYTHONPATH so every + vLLM subprocess picks it up; the fault spec is read from the environment. + """ + site_dir = tmp_path / "ft_inject" + site_dir.mkdir() + (site_dir / "sitecustomize.py").write_text(_FAULT_INJECT_SITECUSTOMIZE) + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv( + "PYTHONPATH", + str(site_dir) + (os.pathsep + existing if existing else ""), + ) + monkeypatch.setenv("VLLM_FT_TEST_INJECT_FAULT", f"rank={rank},step={step}") + + +def _ft_server_args() -> list[str]: + return [ + "--enforce-eager", + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "128", + "--enable-expert-parallel", + "--all2all-backend", + "nixl_ep", + "--enable-fault-tolerance", + "--cpu-distributed-timeout-seconds", + str(CPU_DISTRIBUTED_TIMEOUT_S), + "--fault-tolerance-config", + '{"engine_recovery_timeout_sec": 120}', + ] + + +def _ft_manager(): + """Build the shared DP+EP fault-tolerant server topology (one engine/server).""" + from tests.v1.distributed.test_external_lb_dp import ExternalLBServerManager + + return ExternalLBServerManager( + MODEL_NAME, + DP_SIZE, + api_server_count=1, # FT requires a single API server per engine + base_server_args=_ft_server_args(), + tp_size=1, + ) + + +def _server_for_rank(servers, rank: int): + """Locate the server for a DP rank.""" + for server, sargs in servers: + if "--data-parallel-rank" in sargs: + idx = sargs.index("--data-parallel-rank") + if int(sargs[idx + 1]) == rank: + return server + raise AssertionError(f"no server found for DP rank {rank}") + + +def _complete(client): + """Issue the one standard completion the tests use everywhere.""" + return client.completions.create( + model=MODEL_NAME, + prompt="Hello, my name is", + max_tokens=5, + temperature=0.0, + timeout=10.0, + ) + + +def _in_parallel(fn, servers) -> list: + """Run ``fn(server)`` for all servers concurrently; return results in order.""" + with ThreadPoolExecutor(max_workers=len(servers)) as ex: + return list(ex.map(fn, servers)) + + +def _get_ft_status(server) -> dict: + resp = requests.get(server.url_for("fault_tolerance/status"), timeout=10) + resp.raise_for_status() + return resp.json() + + +def _assert_serving_and_healthy(servers) -> None: + """Wait until every engine is healthy, then serve one request per server.""" + healthy = _wait_for_engines( + list(servers), match_key="status", match_values={"healthy"} + ) + assert all(healthy), healthy + _in_parallel(lambda s: _complete(s.get_client()), servers) + + +def _apply_ft(server, instruction: str, params: dict | None = None) -> dict: + """POST an FT instruction; assert it is accepted (202) and return the body.""" + resp = requests.post( + server.url_for("fault_tolerance/apply"), + json={"instruction": instruction, "params": params or {}}, + timeout=10, + ) + assert resp.status_code == 202, resp.text + return resp.json() + + +def _kill_worker_process(server) -> None: + """SIGKILL only the worker proc, leaving EngineCore and API server alive.""" + workers = [ + p + for p in psutil.Process(server.proc.pid).children(recursive=True) + if "Worker" in " ".join(p.cmdline()) + ] + assert len(workers) == 1, f"expected 1 worker proc, found: {workers}" + workers[0].kill() + + +def _wait_for_engines( + servers: list[RemoteOpenAIServer], + match_key: str, + match_values: set[str], + deadline_s: int = FAULT_DETECTION_DEADLINE_S, +) -> list[dict[str, Any] | None]: + """Poll ``/fault_tolerance/status`` until each server's engine status matches. + + A server matches when its engine-status dict has ``match_key`` equal to + one of ``match_values``. Returns one engine-status dict per server. Servers still + unmatched after ``deadline_s`` get None. + """ + results: dict[int, dict[str, Any]] = {} + pending = dict(enumerate(servers)) + start = time.time() + while pending and time.time() - start < deadline_s: + for i, server in list(pending.items()): + with contextlib.suppress(Exception): + for engine_status in _get_ft_status(server)["engines"]: + if engine_status.get(match_key) in match_values: + results[i] = engine_status + del pending[i] + break + if pending: + time.sleep(1.0) + return [results.get(i) for i in range(len(servers))] + + +@contextlib.contextmanager +def _driving(*servers): + """Pump completions at each server in the background for the block's duration. + + Keeps every engine stepping into its failed component so a fault surfaces. + Errors are expected once faulted and are ignored. + """ + stop = threading.Event() + + def _drive(server): + client = server.get_client() + while not stop.is_set(): + with contextlib.suppress(Exception): + _complete(client) + time.sleep(0.2) + + threads = [threading.Thread(target=_drive, args=(s,), daemon=True) for s in servers] + for t in threads: + t.start() + try: + yield + finally: + stop.set() + for t in threads: + t.join(timeout=2) + + +def _wait_for_ft_apply_outcome(server, request_id: str, deadline_s: int) -> str | None: + """Wait until ``/fault_tolerance/status`` records the FT apply outcome.""" + engine_status = _wait_for_engines( + [server], + match_key="last_ft_request_id", + match_values={request_id}, + deadline_s=deadline_s, + )[0] + return engine_status.get("ft_error") if engine_status else None + + +@pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend") +@multi_gpu_test(num_gpus=2) +def test_injected_fault_retry_recovers_all_ranks(monkeypatch, tmp_path): + """An exception injected into the inference path drives full retry recovery. + + Injecting an exception into ``sync_cudagraph_and_dp_padding`` at a chosen + step on rank 1. + + - Rank 1 raises inside the busy loop and goes UNHEALTHY. + - Rank 0 detects the now-absent peer via the communication timeout and also + goes UNHEALTHY. + + Both being UNHEALTHY is the precondition for ``retry``. The fault is patched + into the DP-sync fn from the test (via a generated ``sitecustomize``). + """ + fault_step = int(os.getenv("FT_FAULT_STEP", "50")) + _install_fault_injection(monkeypatch, tmp_path, rank=1, step=fault_step) + + with _ft_manager() as servers: + assert len(servers) == DP_SIZE + rank0 = _server_for_rank(servers, 0) + rank1 = _server_for_rank(servers, 1) + + # 1. Both engines healthy and serving. + _assert_serving_and_healthy((rank0, rank1)) + + # 2. Drive both ranks so rank 1 accumulates execute_model steps and trips + # the injected fault; rank 0 then times out on the DP allreduce. + with _driving(rank0, rank1): + faulted = _wait_for_engines( + [rank0, rank1], match_key="status", match_values={"unhealthy"} + ) + + for rank, engine_status in enumerate(faulted): + assert engine_status is not None, ( + f"rank {rank} did not report UNHEALTHY within " + f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung" + ) + # The rank that raised carries the fault info from its own exception. + assert faulted[1] is not None + assert faulted[1].get("fault_info"), faulted[1] + + # 3. retry both engines. + for server in (rank0, rank1): + _apply_ft(server, "retry") + + # 4. Recovery completes: both engines return to healthy and serve again. + _assert_serving_and_healthy((rank0, rank1)) + + +@pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend") +@multi_gpu_test(num_gpus=2) +def test_worker_kill_survivor_unhealthy_and_dead_rejects_retry(): + """One worker kill surfaces two status transitions at once. + + SIGKILLing only rank 1's worker leaves both EngineCores alive, so the same + fault is seen two ways: + + - Survivor (rank 0): detects the dead peer via Gloo allreduce / nixl_ep + kernel timeout. Its own executor is fine, so ``on_fault`` marks it + UNHEALTHY with a ``fault_info``. + - Victim (rank 1): detects its own executor failure and marks itself DEAD. + + Recovery is gated on UNHEALTHY: the DEAD engine accepts ``retry`` at the + HTTP layer (202 = background dispatch) but rejects it in the engine, + recording the reason as ``ft_error``. + """ + with _ft_manager() as servers: + assert len(servers) == DP_SIZE + survivor = _server_for_rank(servers, 0) + victim = _server_for_rank(servers, 1) + + # 1. Confirm both engines are healthy and serving. + _assert_serving_and_healthy((survivor, victim)) + + # 2. Kill only the victim's worker; both EngineCores stay alive. + _kill_worker_process(victim) + + # 3. Drive both engines so each keeps stepping into the failed component. + with _driving(survivor, victim): + survivor_faulted, victim_faulted = _wait_for_engines( + [survivor, victim], + match_key="status", + match_values={"dead", "unhealthy"}, + ) + + assert survivor_faulted is not None, ( + "survivor did not report the peer fault within " + f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung" + ) + # The survivor's own executor is fine, so it must be UNHEALTHY, not DEAD. + assert survivor_faulted["status"] == "unhealthy", survivor_faulted + assert survivor_faulted.get("fault_info"), survivor_faulted + + assert victim_faulted is not None, ( + "victim did not report its worker's death within " + f"{FAULT_DETECTION_DEADLINE_S}s" + ) + assert victim_faulted["status"] == "dead", victim_faulted + + # 4. retry is accepted at the HTTP layer (202 = background dispatch)... + request_id = _apply_ft(victim, "retry")["request_id"] + + # 5. ...but the DEAD engine must reject it: recovery requires UNHEALTHY. + ft_error = _wait_for_ft_apply_outcome( + victim, request_id, FAULT_DETECTION_DEADLINE_S + ) + assert ft_error is not None, ( + "rejection was never recorded in /fault_tolerance/status" + ) + assert "status is DEAD" in ft_error, ft_error diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index acc7ce312f54..8acc96413327 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -25,6 +25,7 @@ dp_ep_configs=( hybrid_ssm_configs=( "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" + "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" # GDN (Qwen3.5) "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index 77235582fe13..fc5c04a1ad07 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -83,6 +83,11 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128} ENFORCE_EAGER=${ENFORCE_EAGER:-1} # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048) VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} +# Pin concurrent prefiller and non-DP decoder engines to separate internal +# port windows. DP decoder ranks retain their existing internal port selection. +PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000} +DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000} +INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100} # Resolve the repository root from the script location instead of `.git`. # The ROCm CI image copies `/vllm-workspace` without the Git metadata, so @@ -93,7 +98,7 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") # Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT # Waits for vLLM to start. wait_for_server() { @@ -105,10 +110,20 @@ wait_for_server() { } # Function to clean up previous instances +wait_for_gpu_memory_release() { + if [[ "$SMI_BIN" == *"rocm"* ]]; then + PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()" + fi +} + cleanup_instances() { echo "Cleaning up any running vLLM instances..." - pkill -f "vllm serve" || true + pkill -f "toy_proxy_server.py" || true + pkill -TERM -f "vllm serve" || true + sleep 3 + pkill -9 -f "vllm serve" || true sleep 2 + wait_for_gpu_memory_release } get_num_gpus() { @@ -127,6 +142,7 @@ get_num_gpus() { # Function to run tests for a specific model run_tests_for_model() { local model_name=$1 + cleanup_instances echo "================================" echo "Testing model: $model_name" echo "================================" @@ -154,12 +170,14 @@ run_tests_for_model() { PORT=$((8100 + i)) # Calculate side channel port. Avoid clash with with TP workers. SIDE_CHANNEL_PORT=$((5559 + i)) + INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT" # Build the command with or without model-specific args BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_PORT=$INTERNAL_PORT \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ vllm serve $model_name \ @@ -208,12 +226,21 @@ run_tests_for_model() { PORT=$((8200 + i)) # Calculate side channel port SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE)) + INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) + # For non-DP mode, set VLLM_PORT to pin the internal port; + # For DP mode, set VLLM_DP_MASTER_PORT instead to avoid race condition. + if [[ -z "${DP_EP:-}" ]]; then + DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT" + else + DECODER_INTERNAL_PORT_ENV="VLLM_DP_MASTER_PORT=$INTERNAL_PORT" + fi echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" # Build the command with or without model-specific args BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \ + $DECODER_INTERNAL_PORT_ENV \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ vllm serve $model_name \ @@ -290,7 +317,6 @@ run_tests_for_model() { # Clean up before running next model cleanup_instances - sleep 3 } # Run tests for each model diff --git a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh index 9d8e4df8c539..c3240ab5c179 100755 --- a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh @@ -3,8 +3,8 @@ set -xe # Parse command line arguments KV_BUFFER_DEVICE="cuda" # Default to cuda -PREFILL_GPU_ID=4 # Default GPU IDs -DECODE_GPU_ID=5 +PREFILL_GPU_ID="${PREFILL_GPU_ID:-4}" # Default GPU IDs +DECODE_GPU_ID="${DECODE_GPU_ID:-5}" while [[ $# -gt 0 ]]; do case $1 in --kv_buffer_device) @@ -70,6 +70,7 @@ run_tests_for_model() { --port $PREFILL_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -84,6 +85,7 @@ run_tests_for_model() { --port $DECODE_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -98,7 +100,7 @@ run_tests_for_model() { # Build the command for the proxy server with all the hosts and ports PROXY_PORT=8192 - PROXY_CMD="python ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" + PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" PROXY_CMD+=" --prefiller-ports ${PREFILL_PORT}" PROXY_CMD+=" --decoder-ports ${DECODE_PORT}" # Start the proxy server @@ -110,7 +112,7 @@ run_tests_for_model() { # Run lm eval for this model echo "Running tests for $model_name" - PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py + PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python3 -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py # Clean up before running next model cleanup_instances diff --git a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py index 4422f45847bd..4c1489d3f865 100644 --- a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py +++ b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py @@ -61,7 +61,17 @@ def test_nixl_and_nixl_ep_imports() -> None: importlib.import_module("nixl._bindings") # Exercise the NIXL EP extension used by fused MoE expert parallelism. - nixl_ep = importlib.import_module("nixl_ep") + try: + nixl_ep = importlib.import_module("nixl_ep") + except ImportError as e: + if "materialize_cow_storage" in str(e) or "undefined symbol" in str(e): + pytest.xfail( + "nixl_ep prebuilt extension is ABI-incompatible with this torch " + "(undefined symbol c10::impl::cow::materialize_cow_storage); " + "needs a nixl rebuild against torch 2.13. " + "See pytorch/pytorch#187727 and ai-dynamo/nixl#1798." + ) + raise print(f"nixl_ep: {nixl_ep.__file__}") assert nixl_ep.__file__ is not None diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py new file mode 100644 index 000000000000..aefb536c4c4f --- /dev/null +++ b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import ( + _layer_mapping, + _opaque_fallback_mapping, + _RankContext, + _verify_tiling, + derive_canonical_mappings, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + KVCacheSpec, + KVQuantMode, + MLAAttentionSpec, +) +from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun + +NUM_BLOCKS = 3 + + +def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2): + return _RankContext( + tp_size=tp, + dcp_size=dcp, + pcp_size=pcp, + interleave=interleave, + total_kv_heads=heads * tp if total is None else total, + rank=rank, + ) + + +def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec: + # block_size=4, head_dim=64, int8 + return FullAttentionSpec( + block_size=4, + num_kv_heads=num_kv_heads, + head_size=64, + dtype=torch.int8, + **kwargs, + ) + + +def _mla_spec(**kwargs) -> MLAAttentionSpec: + # page = 4 * 64 = 256B, one 64B latent row per token + return MLAAttentionSpec( + block_size=4, num_kv_heads=1, head_size=64, dtype=torch.int8, **kwargs + ) + + +def _split_nhd_cache(spec) -> torch.Tensor: + return torch.zeros( + NUM_BLOCKS, + 2, + spec.block_size, + spec.num_kv_heads, + spec.head_size, + dtype=torch.int8, + ) + + +def _split_hnd_cache(spec) -> torch.Tensor: + return torch.zeros( + NUM_BLOCKS, + 2, + spec.num_kv_heads, + spec.block_size, + spec.head_size, + dtype=torch.int8, + ).permute(0, 1, 3, 2, 4) + + +def _packed_nhd_cache(spec) -> torch.Tensor: + """Logical (num_blocks, heads, block_size, 2 * head_size) over an NHD + physical layout — the FlashAttention/FlashInfer/Triton/Flex form.""" + return torch.zeros( + NUM_BLOCKS, + spec.block_size, + spec.num_kv_heads, + 2 * spec.head_size, + dtype=torch.int8, + ).permute(0, 2, 1, 3) + + +def _packed_hnd_cache(spec) -> torch.Tensor: + return torch.zeros( + NUM_BLOCKS, + spec.num_kv_heads, + spec.block_size, + 2 * spec.head_size, + dtype=torch.int8, + ) + + +CACHE_BUILDERS = { + "split_nhd": _split_nhd_cache, + "split_hnd": _split_hnd_cache, + "packed_nhd": _packed_nhd_cache, + "packed_hnd": _packed_hnd_cache, +} + + +def _try_mapping(spec, kv_cache, ctx) -> CanonicalPageMapping | None: + return _layer_mapping(spec, kv_cache, NUM_BLOCKS, ctx) + + +def _mapping(spec, kv_cache, ctx) -> CanonicalPageMapping: + mapping = _try_mapping(spec, kv_cache, ctx) + assert mapping is not None + return mapping + + +def _triples(runs: tuple[CopyRun, ...]) -> list[tuple[int, int, int]]: + """Expand runs to explicit (local_offset, canonical_offset, size) copies.""" + out = [] + for run in runs: + for i in range(run.num_fragments): + out.append( + ( + run.local_offset + i * run.local_stride, + run.canonical_offset + i * run.canonical_stride, + run.fragment_size, + ) + ) + return out + + +# --------------------------------------------------------------------------- +# TP-only placement (byte-compatible with the uniform interleave layout) +# --------------------------------------------------------------------------- + + +def test_split_nhd_placement_rank2_of_4(): + spec = _full_spec() + mapping = _mapping(spec, _split_nhd_cache(spec), _ctx(rank=2, tp=4)) + assert mapping.canonical_page_size_bytes == 4 * 1024 + assert mapping.parallelism_agnostic + k_dst = [256, 768, 1280, 1792] + assert _triples(mapping.runs) == [ + (local, canonical, 128) + for local, canonical in zip( + [0, 128, 256, 384, 512, 640, 768, 896], + k_dst + [2048 + o for o in k_dst], + ) + ] + # Heads are sharded, not replicated: this rank writes every block + assert mapping.num_writers == 1 + + +def test_packed_nhd_placement_rank2_of_4(): + spec = _full_spec() + mapping = _mapping(spec, _packed_nhd_cache(spec), _ctx(rank=2, tp=4)) + assert _triples(mapping.runs) == [ + (0, 512, 256), + (256, 1536, 256), + (512, 2560, 256), + (768, 3584, 256), + ] + + +def test_packed_hnd_placement_rank1_of_4(): + spec = _full_spec() + mapping = _mapping(spec, _packed_hnd_cache(spec), _ctx(rank=1, tp=4)) + # Two heads, each a contiguous canonical head region of 4 tokens x 128B + assert _triples(mapping.runs) == [(0, 1024, 1024)] + + +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS)) +def test_single_rank_coalesces_to_one_run(form): + spec = _full_spec() + mapping = _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank=0)) + assert mapping.canonical_page_size_bytes == 1024 + assert _triples(mapping.runs) == [(0, 0, 1024)] + + +# --------------------------------------------------------------------------- +# Replication and writer election +# --------------------------------------------------------------------------- + + +def test_gqa_replicated_heads_rotate_writer(): + # total 2 KV heads on tp=4: replication factor 2, head shard = rank // 2 + spec = _full_spec(num_kv_heads=1) + cache = _split_nhd_cache(spec) + ctx = lambda rank: _ctx(rank, tp=4, total=2) # noqa: E731 + rank2 = _mapping(spec, cache, ctx(2)) + rank3 = _mapping(spec, cache, ctx(3)) + assert rank2.canonical_page_size_bytes == 2 * 512 + # K region: head shard 1 at 64B offsets within 128B token rows + assert _triples(rank2.runs)[:4] == [ + (0, 64, 64), + (64, 192, 64), + (128, 320, 64), + (192, 448, 64), + ] + # Same head shard, so identical bytes: the two take alternate blocks + assert rank3.runs == rank2.runs + assert (rank2.is_writer(0), rank3.is_writer(0)) == (True, False) + assert (rank2.is_writer(1), rank3.is_writer(1)) == (False, True) + _verify_tiling("gqa", [_mapping(spec, cache, ctx(r)) for r in range(4)]) + + +def test_mla_replicas_rotate_writer(): + spec = _mla_spec() + rank0 = _mapping(spec, None, _ctx(rank=0, tp=2)) + rank1 = _mapping(spec, None, _ctx(rank=1, tp=2)) + # Latent pages are stored once per block, not once per rank + assert rank0.canonical_page_size_bytes == 256 + assert _triples(rank0.runs) == [(0, 0, 256)] + assert rank1.runs == rank0.runs + assert (rank0.is_writer(0), rank1.is_writer(0)) == (True, False) + assert (rank0.is_writer(1), rank1.is_writer(1)) == (False, True) + + +# --------------------------------------------------------------------------- +# DCP / PCP token sharding +# --------------------------------------------------------------------------- + + +def test_dcp_interleaves_tokens_within_replicas(): + # tp=4, dcp=2, total 2 KV heads: head shard = rank // 2, cp rank = rank % 2 + spec = _full_spec(num_kv_heads=1) + cache = _split_nhd_cache(spec) + ctx = lambda rank: _ctx(rank, tp=4, dcp=2, total=2) # noqa: E731 + per_rank = [_mapping(spec, cache, ctx(rank)) for rank in range(4)] + assert all(m is not None for m in per_rank) + # 8 canonical tokens x 2 heads x 64B per region + assert per_rank[0].canonical_page_size_bytes == 2048 + assert not per_rank[0].parallelism_agnostic + # rank 2 = head shard 1, cp rank 0: K tokens 0,2,4,6 at head offset 64 + assert _triples(per_rank[2].runs)[:4] == [ + (0, 64, 64), + (64, 320, 64), + (128, 576, 64), + (192, 832, 64), + ] + # every rank contributes (dcp == replication: no residual replicas) + assert all(m.num_writers == 1 for m in per_rank) + _verify_tiling("dcp", per_rank) + + +def test_mla_dcp_shards_latent_tokens(): + spec = _mla_spec() + ctx = lambda rank: _ctx(rank, tp=2, dcp=2) # noqa: E731 + rank0 = _mapping(spec, None, ctx(0)) + rank1 = _mapping(spec, None, ctx(1)) + assert rank0.canonical_page_size_bytes == 512 + assert _triples(rank0.runs) == [(o, 2 * o, 64) for o in (0, 64, 128, 192)] + assert _triples(rank1.runs) == [(o, 2 * o + 64, 64) for o in (0, 64, 128, 192)] + _verify_tiling("mla-dcp", [rank0, rank1]) + + +def test_pcp_tokens_and_tp_heads_compose(): + # tp=2 x pcp=2: rank = pcp_rank * 2 + tp_rank; 4 workers tile the page + spec = _full_spec(num_kv_heads=1) + cache = _packed_nhd_cache(spec) + per_rank = [ + _mapping(spec, cache, _ctx(rank, tp=2, pcp=2, total=2)) for rank in range(4) + ] + assert all(m is not None and m.runs for m in per_rank) + _verify_tiling("pcp", per_rank) + + +def test_interleave_chunks_stay_contiguous(): + # interleave=2: chunks of 2 tokens alternate between the 2 cp ranks and + # coalesce into one contiguous fragment per chunk + spec = _mla_spec() + mapping = _mapping(spec, None, _ctx(rank=0, tp=2, dcp=2, interleave=2)) + assert _triples(mapping.runs) == [(0, 0, 128), (128, 256, 128)] + _verify_tiling( + "interleave", + [ + _mapping(spec, None, _ctx(rank, tp=2, dcp=2, interleave=2)) + for rank in range(2) + ], + ) + + +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS)) +def test_all_ranks_tile_canonical_page(form): + spec = _full_spec() + per_rank = [ + _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank, tp=4)) + for rank in range(4) + ] + _verify_tiling("layer", per_rank) + + +# --------------------------------------------------------------------------- +# Byte-level round trips +# --------------------------------------------------------------------------- + + +def _store_all(mappings, pages, size: int, block_id: int = 0) -> bytes: + buf = bytearray(size) + for mapping, page in zip(mappings, pages): + if not mapping.is_writer(block_id): + continue + for local, canonical, n in _triples(mapping.runs): + buf[canonical : canonical + n] = page[local : local + n] + return bytes(buf) + + +def _load_one(mapping, canonical_bytes: bytes) -> bytes: + page = bytearray(mapping.local_page_size_bytes) + for local, canonical, n in _triples(mapping.runs): + page[local : local + n] = canonical_bytes[canonical : canonical + n] + return bytes(page) + + +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS)) +def test_cross_tp_store_load(form): + """Bytes stored under one TP size are the bytes another TP size loads.""" + total_heads, canonical_size = 8, 4096 + reference = bytes((7 + 31 * i) % 256 for i in range(canonical_size)) + + def mappings_at(tp: int): + spec = _full_spec(num_kv_heads=total_heads // tp) + cache = CACHE_BUILDERS[form](spec) + return [ + _mapping(spec, cache, _ctx(rank, tp=tp, total=total_heads)) + for rank in range(tp) + ] + + for tp in (4, 2, 1): + mappings = mappings_at(tp) + assert all(m is not None for m in mappings) + pages = [_load_one(m, reference) for m in mappings] + assert _store_all(mappings, pages, canonical_size) == reference + + +def test_cp_round_trip(): + # tp=4 / dcp=2 / 2 KV heads: 4 workers jointly hold one canonical page + spec = _full_spec(num_kv_heads=1) + cache = _split_nhd_cache(spec) + mappings = [ + _mapping(spec, cache, _ctx(rank, tp=4, dcp=2, total=2)) for rank in range(4) + ] + reference = bytes((3 + 17 * i) % 256 for i in range(2048)) + pages = [_load_one(m, reference) for m in mappings] + assert _store_all(mappings, pages, 2048) == reference + + +def test_replica_rotation_round_trip(): + """Whichever replica a block elects reproduces the same canonical page.""" + spec = _mla_spec() + mappings = [_mapping(spec, None, _ctx(rank, tp=2)) for rank in range(2)] + reference = bytes((5 + 11 * i) % 256 for i in range(256)) + pages = [_load_one(m, reference) for m in mappings] + for block_id in (0, 1): + assert _store_all(mappings, pages, 256, block_id) == reference + + +# --------------------------------------------------------------------------- +# Fail-closed gates +# --------------------------------------------------------------------------- + + +def test_fail_closed_cases(): + spec = _full_spec() + nhd = _split_nhd_cache(spec) + # Spec heads inconsistent with total heads / tp + assert _try_mapping(spec, nhd, _ctx(0, tp=4, total=2)) is None + # tp not divisible by total KV heads + one_head = _full_spec(num_kv_heads=1) + assert ( + _try_mapping(one_head, _split_nhd_cache(one_head), _ctx(0, tp=3, total=2)) + is None + ) + # DCP wider than the KV replication factor (tokens would shard across + # ranks holding different heads) + assert _try_mapping(spec, nhd, _ctx(0, tp=4, dcp=2)) is None + # Interleave must divide the block size + assert _try_mapping(spec, nhd, _ctx(0, tp=2, dcp=2, interleave=3, total=2)) is None + # Per-token-head scales are packed with the data + quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD) + assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None + # Compressed MLA slots are not 1:1 with tokens + assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None + # Unrecognized physical layouts + swapped = torch.zeros( + NUM_BLOCKS, + spec.block_size, + 2, + spec.num_kv_heads, + spec.head_size, + dtype=torch.int8, + ).permute(0, 2, 1, 3, 4) + assert _try_mapping(spec, swapped, _ctx(0, tp=4)) is None + # Non-attention specs + assert _try_mapping(KVCacheSpec(block_size=4), None, _ctx(0, tp=4, total=8)) is None + + +def test_opaque_fallback_places_page_whole(): + mapping = _opaque_fallback_mapping(1024, 4, 2) + assert mapping.canonical_page_size_bytes == 4096 + assert not mapping.parallelism_agnostic + assert _triples(mapping.runs) == [(0, 2048, 1024)] + assert mapping.num_writers == 1 + _verify_tiling("opaque", [_opaque_fallback_mapping(1024, 4, r) for r in range(4)]) + + +# --------------------------------------------------------------------------- +# derive_canonical_mappings end to end +# --------------------------------------------------------------------------- + + +def _vllm_config(tp=1, dcp=1, pcp=1, pp=1, interleave=1, total_kv_heads=2): + config = MagicMock() + config.parallel_config.tensor_parallel_size = tp + config.parallel_config.decode_context_parallel_size = dcp + config.parallel_config.prefill_context_parallel_size = pcp + config.parallel_config.cp_kv_cache_interleave_size = interleave + config.parallel_config.world_size = pp * tp * pcp + config.parallel_config.rank = 0 + config.model_config.get_total_num_kv_heads.return_value = total_kv_heads + return config + + +def _kv_cache_config(groups): + config = MagicMock() + config.kv_cache_groups = groups + config.num_blocks = NUM_BLOCKS + return config + + +def test_derive_mixed_model_with_dcp(): + attn_spec = _full_spec(num_kv_heads=1) + mla_spec = _mla_spec() + quant_spec = _full_spec( + num_kv_heads=1, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD + ) + kv_cache_config = _kv_cache_config( + [ + KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec), + KVCacheGroupSpec(layer_names=["mla"], kv_cache_spec=mla_spec), + KVCacheGroupSpec(layer_names=["quant"], kv_cache_spec=quant_spec), + ] + ) + kv_caches = { + "attn": _split_nhd_cache(attn_spec), + "quant": _split_nhd_cache(quant_spec), + } + mappings = derive_canonical_mappings( + _vllm_config(tp=4, dcp=2, total_kv_heads=2), kv_cache_config, kv_caches + ) + assert set(mappings) == {"attn", "mla", "quant"} + assert not mappings["attn"].parallelism_agnostic + assert mappings["attn"].runs + # Uncertifiable layers degrade to an opaque page, never disappear + assert not mappings["quant"].parallelism_agnostic + assert ( + mappings["quant"].canonical_page_size_bytes + == 4 * quant_spec.unpadded_page_size_bytes + ) + + +def test_derive_refuses_foreign_worker_groups(): + attn_spec = _full_spec() + kv_cache_config = _kv_cache_config( + [KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec)] + ) + kv_caches = {"attn": _split_nhd_cache(attn_spec)} + assert ( + derive_canonical_mappings( + _vllm_config(tp=2, pp=2, total_kv_heads=4), kv_cache_config, kv_caches + ) + == {} + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py new file mode 100644 index 000000000000..5a66b463e306 --- /dev/null +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for translating vLLM cache metadata to native offloading config.""" + +from typing import Any, cast +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from tests.v1.kv_connector.unit.offloading_connector.utils import MockOffloadingSpec +from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + build_offloading_config, +) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + SchedulerOffloadConfig, + is_store_reachable_swa_chunk, +) +from vllm.platforms import current_platform +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) + + +def _make_vllm_config( + *, + extra_config: dict[str, Any] | None = None, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + prefill_context_parallel_size: int = 1, + decode_context_parallel_size: int = 1, +) -> VllmConfig: + config = MagicMock() + config.cache_config.block_size = 16 + config.cache_config.enable_prefix_caching = True + config.cache_config.prefix_match_unit = None + config.cache_config.cache_dtype = torch.float16 + config.model_config.model = "test-model" + config.model_config.use_mla = False + world_size = ( + tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + ) + with patch.object(current_platform, "device_count", return_value=world_size): + config.parallel_config = ParallelConfig( + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + prefill_context_parallel_size=prefill_context_parallel_size, + decode_context_parallel_size=decode_context_parallel_size, + ) + config.kv_events_config = None + config.use_v2_model_runner = False + config.kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config=dict(extra_config or {}), + ) + return cast(VllmConfig, config) + + +def _make_kv_cache_config() -> KVCacheConfig: + num_blocks = 16 + kv_tensor = KVCacheTensor( + size=num_blocks * 8, + shared_by=["layer"], + block_stride=0, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[kv_tensor], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: + num_blocks = 4 + if packed: + kv_cache_tensors = [ + KVCacheTensor( + size=64, + shared_by=[layer_name], + block_stride=16, + ) + for layer_name in ("layer0", "layer1") + ] + else: + kv_cache_tensors = [ + KVCacheTensor(size=40, shared_by=["layer0"]), + KVCacheTensor(size=24, shared_by=["layer1"]), + ] + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + ) + + +def _mla_spec( + block_size: int = 16, + head_size: int = 512, + dtype: torch.dtype = torch.float32, +) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=head_size, + dtype=dtype, + ) + + +def _make_mla_kv_cache_config( + layer_names: list[str] | None = None, + head_size: int = 512, + dtype: torch.dtype = torch.float32, + num_blocks: int = 4, +) -> KVCacheConfig: + if layer_names is None: + layer_names = ["layer0", "layer1"] + spec = _mla_spec(head_size=head_size, dtype=dtype) + kv_cache_tensors = [ + KVCacheTensor( + size=spec.page_size_bytes * num_blocks, + shared_by=[layer_name], + ) + for layer_name in layer_names + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)], + ) + + +def _make_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor(size=40, shared_by=["full_layer"]), + KVCacheTensor(size=24, shared_by=["mla_layer"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)), + KVCacheGroupSpec(["mla_layer"], _mla_spec()), + ], + ) + + +def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["full_layer"], _full_attention_spec()), + KVCacheGroupSpec( + ["mamba_layer"], + MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + + +def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: + config = _make_vllm_config() + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + ) + return build_offloading_config( + config, kv_cache_config + ).parallel.is_parallelism_agnostic + + +def _replicated_layout( + kv_cache_config: KVCacheConfig, + *, + tensor_parallel_size: int = 4, + pipeline_parallel_size: int = 1, + prefill_context_parallel_size: int = 1, + decode_context_parallel_size: int = 1, + use_mla: bool = True, + use_v2_model_runner: bool = False, + distributed_executor_backend: Any = "mp", + nnodes: int = 1, + world_size: int | None = None, +) -> bool: + config = _make_vllm_config( + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + prefill_context_parallel_size=prefill_context_parallel_size, + decode_context_parallel_size=decode_context_parallel_size, + ) + config.model_config.use_mla = use_mla + config.use_v2_model_runner = use_v2_model_runner + config.parallel_config.distributed_executor_backend = distributed_executor_backend + config.parallel_config.nnodes = nnodes + if world_size is not None: + config.parallel_config.world_size = world_size + return build_offloading_config(config, kv_cache_config).replicated_layout + + +@pytest.mark.parametrize("packed", [False, True]) +def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): + config = _make_vllm_config( + extra_config={"block_size": 32}, + tensor_parallel_size=3, + pipeline_parallel_size=2, + ) + + offloading_config = build_offloading_config( + config, _make_sizing_kv_cache_config(packed) + ) + + assert offloading_config.worker_kv_bytes_per_block == 16 + assert offloading_config.parallel.world_size == 6 + assert offloading_config.cache.blocks_per_chunk == 2 + + +def test_rejects_partially_packed_tensor_layout(): + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + with pytest.raises(AssertionError): + build_offloading_config(_make_vllm_config(), kv_cache_config) + + +def test_zero_blocks_skips_tensor_layout_validation(): + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.num_blocks = 0 + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config) + + assert offloading_config.worker_kv_bytes_per_block == 0 + + +def test_prefill_context_parallelism_does_not_scale_group_blocks(): + config = _make_vllm_config( + extra_config={"block_size": 64}, + prefill_context_parallel_size=2, + ) + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == (16,) + assert offloading_config.cache.tokens_per_hash == 16 + assert offloading_config.cache.blocks_per_chunk == 4 + + +def test_dcp_scales_attention_but_not_mamba_group_blocks(): + config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2) + config.speculative_config = None + + offloading_config = build_offloading_config( + config, _make_mamba_hybrid_kv_cache_config() + ) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 32, + 16, + ) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), + config, + _make_mamba_hybrid_kv_cache_config(), + ) + mamba_group = scheduler_config.kv_group_configs[1] + assert mamba_group.alignment_chunk_count == 2 + assert [ + chunk_idx + for chunk_idx in range(4) + if is_store_reachable_swa_chunk( + chunk_idx, + 4, + mamba_group.alignment_chunk_count, + mamba_group.sliding_window_size_in_chunks, + mamba_group.is_eagle_group, + ) + ] == [1, 3] + + +def test_preserves_data_parallel_index(): + config = _make_vllm_config() + config.parallel_config.data_parallel_index = 2 + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.parallel.data_parallel_index == 2 + + +def test_resolves_heterogeneous_hybrid_block_sizes(): + config = _make_vllm_config() + config.cache_config.block_size = 4 + + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 12, + 16, + ) + assert offloading_config.cache.tokens_per_hash == 4 + assert offloading_config.cache.blocks_per_chunk == 1 + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( + world_size: int, + use_v2_model_runner: bool, +): + assert _replicated_layout( + _make_mla_kv_cache_config(), + tensor_parallel_size=world_size, + use_v2_model_runner=use_v2_model_runner, + ) + + +@pytest.mark.parametrize( + ("kv_cache_config", "case"), + [ + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer"], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=512, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + ), + "sliding-window-mla", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer"], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + HiddenStateCacheSpec( + block_size=16, + num_kv_heads=1, + head_size=512, + dtype=torch.float32, + ), + ) + ], + ), + "hidden-state", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer0"], + ), + KVCacheTensor( + size=_mla_spec(head_size=256).page_size_bytes * 4, + shared_by=["layer1"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "layer0": _mla_spec(), + "layer1": _mla_spec(head_size=256), + }, + ), + ) + ], + ), + "uniform-wrapper", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["mla"], + ), + KVCacheTensor( + size=_full_attention_spec().page_size_bytes * 4, + shared_by=["full"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla"], _mla_spec()), + KVCacheGroupSpec(["full"], _full_attention_spec()), + ], + ), + "mla-full-hybrid", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["mla"], + ), + KVCacheTensor(size=64 * 4, shared_by=["mamba"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla"], _mla_spec()), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((16, 1),), + dtypes=(torch.float32,), + ), + ), + ], + ), + "mla-mamba-hybrid", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer0"], + ), + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer1"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer0"], _mla_spec()), + KVCacheGroupSpec(["layer1"], _mla_spec()), + ], + ), + "multi-group-mla", + ), + ], + ids=[ + "sliding-window-mla", + "hidden-state", + "uniform-wrapper", + "mla-full-hybrid", + "mla-mamba-hybrid", + "multi-group-mla", + ], +) +def test_replicated_layout_excludes_unproven_cache_shapes( + kv_cache_config: KVCacheConfig, + case: str, +): + assert not _replicated_layout(kv_cache_config), case + + +def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting(): + num_blocks = 4 + main_spec = _mla_spec(head_size=512) + indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8) + main_layers = [f"main_{i}" for i in range(61)] + indexer_layers = [f"indexer_{i}" for i in range(61)] + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=main_spec.page_size_bytes * len(main_layers) * num_blocks, + shared_by=main_layers, + ), + KVCacheTensor( + size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks, + shared_by=indexer_layers, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)], + ) + + assert not _replicated_layout(kv_cache_config) + + +@pytest.mark.parametrize( + ("kwargs", "case"), + [ + ({"tensor_parallel_size": 1}, "tp1"), + ({"use_mla": False}, "use-mla-false"), + ({"pipeline_parallel_size": 2, "world_size": 4}, "pp2"), + ({"prefill_context_parallel_size": 2, "world_size": 4}, "pcp2"), + ({"decode_context_parallel_size": 2}, "dcp2"), + ({"world_size": 8}, "world-ne-tp"), + ({"distributed_executor_backend": "ray"}, "ray"), + ({"distributed_executor_backend": "uni"}, "uni"), + ({"distributed_executor_backend": type("DummyExecutor", (), {})}, "class"), + ({"nnodes": 2}, "multi-node"), + ], + ids=[ + "tp1", + "use-mla-false", + "pp2", + "pcp2", + "dcp2", + "world-ne-tp", + "ray", + "uni", + "class", + "multi-node", + ], +) +def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str): + assert not _replicated_layout(_make_mla_kv_cache_config(), **kwargs), case + + +def test_parallelism_agnostic_for_single_full_attention_group(): + assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + + +@pytest.mark.parametrize( + "kv_cache_groups", + [ + [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))], + [ + KVCacheGroupSpec( + ["l0"], + SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + [ + KVCacheGroupSpec(["l0"], _full_attention_spec()), + KVCacheGroupSpec(["l1"], _full_attention_spec()), + ], + ], +) +def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): + assert not _parallelism_agnostic(kv_cache_groups) + + +def test_parallelism_agnostic_disabled_on_v2_model_runner(): + config = _make_vllm_config() + config.use_v2_model_runner = True + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], + ) + + offloading_config = build_offloading_config(config, kv_cache_config) + + assert not offloading_config.parallel.is_parallelism_agnostic + + +def test_accepts_blocks_per_chunk_for_heterogeneous_groups(): + config = _make_vllm_config(extra_config={"blocks_per_chunk": 2}) + + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 12, + 16, + ) + assert offloading_config.cache.blocks_per_chunk == 2 + + +def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): + config = _make_vllm_config(extra_config={"block_size": 64, "blocks_per_chunk": 2}) + + with pytest.raises(ValueError, match="Specify only one"): + build_offloading_config(config, _make_kv_cache_config()) + + +def test_blocks_per_chunk_must_be_positive(): + config = _make_vllm_config(extra_config={"blocks_per_chunk": 0}) + + with pytest.raises(ValueError, match="greater than 0"): + build_offloading_config(config, _make_kv_cache_config()) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index a23fb95b5683..a6bab14e067c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -7,7 +7,11 @@ from tests.v1.kv_connector.unit.utils import create_vllm_config from vllm.config import KVEventsConfig, KVTransferConfig -from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored +from vllm.distributed.kv_events import ( + MEDIUM_CPU, + BlockRemoved, + BlockStored, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, ) @@ -27,6 +31,7 @@ ) from vllm.v1.kv_offload.base import ( Locality, + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -34,7 +39,7 @@ ) from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec -_CPU_MEDIUM = MEDIUM_CPU +_CPU_MEDIUM = Medium.CPU _FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec( kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value, kv_cache_spec_sliding_window=None, @@ -62,8 +67,9 @@ def _wire_hash(block_hash: BlockHash): return maybe_convert_block_hash(block_hash) -def _request(*, block_hashes: list[BlockHash], token_count: int): +def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"): req = MagicMock() + req.request_id = req_id req.block_hashes = block_hashes req.all_token_ids = list(range(1, token_count + 1)) req.lora_request = None @@ -104,10 +110,32 @@ def _record_chunks( return keys +def _record_lookup_chunks( + tracker: OffloadingEventsTracker, + req, + group_config: GroupOffloadConfig, + num_chunks: int, +) -> list[OffloadKey]: + keys: list[OffloadKey] = [] + hbf = group_config.hashes_per_chunk + for chunk_idx in range(num_chunks): + tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1] + assert tail_hash is not None + key = make_offload_key(tail_hash, group_config.group_idx) + tracker.record_lookup( + req, + group_config, + chunk_idx, + key, + ) + keys.append(key) + return keys + + def _stored_event( keys: list[OffloadKey], + medium: Medium = _CPU_MEDIUM, locality: Locality | None = None, - medium: str = _CPU_MEDIUM, ) -> OffloadingEvent: return OffloadingEvent( keys=keys, @@ -119,8 +147,8 @@ def _stored_event( def _removed_event( keys: list[OffloadKey], + medium: Medium = _CPU_MEDIUM, locality: Locality | None = None, - medium: str = _CPU_MEDIUM, ) -> OffloadingEvent: return OffloadingEvent( keys=keys, @@ -130,6 +158,21 @@ def _removed_event( ) +def _lookup_chunk() -> tuple[ + OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey +]: + tracker = _tracker() + req = _request(block_hashes=[_hash(0)], token_count=4) + group_config = _group_config() + key = _record_lookup_chunks( + tracker, + req, + group_config, + num_chunks=1, + )[0] + return tracker, req, group_config, key + + def test_take_events_forwards_locality_to_rich_store(): tracker = _tracker() req = _request(block_hashes=[_hash(0)], token_count=4) @@ -137,7 +180,7 @@ def test_take_events_forwards_locality_to_rich_store(): events = list( tracker.take_events( - [_stored_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + [_stored_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)] ) ) @@ -155,7 +198,7 @@ def test_take_events_forwards_locality_to_placeholder_store(): events = list( tracker.take_events( - [_stored_event([key], locality=Locality.REMOTE, medium=MEDIUM_FS)] + [_stored_event([key], locality=Locality.REMOTE, medium=Medium.STORAGE)] ) ) @@ -172,7 +215,7 @@ def test_take_events_forwards_locality_to_remove(): events = list( tracker.take_events( - [_removed_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + [_removed_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)] ) ) @@ -196,7 +239,7 @@ def test_take_events_publishes_routable_block_stored(): for i, event in enumerate(batch1): assert isinstance(event, BlockStored) - assert event.medium == _CPU_MEDIUM + assert event.medium == _CPU_MEDIUM.value assert event.block_hashes == [_wire_hash(_hash(i))] assert event.block_size == block_size assert event.token_ids == list( @@ -220,18 +263,37 @@ def test_take_events_publishes_routable_block_stored(): assert len(tracker._pending_event_metadata) == 6 -def test_take_events_factor_gt_1_chunk_store_and_remove(): +def test_promotion_emits_full_cpu_stored_event(): + tracker, _, _, key = _lookup_chunk() + + [event] = tracker.take_events([_stored_event([key])]) + + assert isinstance(event, BlockStored) + assert event.medium == MEDIUM_CPU + assert event.block_hashes == [_wire_hash(_hash(0))] + assert event.parent_block_hash is None + assert event.token_ids == [1, 2, 3, 4] + assert event.block_size == 4 + assert event.lora_id is None + assert event.lora_name is None + assert event.extra_keys is None + assert event.group_idx == 0 + assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value + assert event.kv_cache_spec_sliding_window is None + + +def test_lookup_promotion_factor_gt_1_store_and_remove(): block_size = 4 - blocks_per_chunk = 3 + blocks_per_chunk = 2 tracker = _tracker() group_config = _group_config( block_size=block_size, blocks_per_chunk=blocks_per_chunk ) req = _request( - block_hashes=[_hash(i) for i in range(6)], + block_hashes=[_hash(i) for i in range(4)], token_count=block_size * blocks_per_chunk * 2, ) - keys = _record_chunks(tracker, req, group_config, num_chunks=2) + keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2) stored = list(tracker.take_events([_stored_event(keys)])) assert len(stored) == 2 @@ -261,7 +323,7 @@ def test_take_events_factor_gt_1_chunk_store_and_remove(): assert len(removed) == 1 assert isinstance(removed[0], BlockRemoved) assert removed[0].block_hashes == expected_hashes - assert removed[0].medium == _CPU_MEDIUM + assert removed[0].medium == _CPU_MEDIUM.value assert removed[0].group_idx == 0 assert not tracker._pending_event_metadata @@ -293,6 +355,7 @@ def test_take_events_opt_out_keeps_placeholders(): group_config = _group_config() req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) keys = _record_chunks(tracker, req, group_config, num_chunks=3) + _record_lookup_chunks(tracker, req, group_config, num_chunks=3) assert not tracker.self_describing_enabled assert not tracker._pending_event_metadata @@ -315,11 +378,21 @@ def test_take_events_opt_out_keeps_placeholders(): assert len(events[3].block_hashes) == 3 -def test_record_store_skips_sliding_window_group(): +@pytest.mark.parametrize( + "sliding_window_size_in_chunks", + [1, 2], + ids=["ssm", "sliding-window"], +) +def test_event_metadata_skips_non_full_attention_group( + sliding_window_size_in_chunks: int, +): tracker = _tracker() - group_config = _group_config(sliding_window_size_in_chunks=2) + group_config = _group_config( + sliding_window_size_in_chunks=sliding_window_size_in_chunks + ) req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) keys = _record_chunks(tracker, req, group_config, num_chunks=3) + _record_lookup_chunks(tracker, req, group_config, num_chunks=3) assert not tracker._pending_event_metadata @@ -329,6 +402,56 @@ def test_record_store_skips_sliding_window_group(): assert events[0].block_size == 0 +def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit(): + tracker = _tracker() + block_hashes = [_hash(0), _hash(1)] + req = _request(block_hashes=block_hashes, token_count=8) + group_config = _group_config(blocks_per_chunk=2) + key = _record_chunks(tracker, req, group_config, num_chunks=1)[0] + confirmed_meta = tracker._pending_event_metadata[key] + lookup_req = _request( + block_hashes=block_hashes, + token_count=8, + req_id="new-request", + ) + + tracker.record_lookup( + lookup_req, + group_config, + 0, + key, + ) + assert tracker._pending_event_metadata[key] is confirmed_meta + + removed = list(tracker.take_events([_removed_event([key])])) + assert len(removed) == 1 + assert removed[0].block_hashes == [ + _wire_hash(_hash(0)), + _wire_hash(_hash(1)), + ] + + stored = list(tracker.take_events([_stored_event([key])])) + assert len(stored) == 1 + assert stored[0].block_size == 0 + assert stored[0].token_ids == [] + + tracker.record_lookup(lookup_req, group_config, 0, key) + removed = list(tracker.take_events([_removed_event([key])])) + assert removed[0].block_hashes == [ + _wire_hash(_hash(0)), + _wire_hash(_hash(1)), + ] + + +def test_secondary_stored_event_does_not_mutate_cpu_metadata(): + tracker, _, _, key = _lookup_chunk() + expected_metadata = dict(tracker._pending_event_metadata) + + stored = list(tracker.take_events([_stored_event([key], Medium.STORAGE)])) + assert stored[0].token_ids == [1, 2, 3, 4] + assert tracker._pending_event_metadata == expected_metadata + + def test_take_events_groups_removed_hashes_by_kv_group(): tracker = _tracker() group0_config = _group_config(group_idx=0, blocks_per_chunk=2) @@ -378,7 +501,7 @@ def test_reset_cache_clears_side_table(): tracker = _tracker() group_config = _group_config() req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) - _record_chunks(tracker, req, group_config, num_chunks=3) + _record_lookup_chunks(tracker, req, group_config, num_chunks=3) assert tracker._pending_event_metadata @@ -387,7 +510,7 @@ def test_reset_cache_clears_side_table(): assert not tracker._pending_event_metadata -def test_tiering_rejects_self_describing_kv_events(): +def test_tiering_accepts_self_describing_kv_events(): vllm_config = create_vllm_config( block_size=4, max_num_batched_tokens=16, @@ -423,5 +546,9 @@ def test_tiering_rejects_self_describing_kv_events(): ], ) - with pytest.raises(ValueError, match="TieringOffloadingSpec"): - TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + tracker = OffloadingEventsTracker(spec.kv_events_config) + + assert spec.kv_events_config.enable_kv_cache_events + assert spec.kv_events_config.self_describing_kv_events + assert tracker.self_describing_enabled diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index bc5190db00f5..661078041f10 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest import torch @@ -11,6 +11,11 @@ to_keys, ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID +from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( + OffloadingConnectorMetadata, + OffloadingWorkerMetadata, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, _ConnectorMetricName, @@ -18,7 +23,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, + is_store_reachable_swa_chunk, ) +from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -26,6 +33,8 @@ ) from vllm.v1.kv_offload.base import ( LookupResult, + Medium, + OffloadingEvent, OffloadingManager, OffloadPolicy, ReqContext, @@ -33,6 +42,7 @@ get_offload_block_hash, make_offload_key, ) +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -61,7 +71,9 @@ def test_scheduler_reports_allocation_failure(request_runner): runner.run(decoded_tokens=[EOS_TOKEN_ID]) reduced = _reduce_kv_connector_stats(runner) - assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 + # Two attempts: once while running (block becomes full during prefill), + # once from finished_req_ids on the next step. + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2 @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -69,7 +81,7 @@ def test_scheduler_reports_allocation_failure(request_runner): def test_last_block_offloaded_at_request_finish( request_runner, async_scheduling: bool, prompt_offset: int ): - """EOS fills the last block at request finish — verify req_status is kept alive. + """EOS fills the last block at request finish - verify the final block is stored. prompt = block_size + prompt_offset tokens → not a full block at schedule time, so _build_store_jobs creates no store job. After EOS, request_finished @@ -91,19 +103,51 @@ def test_last_block_offloaded_at_request_finish( generate_store_output(list(keys)) ) - # Run with one step (EOS) - runner.run( - decoded_tokens=[EOS_TOKEN_ID], - ) + if prompt_offset == -1: + # EOS fills the block, so a store job is created for block 0. + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,)) + else: + # Block remains partial, so no store job is created. + runner.run(decoded_tokens=[EOS_TOKEN_ID]) cs = runner.connector_scheduler - # Verify req_status is kept alive for _build_store_jobs to process - # regardless of whether there are storable blocks - assert "0" in cs._req_status, ( - "req_status was deleted but should be kept alive " - "for _build_store_jobs to process finished_req_ids." + # After the full run completes, req_status is cleaned up. + assert "0" not in cs._req_status + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_abort_queued_request_does_not_build_store_job( + request_runner, async_scheduling: bool +): + """Aborting a never-scheduled request must not store unallocated KV.""" + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=8, + async_scheduling=async_scheduling, + ) + + runner.new_request(token_ids=[0] * (block_size * 4)) + runner.scheduler.schedule() + + runner.new_request(token_ids=[1] * (block_size * 4)) + queued_req_id = str(runner.req_id) + assert any( + request.request_id == queued_req_id for request in runner.scheduler.waiting ) + runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED) + req_status = runner.connector_scheduler._req_status[queued_req_id] + assert all(group_state.offload_keys for group_state in req_status.group_states) + assert all(not group_state.block_ids for group_state in req_status.group_states) + + scheduler_output = runner.scheduler.schedule() + + metadata = scheduler_output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values()) + assert queued_req_id not in runner.connector_scheduler._req_status + def test_scheduler_reports_lookup_sync_delay(request_runner): runner = request_runner( @@ -123,6 +167,55 @@ def test_scheduler_reports_lookup_sync_delay(request_runner): assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0 +@pytest.mark.parametrize( + ( + "absolute_chunk_index", + "storable_chunk_count", + "alignment_chunk_count", + "sliding_window_chunks", + "is_eagle_group", + "expected", + ), + [ + # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61. + (61, 64, 64, 2, False, False), + (62, 64, 64, 2, False, True), + (60, 64, 64, 2, True, False), + (61, 64, 64, 2, True, True), + # Partial 48-of-64 segment: the reachable tail ends at chunk 47. + (45, 48, 64, 2, False, False), + (46, 48, 64, 2, False, True), + (44, 48, 64, 2, True, False), + (45, 48, 64, 2, True, True), + # A later partial segment uses its own actual end (chunks 64-79). + (76, 80, 64, 3, False, False), + (77, 80, 64, 3, False, True), + # No alignment means no store-pruning optimization. + (0, 1, None, None, False, True), + # A tail at least as large as the segment keeps every chunk. + (0, 2, 64, 2, False, True), + ], +) +def test_is_store_reachable_swa_chunk( + absolute_chunk_index: int, + storable_chunk_count: int, + alignment_chunk_count: int | None, + sliding_window_chunks: int | None, + is_eagle_group: bool, + expected: bool, +): + assert ( + is_store_reachable_swa_chunk( + absolute_chunk_index, + storable_chunk_count, + alignment_chunk_count, + sliding_window_chunks, + is_eagle_group, + ) + is expected + ) + + def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): """A deferred lookup reports its async delay once it resolves.""" runner = request_runner( @@ -143,6 +236,159 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0 +def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner): + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + runner.manager.lookup.return_value = LookupResult.RETRY + runner.manager.take_events.return_value = [] + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.new_request( + token_ids=[1] * 12, + kv_transfer_params={"max_offload_tokens": 0}, + ) + runner.run(decoded_tokens=[]) + + tracker = runner.connector_scheduler._events_tracker + assert runner.manager.lookup.call_count == 3 + assert not tracker._pending_event_metadata + assert list(runner.connector_scheduler.take_events()) == [] + + runner.manager.lookup.return_value = LookupResult.MISS + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + assert not tracker._pending_event_metadata + assert list(runner.connector_scheduler.take_events()) == [] + + +def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal( + request_runner, +): + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + blocks_per_chunk=2, + ) + raw_events: list[OffloadingEvent] = [] + + def take_raw_events(): + yield from raw_events + raw_events.clear() + + runner.manager.lookup.return_value = LookupResult.RETRY + runner.manager.take_events.side_effect = take_raw_events + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + + runner.new_request(token_ids=[1] * 8) + runner.run(decoded_tokens=[]) + + tracker = runner.connector_scheduler._events_tracker + assert not tracker._pending_event_metadata + key = runner.manager.lookup.call_args.args[0] + req_id = str(runner.req_id) + req_status = runner.connector_scheduler._req_status[req_id] + + runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED) + + assert not tracker._pending_event_metadata + + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False)) + events = list(runner.connector_scheduler.take_events()) + assert len(events) == 1 + assert isinstance(events[0], BlockStored) + assert events[0].block_size == 0 + assert events[0].token_ids == [] + + runner.manager.lookup.return_value = LookupResult.HIT + group_config = runner.connector_scheduler.config.kv_group_configs[0] + assert ( + runner.connector_scheduler._maximal_prefix_lookup( + [key], + req_status.req_context, + req_status.req, + group_config, + 0, + ) + == 1 + ) + assert key in tracker._pending_event_metadata + + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=True)) + [event] = runner.connector_scheduler.take_events() + assert isinstance(event, BlockRemoved) + assert event.medium == MEDIUM_CPU + assert len(event.block_hashes) == 2 + assert key not in tracker._pending_event_metadata + + +@pytest.mark.parametrize("blocks_per_chunk", [1, 2]) +def test_promotion_hit_precedes_stored_event_translation( + request_runner, + blocks_per_chunk: int, +): + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + blocks_per_chunk=blocks_per_chunk, + ) + token_ids = [1] * 4 * blocks_per_chunk + + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.new_request(token_ids=token_ids) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=tuple(range(blocks_per_chunk)), + ) + runner.scheduler.reset_prefix_cache() + runner.connector_scheduler._events_tracker.reset() + + raw_events: list[OffloadingEvent] = [] + + def lookup(key, req_context): + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False)) + return LookupResult.HIT + + def take_raw_events(): + yield from raw_events + raw_events.clear() + + runner.manager.lookup.side_effect = lookup + runner.manager.take_events.side_effect = take_raw_events + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + publisher = MagicMock() + runner.scheduler.kv_event_publisher = publisher + + runner.new_request(token_ids=token_ids) + runner.run( + decoded_tokens=[], + expected_loaded=tuple(range(blocks_per_chunk)), + ) + + events = [ + event + for publish_call in publisher.publish.call_args_list + for event in publish_call.args[0].events + if isinstance(event, BlockStored) and event.medium == MEDIUM_CPU + ] + assert len(events) == 1 + assert len(events[0].block_hashes) == blocks_per_chunk + assert events[0].block_size == 4 + assert events[0].token_ids == token_ids + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_offloading_connector(request_runner, async_scheduling: bool): block_size = 4 @@ -241,7 +487,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2)) # single block lookup with a hit in a middle block @@ -249,7 +495,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5)) @@ -307,7 +553,7 @@ def test_request_preemption(request_runner, async_scheduling: bool): # request should now return from preemption # re-load [0, ..., 8] from the CPU and store [9, 10, 11] - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -326,14 +572,14 @@ def test_request_preemption(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_on_request_finished_is_not_deferred_until_store_completion( +def test_on_request_finished_not_deferred_until_store_completion( request_runner, async_scheduling: bool ): - """on_request_finished fires when no more stores will be submitted. + """on_request_finished fires after the last prepare_store is submitted. - A request can finish while its GPU->primary store is still in flight. The - manager-level hook should not wait for that completion; complete_store may - still arrive afterward for already-submitted transfer jobs. + The manager contract guarantees no more submit-side calls (prepare_store) + after on_request_finished. However, complete_store callbacks for + already-submitted transfers may still arrive afterward. """ block_size = 4 blocks_per_chunk = 3 @@ -370,8 +616,9 @@ def test_on_request_finished_is_not_deferred_until_store_completion( complete_transfers=False, ) - # Finish the request while its stores are still in flight. The hook should - # fire immediately even though no complete_store has arrived yet. + # Finish the request while its stores are still in flight. The hook fires + # once the last prepare_store is issued (on the next schedule step), even + # though complete_store has not yet been called. runner.run( decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False, @@ -381,8 +628,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion( assert calls == [("on_request_finished", req_id)], calls - # Drain the stores afterward. The already-submitted complete_store calls - # are allowed to arrive after on_request_finished. + # Drain the stores afterward. complete_store is allowed after the hook. runner.run( decoded_tokens=[], complete_transfers=True, @@ -395,11 +641,50 @@ def test_on_request_finished_is_not_deferred_until_store_completion( finished_idx = calls.index(("on_request_finished", req_id)) store_indices = [i for i, c in enumerate(calls) if c == ("complete_store", req_id)] - # The request-level hook no longer waits for already-submitted transfers. + # complete_store arrives after on_request_finished, as allowed by the contract. assert store_indices, calls assert finished_idx < min(store_indices), calls +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_on_request_finished_fires_after_final_block_store( + request_runner, async_scheduling: bool +): + """on_request_finished fires after the final-block prepare_store at EOS. + + When EOS fills a partial block, request_finished() keeps req_status alive + so _build_store_jobs can create a store job for it on the next step. + """ + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=10, + async_scheduling=async_scheduling, + ) + + calls: list[tuple[str, str]] = [] + runner.manager.on_request_finished.side_effect = lambda req_context: calls.append( + ("on_request_finished", req_context.req_id) + ) + + def prepare_store(keys, req_context): + calls.append(("prepare_store", req_context.req_id)) + return generate_store_output(keys) + + runner.manager.prepare_store.side_effect = prepare_store + + runner.new_request(token_ids=[0] * (block_size - 1)) + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,)) + + req_id = str(runner.req_id) + assert calls.count(("on_request_finished", req_id)) == 1, calls + + finished_idx = calls.index(("on_request_finished", req_id)) + prepare_indices = [i for i, c in enumerate(calls) if c == ("prepare_store", req_id)] + assert prepare_indices, calls + assert finished_idx > max(prepare_indices), calls + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool): block_size = 4 @@ -427,7 +712,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -439,7 +724,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a new request to load the same first block runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -491,7 +776,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -588,7 +873,10 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo touch_calls = runner.manager.touch.call_args_list assert len(touch_calls) == 6 - runner.run(decoded_tokens=[EOS_TOKEN_ID]) + # EOS fills the 7th block (offset 6). The extra schedule step processes + # finished_req_ids and stores block 6 for both groups before the request's + # GPU blocks are freed. + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(6,)) runner.scheduler.reset_prefix_cache() @@ -601,19 +889,11 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # Group 1 (sliding window, window=2): only the last 2 blocks # are within the window → loads blocks 1,2 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), - # The deferred store from the previous request's last block - # completes during this step, and its blocks are flushed because - # they were reallocated to the new request. - # Only block 1 (sliding window group) is stored — block 0's - # deferred store is flushed because it was reallocated. - expected_stored=((0, 1),), - expected_flushed=((0, 1),), ) - # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups) - # + 2 from _get_reqs_to_store (2 groups) + # 2 touch calls from get_num_new_matched_tokens (2 groups) touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 + assert len(touch_calls) == 2 # full attention group touched all 3 blocks assert len(touch_calls[0].args[0]) == 3 # sliding window group touched just the last 2 blocks @@ -793,73 +1073,144 @@ def _make_scheduler_with_lookup( scheduler = object.__new__(OffloadingConnectorScheduler) scheduler.manager = manager + scheduler._events_tracker = MagicMock() return scheduler _EMPTY_REQ_CTX = ReqContext(req_id="") +_LOOKUP_REQ = MagicMock() +_LOOKUP_REQ.request_id = "req" +_LOOKUP_GROUP_CONFIG = MagicMock() + + +def _maximal_lookup(sched, keys, start_chunk_idx: int = 0): + return sched._maximal_prefix_lookup( + keys, + _EMPTY_REQ_CTX, + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + start_chunk_idx, + ) class TestMaximalPrefixLookup: def test_all_hit(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 + assert _maximal_lookup(sched, to_keys([1, 2])) == 2 + + def test_records_absolute_chunk_indices(self): + keys = to_keys([1, 2]) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) + + assert _maximal_lookup(sched, keys, start_chunk_idx=3) == 2 + assert sched._events_tracker.record_lookup.call_args_list == [ + call( + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + 3, + keys[0], + ), + call( + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + 4, + keys[1], + ), + ] def test_all_miss(self): sched = _make_scheduler_with_lookup({}) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + assert _maximal_lookup(sched, to_keys([1, 2])) == 0 + sched._events_tracker.record_lookup.assert_not_called() def test_partial_prefix(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 + assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2 def test_miss_then_hit(self): sched = _make_scheduler_with_lookup({2: LookupResult.HIT}) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + assert _maximal_lookup(sched, to_keys([1, 2])) == 0 def test_single_hit(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT}) - assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 + assert _maximal_lookup(sched, to_keys([1])) == 1 def test_empty(self): sched = _make_scheduler_with_lookup({}) - assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0 + assert _maximal_lookup(sched, []) == 0 + + @pytest.mark.parametrize( + "pending_result", + [LookupResult.RETRY, LookupResult.HIT_PENDING], + ) + def test_pending_result_is_not_recorded( + self, + pending_result: LookupResult, + ): + sched = _make_scheduler_with_lookup({1: pending_result}) + + assert _maximal_lookup(sched, to_keys([1])) is None + sched._events_tracker.record_lookup.assert_not_called() def test_retry_defers(self): + keys = to_keys([1, 2]) sched = _make_scheduler_with_lookup( {1: LookupResult.RETRY, 2: LookupResult.HIT} ) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert _maximal_lookup(sched, keys) is None assert sched.manager.lookup.call_count == 2 + sched._events_tracker.record_lookup.assert_called_once_with( + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + 1, + keys[1], + ) def test_retry_after_hit_defers(self): + keys = to_keys([1, 2]) sched = _make_scheduler_with_lookup( {1: LookupResult.HIT, 2: LookupResult.RETRY} ) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert _maximal_lookup(sched, keys) is None + sched._events_tracker.record_lookup.assert_called_once_with( + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + 0, + keys[0], + ) def test_hit_pending_defers(self): + keys = to_keys([1, 2]) sched = _make_scheduler_with_lookup( {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT} ) - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert _maximal_lookup(sched, keys) is None assert sched.manager.lookup.call_count == 2 + sched._events_tracker.record_lookup.assert_called_once_with( + _LOOKUP_REQ, + _LOOKUP_GROUP_CONFIG, + 1, + keys[1], + ) def test_hit_pending_does_not_stop_scan(self): """HIT_PENDING defers but does not break — scan continues until miss.""" sched = _make_scheduler_with_lookup( {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT} ) - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None + assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None assert sched.manager.lookup.call_count == 2 + sched._events_tracker.record_lookup.assert_not_called() def test_retry_stops_at_miss(self): """RETRY is treated as hit for iteration, but miss stops the scan.""" sched = _make_scheduler_with_lookup( {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT} ) - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None + assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None # lookup should have been called for blocks 1 and 2 (stops at miss) assert sched.manager.lookup.call_count == 2 + sched._events_tracker.record_lookup.assert_not_called() class TestSlidingWindowLookup: @@ -1011,7 +1362,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling # New request with 2 offloaded chunks; first matches what's in CPU. runner.new_request(token_ids=[0] * tokens_per_chunk * 2) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -1042,7 +1393,7 @@ def test_loads_do_not_populate_fence_index(request_runner): async_scheduling=False, ) runner.new_request(token_ids=[0] * 12) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.run(decoded_tokens=[], complete_transfers=False) assert runner.connector_scheduler._block_id_to_pending_jobs == {} @@ -1088,7 +1439,7 @@ def capture_fence(): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * 4) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1139,7 +1490,7 @@ def capture_fence(): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[1] * 4) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1191,6 +1542,51 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool): assert runner.manager.complete_store.call_count == 0 +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_complete_store_waits_for_all_worker_acks( + request_runner, async_scheduling: bool +): + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk + runner = request_runner( + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + worker_count=3, + ) + runner.new_request(token_ids=[0] * tokens_per_chunk) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0, 0], complete_transfers=False) + assert len(runner.connector_scheduler._jobs) == 1 + job_id = next(iter(runner.connector_scheduler._jobs)) + assert runner.connector_scheduler._jobs[job_id].pending_count == 3 + runner.manager.complete_store.reset_mock() + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ) + ) + ) + assert runner.manager.complete_store.call_count == 0 + assert runner.connector_scheduler._jobs[job_id].pending_count == 2 + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 2} + ) + ) + ) + assert runner.manager.complete_store.call_count == 1 + assert job_id not in runner.connector_scheduler._jobs + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): """Validates max_offload_tokens: type coercion, boundary values, and capping. @@ -1365,7 +1761,7 @@ def test_reset_cache(request_runner, async_scheduling: bool): # Leave the load in-flight so that reset_cache must flush it. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1419,8 +1815,8 @@ def test_reset_cache(request_runner, async_scheduling: bool): def test_reset_cache_finalizes_finished_request_with_pending_store( request_runner, async_scheduling: bool ): - """reset_cache drops a finished request whose in-flight stores it discards - without calling on_request_finished twice. + """reset_cache fires on_request_finished for a finished request whose + in-flight stores it discards, exactly once. """ block_size = 4 blocks_per_chunk = 3 @@ -1456,16 +1852,15 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( assert req_status.transfer_jobs, "expected an in-flight store before finish" assert any(job.is_store for job in cs._jobs.values()) - # Finish the request while its store is still in flight. request_finished - # fires the hook eagerly, but the entry stays tracked so later completions - # can still call complete_store(). + # Finish the request while its store is still in flight. The manager hook + # is deferred because the final store decision has not happened yet. req_status.req.status = RequestStatus.FINISHED_STOPPED cs.request_finished(req_status.req) - assert finalized == [req_id] + assert finalized == [] assert req_id in cs._req_status - # reset_cache discards the in-flight store and drops the state without a - # duplicate on_request_finished call. + # reset_cache discards both the in-flight and not-yet-prepared final stores, + # so it issues the deferred notification before dropping the state. cs.reset_cache() assert finalized == [req_id] assert req_id not in cs._req_status @@ -1553,9 +1948,7 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner # preemption batch's ModelRunnerOutput is consumed by update_from_output(). free_block_queue.num_free_blocks = num_free_blocks_empty assert runner.scheduler.reset_prefix_cache() - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len( - key - ) + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys) readmit_output = runner.scheduler.schedule() @@ -1669,7 +2062,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) runner.manager.lookup.return_value = LookupResult.HIT - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2 runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0: full prefix lookup hits 2 offloaded chunks @@ -1840,6 +2233,13 @@ def _make_req_status( req.request_id = "test-req" req.num_tokens = num_tokens req.kv_transfer_params = None + num_hash_blocks = max( + len(hashes) * scheduler.config.kv_group_configs[idx].hashes_per_chunk + for idx, hashes in enumerate(offload_keys_per_group) + ) + req.block_hashes = [BlockHash(str(i).encode()) for i in range(num_hash_blocks)] + req.all_token_ids = list(range(num_tokens)) + req.lora_request = None state = RequestOffloadState( config=scheduler.config, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 25bb664ff402..b65cfe8150e4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -6,6 +6,10 @@ import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( + OffloadingConnectorMetadata, + TransferJob, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -22,7 +26,17 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, CanonicalKVCaches, + GPULoadStoreSpec, + LoadStoreSpec, + OffloadingManager, OffloadingSpec, + OffloadingWorker, +) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingModelConfig, + OffloadingParallelConfig, ) NUM_BLOCKS = 10 @@ -87,7 +101,26 @@ def _allocate_and_reshape_kv_caches( set_kv_cache_layout(None) -def _make_worker(kv_cache_config: KVCacheConfig): +def _single_rank_vllm_config(total_kv_heads: int): + """A one-rank (TP=1) parallel config, as canonical mappings are derived + from it.""" + vllm_config = MagicMock() + parallel_config = vllm_config.parallel_config + parallel_config.tensor_parallel_size = 1 + parallel_config.decode_context_parallel_size = 1 + parallel_config.prefill_context_parallel_size = 1 + parallel_config.cp_kv_cache_interleave_size = 1 + parallel_config.world_size = 1 + parallel_config.rank = 0 + vllm_config.model_config.get_total_num_kv_heads.return_value = total_kv_heads + return vllm_config + + +def _make_worker( + kv_cache_config: KVCacheConfig, + replicated_layout: bool = False, + rank: int = 0, +): """ Create an OffloadingConnectorWorker with mocked dependencies. """ @@ -96,10 +129,14 @@ def _make_worker(kv_cache_config: KVCacheConfig): ) spec = MagicMock(spec=OffloadingSpec) + spec.replicated_layout = replicated_layout + spec.config = MagicMock() + spec.config.parallel.rank = rank spec.get_worker.return_value = MagicMock() worker = OffloadingConnectorWorker( spec=spec, + vllm_config=_single_rank_vllm_config(NUM_KV_HEADS), kv_cache_config=kv_cache_config, ) worker.worker = MagicMock() @@ -107,11 +144,174 @@ def _make_worker(kv_cache_config: KVCacheConfig): return worker, spec +def _store_metadata(job_id: int) -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata( + load_jobs={}, + store_jobs={ + job_id: TransferJob( + req_id="req", + src_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)), + dst_spec=LoadStoreSpec(), + ) + }, + ) + + +def _load_metadata(job_id: int) -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata( + load_jobs={ + job_id: TransferJob( + req_id="req", + src_spec=LoadStoreSpec(), + dst_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)), + ) + }, + store_jobs={}, + ) + + +def _empty_metadata() -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata(load_jobs={}, store_jobs={}) + + +def _offloading_config(rank: int = 0) -> OffloadingConfig: + return OffloadingConfig( + groups=(), + worker_kv_bytes_per_block=0, + enable_kv_cache_events=False, + extra_config={}, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float16"), + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), + parallel=OffloadingParallelConfig( + rank=rank, + world_size=2, + tp_size=2, + pp_size=1, + pcp_size=1, + dcp_size=1, + data_parallel_index=0, + is_parallelism_agnostic=False, + ), + ) + + +class BareExternalOffloadingSpec(OffloadingSpec): + def get_manager(self) -> OffloadingManager: + raise NotImplementedError + + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: + raise NotImplementedError + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- +def test_prepare_store_kv_non_writer_marks_completed_without_submit(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + + worker.prepare_store_kv(_store_metadata(7)) + worker.start_kv_transfers(_empty_metadata()) + + assert worker._unsubmitted_store_jobs == [] + assert worker.worker is not None + worker.worker.submit_store.assert_not_called() + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {7: 1} + + +def test_prepare_store_kv_writer_submits_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=0, + ) + + worker.prepare_store_kv(_store_metadata(8)) + assert worker.build_connector_worker_meta() is None + worker.start_kv_transfers(_empty_metadata()) + + assert worker.worker is not None + worker.worker.submit_store.assert_called_once() + + +def test_prepare_store_kv_non_replicated_rank_gt_zero_queues_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=False, + rank=1, + ) + + worker.prepare_store_kv(_store_metadata(9)) + assert worker.build_connector_worker_meta() is None + assert len(worker._unsubmitted_store_jobs) == 1 + + worker.start_kv_transfers(_empty_metadata()) + + assert worker.worker is not None + worker.worker.submit_store.assert_called_once() + assert worker._unsubmitted_store_jobs == [] + + +def test_handle_preemptions_non_writer_acks_flushed_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + metadata = _store_metadata(10) + metadata.jobs_to_flush = {10} + + worker.handle_preemptions(metadata) + + assert worker.worker is not None + worker.worker.submit_store.assert_not_called() + worker.worker.wait.assert_called_once_with({10}) + assert metadata.store_jobs == {} + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {10: 1} + + +def test_start_kv_transfers_non_writer_still_submits_load(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + + worker.start_kv_transfers(_load_metadata(10)) + + assert worker.worker is not None + worker.worker.submit_load.assert_called_once() + assert worker.build_connector_worker_meta() is None + + +def test_offloading_connector_worker_accepts_plugin_spec_default_layout(): + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( + OffloadingConnectorWorker, + ) + + spec = BareExternalOffloadingSpec(_offloading_config(rank=1)) + + OffloadingConnectorWorker( + spec=spec, + vllm_config=_single_rank_vllm_config(NUM_KV_HEADS), + kv_cache_config=KVCacheConfig( + num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[] + ), + ) + + assert spec.replicated_layout is False + + @pytest.mark.parametrize("backend", ATTN_BACKENDS) def test_register_kv_caches(backend): """Test register_kv_caches with multiple groups covering all layer types. @@ -125,8 +325,7 @@ def test_register_kv_caches(backend): own dedicated tensors. Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce - kv_caches, which automatically applies - _update_hybrid_attention_mamba_layout for hybrid models. + the raw per-layer kv_caches registered by the connector. Verifies that the canonicalized CanonicalKVCaches has the correct block tensors, tensor_idx references, and page sizes across all groups. @@ -340,6 +539,8 @@ def test_register_kv_caches(backend): for actual, expected in zip(actual_refs, exp_refs): assert actual.tensor_idx == expected.tensor_idx assert actual.page_size_bytes == expected.page_size_bytes + # Every layer gets a canonical mapping, certified or opaque + assert actual.mapping is not None @pytest.mark.parametrize("backend", ATTN_BACKENDS) @@ -440,9 +641,15 @@ def test_register_kv_caches_uniform_type(backend): assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes) assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes) - assert group_refs[0] == CanonicalKVCacheRef( - tensor_idx=0, page_size_bytes=spec_a.page_size_bytes - ) - assert group_refs[1] == CanonicalKVCacheRef( - tensor_idx=1, page_size_bytes=spec_b.page_size_bytes - ) + for ref, expected_tensor_idx, expected_spec in ( + (group_refs[0], 0, spec_a), + (group_refs[1], 1, spec_b), + ): + assert ref.tensor_idx == expected_tensor_idx + assert ref.page_size_bytes == expected_spec.page_size_bytes + assert ref.mapping is not None + + # Only layer_a matches the model's total KV head count, so layer_b gets an + # opaque mapping rather than a certified, parallelism-agnostic one + assert group_refs[0].mapping.parallelism_agnostic + assert not group_refs[1].mapping.parallelism_agnostic diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index b878e294a6ef..9acf5f5d49fd 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -179,6 +179,7 @@ def __init__( async_scheduling: bool = True, kv_cache_groups: list[KVCacheGroupSpec] | None = None, extra_config_overrides: dict[str, Any] | None = None, + worker_count: int = 1, ): assert blocks_per_chunk == 1 or kv_cache_groups is None, ( "blocks_per_chunk > 1 requires all groups to have the same " @@ -198,6 +199,7 @@ def __init__( disable_hybrid_kv_cache_manager=False, ) vllm_config.scheduler_config.async_scheduling = async_scheduling + vllm_config.parallel_config.world_size = worker_count extra_config: dict[str, Any] = { "spec_name": "MockOffloadingSpec", @@ -482,8 +484,13 @@ def _run( # Strict-always-False frees the request immediately on EOS, but # the worker may still have a deferred store queued. In production # the next request's step drains it; in single-request tests we - # must keep stepping until the scheduler sees no in-flight jobs. - if not self.scheduler.requests and not self.connector_scheduler._jobs: + # must keep stepping until the scheduler sees no in-flight jobs + # and no pending finished_req_ids awaiting build_connector_meta. + if ( + not self.scheduler.requests + and not self.connector_scheduler._jobs + and not self.scheduler.finished_req_ids + ): break scheduler_output = self.scheduler.schedule() @@ -544,19 +551,30 @@ def _run( if ( prev_token_id == EOS_TOKEN_ID and prev_token_id != token_id - and (self.scheduler.requests or self.connector_scheduler._jobs) + and ( + self.scheduler.requests + or self.connector_scheduler._jobs + or self.scheduler.finished_req_ids + ) ): # continue for one more step to allow offloading to kick off continue if token_id is None: if self.async_scheduling: - # sample last token + # Flush the previous step's output. engine_outputs = self.scheduler.update_from_output( prev_scheduler_output, prev_model_runner_output ) self._record_kv_connector_stats(engine_outputs) - break + prev_model_runner_output = None + if self.scheduler.requests: + # Request still running, just exhausted decoded_tokens. + break + if not self.scheduler.finished_req_ids and ( + not complete_transfers or not self.connector_scheduler._jobs + ): + break self._parse_transfers() @@ -652,6 +670,7 @@ def runner_factory( blocks_per_chunk=1, kv_cache_groups=None, extra_config_overrides=None, + worker_count=1, ): runner = RequestRunner( block_size=block_size, @@ -660,6 +679,7 @@ def runner_factory( async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, extra_config_overrides=extra_config_overrides, + worker_count=worker_count, ) runners.append(runner) return runner diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 951b447fd6b9..4dcddcaa30fa 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -410,7 +410,7 @@ def test_lookup_key_client_lookup_prepends_typed_tag(): # Blocking lookup (non_block defaults to False) runs on the executor and # returns the resolved hit length. - assert client.lookup("req0", token_len=128, block_hashes=[]) == 5 + assert client.lookup("req0", num_tokens=128, block_hashes=[]) == 5 sent_frames = fake_socket.send_multipart.call_args[0][0] assert sent_frames[0] == protocol.LOOKUP_MSG @@ -439,11 +439,11 @@ def test_lookup_key_client_reset_uses_typed_protocol(): assert client.reset() is False -def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0): +def _poll_lookup(client, req_id, num_tokens=128, block_hashes=(), timeout=5.0): """Drive non-blocking lookup until the executor completes it.""" deadline = time.monotonic() + timeout while time.monotonic() < deadline: - result = client.lookup(req_id, token_len, list(block_hashes), non_block=True) + result = client.lookup(req_id, num_tokens, list(block_hashes), non_block=True) if result is not None: return result time.sleep(0.005) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 6e003798c7a4..7460e45d17ed 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -3,6 +3,8 @@ from math import lcm +import torch + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 ExternalCachedBlockPool, MooncakeStoreCoordinator, @@ -14,10 +16,20 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, + MambaSpec, SlidingWindowSpec, ) +def _mamba_align(block_size=32): + return MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for @@ -196,6 +208,49 @@ def test_coordinator_group_block_size_double_hash(): assert hit % 32 == 0 +# ----- Fine-grained partial hits (full attention + mamba "align") ----- + + +def test_coordinator_fine_grained_partial_tail_hit(): + """K3 shape: FA + mamba-align, block_size=32 over hash_block_size=16. When + both groups have the sub-block boundary hash, the reconciled hit lands on + the hash boundary (48), not the block boundary (32).""" + groups = [ + KVCacheGroupSpec(["L0"], _full(32)), + KVCacheGroupSpec(["L1"], _mamba_align(32)), + ] + coord = _make_coord(groups, hash_block_size=16) + assert coord.enable_partial_hash_hits + hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc. + # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary + # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full. + exists = {(g, bytes(h)) for g in (0, 1) for h in (hs[1], hs[2])} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + assert hit == 48 + + +def test_coordinator_fine_grained_clips_when_one_group_missing_tail(): + """If only one group has the sub-block boundary, min-convergence clips the + reconciled hit back to the block boundary (32).""" + groups = [ + KVCacheGroupSpec(["L0"], _full(32)), + KVCacheGroupSpec(["L1"], _mamba_align(32)), + ] + coord = _make_coord(groups, hash_block_size=16) + hs = _hashes(4) + # Full block 0 for both; partial boundary hs[2] only for FA (group 0). + exists = {(g, bytes(hs[1])) for g in (0, 1)} + exists |= {(0, bytes(hs[2]))} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + assert hit == 32 + + # ----- store_mask ----- @@ -476,3 +531,119 @@ def test_load_mask_without_eagle_unchanged(): assert hit == 64 masks = coord.load_mask(hs, token_len=hit) assert masks[0] == [True, True, True, True] + + +def _mamba(block_size=16): + return MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + +def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun(): + """Full+Mamba with eagle must not overrun the attention-verified hit. + + ``MambaManager`` ignores ``drop_eagle_block`` (a Mamba block at position + p IS the recurrent state after (p + 1) * block_size tokens; there is + nothing to recompute), so granting the Mamba group the one-block eagle + peek margin lets it match one block PAST the eagle-pruned full-attention + hit, and adopting that length resumes the recurrent state ahead of the + verified token prefix (#43559). Gating the margin on ``not + isinstance(spec, MambaSpec)`` pins the hit to the attention-verified 48. + """ + groups = [ + KVCacheGroupSpec(["L0"], _full(16)), + KVCacheGroupSpec(["L1"], _mamba(16)), + ] + coord = _make_coord(groups, hash_block_size=16, use_eagle=True) + hs = _hashes(4) + exists = {(g, bytes(h)) for g in (0, 1) for h in hs} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + # FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The + # Mamba group must serve its state@48 snapshot, not peek to state@64. + assert hit == 48 + + +def test_eagle_flag_propagates_to_all_merged_swa_groups(): + """Regression for MTP x PD external-store 0% prefix hit. + + DSV4 splits SWA layers into several KV cache groups sharing one spec, and + only the group containing the MTP layer is annotated ``is_eagle_group``. + The lookup merges equal-spec groups, applies the eagle drop to the merged + group, and requires each chunk hash to exist in EVERY member group — so + the save-side masks must eagle-shift every member, not just the annotated + one. Without propagation the non-annotated groups never store the eagle + proof-run chunks and every external lookup returns 0. + """ + swa = _swa(block_size=16, sliding_window=32) + groups = [ + KVCacheGroupSpec(["L0"], _full(64)), + KVCacheGroupSpec(["L1"], swa), + KVCacheGroupSpec(["L2"], swa, is_eagle_group=True), + ] + coord = _make_coord(groups, hash_block_size=16, use_eagle=True) + assert coord.eagle_group_ids == {1, 2} + + # Save side: both SWA groups must produce identical (eagle-shifted) masks. + masks = coord.store_mask(128, num_prompt_tokens=130) + assert masks[1] == masks[2] + + # Round trip: everything store_mask kept is in the store; the eagle + # lookup must then serve a non-zero hit (it was 0 before the fix). + hs = _hashes(128 // 16) + exists = set() + for g_idx, g in enumerate(groups): + ghashes = chunk_hashes_for_block_size(hs, 16, g.kv_cache_spec.block_size) + mask = masks[g_idx] + for i in range(128 // g.kv_cache_spec.block_size): + if mask is None or mask[i]: + exists.add((g_idx, bytes(ghashes[i]))) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=128, cached_block_pool=ExternalCachedBlockPool(16, exists) + ) + assert hit == 64 + + +def test_dsv4_five_group_eagle_store_lookup_round_trip(): + """Cover the five KV groups observed with DeepSeek-V4-Flash + MTP. + + The two 64-token SWA groups have identical specs, but only one owns the + EAGLE layer. Saving each group through its store mask must still leave a + prefix that the merged SWA lookup can consume. + """ + swa_64_sw128 = _swa(block_size=64, sliding_window=128) + groups = [ + KVCacheGroupSpec(["full_mla"], _full(block_size=256)), + KVCacheGroupSpec(["swa"], swa_64_sw128), + KVCacheGroupSpec(["mtp"], swa_64_sw128, is_eagle_group=True), + KVCacheGroupSpec(["c4_state"], _swa(block_size=4, sliding_window=8)), + KVCacheGroupSpec(["c128_state"], _swa(block_size=8, sliding_window=128)), + ] + coord = _make_coord(groups, hash_block_size=4, use_eagle=True) + token_len = 768 + hashes = _hashes(token_len // coord.hash_block_size) + + # Mirror MooncakeStoreWorker's aligned save: only keys selected by each + # group's store mask are visible to the external lookup. + exists: set[tuple[int, bytes]] = set() + store_masks = coord.store_mask(token_len) + for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)): + group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec) + for chunk_id, block_hash in enumerate(group_hashes): + if mask is None or mask[chunk_id]: + exists.add((gid, bytes(block_hash))) + + _masks, hit = coord.find_longest_cache_hit( + hashes, + max_length=token_len, + cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists), + ) + + # The final 256-token segment has no lookahead block, so EAGLE falls back + # to the previous aligned boundary instead of consuming all 768 tokens. + assert hit == 512 diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index f09e0a24729b..6c8af7fb73c1 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -32,6 +32,7 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MambaSpec, SlidingWindowSpec, ) @@ -173,7 +174,6 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): worker = _build_worker_with_dict_store(vllm_config, cfg, store) worker.tp_size = 1 worker.pp_size = 1 - worker.put_step = 1 worker.num_kv_head = 8 # Register kv_caches using mocked thread classes so register_kv_caches @@ -215,7 +215,7 @@ def _fake_thread_init(*args, **kwargs): block_size=worker.block_size, coord=worker.coord, tp_rank=worker.tp_rank, - put_step=worker.put_step, + group_put_steps=worker._group_tp_replication_factors, kv_role=worker.kv_role, ready_event=ready, enable_kv_event=False, @@ -240,7 +240,10 @@ def _fake_thread_init(*args, **kwargs): worker.store = store # Both groups stored all 4 blocks -> full hit. - assert worker.lookup(token_len=64, block_hashes=hs) == 64 + assert worker.lookup(num_tokens=65, block_hashes=hs) == 64 + # Exact-multiple prompt: the full hit is re-derived one block lower, + # where both groups' stored blocks still cover the SWA window. + assert worker.lookup(num_tokens=64, block_hashes=hs) == 48 # Evict SWA's first two blocks (outside its window of 32 tokens = 2 blocks). swa_keys_outside_window = [ @@ -253,7 +256,12 @@ def _fake_thread_init(*args, **kwargs): # SWA window=32 -> only last 2 blocks must be present in SWA group. # Full has all 4. Coordinator should still return 64. - assert worker.lookup(token_len=64, block_hashes=hs) == 64 + assert worker.lookup(num_tokens=65, block_hashes=hs) == 64 + # Exact-multiple prompt after eviction: the boundary one block lower + # needs SWA block 1, which is gone — no usable stored boundary remains + # (the pre-fix arithmetic clamp would have returned 48 and livelocked + # on load failure -> recompute -> same lookup). + assert worker.lookup(num_tokens=64, block_hashes=hs) == 0 def test_recv_skips_swa_blocks_before_window(): @@ -326,14 +334,14 @@ def batch_get_into_multi_buffers(self, keys, addrs, sizes): def test_chunked_token_database_hash_block_size_smaller_than_block_size(): """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens - keys each 16-token chunk by its last fine hash, keeping the Mooncake key - at one digest instead of concatenating all 4 fine hashes.""" + keys each chunk by its ending fine hash, including a partial tail.""" md = KeyMetadata("m", 0, 0, 0, 0, group_id=3) db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4) db.set_kv_caches_base_addr([0]) db.set_block_len([512]) - # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks. fine_hashes = [BlockHash(bytes([i + 1]) * 4) for i in range(8)] + + # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks. out = list(db.process_tokens(token_len=32, block_hashes=fine_hashes)) assert len(out) == 2 assert out[0][0] == 0 and out[0][1] == 16 @@ -342,3 +350,243 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): # prior three. assert out[0][2].hex() == fine_hashes[3].hex() assert out[1][2].hex() == fine_hashes[7].hex() + + # Sub-block hit: emit the partial chunk under its ending fine hash. + out = list(db.process_tokens(token_len=12, block_hashes=fine_hashes[:3])) + assert [(s, e) for s, e, _ in out] == [(0, 12)] + assert out[0][2].hex() == fine_hashes[2].hex() + + # Cross-block hit: emit both the full chunk and its partial tail. + out = list(db.process_tokens(token_len=28, block_hashes=fine_hashes[:7])) + assert [(s, e) for s, e, _ in out] == [(0, 16), (16, 28)] + assert out[0][2].hex() == fine_hashes[3].hex() + assert out[1][2].hex() == fine_hashes[6].hex() + + +def test_sub_block_partial_tail_offload_reads_cow_block(): + """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the + partial tail is offloaded for both groups under the boundary sub-hash. The + full-attention block is read from the request block table; the mamba block + is the core-provided CoW target, not block_ids.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + assert coord.enable_partial_hash_hits + + class _RecordingStore(_DictStore): + def __init__(self): + super().__init__() + self.puts: dict[str, list[int]] = {} + + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + for key, addr in zip(keys, addrs): + self.puts[key] = addr + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _RecordingStore() + token_dbs = [] + for g_idx in range(2): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=16, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + # The surrounding metadata may describe a longer resumed replay, but the + # handoff identifies the exact state boundary to persist. + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)] + mamba_cow_block = 7 + req = ReqMeta( + req_id="r0", + token_len_chunk=0, + block_ids=([1], [2]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=20, + partial_tail_offloads=[(1, mamba_cow_block, 12)], + ) + + send._maybe_offload_partial_tail(req) + + # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2]. + partial_hash = hs[2] + fa_key = token_dbs[0].key_for(partial_hash) + mamba_key = token_dbs[1].key_for(partial_hash) + assert set(store.puts) == {fa_key, mamba_key} + # FA reads block_ids[0][0] = block 1: addr = base(0) + 1 * 512. + assert store.puts[fa_key] == [512] + # Mamba reads the CoW block 7, not block_ids[1][0]=2. + assert store.puts[mamba_key] == [10_000 + mamba_cow_block * 512] + + +def test_offload_syncs_event_before_put(): + """An offload-carrying meta synchronizes its CoW-fence event before the + store put reads the blocks, then completes in one pass and drains the + completion counter.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + event = MagicMock() + + class _FencedStore(_DictStore): + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + assert event.synchronize.called, "put must run after the event sync" + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _FencedStore() + token_dbs = [] + for g_idx in range(2): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=16, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)] + req = ReqMeta( + req_id="r1", + token_len_chunk=0, + block_ids=([1], [2]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=12, + partial_tail_offloads=[(1, 7, 12)], + ) + req.current_event = event + send.add_stored_request("r1") + + send.request_queue.put(req) + send._handle_request(send.request_queue.get()) + assert send.request_queue.qsize() == 0 + assert store._data + assert send.stored_requests["r1"] == 0 + event.synchronize.assert_called_once() + + +def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): + """The K3-shaped 900/128/1536 scenario scaled to 12/4/16, with a + full-attention group whose block (4) is smaller than the lcm (16): the + offload must persist every FA block up to the boundary — the normal save + floors to the lcm, so those blocks are otherwise never written and the + consumer's per-group lookup would miss. The mamba boundary block still + reads the core-provided CoW target.""" + full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + assert coord.enable_partial_hash_hits + + class _RecordingStore(_DictStore): + def __init__(self): + super().__init__() + self.puts: dict[str, list[int]] = {} + + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + for key, addr in zip(keys, addrs): + self.puts[key] = addr + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _RecordingStore() + token_dbs = [] + for g_idx, block_size in enumerate([4, 16]): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=block_size, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)] # 3 hash units = 12 tok + mamba_cow_block = 7 + req = ReqMeta( + req_id="r2", + token_len_chunk=0, + block_ids=([1, 2, 3], [4]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=12, + partial_tail_offloads=[(1, mamba_cow_block, 12)], + ) + + send._maybe_offload_partial_tail(req) + + # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal + # block-end hashes; mamba (block 16): the partial boundary block under + # the boundary sub-hash, read from the CoW target. + expected = { + token_dbs[0].key_for(hs[0]): [1 * 512], + token_dbs[0].key_for(hs[1]): [2 * 512], + token_dbs[0].key_for(hs[2]): [3 * 512], + token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512], + } + assert store.puts == expected diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py b/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py new file mode 100644 index 000000000000..77cae42ffcc1 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for ChunkedTokenDatabase.prepare_values.""" + +import random + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + ChunkedTokenDatabase, + KeyMetadata, +) +from vllm.utils.math_utils import cdiv + +BLOCK_SIZE = 128 + + +def _reference_prepare_value( + db: ChunkedTokenDatabase, start: int, end: int, block_ids: list[int] +) -> tuple[list[int], list[int], int]: + """Compute a token range with the original scalar implementation.""" + addr_list = [] + size_list = [] + block_id = block_ids[start // db.block_size] + length = len(db.block_len) + for index, base_addr in enumerate(db.kv_caches_base_addr): + addr = base_addr + block_id * db.block_len[index % length] + assert (end - start) % db.block_size == 0 + size = db.block_len[index % length] * cdiv(end - start, db.block_size) + addr_list.append(addr) + size_list.append(size) + return addr_list, size_list, block_id + + +def _make_db(num_regions: int, num_block_lens: int) -> ChunkedTokenDatabase: + md = KeyMetadata(model_name="t", tp_rank=1, pcp_rank=0, dcp_rank=0, pp_rank=0) + db = ChunkedTokenDatabase(md, BLOCK_SIZE) + db.set_kv_caches_base_addr( + [0x7F00_0000_0000 + i * (1 << 30) for i in range(num_regions)] + ) + # Exercise repeated block lengths when there are more cache regions. + db.set_block_len([30_208 + 512 * i for i in range(num_block_lens)]) + return db + + +@pytest.mark.parametrize("num_regions,num_block_lens", [(96, 96), (96, 2), (1, 1)]) +def test_prepare_values_matches_reference(num_regions: int, num_block_lens: int): + db = _make_db(num_regions, num_block_lens) + rng = random.Random(0) + n_blocks = 300 + block_ids = [rng.randrange(0, 1 << 20) for _ in range(n_blocks)] + chunks = [] + b = 0 + while b < n_blocks - 4: + span = rng.choice([1, 1, 1, 2, 4]) + chunks.append((b * BLOCK_SIZE, (b + span) * BLOCK_SIZE)) + b += span + rng.choice([0, 1]) + + addrs, sizes, bids = db.prepare_values(chunks, block_ids) + assert len(addrs) == len(sizes) == len(bids) == len(chunks) + for (start, end), addr, size, bid in zip(chunks, addrs, sizes, bids): + ref_addr, ref_size, ref_bid = _reference_prepare_value( + db, start, end, block_ids + ) + assert addr == ref_addr + assert size == ref_size + assert bid == ref_bid + # Native bindings require Python ints rather than numpy scalars. + assert all(type(a) is int for a in addr) + assert type(bid) is int + + +def test_prepare_value_single_matches_reference(): + db = _make_db(8, 8) + block_ids = list(range(64)) + got = db.prepare_value(5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids) + assert got == _reference_prepare_value( + db, 5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids + ) + + +def test_prepare_values_empty(): + db = _make_db(4, 4) + assert db.prepare_values([], [1, 2, 3]) == ([], [], []) + + +def test_prepare_values_rejects_unaligned_chunk(): + db = _make_db(4, 4) + with pytest.raises(AssertionError): + db.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1]) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 7e2919629876..13960c40340e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -13,11 +13,16 @@ ) -def _make_bare_scheduler() -> MooncakeStoreScheduler: +def _make_bare_scheduler( + *, hash_block_size: int = 16, enable_partial_hash_hits: bool = False +) -> MooncakeStoreScheduler: scheduler = object.__new__(MooncakeStoreScheduler) scheduler.kv_role = "kv_both" scheduler.lookup_async = False + scheduler.enable_lookup = True scheduler._block_size = 16 + scheduler._hash_block_size = hash_block_size + scheduler.enable_partial_hash_hits = enable_partial_hash_hits scheduler.load_specs = {} scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} @@ -136,6 +141,7 @@ def test_preemption_resets_tracker_before_request_finished(): block_hashes=[b"h0", b"h1"], prefill_end_tokens=48, ) + scheduler._request_trackers["req-0"].has_pending_offload = True scheduler.build_connector_meta(_make_preemption_scheduler_output()) @@ -144,6 +150,7 @@ def test_preemption_resets_tracker_before_request_finished(): assert tracker.allocated_block_ids == () assert tracker.num_saved_tokens == 0 assert tracker.token_ids is None + assert tracker.has_pending_offload is False assert tracker.prefill_end_tokens == 0 request = SimpleNamespace(request_id="req-0") assert scheduler.request_finished(request, ([0, 1],)) == (False, None) @@ -473,27 +480,25 @@ def test_from_request_tracker_no_load_saves_normally(): class _StubLookupClient: def __init__(self, hit_tokens: int) -> None: self._hit_tokens = hit_tokens + self.num_tokens: list[int] = [] def lookup( self, req_id: str, - token_len: int, + num_tokens: int, block_hashes: list[bytes], non_block: bool = False, ) -> int: + self.num_tokens.append(num_tokens) return self._hit_tokens def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned(): - # When the external store hits the entire prompt, scheduler must leave at - # least one token uncomputed for sampling but stay on a block boundary. - # Otherwise the recv-side load mask floors token_len to - # (num_tokens-1)//block_size, the tail partial chunk is dropped, and -- if - # the local cache covers the aligned prefix -- key_list ends up empty - # (ZeroDivisionError in the recv thread's `tp_rank % len(key_list)`). + # The worker re-derives a full external hit below the request end on an + # existing boundary, so the scheduler receives the usable aligned hit. scheduler = _make_bare_scheduler() scheduler.load_async = True - scheduler.client = _StubLookupClient(hit_tokens=48) # full hit on 48-token prompt + scheduler.client = _StubLookupClient(hit_tokens=32) request = SimpleNamespace( request_id="req-0", @@ -510,6 +515,7 @@ def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned(): assert need_to_allocate == 16 assert load_async is True load_spec = scheduler.load_specs["req-0"] + assert scheduler.client.num_tokens == [48] assert load_spec.vllm_cached_tokens == 16 assert load_spec.kvpool_cached_tokens == 32 assert load_spec.kvpool_cached_tokens % 16 == 0 @@ -522,7 +528,7 @@ def test_full_external_hit_with_full_local_hit_skips_load(): # into any block-aligned key. scheduler = _make_bare_scheduler() scheduler.load_async = True - scheduler.client = _StubLookupClient(hit_tokens=48) + scheduler.client = _StubLookupClient(hit_tokens=32) request = SimpleNamespace( request_id="req-0", @@ -537,3 +543,251 @@ def test_full_external_hit_with_full_local_hit_skips_load(): assert need_to_allocate == 0 assert load_async is False assert "req-0" not in scheduler.load_specs + + +def test_partial_hash_hit_block_aligned_local_loads_partial_tail(): + # Fine-grained on (hash=4, block=16): a block-aligned local hit can pull a + # sub-block remote hit (24 = a hash boundary inside block 1). Loads [16, 24). + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=24) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=32, + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=16 + ) + + assert need_to_allocate == 8 + assert load_async is True + load_spec = scheduler.load_specs["req-0"] + assert load_spec.vllm_cached_tokens == 16 + assert load_spec.kvpool_cached_tokens == 24 + + +def test_partial_hash_hit_no_remote_gain_skips_load(): + # Core always presents a block-aligned local hit (it floors a sub-block + # tail before calling the connector). When the remote hit does not exceed + # that block-aligned local hit, nothing is loaded. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=16) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=32, + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=16 + ) + + assert need_to_allocate == 0 + assert load_async is False + assert "req-0" not in scheduler.load_specs + + +def test_sub_block_prompt_looks_up_with_fine_grained(): + # A prompt smaller than one block (12 < block 16). With fine-grained partial + # hits the sub-block prefix is worth looking up (floor is the hash unit 4, + # not a full block), so a remote partial hit is loaded. Pre-change the + # block-size floor returned (0, False) for such prompts. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=8) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert need_to_allocate == 8 + assert load_async is True + assert scheduler.load_specs["req-0"].kvpool_cached_tokens == 8 + + +def test_sub_block_prompt_not_looked_up_without_fine_grained(): + # Without fine-grained partial hits, sub-block prompts still skip the lookup + # (there is no full block, and no sub-block key granularity). + scheduler = _make_bare_scheduler() + scheduler.client = _StubLookupClient(hit_tokens=8) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert need_to_allocate == 0 + assert load_async is False + assert "req-0" not in scheduler.load_specs + + +def test_disabled_lookup_reports_no_hit_without_querying_client(): + # With enable_lookup=False the connector reports no external hit without + # consulting the lookup client, so admission is never deferred on a store + # lookup. Used by instances that only contribute store capacity. + scheduler = _make_bare_scheduler() + scheduler.enable_lookup = False + scheduler.client = _StubLookupClient(hit_tokens=32) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=48, + block_hashes=[b"h0", b"h1", b"h2"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert need_to_allocate == 0 + assert load_async is False + assert scheduler.client.num_tokens == [] + assert scheduler.load_specs == {} + + +def test_pending_partial_tail_emits_offload_only_reqmeta(): + # A sub-block prompt never produces a block-aligned save, so the partial- + # tail offload arriving this step is emitted as an offload-only ReqMeta + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so + # the worker skips the normal save). Pending-offload state delays the free + # without advancing the normal-save watermark before the put succeeds. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(12)), + block_hashes=[b"h0", b"h1", b"h2"], + num_output_placeholders=0, + num_prompt_tokens=12, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=12, + allocated_block_ids=([0],), + num_saved_tokens=0, + token_ids=list(range(12)), + prefill_end_tokens=12, + ) + + out = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + new_block_ids=[], + num_computed_tokens=[], + resumed_req_ids=set(), + ), + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + partial_tail_offloads={"req-0": [(1, 7, 12)]}, + ) + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + req_meta = meta.requests[0] + assert req_meta.req_id == "req-0" + assert req_meta.can_save is True + assert req_meta.token_len_chunk == 0 + assert req_meta.partial_tail_offloads == [(1, 7, 12)] + assert req_meta.num_prompt_tokens == 12 + assert req_meta.block_ids == ([0],) + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 0 + assert tracker.has_pending_offload is True + request = SimpleNamespace(request_id="req-0") + assert scheduler.request_finished(request, ([0],)) == (True, None) + + +def test_resumed_partial_tail_uses_handoff_boundary(): + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(20)), + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"], + num_output_placeholders=0, + num_prompt_tokens=12, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=20, + allocated_block_ids=([0, 1],), + num_saved_tokens=0, + token_ids=list(range(20)), + # Resumption replays prompt + previously generated tokens. + prefill_end_tokens=20, + ) + + out = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + new_block_ids=[], + num_computed_tokens=[], + resumed_req_ids=set(), + ), + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + partial_tail_offloads={"req-0": [(1, 7, 12)]}, + ) + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + assert meta.requests[0].partial_tail_offloads == [(1, 7, 12)] + # Ordinary metadata retains the full resumed prefill range. + assert meta.requests[0].num_prompt_tokens == 20 + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 0 + assert tracker.has_pending_offload is True + + +def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary(): + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + num_output_placeholders=0, + num_prompt_tokens=36, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=44, + allocated_block_ids=([0, 1],), + num_saved_tokens=32, + token_ids=list(range(44)), + prefill_end_tokens=48, + ) + out = _make_scheduler_output(scheduled_spec_tokens=None) + out.partial_tail_offloads = {"req-0": [(0, 7, 36)]} + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + assert meta.requests[0].can_save is True + assert meta.requests[0].partial_tail_offloads == [(0, 7, 36)] + assert meta.requests[0].num_prompt_tokens == 48 + # Ordinary saving still covers the full resumed prefill range. + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 48 + assert tracker.has_pending_offload is True diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index f8f43662f26a..0403affea348 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -90,7 +90,7 @@ def _make_store_sending_thread( block_size=block_size, coord=coord, tp_rank=tp_rank, - put_step=put_step, + group_put_steps=[put_step] * len(token_databases), kv_role="kv_producer", ready_event=threading.Event(), replicate_config=replicate_config, @@ -527,6 +527,29 @@ def test_store_sending_thread_delta_strides_with_local_phase(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def test_tp_sharded_group_saves_every_block_on_every_rank(): + """Sharded ranks must write every block because peers hold different bytes.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) + thread.group_put_steps = [1] + + thread.add_stored_request("req-a") + thread._handle_request( + ReqMeta( + req_id="req-a", + token_len_chunk=64, + block_ids=([0, 1, 2, 3],), + block_hashes=[b"a0", b"a1", b"a2", b"a3"], + can_save=True, + ) + ) + + keys = store.batch_is_exist.call_args.args[0] + assert len(keys) == 4 + + def test_store_sending_thread_retries_skipped_range_after_pressure(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -576,6 +599,88 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def _make_partial_tail_send_thread(store): + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + ) + db = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), + block_size=4, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([0x1000]) + db.set_block_len([256]) + return _make_store_sending_thread( + store, + coord=coord, + token_databases=[db], + ) + + +def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta: + return ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=(block_ids,), + block_hashes=[b"a0", b"a1", b"a2"], + can_save=True, + partial_tail_offloads=[(1, 7, 12)], + ) + + +def test_partial_tail_offload_skips_null_source_blocks(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, 256] + thread = _make_partial_tail_send_thread(store) + + assert thread._maybe_offload_partial_tail(_make_partial_tail_req([0, 2, 3])) + + keys, addrs, _sizes, _replicate_config = ( + store.batch_put_from_multi_buffers.call_args.args + ) + assert keys == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132", + ] + assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]] + + +def test_partial_tail_offload_honors_active_pressure_gate(): + store = MagicMock() + thread = _make_partial_tail_send_thread(store) + thread._store_pressure_active = True + thread._skip_store_requests.add("req-a") + thread.add_stored_request("req-a") + + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + + store.batch_is_exist.assert_not_called() + store.batch_put_from_multi_buffers.assert_not_called() + assert thread.stored_requests["req-a"] == 0 + + +def test_partial_tail_put_failure_activates_pressure_gate(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, -200, 256] + thread = _make_partial_tail_send_thread(store) + thread.add_stored_request("req-a") + + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + + assert thread._store_pressure_active is True + assert thread._skip_store_requests == {"req-a"} + assert thread._saved_offset.get("req-a", 0) == 0 + assert thread.stored_requests["req-a"] == 0 + + thread.add_stored_request("req-a") + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + assert store.batch_put_from_multi_buffers.call_count == 1 + + def test_store_sending_thread_delta_start_rank_saves_second_local_chunk(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -655,6 +760,68 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): assert masked_hashes == [b"a2".hex()] +def test_store_sending_thread_prepares_missing_chunks_once_per_group(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 1, 0, 1, 0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256, 512, 512] + coord = SimpleNamespace( + lcm_block_size=16, + store_mask=lambda token_len, start_token, num_prompt_tokens=None: ( + None, + None, + ), + ) + + db0 = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), + block_size=16, + ) + db0.set_kv_caches_base_addr([0x1000]) + db0.set_block_len([256]) + db0.prepare_values = MagicMock(wraps=db0.prepare_values) + db0.prepare_value = MagicMock(side_effect=AssertionError("scalar path called")) + + db1 = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=16, + ) + db1.set_kv_caches_base_addr([0x2000]) + db1.set_block_len([512]) + db1.prepare_values = MagicMock(wraps=db1.prepare_values) + db1.prepare_value = MagicMock(side_effect=AssertionError("scalar path called")) + + thread = _make_store_sending_thread( + store, + coord=coord, + token_databases=[db0, db1], + ) + thread.add_stored_request("req-a") + thread._handle_request( + ReqMeta( + req_id="req-a", + token_len_chunk=48, + block_ids=([0, 1, 2], [2, 1, 0]), + block_hashes=[b"a0", b"a1", b"a2"], + can_save=True, + ) + ) + + db0.prepare_value.assert_not_called() + db1.prepare_value.assert_not_called() + db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [0, 1, 2]) + db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [2, 1, 0]) + + keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args + assert [key.rsplit("@", 1)[-1] for key in keys] == [ + "6130", + "6132", + "6131", + "6132", + ] + assert addrs == [[0x1000], [0x1200], [0x2200], [0x2000]] + assert sizes == [[256], [256], [512], [512]] + + def test_store_sending_thread_only_skips_on_no_available_handle(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -1212,7 +1379,8 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ] assert len(keys) == len(block_hashes) # PUT side: mirrors KVCacheStoreSendingThread's striding slice. - put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step]) + put_step = w._group_tp_replication_factors[0] + put_keys.update(keys[w.tp_rank % put_step :: put_step]) # GET side: KVCacheStoreRecvingThread fetches every key. get_keys_per_rank[tp_rank] = set(keys) @@ -1551,6 +1719,15 @@ def _register_with_mocked_threads( worker.register_kv_caches(kv_caches) +def _refresh_group_tp_replication_factors( + worker: mooncake_store_worker.MooncakeStoreWorker, +) -> None: + worker._group_tp_replication_factors = ( + worker._compute_group_tp_replication_factors() + ) + worker._init_lookup_key_prefixes() + + def _make_bare_worker( *, num_gpu_blocks: int = 10, @@ -1568,11 +1745,10 @@ def _make_bare_worker( worker.cache_config.num_gpu_blocks = num_gpu_blocks worker.store = MagicMock() worker.store.register_buffer.return_value = 0 - worker.use_mla = False worker.kv_role = kv_role + worker._capacity_only = False worker.block_size = block_size worker.tp_rank = 0 - worker.put_step = 1 worker.enable_kv_events = False worker.kv_send_thread = None worker.kv_recv_threads = [] @@ -1602,7 +1778,6 @@ def _make_bare_worker( worker.pcp_size = 1 worker.dcp_size = 1 worker.hash_block_size = block_size - worker.metadata = KeyMetadata("test-model", 0, 0, 0, 0) # Pre-build a single-group token_dbs so lookup-only tests don't have to # call register_kv_caches. worker.token_dbs = [ @@ -1617,7 +1792,7 @@ def _make_bare_worker( scheduler_block_size=block_size, hash_block_size=block_size, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) return worker @@ -1626,9 +1801,8 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) - assert worker._lookup_expected_per_key == 4 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0", @@ -1643,21 +1817,194 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): worker.num_kv_head = 1 worker.pcp_size = 2 worker.dcp_size = 1 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) - assert worker._lookup_expected_per_key == 2 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0", ) +def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): + """Replicated attention needs one namespace; sharded Mamba needs every rank.""" + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + _refresh_group_tp_replication_factors(worker) + + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + + +def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + MLAAttentionSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + mla = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=None) + gqa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], mla), + KVCacheGroupSpec(["l1"], gqa), + KVCacheGroupSpec(["l2"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), block_size=16 + ) + for g_idx in range(3) + ] + + _refresh_group_tp_replication_factors(worker) + assert worker._group_tp_replication_factors == (4, 2, 1) + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + assert worker._lookup_key_prefixes[2] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:3@pcp0@dcp0@pp_rank:0@group:2", + ) + + +@pytest.mark.parametrize("spec_order", [("mla", "gqa"), ("gqa", "mla")]) +def test_uniform_group_uses_common_inner_replication_factor(spec_order): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + specs_by_name = { + "mla": MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + "gqa": FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + } + inner_specs = {name: specs_by_name[name] for name in spec_order} + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs=inner_specs, + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(list(inner_specs), uniform_spec), + ] + + _refresh_group_tp_replication_factors(worker) + + assert worker._group_tp_replication_factors == (2,) + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:0", + ) + + +def test_lookup_rejects_boundary_missing_one_mamba_shard(): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=16, + ) + _refresh_group_tp_replication_factors(worker) + + # 33 tokens for two 16-token blocks: the hit stops below the request end, + # so the full-hit re-derivation stays out of the shard accounting. + worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys) + assert worker.lookup(33, [b"h0", b"h1"]) == 32 + + worker.store.batch_is_exist.side_effect = lambda keys: [ + 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys + ] + assert worker.lookup(33, [b"h0", b"h1"]) == 0 + + def test_lookup_requires_all_dcp_rank_namespaces(): worker = _make_bare_worker(block_size=16) worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [1, 1, 0, 1] assert worker.lookup(16, [b"a0"]) == 0 @@ -1675,6 +2022,101 @@ def test_lookup_partial_prefix_returns_first_hit_length(): assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32 +def test_lookup_partial_tail_uses_hash_alignment(): + """A stored sub-block tail can serve a request extending past it.""" + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["full"], full), + KVCacheGroupSpec(["mamba"], mamba), + ] + worker.hash_block_size = 4 + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id), + block_size=16, + hash_block_size=4, + ) + for group_id in range(2) + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=4, + ) + _refresh_group_tp_replication_factors(worker) + worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1] + + assert worker.lookup(13, [b"h0", b"h1", b"h2"]) == 12 + + +def test_lookup_full_hit_reuses_existing_boundary(): + """A full hit is re-derived below the request end without another RPC.""" + worker = _make_bare_worker(block_size=16) + worker.store.batch_is_exist.return_value = [1, 1] + + assert worker.lookup(32, [b"h0", b"h1"]) == 16 + assert worker.store.batch_is_exist.call_count == 1 + + +def test_lookup_full_hit_with_eagle_pops_once_not_twice(): + """Eagle already leaves the last block for the drafter, so a + full-prompt re-derivation must never fire for eagle-governed hits: + firing would anchor the search one block lower and pop a second + block, regressing the hit by an extra producer boundary.""" + worker = _make_bare_worker(block_size=16) + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=16, + use_eagle=True, + ) + worker.store.batch_is_exist.return_value = [1, 1, 1, 1] + + # 64-token exact-multiple prompt, all 4 blocks stored: one eagle pop + # gives 48; a spurious re-derivation (anchored at 48) would pop again + # and return 32. + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 48 + assert worker.store.batch_is_exist.call_count == 1 + + +def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable(): + """The motivating livelock: the producer of a 64-token prompt stored + only its SWA tail window (blocks 2-3). The old arithmetic clamp turned + the full hit into 48, whose SWA window needs the never-written block 1, + so every load failed and the recompute re-entered the same lookup. The + re-derivation must report that no stored boundary below the request end + is usable.""" + from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec + + worker = _make_bare_worker(block_size=16) + swa = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=None, sliding_window=32 + ) + worker._kv_cache_groups = [KVCacheGroupSpec(["layer0"], swa)] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=worker.hash_block_size, + hash_block_size=worker.hash_block_size, + ) + worker.store.batch_is_exist.return_value = [0, 0, 1, 1] + + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 0 + assert worker.store.batch_is_exist.call_count == 1 + + def test_lookup_swa_single_group_returns_full_when_tail_window_present(): """Single-SWA, sliding_window=32 (= 2 blocks): producer stored only the tail. Coordinator-driven lookup returns full prefix even though the @@ -1692,7 +2134,7 @@ def test_lookup_swa_single_group_returns_full_when_tail_window_present(): hash_block_size=worker.hash_block_size, ) worker.store.batch_is_exist.return_value = [0, 0, 1, 1] - assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 64 + assert worker.lookup(65, [b"h0", b"h1", b"h2", b"h3"]) == 64 def test_lookup_checks_all_potential_swa_hit_boundaries(): @@ -1733,7 +2175,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) # Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11. # Only the first full chunk and the SWA chunk ending at token 32 exist, so # lookup should recover a 32-token external prefix hit. A sparse @@ -1794,7 +2236,7 @@ def test_lookup_applies_swa_mask_before_accessing_hashes(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)]) accessed_before_rpc: list[int] = [] @@ -2157,7 +2599,7 @@ def test_lookup_records_mooncake_metrics(): worker = _make_bare_worker() worker.store.batch_is_exist.return_value = [1, 1] - result = worker.lookup(32, [b"a0", b"a1"]) + result = worker.lookup(33, [b"a0", b"a1"]) stats = worker.get_kv_connector_stats() assert result == 32 diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py new file mode 100644 index 000000000000..e38ec3f7e9be --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hardware-fair KV-read routing across MoRIIO heterogeneous P/D configs. + +Dependency-light (no GPU / ROCm / mori): builds bare ``MoRIIOConnectorWorker`` +instances via ``object.__new__`` and drives the REAL read-source routing +decision (``_resolve_read_source`` / ``_next_flex_tp_rank``) directly -- no +routing logic is re-implemented here, so a future change to those functions is +what these tests exercise. + +Scope -- the CONNECTOR (decode-side) read routing. The connector never selects +the prefill instance; the proxy does and hands the connector a ``remote_host`` + +``remote_dp_rank`` per request. The connector's only fairness lever is WHICH +prefill (dp, tp) rank each read targets, and the RFC's deployments collapse to +three real connector behaviours: + + symmetric TP (TP prefill + TP decode) -> decode tp_k reads prefill tp_k + flexible (TP prefill + DP decode) -> round-robin over prefill tp0..N-1 + owner DP (DP prefill, any decode) -> read the owner rank, tp0 + +These assume the proxy delivers a FAIR request stream (each prefill instance + +owner dp-rank an equal share) and verify the connector never re-introduces a +bottleneck. That proxy contract is checked separately against the real +``flat_interleaved_dp_route`` in the toy-proxy tests (PR #46115). + +A node runs one of two modes (8 GPUs each): + * TP8 -> (dp_size, tp_size) = (1, 8); MLA latent KV REPLICATED on all ranks. + * DP8EP -> (dp_size, tp_size) = (8, 1); KV PARTITIONED, one owner rank/request. +""" + +from collections import Counter +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ReqMeta, + get_port_offset, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, +) + +MODE_DIMS = {"TP8": (1, 8), "DP8EP": (8, 1)} # (dp_size, tp_size), 8 GPUs/node + + +@dataclass(frozen=True) +class PDConfig: + name: str + p_mode: str + d_mode: str + + @property + def p_dp(self) -> int: + return MODE_DIMS[self.p_mode][0] + + @property + def p_tp(self) -> int: + return MODE_DIMS[self.p_mode][1] + + @property + def d_dp(self) -> int: + return MODE_DIMS[self.d_mode][0] + + @property + def d_tp(self) -> int: + return MODE_DIMS[self.d_mode][1] + + @property + def n_prefill_gpus(self) -> int: + return self.p_dp * self.p_tp + + +CONFIGS = [ + PDConfig("1P_TP8:1D_TP8", "TP8", "TP8"), + PDConfig("2P_TP8:1D_DP8EP", "TP8", "DP8EP"), + PDConfig("2P_TP8:2D_TP8", "TP8", "TP8"), + PDConfig("2P_DP8EP:3D_DP8EP", "DP8EP", "DP8EP"), + PDConfig("2P_DP8EP:4D_TP8", "DP8EP", "TP8"), +] +CONFIG_IDS = [c.name for c in CONFIGS] + + +def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int): + w = object.__new__(MoRIIOConnectorWorker) + w.world_size = world_size + w.tp_rank = tp_rank + w.dp_rank = dp_rank + w.use_mla = True + return w + + +def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0"): + return ReqMeta( + transfer_id="t", + local_block_ids=[1], + remote_block_ids=[2], + remote_host=host, + remote_port=1234, + remote_handshake_port=6301, + remote_notify_port=61005, + remote_engine_id=f"{host}:6301", + tp_size=p_tp, + remote_dp_size=p_dp, + remote_dp_rank=remote_dp_rank, + ) + + +def build_decode_workers(cfg: PDConfig) -> list: + """Decode workers that issue reads for one prefill instance. A TP decode + instance reads from every tp rank; a DP decode instance from each dp-rank + worker. Reused across requests so per-worker round-robin state advances.""" + if cfg.d_tp > 1: + return [ + make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0) + for r in range(cfg.d_tp) + ] + return [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d) for d in range(cfg.d_dp) + ] + + +def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter: + """Drive the REAL _resolve_read_source over a fair input stream; tally which + prefill GPU each read targets. The only modelled assumption is the proxy + contract -- each owner dp-rank delivered equally (``for owner_dp in + range(p_dp)``); no proxy algorithm is reproduced.""" + workers = build_decode_workers(cfg) + hits: Counter = Counter() + for _ in range(rounds): + for owner_dp in range(cfg.p_dp): + for worker in workers: + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=owner_dp) + chosen_tp, _flexible = worker._resolve_read_source(meta) + target = get_port_offset(owner_dp, chosen_tp, cfg.p_tp) + assert 0 <= target < cfg.n_prefill_gpus + hits[target] += 1 + return hits + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_kv_read_load_is_hardware_fair(cfg: PDConfig) -> None: + # rounds a multiple of p_tp so the round-robin closes on an exact cycle. + hits = prefill_target_multiset(cfg, rounds=16) + gpus = set(range(cfg.n_prefill_gpus)) + assert set(hits) == gpus, f"unused prefill GPUs: {sorted(gpus - set(hits))}" + assert max(hits.values()) == min(hits.values()), dict(sorted(hits.items())) + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_flexible_gate_fires_only_for_tp_prefill_dp_decode(cfg: PDConfig) -> None: + flags = set() + for worker in build_decode_workers(cfg): + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=0) + _chosen, flexible = worker._resolve_read_source(meta) + flags.add(flexible) + is_mirror = cfg.d_tp == 1 and cfg.p_dp == 1 and cfg.p_tp > 1 + assert flags == {is_mirror} + + +def test_symmetric_tp_is_a_bijection() -> None: + targets = [] + for tp_rank in range(8): + worker = make_decode_worker(world_size=8, tp_rank=tp_rank, dp_rank=0) + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=8, p_dp=1, remote_dp_rank=0) + ) + assert not flexible + targets.append(chosen_tp) + assert sorted(targets) == list(range(8)) + + +def test_owner_dp_read_is_faithful_and_covers_every_rank() -> None: + worker = make_decode_worker(world_size=8, tp_rank=3, dp_rank=0) + targets = [] + for owner_dp in range(8): + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=1, p_dp=8, remote_dp_rank=owner_dp) + ) + assert not flexible + assert chosen_tp == 0 # p_tp == 1, only tp0 exists + targets.append(get_port_offset(owner_dp, chosen_tp, 1)) + assert sorted(targets) == list(range(8)) + + +def test_flexible_round_robin_is_deterministic_uniform_and_staggered() -> None: + w0 = make_decode_worker(world_size=1, tp_rank=0, dp_rank=0) + seq = [w0._next_flex_tp_rank(8) for _ in range(64)] + assert seq[:8] == list(range(8)) + assert Counter(seq) == Counter({t: 8 for t in range(8)}) + + first_pick = [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d)._next_flex_tp_rank(8) + for d in range(8) + ] + assert sorted(first_pick) == list(range(8)) + + +def test_read_blocks_for_req_threads_chosen_tp() -> None: + # The resolved (chosen_tp, flexible) must reach _read_blocks, which keys the + # session AND the notify port off that single value -- so a read and its + # completion notify address the same prefill rank. + worker = make_decode_worker(world_size=1, tp_rank=0, dp_rank=3) + worker._read_blocks = MagicMock() + worker._read_blocks_for_req("r", make_meta(p_tp=8, p_dp=1, remote_dp_rank=0)) + kw = worker._read_blocks.call_args.kwargs + assert kw["flexible"] is True + assert kw["chosen_tp"] == 3 # first flexible pick = dp_rank seed + assert get_port_offset(0, kw["chosen_tp"], 8) == 3 diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 31d34c69b3fb..ba9f73f24058 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -204,13 +204,12 @@ def get_xfer_telemetry(self, handle: int) -> dict: def _make_fake_nixl_pkg(): """Context manager that creates a temporary package making `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper. - Also creates rixl package for ROCm compatibility. + Also creates the ROCm NIXL packages. Automatically cleans up the temporary directory when done. """ with tempfile.TemporaryDirectory() as td: - # Create both nixl and rixl packages for cross-platform compatibility - for pkg_name in ["nixl", "rixl"]: + for pkg_name in ["nixl", "nixl_rocm"]: pkg_root = os.path.join(td, pkg_name, "_api") os.makedirs(pkg_root, exist_ok=True) @@ -480,8 +479,9 @@ def __init__( super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout - # Mock register_kv_caches attribute needed for tests that do not call it. + # Mock register_kv_caches attributes needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} + self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) @@ -766,8 +766,9 @@ def check_handshake(remote_tp_size: int): assert remote_info.remote_tp_size == remote_tp_size assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size) # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks - assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio - assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio + split_key = (-tp_ratio, worker.block_size) + assert split_key in worker.src_xfer_handles_by_tp_ratio + assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio assert remote_engine_id in worker.dst_xfer_side_handles assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set( range(tp_ratio) @@ -1116,18 +1117,6 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices( block_lens=[remote_block_len], ) - assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == ( - local_block_len - ) - assert ( - worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True) - == worker._mamba_ssm_size[0] - ) - assert ( - worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True) - == worker._mamba_ssm_size[1] - ) - assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ [0x1000 + local_block_len, local_block_len, 0] ] @@ -1409,6 +1398,53 @@ def test_kv_connector_stats(default_vllm_config, dist_init): assert stats_after_reset is None +def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init): + """reqs_to_send deadlines are stamped with the scheduler process's + perf_counter, whose epoch differs across processes and (by boot-time + deltas) across nodes. Without rebasing, a P worker on a node whose + monotonic clock is ahead of the scheduler's by more than the TTL + expires the lease on arrival and reports done_sending before D has + read the blocks — the freed blocks can then be reallocated and the + remote read pulls another request's data (silent accuracy corruption). + The worker must anchor the remaining TTL to its own clock. + """ + vllm_config = create_vllm_config() + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + + req_id = "req-lease-clock" + ttl = 480.0 + # Simulate a scheduler whose monotonic clock is 10,000 s behind this + # worker's (e.g. its node booted much later): the raw deadline is + # then already far in the past in this worker's clock domain. + scheduler_clock = time.perf_counter() - 10_000.0 + + metadata = NixlConnectorMetadata() + metadata.reqs_in_batch = {req_id} + metadata.reqs_to_send = {req_id: scheduler_clock + ttl} + metadata.scheduler_clock = scheduler_clock + connector.bind_connector_metadata(metadata) + dummy_ctx = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + ) + connector.start_load_kv(dummy_ctx) + + remaining = worker._reqs_to_send[req_id] - time.perf_counter() + assert ttl - 5.0 < remaining <= ttl + 5.0 + + # The expiry sweep must not release the request. + done_sending, _ = worker.get_finished() + assert req_id not in done_sending + assert req_id in worker._reqs_to_process + + def test_kv_connector_stats_aggregation(): """ Test KV transfer stats aggregation across TP ranks using @@ -2104,7 +2140,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: @@ -3130,10 +3166,12 @@ def test_mla_broadcast_notif_uses_remote_request_id( ) +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) def test_kv_both_deprecation_warning(default_vllm_config, dist_init): """kv_role='kv_both' should emit a deprecation log warning.""" - from unittest.mock import patch - from vllm.logger import _print_warning_once _print_warning_once.cache_clear() @@ -3156,10 +3194,12 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): assert "deprecated" in msg +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init): """kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning.""" - from unittest.mock import patch - for role in ("kv_consumer", "kv_producer"): vllm_config = create_vllm_config(kv_role=role) with patch( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 4945942ba3a8..c3058eaddbad 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -365,6 +365,28 @@ def test_apply_prefix_caching_mamba_hybrid( [[6, 7, 8, 9], [99]], id="fa_prefix_hit_and_ssm_trim", ), + # Multi-slot SSM ("all" mode): a local prefix hit leaves fewer local + # slots; the earlier remote slots are covered locally → remote tail. + pytest.param( + 10, + 10, + [list(range(10)), [5, 6]], + [list(range(10)), [1, 2, 3]], + [list(range(10)), [5, 6]], + [list(range(10)), [2, 3]], + id="ssm_multi_block_local_hit_tail", + ), + # Multi-slot SSM ("all" mode): the one trailing local position holds + # the token D recomputes itself → local head-clip. + pytest.param( + 10, + 10, + [list(range(10)), [4, 5, 6]], + [list(range(10)), [8, 9]], + [list(range(10)), [4, 5]], + [list(range(10)), [8, 9]], + id="ssm_multi_block_local_extra_head_clip", + ), ], ) def test_apply_prefix_caching_ssm_prefix_cache_hit( @@ -402,6 +424,28 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit( ) +@pytest.mark.cpu_test +def test_apply_prefix_caching_ssm_unpairable_slots_rejected(): + """Local SSM slots can only exceed the remote ones by the position D + recomputes itself. A larger excess means the lists aren't + position-aligned: fail loudly rather than transfer into wrong slots.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = object.__new__(NixlConnectorWorker) + worker._has_mamba = True + worker._physical_blocks_per_logical_kv_block = 10 + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True) + + with pytest.raises(AssertionError, match="unpairable SSM state slots"): + worker._apply_prefix_caching( + [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10 + ) + + @pytest.mark.cpu_test @pytest.mark.parametrize( "local_physical_per_logical,remote_physical_per_logical," @@ -708,6 +752,120 @@ def test_get_block_descs_ids_kernel_block_mismatch(): assert list(result) == expected, f"Expected {expected}, got {list(result)}" +@pytest.mark.cpu_test +def test_get_block_descs_ids_hetero_block_size_hybrid(): + """With a block-size ratio, FA desc ids are ratio-expanded while SSM + desc ids keep the unexpanded logical stride (state blocks are never + sub-split).""" + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = _make_mock_worker_for_desc_ids( + num_regions=2, + has_mamba=True, + group_spec_types=(FullAttentionSpec, MambaSpec), + block_len_per_layer=[100], + ) + + ratio = 4 + # FA ids are already remote-granularity (expanded) sub-block ids. + fa_sub_blocks = [3, 5] + ssm_blocks = [1] + result = worker._compute_desc_ids( + block_ids=(fa_sub_blocks, ssm_blocks), + dst_num_blocks=100, + block_size_ratio=ratio, + physical_blocks_per_logical=1, + ) + + # FA regions have 100*4 entries each; SSM regions (4 per layer) start at + # 2*400 and stride by the unexpanded 100 logical blocks. + expected = [3, 5, 403, 405, 801, 901, 1001, 1101] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +def _bind_worker_method(worker, name): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + method = getattr(NixlConnectorWorker, name) + setattr(worker, name, method.__get__(worker, NixlConnectorWorker)) + + +@pytest.mark.cpu_test +def test_map_block_ids_for_block_size_ratio_hybrid(): + """Attention groups expand to remote granularity and clip to the remote + coverage; mamba state blocks pass through 1:1.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + _bind_worker_method(worker, "get_mapped_blocks") + _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio") + + local, remote = worker._map_block_ids_for_block_size_ratio( + [[1, 2, 3], [7]], + [list(range(30, 40)), [42]], + 4, + ) + # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks. + assert local == [list(range(4, 14)), [7]] + assert remote == [list(range(30, 40)), [42]] + + # Attention-only full prefix hit: empty local list is preserved. + worker._group_spec_types = (FullAttentionSpec,) + local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4) + assert local == [] + + +@pytest.mark.cpu_test +def test_post_process_zeroes_untransferred_tail(): + """The untransferred sub-blocks of the last local block are zeroed on + receive; mamba state caches are untouched by the attention permute.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + ratio = 4 + block_tokens = 8 # 2 tokens per remote sub-block + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.device_type = "cpu" + worker.enable_permute_local_kv = False + attn_cache = torch.ones(6, block_tokens, 2, 4) + mamba_cache = torch.ones(6, 16) + worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + fa_group = MagicMock(layer_names=["attn.0"]) + ssm_group = MagicMock(layer_names=["mamba.0"]) + worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group]) + # The cached property filters mamba layers out of the permuted caches. + attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) + assert len(attn_caches) == 1 and attn_caches[0] is attn_cache + worker._attention_kv_caches = attn_caches + _bind_worker_method(worker, "post_process_device_kv_on_receive") + + # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + + # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). + assert torch.all(attn_cache[2] == 1) + assert torch.all(attn_cache[3, :4] == 1) + assert torch.all(attn_cache[3, 4:] == 0) + # Untouched blocks and the mamba cache keep their content. + assert torch.all(attn_cache[4] == 1) + assert torch.all(mamba_cache == 1) + + @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM @@ -1266,3 +1424,240 @@ def test_logical_to_kernel_block_ids_with_remote_ratio( assert list(result) == expected_kernel_block_ids, ( f"Expected {expected_kernel_block_ids}, got {result}" ) + + +@pytest.mark.cpu_test +def test_exchange_clipped_blocks_ssm_single_state(): + """In single-state cache modes, SSM lists are reduced to the running + state slot: speculative scratch slots, null placeholders and the previous + step's state carry nothing. Attention groups pass through untouched.""" + sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) + sched.blocks_per_sw = [0, 0] + sched._ssm_spec_blocks = [None, 2] + sched._ssm_state_slots_are_positional = False + + # Align-mode list: null placeholders, state block, 2 speculative slots. + clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 0, 7, 8, 9])) + assert clipped == ([1, 2, 3], [7]) + + # Same, still holding the previous step's state block (freed a step later). + assert sched.get_exchange_clipped_blocks(([1], [0, 6, 7, 8, 9]))[1] == [7] + + # Default (mamba_block_size=max_model_len): state block, 2 scratch slots. + assert sched.get_exchange_clipped_blocks(([1], [7, 8, 9]))[1] == [7] + + # Scratch slots not allocated: the state slot still survives. + assert sched.get_exchange_clipped_blocks(([1], [5]))[1] == [5] + + # Non-mamba models pass through unchanged. + fa_sched = make_nixl_scheduler(has_mamba=False) + assert fa_sched.get_exchange_clipped_blocks(([1, 2],)) == ([1, 2],) + + +@pytest.mark.cpu_test +def test_exchange_clipped_blocks_ssm_positional_states(): + """In "all" mode every position holds a state, so only the speculative + slots go; placeholders stay to keep the list position-indexed.""" + sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) + sched.blocks_per_sw = [0, 0] + sched._ssm_spec_blocks = [None, 2] + sched._ssm_state_slots_are_positional = True + + clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 5, 6, 7, 8, 9])) + assert clipped == ([1, 2, 3], [0, 5, 6, 7]) + + +# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ───────────────────── + + +def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): + """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed + MambaSpec) groups whose layers share the same HMA tensors, with a + mamba-aligned unified page and an MLA kernel block smaller than the + logical block.""" + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + # 12-token logical blocks over a 4-token MLA kernel block. + mla_spec = MLAAttentionSpec( + block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16 + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=12, + # GDN-decomposable conv (Q|K|V = 2|2|4 cols x 3 rows) + fp32 temporal. + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + assert kda_spec.page_size_bytes == unified_page + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + +@pytest.mark.cpu_test +def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): + """Hybrid MLA+KDA registration: HMA tensors shared by both layer types + must be flagged as MLA regions even when a KDA layer registers them + first, expose TP-independent kernel-granularity block lens, and build + FA + mamba descriptors for every region.""" + from unittest.mock import MagicMock + + from vllm.config import set_current_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + kv_cache_config = _make_hybrid_mla_kv_cache_config() + unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes + vllm_config = create_vllm_config(block_size=12) + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared tensors would not be + # deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [4] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + with ( + patch.object(bw, "NixlWrapper"), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config) + worker.use_mla = True # opt-125m test config is not MLA; force the flag + worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata" + + tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)] + # KDA layer first per tensor: exercises the dual-purpose flag merge. + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + + # 12-token logical blocks over the 4-token MLA kernel block. + assert worker._physical_blocks_per_logical_kv_block == 3 + assert worker.block_size == 4 and worker.num_blocks == 12 + # Both shared tensors are dual-purpose: their FA view is MLA even though + # a KDA layer registered them first. + assert worker._region_is_mla == [True, True] + assert worker.num_regions == 2 and worker.num_descs == 24 + # Kernel-granularity block lens; TP-independent for MLA hybrids. + assert worker.block_len_per_layer == [unified_page // 3] * 2 + # Split handles must replicate every FA descriptor (MLA isn't head-sharded). + assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24 + # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page. + # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks. + assert worker.src_blocks_data.shape == (24 + 32, 3) + fa_descs = worker.src_blocks_data[:24] + assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3 + assert all(size == unified_page // 3 for size in fa_descs[:, 1]) + + +@pytest.mark.cpu_test +def test_push_write_hybrid_mla_replicates_attention(): + """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be + written to every covered D rank (replicated MLA latent) while SSM state + is written per-rank through the split handles.""" + import threading + from collections import defaultdict + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + ) + from vllm.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec + + worker = object.__new__(NixlPushConnectorWorker) + worker.shutdown = lambda: None # skeleton worker: silence __del__ + worker.use_mla = True + worker._has_mamba = True + worker._group_spec_types = (MLAAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.transfer_topo.tp_ratio.return_value = -2 + remote_info = MagicMock() + remote_info.remote_physical_blocks_per_logical = 1 + remote_info.remote_block_size = 4 + worker.transfer_topo.get_engine_info.return_value = remote_info + + engine_id = "remote-engine" + # Read-oriented mapping collapses the replicated attention group to one + # source rank; the SSM state is sharded across both covered D ranks. + worker.tp_mappings = { + engine_id: TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + } + worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}} + worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]} + worker.src_xfer_handles_by_block_size = {4: 300} + worker._sending_transfers = defaultdict(list) + worker._sending_transfers_lock = threading.Lock() + worker.kv_cache_config = _make_hybrid_mla_kv_cache_config() + worker._xfer_blocks = MagicMock(return_value=1) + + meta = MagicMock() + meta.remote.engine_id = engine_id + meta.remote.block_ids = [[7, 8], [3]] + meta.local_physical_block_ids = [[1, 2], [5]] + + worker._xfer_blocks_for_req("req-1", meta) + + calls = worker._xfer_blocks.call_args_list + assert len(calls) == 2 + for call, rank, local_handle, remote_handle in zip( + calls, (0, 1), (200, 201), (100, 101) + ): + spec = call.kwargs["read_spec"] + assert spec.remote_rank == rank + # Attention group replicated to every rank, SSM by membership. + assert spec.local_block_ids == [[1, 2], [5]] + assert spec.remote_block_ids == [[7, 8], [3]] + assert call.kwargs["local_xfer_side_handle"] == local_handle + assert call.kwargs["remote_xfer_side_handle"] == remote_handle diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py new file mode 100644 index 000000000000..6ac9ce030d05 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models +under heterogeneous P/D block geometry (TP-sharded KDA-style state, so +the mamba-aligned logical block size differs between P and D while the +kernel-granularity pages stay equal). + +The invariant under test: every LOCAL byte range a request's READ transfers +into must lie within that request's own blocks. A violation means an +incoming transfer can overwrite a co-resident request's KV or mamba state +mid-decode (silent corruption of an unrelated request). +""" + +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from .utils import create_vllm_config + + +class _RecordingNixl: + """Minimal NIXL wrapper stand-in that records descriptor lists and + prepared transfers so tests can resolve desc ids to byte ranges.""" + + def __init__(self, *args, **kwargs): + self.dlists: dict[int, np.ndarray] = {} + self.xfers: list[tuple] = [] + self._next_handle = 1 + + def get_reg_descs(self, caches_data, mem_type): + return caches_data + + def register_memory(self, descs, backends=None): + pass + + def deregister_memory(self, descs): + pass + + def get_agent_metadata(self): + return b"agent-meta" + + def get_xfer_descs(self, blocks_data, mem_type): + return blocks_data + + def prep_xfer_dlist(self, agent, descs): + handle = self._next_handle + self._next_handle += 1 + self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3) + return handle + + def add_remote_agent(self, metadata): + return "remote-agent" + + def make_prepped_xfer( + self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None + ): + handle = self._next_handle + self._next_handle += 1 + self.xfers.append( + ( + op, + local_handle, + np.asarray(local_ids), + remote_handle, + np.asarray(remote_ids), + ) + ) + return handle + + def transfer(self, handle): + pass + + def check_xfer_state(self, handle): + return "DONE" + + def get_xfer_telemetry(self, handle): + from types import SimpleNamespace + + return SimpleNamespace( + xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1 + ) + + def release_xfer_handle(self, handle): + pass + + def release_dlist_handle(self, handle): + pass + + def send_notif(self, agent, notif_msg=None): + pass + + def get_new_notifs(self): + return {} + + def remove_remote_agent(self, agent): + pass + + +def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks): + """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( + base_worker as bw, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + mla_spec = MLAAttentionSpec( + block_size=local_block_size, + num_kv_heads=1, + head_size=6, + dtype=torch.float16, + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=local_block_size, + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_logical_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_logical_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + vllm_config = create_vllm_config(block_size=local_block_size) + vllm_config.cache_config.enable_prefix_caching = False + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared-tensor regions this test builds + # would not be deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + from unittest.mock import MagicMock + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + from vllm.config import set_current_vllm_config + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + worker.use_mla = True + + # Attention caches are kernel-block granular on dim 0, as the + # receive post-process assumes. + ppl = local_block_size // kernel_block_size + tensors = [ + torch.zeros( + num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8 + ) + for _ in range(2) + ] + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + # Keep tensors alive alongside the worker; flat views for byte checks. + worker._test_tensors = [t.view(-1) for t in tensors] + worker._test_tensors_2d = tensors + worker._test_unified_page = unified_page + return worker + + +def _make_remote_meta( + worker, + remote_block_size, + remote_kernel_block_size, + remote_num_logical, + remote_ssm_sizes, +): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, + ) + + remote_ppl = remote_block_size // remote_kernel_block_size + # Kernel-granularity pages are TP-independent for MLA hybrids and must + # match the local ones for the handshake to pass, scaled down by the + # block-size ratio when the remote's kernel block is smaller. + block_size_ratio = worker.block_size // remote_kernel_block_size + kernel_page = worker.block_len_per_layer[0] // block_size_ratio + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"remote-agent-meta", + device_id=0, + kv_caches_base_addr=[0x10_000_000, 0x20_000_000], + num_blocks=remote_num_logical * remote_ppl, + block_lens=[kernel_page, kernel_page], + kv_cache_layout=worker.kv_cache_layout, + block_size=remote_kernel_block_size, + ssm_sizes=remote_ssm_sizes, + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=remote_ppl, + ) + + +def _owned_byte_ranges(worker, group_logical_ids): + """Byte ranges owned by a request: for each HMA region tensor, every + logical block id of every group maps to one unified page.""" + unified_page = worker._test_unified_page + bases = [t.data_ptr() for t in worker._test_tensors] + owned = [] + for base in bases: + for ids in group_logical_ids: + for b in ids: + owned.append((base + b * unified_page, base + (b + 1) * unified_page)) + return owned + + +def _assert_local_writes_within(worker, owned_ranges): + nixl = worker.nixl_wrapper + assert nixl.xfers, "no transfers were posted" + violations = [] + total_descs = 0 + for op, local_handle, local_ids, _, remote_ids in nixl.xfers: + assert len(local_ids) == len(remote_ids) + desc_arr = nixl.dlists[local_handle] + for i in local_ids: + addr, length, _dev = desc_arr[int(i)] + addr, length = int(addr), int(length) + total_descs += 1 + if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges): + violations.append((int(i), hex(addr), length)) + assert not violations, ( + f"{len(violations)}/{total_descs} local descriptors write outside " + f"the request's own blocks: {violations[:10]}" + ) + return total_descs + + +@pytest.mark.cpu_test +def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): + """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens + (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), + equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and + TP-sharded KDA state. Every local descriptor of the request's reads must + stay within its own blocks.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + assert worker._physical_blocks_per_logical_kv_block == 3 + + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + for rank in (0, 1): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + + # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok + # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. + # Sparse, non-contiguous ids so neighbor blocks exist on all sides. + local_ids = ([2, 5], [1], [7]) + remote_ids = [[1, 4], [5], [2]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 2, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + worker._read_blocks_for_req("req-b", meta) + + owned = _owned_byte_ranges(worker, local_ids) + total = _assert_local_writes_within(worker, owned) + # Multi-read: rank 0 carries the replicated MLA + its SSM shard, + # rank 1 carries only its SSM shard. + assert len(worker.nixl_wrapper.xfers) == 2 + assert total > 0 + + +def _resolve( + desc_arr, + idx, + bases, + region_size, + unified_page, + desc_page, + logical_ids_attn, + block_tokens, +): + """Resolve a desc id to (region, kind, token_start) where kind is 'attn' + (desc-page sized, sub-block-aligned, in the request's attention blocks) + or 'mamba'. token_start is the request-relative token offset, so local + and remote are comparable even when their kernel blocks differ in size.""" + addr, length, _ = (int(x) for x in desc_arr[int(idx)]) + for region, base in enumerate(bases): + off = addr - base + if 0 <= off < region_size: + b = off // unified_page + rem = off % unified_page + if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn: + pos = logical_ids_attn.index(b) + tokens_per_desc = block_tokens * desc_page // unified_page + sub = rem // desc_page + return (region, "attn", pos * block_tokens + sub * tokens_per_desc) + return (region, "mamba", None) + raise AssertionError(f"desc {idx} addr {addr:#x} not in any region") + + +def _run_hetero_case( + local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None +): + """Full pull-path run for one geometry; returns pairing records. + + ``remote_kernel`` defaults to the local kernel block size; a smaller + value additionally exercises block_size_ratio > 1. + """ + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + remote_kernel = remote_kernel or kernel + block_size_ratio = kernel // remote_kernel + remote_ppl = remote_block // remote_kernel + matched = num_tokens - 1 # mamba N-1 rule + n_local = -(-num_tokens // local_block) + n_remote = -(-matched // remote_block) + + worker = _make_mla_hybrid_worker( + local_block_size=local_block, + kernel_block_size=kernel, + num_logical_blocks=max(2 * n_local + 4, 8), + ) + # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size + # shards of each. + meta_r = _make_remote_meta( + worker, + remote_block_size=remote_block, + remote_kernel_block_size=remote_kernel, + remote_num_logical=max(2 * n_remote + 4, 8), + remote_ssm_sizes=(48 // tp_size, 64 // tp_size), + ) + for rank in range(tp_size): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + + # Sparse ids so neighbors exist between the request's blocks. + local_attn = [2 * i + 1 for i in range(n_local)] + remote_attn = [2 * i + 2 for i in range(n_remote)] + local_ids = (local_attn, [0], [2 * n_local + 2]) + remote_ids = [remote_attn, [1], [0]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": tp_size, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + # Sentinel-fill the local KV so untouched bytes are detectable. + for t in worker._test_tensors: + t.fill_(0xAA) + + worker._read_blocks_for_req("req-b", meta) + + # Invariant 1: all local writes within the request's own blocks. + owned = _owned_byte_ranges(worker, local_ids) + _assert_local_writes_within(worker, owned) + + # Invariant 2: local<->remote attention pairs are token-aligned. + nixl = worker.nixl_wrapper + local_bases = [t.data_ptr() for t in worker._test_tensors] + remote_bases = [0x10_000_000, 0x20_000_000] + local_unified = worker._test_unified_page + remote_unified = (local_unified // local_block) * remote_block + # With block_size_ratio > 1 the local page is split into ratio sub-descs, + # each the size of a whole remote kernel page. + desc_page = worker.block_len_per_layer[0] // block_size_ratio + meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified + covered_tokens = set() + for op, lh, lids, rh, rids in nixl.xfers: + larr, rarr = nixl.dlists[lh], nixl.dlists[rh] + for li, ri in zip(lids, rids): + lreg, lkind, ltok = _resolve( + larr, + li, + local_bases, + len(worker._test_tensors[0]), + local_unified, + desc_page, + local_attn, + local_block, + ) + rreg, rkind, rtok = _resolve( + rarr, + ri, + remote_bases, + meta_r_num_blocks_bytes, + remote_unified, + desc_page, + remote_attn, + remote_block, + ) + assert lkind == rkind, ( + f"pair kind mismatch: local {lkind} vs remote {rkind} " + f"(local desc {li}, remote desc {ri})" + ) + assert lreg == rreg, ( + f"region mismatch: local {lreg} vs remote {rreg} for " + f"tokens {ltok} vs {rtok}" + ) + if lkind == "attn": + assert ltok == rtok, ( + f"TOKEN MISALIGNMENT: local sub-block holds tokens " + f"[{ltok}..) but receives remote tokens [{rtok}..) " + f"(geometry local_block={local_block}, " + f"remote_block={remote_block}, N={num_tokens})" + ) + covered_tokens.add(ltok) + + # Invariant 3: full coverage of the matched tokens, at the finest + # transfer granularity (the remote kernel block). + needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)} + missing = needed - covered_tokens + assert not missing, ( + f"tokens never transferred: {sorted(missing)[:8]} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + # Invariant 4: no stale bytes after receive completion. The scheduler + # excludes the blocks covering the matched tokens from alloc-time KV + # zeroing (the zeroing would race the RDMA write), so every byte of + # those blocks must be either written by the transfer or zeroed by the + # receive post-process. Stale bytes surface as mid-response garbage + # once decode grows into the untransferred tail. + for op, lh, lids, rh, rids in nixl.xfers: + larr = nixl.dlists[lh] + for li in lids: + addr, length, _ = (int(x) for x in larr[int(li)]) + for t in worker._test_tensors: + off = addr - t.data_ptr() + if 0 <= off < t.numel(): + t[off : off + length] = 0 # simulate the RDMA write + break + done_sending, done_recving = worker.get_finished() + assert "req-b" in done_recving + n_excluded = -(-matched // local_block) + stale = [] + for b in local_attn[:n_excluded]: + for region, t in enumerate(worker._test_tensors): + page = t[b * local_unified : (b + 1) * local_unified] + n_stale = int((page == 0xAA).sum()) + if n_stale: + stale.append((region, b, n_stale)) + assert not stale, ( + f"stale (unzeroed, untransferred) bytes in matched-range attention " + f"blocks (region, block, bytes): {stale} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "local_block,remote_block", + [ + (12, 8), # ppl 3 vs 2 + (36, 8), # ppl 9 vs 2 (large ppl asymmetry, scaled) + (24, 4), # ppl 6 vs 1 + (16, 24), # remote larger than local (D_TP > P_TP direction) + ], +) +@pytest.mark.parametrize("num_tokens", list(range(2, 40))) +def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): + """Sweep prompt lengths across block-boundary residues for several + hetero-ppl geometries; assert neighbor-safety, token alignment, and + coverage of every transferred kernel block.""" + _run_hetero_case( + local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around the remote kernel block (4), the local kernel block + # (8), the remote logical block (8) and the local logical block (24). + [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49], +) +def test_hetero_ppl_with_block_size_ratio(num_tokens): + """Both hetero regimes at once: kernel blocks differ (local 8 / remote + 4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs + 2). The transfer is clipped at remote sub-block granularity by the + pairing and front-trimmed by _apply_prefix_caching, so the + untransferred tail can span both a partial block and whole blocks — + the case each of the two former zeroing paths handled only half of.""" + _run_hetero_case( + local_block=24, + kernel=8, + remote_block=8, + remote_kernel=4, + num_tokens=num_tokens, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around every geometric boundary: kernel block (64), remote + # logical block (768), local logical block (5760), plus odd offsets. + [ + 2, + 63, + 64, + 65, + 127, + 128, + 300, + 640, + 767, + 768, + 769, + 831, + 832, + 1000, + 1535, + 1536, + 1537, + 2303, + 2304, + 2305, + 3001, + 5759, + 5760, + 5761, + 5824, + 6528, + 6529, + ], +) +def test_mla_hybrid_large_ppl_geometry(num_tokens): + """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): + decode (local) logical block 5760 / kernel 64 (ppl=90), prefill + (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with + replicated MLA and 8-way TP-sharded KDA state.""" + _run_hetero_case( + local_block=5760, + kernel=64, + remote_block=768, + num_tokens=num_tokens, + tp_size=8, + ) + + +@pytest.mark.cpu_test +def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): + """The MLA per-token page is TP-independent, so kernel block lengths + differing by anything other than the block-size ratio must fail the + handshake loudly rather than transfer at mismatched geometry.""" + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + # Equal kernel block sizes (ratio 1), but a half-sized per-token page. + meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer] + with pytest.raises((AssertionError, RuntimeError)): + worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index ad5ccbbbe86c..69887de4a015 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -25,6 +25,7 @@ import threading import time from collections import defaultdict +from concurrent.futures import Future from typing import Any from unittest.mock import MagicMock, patch @@ -42,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) +from vllm.v1.kv_cache_interface import FullAttentionSpec from vllm.v1.outputs import KVConnectorOutput from .utils import make_nixl_push_scheduler @@ -104,9 +106,9 @@ def get_unhashed_block_ids_all_groups(self) -> tuple[list[int], ...]: def _stub_sw_clipping(scheduler) -> None: - """Make ``get_sw_clipped_blocks`` a passthrough so tests don't need - the full sliding-window machinery.""" - scheduler.get_sw_clipped_blocks = lambda block_ids: block_ids + """Make ``get_exchange_clipped_blocks`` a passthrough so tests don't + need the full sliding-window machinery.""" + scheduler.get_exchange_clipped_blocks = lambda block_ids, clip_ssm=True: block_ids # ----------------------------------------------------------------- # @@ -320,6 +322,7 @@ def fresh(cls) -> _StubWriterWorker: w._finished_blocks_inbox = queue.Queue() w._pending_completion_notifs = queue.Queue() w._evict_finished_inbox = queue.Queue() + w._deferred_push_inbox = queue.Queue() w._push_writer_wake = threading.Event() w._push_writer_stop = threading.Event() w._push_writer_thread = None @@ -334,6 +337,10 @@ def fresh(cls) -> _StubWriterWorker: w.world_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} + w._physical_blocks_per_logical_kv_block = 1 + # Single non-hybrid attention group, matching the stub block id lists. + w._has_mamba = False + w._group_spec_types = (FullAttentionSpec,) # Track _do_start_push_kv invocations. calls: list[tuple[str, Any, dict[str, Any]]] = [] @@ -506,6 +513,112 @@ def test_start_load_kv_enqueues_to_writer(self): assert w.start_push_calls == [] +# The P→D handshake must run on the base worker's background executor, never +# blocking the writer thread: ``_do_start_push_kv`` defers the WRITE until the +# handshake resolves, then re-drives via ``_deferred_push_inbox``. These call +# the *real* ``_do_start_push_kv`` (the stub overrides it for matching tests). +def _real_do_start_push_kv(w, *args): + return NixlPushConnectorWorker._do_start_push_kv(w, *args) + + +def test_do_start_push_kv_defers_then_writes_when_handshake_ready(): + """Full happy-path lifecycle: an in-flight handshake defers the WRITE (no + NIXL op from the writer or the executor callback); once it resolves the + request is re-queued on ``_deferred_push_inbox`` with the wake set; and on + re-drive with the handshake ready the WRITE is issued with a correct + ReqMeta.""" + w = _StubWriterWorker.fresh() + w._logical_to_kernel_block_ids = lambda x, ratio: x + xfer_calls: list[dict[str, Any]] = [] + w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw) + + fut: Future = Future() + w._ensure_handshake = lambda *a, **k: fut + + rd = _registration_data("req-hs", decode_engine_id="decode-engine") + _real_do_start_push_kv(w, "req-hs", ([1, 2, 3],), rd) + + # Handshake pending -> nothing issued, nothing queued, no wake. + assert xfer_calls == [] + assert w._deferred_push_inbox.qsize() == 0 + assert not w._push_writer_wake.is_set() + + # Handshake completes: request re-queued for the writer, wake set, but no + # *direct* WRITE from the callback (wrong thread for NIXL ops). + fut.set_result(({(0, 0): "agent"}, 0.0)) + assert xfer_calls == [] + assert w._push_writer_wake.is_set() + rid, blocks, reg = w._deferred_push_inbox.get_nowait() + assert (rid, blocks, reg) == ("req-hs", ([1, 2, 3],), rd) + + # Re-drive on the writer with the handshake now ready -> WRITE issued. + w._ensure_handshake = lambda *a, **k: None + _real_do_start_push_kv(w, rid, blocks, reg) + assert len(xfer_calls) == 1 + assert xfer_calls[0]["req_id"] == "req-hs" + meta = xfer_calls[0]["meta"] + assert meta.remote is not None + assert meta.remote.engine_id == "decode-engine" + # RemoteMeta.request_id is D's request id from the registration. + assert meta.remote.request_id == "req-hs" + + +def test_do_start_push_kv_drops_request_on_handshake_failure(): + """Handshake raises: the request is dropped (not re-queued, no WRITE) and + the failure is logged. Blocks are reclaimed by the lease/watchdog, matching + the old blocking behaviour.""" + w = _StubWriterWorker.fresh() + w._logical_to_kernel_block_ids = lambda x: x + xfer_calls: list[dict[str, Any]] = [] + w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw) + failures: list[dict[str, Any]] = [] + w._log_failure = lambda **kw: failures.append(kw) + + fut: Future = Future() + w._ensure_handshake = lambda *a, **k: fut + + _real_do_start_push_kv(w, "req-fail", ([9],), _registration_data("req-fail")) + fut.set_exception(RuntimeError("handshake boom")) + + assert w._deferred_push_inbox.qsize() == 0 + assert xfer_calls == [] + assert len(failures) == 1 + assert failures[0]["failure_type"] == "push_handshake_failed" + + +def test_writer_loop_drains_deferred_push_inbox(): + """The writer loop drains ``_deferred_push_inbox`` and re-drives + ``_do_start_push_kv`` for each entry (event-driven, no polling).""" + w = _StubWriterWorker.fresh() + w.nixl_wrapper = MagicMock() + w.nixl_wrapper.get_new_notifs.return_value = {} + + processed = threading.Event() + + def _tracked(rid, blocks, rd): + w.start_push_calls.append((rid, blocks, rd)) + processed.set() + + w._do_start_push_kv = _tracked # type: ignore[method-assign] + + w._deferred_push_inbox.put( + ("req-retry", ([1, 2],), _registration_data("req-retry")) + ) + w._push_writer_wake.set() + + t = threading.Thread(target=w._push_writer_loop, daemon=True) + t.start() + try: + assert processed.wait(timeout=2.0), "writer did not drain deferred inbox" + finally: + w._push_writer_stop.set() + w._push_writer_wake.set() + t.join(timeout=2) + + assert len(w.start_push_calls) == 1 + assert w.start_push_calls[0][0] == "req-retry" + + class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained diff --git a/tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py similarity index 94% rename from tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py rename to tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py index 2371b5555bce..da8be07157d8 100644 --- a/tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py +++ b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm. +"""Verify that GPU memory is released after NixlConnector shutdown on ROCm. Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing @@ -62,7 +62,7 @@ def _full_gpu_cleanup(): @pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)]) -def test_gpu_memory_rixl_hma(model_name, sw_size): +def test_gpu_memory_nixl_hma(model_name, sw_size): """Track GPU memory through NixlConnector create/infer/shutdown cycle.""" from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig @@ -84,7 +84,7 @@ def test_gpu_memory_rixl_hma(model_name, sw_size): } print("\n" + "=" * 90) - print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)") + print("GPU MEMORY -- NIXL NixlConnector HMA (ROCm)") print("=" * 90) gc.collect() torch.accelerator.empty_cache() @@ -169,14 +169,14 @@ def test_gpu_memory_rixl_hma(model_name, sw_size): @pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"]) -def test_gpu_memory_no_rixl_baseline(model_name): +def test_gpu_memory_no_nixl_baseline(model_name): """Same workload without NixlConnector. Comparing driver-level memory - between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact.""" + between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact.""" from vllm import LLM, SamplingParams from vllm.distributed.parallel_state import cleanup_dist_env_and_memory print("\n" + "=" * 90) - print("CONTROL -- same model, no RIXL connector") + print("CONTROL -- same model, no NIXL connector") print("=" * 90) gc.collect() torch.accelerator.empty_cache() @@ -209,7 +209,7 @@ def test_gpu_memory_no_rixl_baseline(model_name): drv_base = snap0["drv_used_mb"] drv_leaked = snap_final["drv_used_mb"] - drv_base drv_peak = snap_peak["drv_used_mb"] - drv_base - print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB") + print(f"\n Driver leaked (no NIXL): {drv_leaked:.0f} MB") print("=" * 90) leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0 diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 7c735098230e..d72bbc9ec2d4 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -161,3 +161,57 @@ def test_fa_and_ssm_different_split_factors(self): # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0) # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0) assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)] + + def test_hetero_block_size_splits(self): + """With a block-size ratio, single-source FA sub-block descs pass + through whole; SSM descs are unexpanded and split per source.""" + plan = TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded). + src_blocks_data = np.array( + [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + (3000, 400, 0), + ], + dtype=np.uint64, + ) + + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2)) + + assert len(splits) == 2 + fa_passthrough = [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + ] + assert splits[0] == fa_passthrough + [(3000, 200, 0)] + assert splits[1] == fa_passthrough + [(3200, 200, 0)] + + def test_hetero_block_size_head_sharded_asserts(self): + """Head-sharded FA reads (multiple FA sources) are incompatible with + a block-size mismatch and must fail loudly.""" + plan = TPMapping( + source_ranks_per_group=((0, 1), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 1}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + src_blocks_data = np.array( + [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)], + dtype=np.uint64, + ) + + with pytest.raises(AssertionError, match="Head-sharded"): + list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2)) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 7df9e20e6a58..7d4b17ab399d 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -362,6 +362,7 @@ def wrapper(*args, **kwargs): class MockKVConfig: matched_tokens: int = 0 is_async: bool = False + num_defers_before_matching: int = 0 class MockKVConnectorMetadata(KVConnectorMetadata): @@ -384,6 +385,12 @@ def __init__( self.config = MockKVConfig( matched_tokens=extra_config["matched_tokens"], is_async=extra_config["is_async"], + num_defers_before_matching=extra_config.get( + "num_defers_before_matching", 0 + ), + ) + self._defers_left: defaultdict[str, int] = defaultdict( + lambda: self.config.num_defers_before_matching ) def get_num_new_matched_tokens( @@ -391,6 +398,9 @@ def get_num_new_matched_tokens( request: Request, num_computed_tokens: int, ) -> tuple[int | None, bool]: + if self._defers_left[request.request_id] > 0: + self._defers_left[request.request_id] -= 1 + return (None, False) return (self.config.matched_tokens, self.config.is_async) def update_state_after_alloc( diff --git a/tests/v1/kv_offload/cpu/__init__.py b/tests/v1/kv_offload/cpu/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/kv_offload/cpu/policies/__init__.py b/tests/v1/kv_offload/cpu/policies/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py new file mode 100644 index 000000000000..14ccf8b67e7b --- /dev/null +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable + +import pytest + +from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager +from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy + + +class _DummyCachePolicy(CachePolicy): + """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded + by module path, so it must be importable at module scope (mirrors + tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec).""" + + def __init__(self, cache_capacity: int) -> None: + self.cache_capacity = cache_capacity + + def get(self, key: OffloadKey) -> BlockStatus | None: + return None + + def insert(self, key: OffloadKey, block: BlockStatus) -> None: + pass + + def remove(self, key: OffloadKey) -> None: + pass + + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + pass + + def evict( + self, n: int, protected: set[OffloadKey] + ) -> list[tuple[OffloadKey, BlockStatus]] | None: + return None + + def clear(self) -> None: + pass + + +@pytest.fixture(autouse=True) +def restore_cache_policy_registry(): + """Save and restore CachePolicyFactory._registry between tests.""" + original = dict(CachePolicyFactory._registry) + yield + CachePolicyFactory._registry = original + + +class TestCachePolicyFactory: + """Unit tests for CachePolicyFactory (registration/resolution by name).""" + + def test_pre_registered_policies_can_be_imported(self): + """If someone moves a policy module but forgets to update + factory.py, CI fails.""" + for name in CachePolicyFactory._registry: + cls = CachePolicyFactory._registry[name]() + assert issubclass(cls, CachePolicy) + + def test_lru_and_arc_registered(self): + assert CachePolicyFactory.get_cache_policy_cls("lru") is LRUCachePolicy + assert CachePolicyFactory.get_cache_policy_cls("arc") is ARCCachePolicy + + def test_register_and_resolve_custom_policy(self): + CachePolicyFactory.register_cache_policy( + "dummy", + "tests.v1.kv_offload.cpu.policies.test_factory", + "_DummyCachePolicy", + ) + policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy") + assert policy_cls is _DummyCachePolicy + + manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy") + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_raises(self): + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent") + + def test_duplicate_registration_raises(self): + with pytest.raises(ValueError, match="is already registered"): + CachePolicyFactory.register_cache_policy("lru", "some.module", "SomeClass") + + def test_dynamic_load_via_cache_policy_module_path(self): + """Out-of-tree policy loaded via cache_policy_module_path, no + register_cache_policy() call -- this is how external projects + integrate a custom CachePolicy without forking/patching vLLM. + Mirrors tests/v1/kv_offload/test_factory.py's + test_dynamic_load_via_spec_module_path.""" + policy_cls = CachePolicyFactory.get_cache_policy_cls( + "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory" + ) + assert policy_cls is _DummyCachePolicy + + def test_manager_resolves_policy_via_module_path(self): + """End-to-end: CPUOffloadingManager resolves an unregistered policy + purely from cache_policy_module_path.""" + manager = CPUOffloadingManager( + num_blocks=4, + cache_policy="_DummyCachePolicy", + cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory", + ) + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_without_module_path_raises(self): + """eviction_policy not in registry + no cache_policy_module_path -> + ValueError, same shape as the OffloadingSpecFactory error path.""" + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent", None) diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 2f1ce67e9c41..bfbb18251c66 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -7,6 +7,7 @@ import pytest import torch +from vllm import _custom_ops as ops from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import set_random_seed @@ -16,6 +17,7 @@ CanonicalKVCacheTensor, GPULoadStoreSpec, ) +from vllm.v1.kv_offload.cpu import gpu_worker from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion @@ -32,6 +34,18 @@ NUM_MAPPINGS_PER_GROUP = [2] +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gpu_worker, "HAS_TRITON", True) + monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: True) + + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]] + assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is ( + ops.swap_blocks_batch + ) + + @pytest.mark.parametrize("gpu_to_cpu", [True, False]) @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS) @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) @@ -41,7 +55,10 @@ @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", DEVICES) -@pytest.mark.parametrize("use_shared_memory", [False, True]) +@pytest.mark.parametrize( + ("use_shared_memory", "replicated_layout"), + [(False, False), (True, False), (True, True)], +) @torch.inference_mode() def test_transfer( default_vllm_config, @@ -55,6 +72,7 @@ def test_transfer( seed: int, device: str, use_shared_memory: bool, + replicated_layout: bool, ) -> None: set_random_seed(seed) @@ -95,11 +113,15 @@ def test_transfer( gpu_page_size_bytes * num_tensors * blocks_per_chunk, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, ) + simulated_world_size = 2 + kv_bytes_per_block = ( + cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size + ) mmap_region = SharedOffloadRegion( engine_id=str(uuid.uuid4()), num_blocks=num_cpu_blocks, rank=0, - kv_bytes_per_block=cpu_page_size, + kv_bytes_per_block=kv_bytes_per_block, cpu_page_size=cpu_page_size, ) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 97ed5168f744..5a36be16bf23 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -6,10 +6,10 @@ import numpy as np import pytest -from vllm.distributed.kv_events import MEDIUM_CPU from vllm.v1.kv_offload.base import ( LoadStoreSpec, LookupResult, + Medium, OffloadingEvent, OffloadKey, PrepareStoreOutput, @@ -37,6 +37,7 @@ def make_req_context( def make_cpu_manager( num_blocks: int = 4, cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, store_threshold: int = 0, max_tracker_size: int = 64_000, @@ -44,6 +45,7 @@ def make_cpu_manager( return CPUOffloadingManager( num_blocks=num_blocks, cache_policy=cache_policy, + cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, @@ -115,7 +117,7 @@ def verify_events( stores: list[set[OffloadKey]] = [] evictions: list[set[OffloadKey]] = [] for event in events: - assert event.medium == MEDIUM_CPU + assert event.medium == Medium.CPU if event.removed: evictions.append(set(event.keys)) else: diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index d08dbee57655..ecd706828c40 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -1,35 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for OffloadingSpecFactory. +"""Unit tests for native offloading specs and their factory.""" -These tests verify: -1. Pre-registration integrity — registered module paths can actually import - and yield correct OffloadingSpec subclasses (CI sentinel against file moves). -2. End-to-end factory → spec construction with real configs. -3. Downstream collaboration — build_metric_definitions delegation. -4. Error paths — unregistered specs, missing config, duplicate registration. -""" - -from typing import cast -from unittest.mock import MagicMock, patch +from typing import Any +from unittest.mock import MagicMock import pytest -import torch -from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( - build_offloading_config, -) -from vllm.platforms import current_platform -from vllm.v1.kv_cache_interface import ( - FullAttentionSpec, - KVCacheConfig, - KVCacheGroupSpec, - KVCacheTensor, - MLAAttentionSpec, - SlidingWindowSpec, -) from vllm.v1.kv_offload.base import ( CanonicalKVCaches, OffloadingHistogramMetadata, @@ -37,220 +14,81 @@ OffloadingSpec, OffloadingWorker, ) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingGroupConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec from vllm.v1.kv_offload.factory import OffloadingSpecFactory from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - @pytest.fixture(autouse=True) def restore_registry(): - """Save and restore OffloadingSpecFactory._registry between tests.""" original = dict(OffloadingSpecFactory._registry) yield OffloadingSpecFactory._registry = original -def _get_extra_config(config: VllmConfig) -> dict: - assert config.kv_transfer_config is not None - return config.kv_transfer_config.kv_connector_extra_config - - -def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec: - return OffloadingSpecFactory.create_spec( - build_offloading_config(config, kv_cache_config) - ) - - -def _make_vllm_config( +def _make_offloading_config( + *, spec_name: str | None = "CPUOffloadingSpec", - cpu_bytes_to_use: int | None = None, - store_threshold: int = 0, - extra_config: dict | None = None, -): - """Build a real VllmConfig with kv_transfer_config set for offloading.""" - from vllm.config import ( - CacheConfig, - DeviceConfig, - ModelConfig, - SchedulerConfig, - VllmConfig, - ) - - model_config = ModelConfig( - model="facebook/opt-125m", - trust_remote_code=True, - dtype="float16", - seed=42, - ) - scheduler_config = SchedulerConfig( - max_num_seqs=16, - max_num_batched_tokens=64, - max_model_len=10000, - enable_chunked_prefill=True, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - cache_config = CacheConfig( - block_size=16, - gpu_memory_utilization=0.9, - cache_dtype="auto", - enable_prefix_caching=True, - ) - - cfg = extra_config or {} - if cpu_bytes_to_use is not None: - cfg["cpu_bytes_to_use"] = cpu_bytes_to_use - cfg["spec_name"] = spec_name - if store_threshold > 0: - cfg["store_threshold"] = store_threshold - - kv_transfer_config = KVTransferConfig( - kv_connector="OffloadingConnector", - kv_role="kv_both", - kv_connector_extra_config=cfg, - ) - return VllmConfig( - scheduler_config=scheduler_config, - model_config=model_config, - cache_config=cache_config, - kv_transfer_config=kv_transfer_config, - device_config=DeviceConfig("cpu"), - ) - - -def _make_layout_vllm_config( - spec_name: str = "CPUOffloadingSpec", - cpu_bytes_to_use: int | None = None, - extra_config: dict | None = None, - tensor_parallel_size: int = 1, - pipeline_parallel_size: int = 1, - prefill_context_parallel_size: int = 1, - decode_context_parallel_size: int = 1, -) -> VllmConfig: - config = MagicMock() - config.cache_config.block_size = 16 - config.cache_config.enable_prefix_caching = True - config.cache_config.prefix_match_unit = None - config.cache_config.cache_dtype = torch.float16 - config.model_config.model = "test-model" - world_size = ( - tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size - ) - with patch.object(current_platform, "device_count", return_value=world_size): - config.parallel_config = ParallelConfig( - tensor_parallel_size=tensor_parallel_size, - pipeline_parallel_size=pipeline_parallel_size, - prefill_context_parallel_size=prefill_context_parallel_size, - decode_context_parallel_size=decode_context_parallel_size, - ) - config.kv_events_config = None - config.use_v2_model_runner = False - - connector_extra_config = dict(extra_config or {}) - connector_extra_config["spec_name"] = spec_name + cpu_bytes_to_use: int | None = 65536, + worker_kv_bytes_per_block: int = 8, + groups: tuple[OffloadingGroupConfig, ...] | None = None, + tokens_per_hash: int = 16, + blocks_per_chunk: int = 1, + rank: int = 0, + world_size: int = 1, + tp_size: int | None = None, + pp_size: int = 1, + pcp_size: int = 1, + dcp_size: int = 1, + data_parallel_index: int = 0, + is_parallelism_agnostic: bool = False, + replicated_layout: bool = False, + extra_config: dict[str, Any] | None = None, +) -> OffloadingConfig: + normalized_extra_config = dict(extra_config or {}) + if spec_name is not None: + normalized_extra_config["spec_name"] = spec_name if cpu_bytes_to_use is not None: - connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use - config.kv_transfer_config = KVTransferConfig( - kv_connector="OffloadingConnector", - kv_role="kv_both", - kv_connector_extra_config=connector_extra_config, - ) - return cast(VllmConfig, config) - - -def _make_kv_cache_config(): - """Build a minimal KVCacheConfig with one KV cache tensor.""" - num_blocks = 16 - num_kv_heads = 1 - head_size = 1 - dtype = torch.float32 - page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8 - kv_tensor = KVCacheTensor( - size=num_blocks * page_size, shared_by=["layer"], block_stride=0 - ) - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[kv_tensor], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=16, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - ), - ) - ], - ) - - -def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: - num_blocks = 4 - if packed: - kv_cache_tensors = [ - KVCacheTensor( - size=64, - shared_by=[layer_name], - block_stride=16, - ) - for layer_name in ("layer0", "layer1") - ] - else: - kv_cache_tensors = [ - KVCacheTensor(size=40, shared_by=["layer0"]), - KVCacheTensor(size=24, shared_by=["layer1"]), - ] - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer0", "layer1"], - FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], + normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use + + if groups is None: + groups = (OffloadingGroupConfig(16, ("layer",)),) + + return OffloadingConfig( + groups=groups, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + enable_kv_cache_events=False, + extra_config=normalized_extra_config, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float16"), + cache=OffloadingCacheConfig( + tokens_per_hash=tokens_per_hash, + blocks_per_chunk=blocks_per_chunk, + ), + parallel=OffloadingParallelConfig( + rank=rank, + world_size=world_size, + tp_size=world_size if tp_size is None else tp_size, + pp_size=pp_size, + pcp_size=pcp_size, + dcp_size=dcp_size, + data_parallel_index=data_parallel_index, + is_parallelism_agnostic=is_parallelism_agnostic, + ), + replicated_layout=replicated_layout, ) -def _make_hybrid_kv_cache_config() -> KVCacheConfig: - return KVCacheConfig( - num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor(size=40, shared_by=["full_layer"]), - KVCacheTensor(size=24, shared_by=["mla_layer"]), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full_layer"], - FullAttentionSpec( - block_size=12, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mla_layer"], - MLAAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=576, - dtype=torch.float32, - ), - ), - ], - ) +def _create_spec(**kwargs: Any) -> OffloadingSpec: + return OffloadingSpecFactory.create_spec(_make_offloading_config(**kwargs)) class SingleArgExternalOffloadingSpec(OffloadingSpec): @@ -261,104 +99,64 @@ def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: raise NotImplementedError -# --------------------------------------------------------------------------- -# Pre-registration integrity (CI sentinel) -# --------------------------------------------------------------------------- - - def test_pre_registered_specs_can_be_imported(): - """If someone moves cpu/spec.py but forgets to update factory.py, CI fails.""" for name in OffloadingSpecFactory._registry: cls = OffloadingSpecFactory._registry[name]() assert issubclass(cls, OffloadingSpec) def test_cpu_spec_registered(): - """CPUOffloadingSpec is registered and importable.""" cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]() assert cls is CPUOffloadingSpec def test_tiering_spec_registered(): - """TieringOffloadingSpec is registered and importable.""" cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]() assert cls is TieringOffloadingSpec -# --------------------------------------------------------------------------- -# Normal path — get_spec_cls -# --------------------------------------------------------------------------- - - def test_get_spec_cls_returns_registered_class(): - """Registered spec_name returns correct class.""" - config = _make_vllm_config(spec_name="CPUOffloadingSpec") - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) + spec_cls = OffloadingSpecFactory.get_spec_cls( + _make_offloading_config().extra_config + ) assert spec_cls is CPUOffloadingSpec -def test_get_spec_cls_default_to_cpu(): - """Default spec_name (absent from config) resolves to CPUOffloadingSpec.""" - config = _make_vllm_config(spec_name=None) - config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) +def test_get_spec_cls_defaults_to_cpu(): + spec_cls = OffloadingSpecFactory.get_spec_cls( + _make_offloading_config(spec_name=None).extra_config + ) assert spec_cls is CPUOffloadingSpec -# --------------------------------------------------------------------------- -# End-to-end — create_spec -# --------------------------------------------------------------------------- - - -def test_create_cpu_offloading_spec_end_to_end(): - """Full factory → spec construction with real VllmConfig/KVCacheConfig. - - Verifies: - - cpu_bytes_to_use validation and num_blocks calculation - - block_size % tokens_per_hash assertion - - spec instance is CPUOffloadingSpec - """ - config = _make_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_kv_cache_config() - spec = _create_spec(config, kv_cache_config) +def test_create_cpu_offloading_spec(): + spec = _create_spec() assert isinstance(spec, CPUOffloadingSpec) assert spec.num_blocks > 0 -@pytest.mark.parametrize("packed", [False, True]) -def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool): - cpu_bytes_to_use = 1920 - config = _make_layout_vllm_config( - cpu_bytes_to_use=cpu_bytes_to_use, - extra_config={"block_size": 32}, - tensor_parallel_size=3, - pipeline_parallel_size=2, +def test_cpu_spec_sizes_normalized_worker_layout(): + # The CPU spec now rounds the offloaded row up to the mmap page size + # (matching the shared region), so kv_bytes_per_chunk picks up padding + # while cpu_page_size_per_worker stays the un-padded per-worker slot. + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=alignment * 3, + worker_kv_bytes_per_block=16, + blocks_per_chunk=2, + world_size=6, + tp_size=3, + pp_size=2, ) - spec = _create_spec(config, _make_sizing_kv_cache_config(packed)) - assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 - assert spec.kv_bytes_per_chunk == 192 - assert spec.num_blocks == cpu_bytes_to_use // 192 - - -def test_cpu_spec_rejects_partially_packed_tensor_layout(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_sizing_kv_cache_config(packed=False) - kv_cache_config.kv_cache_tensors[0].block_stride = 16 - - with pytest.raises(AssertionError): - _create_spec(config, kv_cache_config) - + assert spec.kv_bytes_per_chunk == alignment + assert spec.num_blocks == 3 -def test_cpu_spec_zero_blocks_skips_tensor_layout_validation(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_sizing_kv_cache_config(packed=False) - kv_cache_config.num_blocks = 0 - kv_cache_config.kv_cache_tensors[0].block_stride = 16 - spec = _create_spec(config, kv_cache_config) +def test_cpu_spec_zero_worker_bytes_produces_empty_cache(): + spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4) assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 0 @@ -368,225 +166,405 @@ def test_cpu_spec_zero_blocks_skips_tensor_layout_validation(): def test_tiering_spec_aligns_row_size(): alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT - cpu_bytes_to_use = alignment * 3 - config = _make_layout_vllm_config( + spec = _create_spec( spec_name="TieringOffloadingSpec", - cpu_bytes_to_use=cpu_bytes_to_use, - extra_config={"block_size": 32}, - tensor_parallel_size=3, - pipeline_parallel_size=2, + cpu_bytes_to_use=alignment * 3, + worker_kv_bytes_per_block=16, + blocks_per_chunk=2, + world_size=6, + tp_size=3, + pp_size=2, ) - spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False)) - assert isinstance(spec, TieringOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 assert spec.kv_bytes_per_chunk == alignment - assert spec.num_blocks == cpu_bytes_to_use // alignment + assert spec.num_blocks == 3 -def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={"block_size": 64}, - prefill_context_parallel_size=2, +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int): + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + spec_name="TieringOffloadingSpec", + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=True, ) - spec = _create_spec(config, _make_kv_cache_config()) + assert isinstance(spec, TieringOffloadingSpec) + assert spec.replicated_layout is True + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block + assert spec.num_blocks == 8 + - assert spec.tokens_per_block == (16,) - assert spec.tokens_per_hash == 16 - assert spec.blocks_per_chunk == 4 +def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch): + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + spec_name="TieringOffloadingSpec", + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=4, + replicated_layout=True, + ) + assert isinstance(spec, TieringOffloadingSpec) -def test_offloading_config_preserves_data_parallel_index(): - config = _make_layout_vllm_config() - config.parallel_config.data_parallel_index = 2 + region = MagicMock() + region_calls: list[dict[str, Any]] = [] + worker_calls: list[dict[str, Any]] = [] - offloading_config = build_offloading_config(config, _make_kv_cache_config()) + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return region - assert offloading_config.parallel.data_parallel_index == 2 + def fake_worker_ctor(**kwargs): + worker_calls.append(kwargs) + return MagicMock() + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor) + monkeypatch.setattr( + tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5 + ) -def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - config.cache_config.block_size = 4 + kv_caches = MagicMock() + spec.create_worker(kv_caches) - spec = _create_spec(config, _make_hybrid_kv_cache_config()) + assert region_calls[0]["rank"] == 0 + assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block + assert worker_calls[0]["kv_caches"] is kv_caches + assert worker_calls[0]["mmap_region"] is region - assert spec.tokens_per_block == (12, 16) - assert spec.tokens_per_hash == 4 - assert spec.blocks_per_chunk == 1 +def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch): + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module -def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec: - return FullAttentionSpec( - block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32 + spec = _create_spec( + spec_name="TieringOffloadingSpec", + worker_kv_bytes_per_block=4096, + world_size=4, ) + assert isinstance(spec, TieringOffloadingSpec) + + region_calls: list[dict[str, Any]] = [] + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return MagicMock() -def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: - config = _make_layout_vllm_config() - kv_cache_config = KVCacheConfig( - num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock()) + monkeypatch.setattr( + tiering_spec_module.torch.accelerator, + "current_device_index", + lambda: 5, ) - offloading_config = build_offloading_config(config, kv_cache_config) - return offloading_config.parallel.is_parallelism_agnostic + spec.create_worker(MagicMock()) -def test_parallelism_agnostic_for_single_full_attention_group(): - assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + assert region_calls[0]["rank"] == 1 + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: int): + # On shared-region (CUDA-alike) platforms the default spec now honors + # replicated layout: a single MLA copy (num_copies=1), matching tiering. + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=True, + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.replicated_layout is True + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block + assert spec.num_blocks == 8 + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_cpu_spec_replicated_disabled_without_shared_region( + monkeypatch, world_size: int +): + # Data-loss guard: non-CUDA-alike platforms keep a per-rank private pinned + # tensor (no shared medium), so replicated layout MUST stay off. Otherwise + # the rank-0 writer gate would ack rank>0 stores without writing, leaving + # those private buffers empty and corrupting subsequent loads. + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + monkeypatch.setattr( + cpu_spec_module.current_platform, "is_cuda_alike", lambda: False + ) + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=True, + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.replicated_layout is False + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size + assert spec.num_blocks == 2 + + +@pytest.mark.parametrize("config_replicated", [True, False]) +@pytest.mark.parametrize("cuda_alike", [True, False]) +def test_cpu_spec_replicated_layout_truth_matrix( + monkeypatch, cuda_alike: bool, config_replicated: bool +): + # replicated_layout is enabled iff the config gate passes AND the deployment + # actually allocates on the shared region (CUDA-alike). + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + monkeypatch.setattr( + cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike + ) + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=4, + replicated_layout=config_replicated, + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.replicated_layout is (cuda_alike and config_replicated) + + +def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch): + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=4, + ) + assert isinstance(spec, CPUOffloadingSpec) + + region = MagicMock() + region_calls: list[dict[str, Any]] = [] + worker_calls: list[dict[str, Any]] = [] + + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return region + + def fake_worker_ctor(**kwargs): + worker_calls.append(kwargs) + return MagicMock() + + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) + monkeypatch.setattr( + cpu_spec_module.torch.accelerator, "current_device_index", lambda: 5 + ) + + kv_caches = MagicMock() + spec.create_worker(kv_caches) + + assert region_calls[0]["engine_id"] == "test-engine" + assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4 + assert worker_calls[0]["kv_caches"] is kv_caches + assert worker_calls[0]["mmap_region"] is region + + +def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch): + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4) + assert isinstance(spec, CPUOffloadingSpec) + + region_calls: list[dict[str, Any]] = [] + worker_calls: list[dict[str, Any]] = [] + + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return MagicMock() + + def fake_worker_ctor(**kwargs): + worker_calls.append(kwargs) + return MagicMock() + + monkeypatch.setattr( + cpu_spec_module.current_platform, "is_cuda_alike", lambda: False + ) + monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) + + spec.create_worker(MagicMock()) + + # Non-CUDA-alike platforms keep the per-rank pinned-tensor path. + assert region_calls == [] + assert worker_calls[0]["mmap_region"] is None + + +def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + # worker_kv_bytes_per_block=0 yields num_blocks=0; a zero-byte region cannot + # be mmap'd, so even on CUDA-alike this must fall back to the tensor path. + spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4) + assert isinstance(spec, CPUOffloadingSpec) + assert spec.num_blocks == 0 + + region_calls: list[dict[str, Any]] = [] + worker_calls: list[dict[str, Any]] = [] + + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr( + cpu_spec_module, + "SharedOffloadRegion", + lambda **kwargs: region_calls.append(kwargs), + ) + monkeypatch.setattr( + cpu_spec_module, + "CPUOffloadingWorker", + lambda **kwargs: worker_calls.append(kwargs), + ) + + spec.create_worker(MagicMock()) + + assert region_calls == [] + assert worker_calls[0]["mmap_region"] is None @pytest.mark.parametrize( - "kv_cache_groups", + ("replicated_layout", "device_index", "world_size", "expected_rank"), [ - # MLA latent KV is replicated per rank, never head-sharded. - [ - KVCacheGroupSpec( - ["l0"], - MLAAttentionSpec( - block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32 - ), - ) - ], - # Sliding window is not full attention. - [ - KVCacheGroupSpec( - ["l0"], - SlidingWindowSpec( - block_size=16, - num_kv_heads=4, - head_size=128, - dtype=torch.float32, - sliding_window=128, - ), - ) - ], - # Hybrid model: more than one KV cache group. - [ - KVCacheGroupSpec(["l0"], _full_attention_spec()), - KVCacheGroupSpec(["l1"], _full_attention_spec()), - ], + (True, 5, 4, 0), # replicated: always slot 0 + (True, 0, 4, 0), # replicated: slot 0 regardless of device + (False, 5, 4, 1), # non-replicated: 5 % 4 == 1 + (False, 7, 4, 3), # non-replicated: 7 % 4 == 3 ], ) -def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): - assert not _parallelism_agnostic(kv_cache_groups) +def test_cpu_spec_create_worker_rank_assignment( + monkeypatch, replicated_layout, device_index, world_size, expected_rank +): + import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=replicated_layout, + ) + + region_calls: list[dict[str, Any]] = [] + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return MagicMock() -def test_parallelism_agnostic_disabled_on_v2_model_runner(): - config = _make_layout_vllm_config() - config.use_v2_model_runner = True - kv_cache_config = KVCacheConfig( - num_blocks=0, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], + monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", MagicMock()) + monkeypatch.setattr( + cpu_spec_module.torch.accelerator, "current_device_index", lambda: device_index ) - offloading_config = build_offloading_config(config, kv_cache_config) - assert not offloading_config.parallel.is_parallelism_agnostic + spec.create_worker(MagicMock()) -def test_create_dynamic_spec_receives_translated_config(): - config = _make_layout_vllm_config( - spec_name="SingleArgExternalOffloadingSpec", - extra_config={ - "spec_module_path": "tests.v1.kv_offload.test_factory", - }, + assert region_calls[0]["rank"] == expected_rank + + +def test_offloading_spec_has_replicated_layout_default(): + spec = SingleArgExternalOffloadingSpec(_make_offloading_config()) + assert spec.replicated_layout is False + + +def test_offloading_spec_uses_normalized_chunk_geometry(): + groups = ( + OffloadingGroupConfig(12, ("full_layer",)), + OffloadingGroupConfig(16, ("mla_layer",)), + ) + spec = _create_spec( + groups=groups, + tokens_per_hash=4, + blocks_per_chunk=2, ) - kv_cache_config = _make_kv_cache_config() - offloading_config = build_offloading_config(config, kv_cache_config) - spec = OffloadingSpecFactory.create_spec(offloading_config) + assert spec.tokens_per_block == (12, 16) + assert spec.tokens_per_hash == 4 + assert spec.blocks_per_chunk == 2 - assert isinstance(spec, SingleArgExternalOffloadingSpec) - assert spec.config is offloading_config +def test_create_dynamic_spec_receives_config(): + config = _make_offloading_config( + spec_name="SingleArgExternalOffloadingSpec", + extra_config={"spec_module_path": "tests.v1.kv_offload.test_factory"}, + ) -# --------------------------------------------------------------------------- -# Dynamic import via spec_module_path -# --------------------------------------------------------------------------- + spec = OffloadingSpecFactory.create_spec(config) + + assert isinstance(spec, SingleArgExternalOffloadingSpec) + assert spec.config is config def test_dynamic_load_via_spec_module_path(): - """External spec loaded via spec_module_path. - - This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec) - integrate with vLLM without being pre-registered in the factory. - The fallback path: registry miss → spec_module_path → importlib.import_module. - """ - config = _make_vllm_config(spec_name="CPUOffloadingSpec") - # Delete from registry to force the dynamic import path del OffloadingSpecFactory._registry["CPUOffloadingSpec"] - # spec_name not in registry → falls through to spec_module_path - config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = ( - "vllm.v1.kv_offload.cpu.spec" + config = _make_offloading_config( + extra_config={"spec_module_path": "vllm.v1.kv_offload.cpu.spec"} ) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - assert spec_cls is CPUOffloadingSpec + spec_cls = OffloadingSpecFactory.get_spec_cls(config.extra_config) -# --------------------------------------------------------------------------- -# Error paths -# --------------------------------------------------------------------------- + assert spec_cls is CPUOffloadingSpec def test_unregistered_spec_without_module_path_raises(): - """spec_name not in registry + no spec_module_path → ValueError.""" - config = _make_vllm_config(spec_name="NonexistentSpec") + config = _make_offloading_config(spec_name="NonexistentSpec") with pytest.raises(ValueError, match="Unsupported spec type"): - OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) + OffloadingSpecFactory.get_spec_cls(config.extra_config) - # create_spec should also fail (calls get_spec_cls internally) - kv_cache_config = _make_kv_cache_config() with pytest.raises(ValueError, match="Unsupported spec type"): - _create_spec(config, kv_cache_config) + OffloadingSpecFactory.create_spec(config) def test_cpu_spec_missing_cpu_bytes_to_use_raises(): - """CPUOffloadingSpec requires cpu_bytes_to_use → Exception.""" - config = _make_vllm_config(cpu_bytes_to_use=None) - config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None) - kv_cache_config = _make_kv_cache_config() with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"): - _create_spec(config, kv_cache_config) + _create_spec(cpu_bytes_to_use=None) def test_duplicate_registration_raises(): - """register_spec with existing name → ValueError.""" with pytest.raises(ValueError, match="is already registered"): OffloadingSpecFactory.register_spec( "CPUOffloadingSpec", "some.module", "SomeClass" ) -# --------------------------------------------------------------------------- -# Downstream collaboration — build_metric_definitions -# --------------------------------------------------------------------------- - - def test_build_metric_definitions_below_threshold(): - """store_threshold < 2 keeps stores_skipped disabled.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=1) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) + extra_config = {"store_threshold": 1} + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions(extra_config) + assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics def test_build_metric_definitions_allocation_size_histogram(): - """CPU allocation size is always reported as a histogram.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=0) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions({}) metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE] + assert isinstance(metadata, OffloadingHistogramMetadata) assert metadata.buckets == ( 1, @@ -603,49 +581,10 @@ def test_build_metric_definitions_allocation_size_histogram(): def test_build_metric_definitions_returns_counter_at_threshold(): - """store_threshold >= 2 → returns stores_skipped counter definition.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=2) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) - assert CPUOffloadingMetrics.STORES_SKIPPED in metrics - + extra_config = {"store_threshold": 2} + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions(extra_config) -def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={"blocks_per_chunk": 2}, - ) - - spec = _create_spec(config, _make_hybrid_kv_cache_config()) - - assert spec.tokens_per_block == (12, 16) - assert spec.blocks_per_chunk == 2 - - -def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={ - "block_size": 64, - "blocks_per_chunk": 2, - }, - ) - - with pytest.raises(ValueError, match="Specify only one"): - _create_spec(config, _make_kv_cache_config()) - - -def test_blocks_per_chunk_must_be_positive(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={ - "blocks_per_chunk": 0, - }, - ) - - with pytest.raises(ValueError, match="greater than 0"): - _create_spec(config, _make_kv_cache_config()) + assert CPUOffloadingMetrics.STORES_SKIPPED in metrics diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 6c11f2d465fb..027f01133b2d 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -51,6 +51,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: data_parallel_index=0, is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), ), + replicated_layout=kwargs.get("replicated_layout", False), ) spec = MagicMock(spec=OffloadingSpec) spec.config = config @@ -85,7 +86,7 @@ def test_get_file_name_full_structure(): path = fm.get_file_name(key) expected_path = ( - "/tmp/cache/test-model_42b94bdc9933_r3/000/10_g2/0001020304050607.bin" + "/tmp/cache/test-model_de3bba26cf36_r3/000/10_g2/0001020304050607.bin" ) assert path == expected_path @@ -119,6 +120,7 @@ def test_get_run_config_fields(): } ], "inference_engine": "vllm", + "parallel_agnostic": False, } @@ -164,6 +166,7 @@ def test_parallel_agnostic_collapses_namespace_when_config_allows(): assert fm.fields["pcp_size"] == 1 assert fm.fields["dcp_size"] == 1 assert fm.rank == 0 + assert "parallel_agnostic" not in fm.fields def test_parallel_agnostic_ignored_when_config_disallows(): @@ -175,6 +178,7 @@ def test_parallel_agnostic_ignored_when_config_disallows(): ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False def test_namespace_kept_without_parallel_agnostic_opt_in(): @@ -186,3 +190,123 @@ def test_namespace_kept_without_parallel_agnostic_opt_in(): ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False + + +def test_parallel_agnostic_separates_persistent_layouts(): + agnostic = make_mapper_from_offloading_spec( + is_parallelism_agnostic=True, + parallel_agnostic=True, + ) + specific = make_mapper_from_offloading_spec( + is_parallelism_agnostic=False, + parallel_agnostic=True, + ) + + assert agnostic.base_path != specific.base_path + assert "parallel_agnostic" not in agnostic.fields + assert specific.fields["parallel_agnostic"] is False + + +# --------------------------------------------------------------------------- +# replicated_layout: OR'd into parallel-agnostic identity for compact rows +# --------------------------------------------------------------------------- + + +def test_replicated_layout_collapses_parallel_identity(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=1, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=3, **shared) + + assert tp2.base_path == tp4.base_path + for fm in (tp2, tp4): + assert fm.fields["tp_size"] == 1 + assert fm.fields["pp_size"] == 1 + assert fm.fields["pcp_size"] == 1 + assert fm.fields["dcp_size"] == 1 + assert fm.rank == 0 + assert "parallel_agnostic" not in fm.fields + assert fm.fields["replicated_layout"] is True + assert fm.get_file_name(make_offload_key(b"\x01" * 8, 0)).startswith( + f"{fm.base_path}_r0/" + ) + + +def test_replicated_layout_requires_caller_opt_in(): + fm = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=True, + parallel_agnostic=False, + ) + assert fm.fields["tp_size"] == 2 + assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False + assert "replicated_layout" not in fm.fields + baseline = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=False, + parallel_agnostic=False, + ) + assert fm.base_path == baseline.base_path + + +def test_non_replicated_keeps_parallel_identity(): + fm = make_mapper_from_offloading_spec( + tp_size=4, + world_size=4, + rank=2, + replicated_layout=False, + is_parallelism_agnostic=False, + parallel_agnostic=True, + ) + assert fm.fields["tp_size"] == 4 + assert fm.rank == 2 + assert fm.fields["parallel_agnostic"] is False + assert fm.get_file_name(make_offload_key(b"\x02" * 8, 0)).startswith( + f"{fm.base_path}_r2/" + ) + + +def test_replicated_and_parallelism_agnostic_separate_layouts(): + shared = dict( + model_name="shared-model", + groups=((16, "layer0"),), + tp_size=2, + world_size=2, + rank=1, + parallel_agnostic=True, + ) + via_agnostic = make_mapper_from_offloading_spec( + is_parallelism_agnostic=True, + replicated_layout=False, + **shared, + ) + via_replicated = make_mapper_from_offloading_spec( + is_parallelism_agnostic=False, + replicated_layout=True, + **shared, + ) + assert via_agnostic.base_path != via_replicated.base_path + assert "replicated_layout" not in via_agnostic.fields + assert via_replicated.fields["replicated_layout"] is True + + +def test_replicated_layout_run_config_tp_invariant(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared) + assert tp2.get_run_config() == tp4.get_run_config() diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py index 2f0df8990233..630fcdef4407 100644 --- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py +++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py @@ -81,6 +81,8 @@ async def lifespan(app: FastAPI): app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size)) mode = "decoder-first" if global_args.decoder_first else "prefiller-first" + if global_args.p2p: + mode += ",p2p" pd_host = global_args.p2p_connector_host pd_port = global_args.p2p_connector_port n_pref = len(app.state.prefill_clients) @@ -145,6 +147,14 @@ def parse_args(): help="Send decode request before prefill so decoder is already " "waiting when KV blocks arrive (decoder-first mode)", ) + p.add_argument( + "--p2p", + action="store_true", + help="P2P mode: do not inject kv_transfer_params on the prefiller; " + "on the decoder, inject kv_transfer_params with a top-level " + "'remote_kv_source' block ({kv_request_id, remote_host, remote_port}) " + "instead of the default 'remote_prefiller' block.", + ) p.add_argument( "--prefiller-dp-size", type=int, @@ -206,11 +216,14 @@ def _auth_headers(request_id: str) -> dict: async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None): """Send a prefill-only request (max_tokens=1) to the prefiller.""" data = req_data.copy() - data["kv_transfer_params"] = { - "decode": { - "kv_request_id": request_id, - }, - } + if global_args.p2p: + data.pop("kv_transfer_params", None) + else: + data["kv_transfer_params"] = { + "remote_decoder": { + "kv_request_id": request_id, + }, + } data["stream"] = False data["max_tokens"] = 1 data.pop("max_completion_tokens", None) @@ -259,8 +272,9 @@ async def _handle_completions(api: str, request: Request): # Inject the prefiller's P2PConnector address so the decoder can pull # KV blocks from it. remote_port = base + prefill_rank targets the # replica that produced the KV (base+0 == base when dp=1). + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" req_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, @@ -311,9 +325,10 @@ async def _handle_completions_decoder_first(api: str, request: Request): prefill_client = _get_next(request.app, "prefill") decode_client = _get_next(request.app, "decode") + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" decode_data = req_data.copy() decode_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 01fd295302a8..c8a9d7f3ea35 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -21,9 +21,11 @@ from vllm.v1.kv_offload.tiering.p2p.manager import ( _UNBOUND_STORE_TIMEOUT_S, P2PSecondaryTierManager, + _annotate_req_context, ) from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, + SessionCloseResult, SessionPollResult, StoreResult, ) @@ -33,15 +35,15 @@ # --------------------------------------------------------------------------- -def _prefill_kv_params( +def _remote_prefiller_kv_params( remote_host: str = "10.0.0.1", remote_port: int = 8000, kv_request_id: str = "req-1", ) -> dict: - """Decoder-side kv_transfer_params: ``prefill`` sub-dict carries + """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries kv_request_id + remote_host + remote_port.""" return { - "prefill": { + "remote_prefiller": { "kv_request_id": kv_request_id, "remote_host": remote_host, "remote_port": remote_port, @@ -49,14 +51,34 @@ def _prefill_kv_params( } -def _decode_kv_params(kv_request_id: str = "req-1") -> dict: - """Prefiller-side kv_transfer_params: ``decode`` sub-dict carries +def _remote_kv_source_kv_params( + remote_host: str = "10.0.0.1", + remote_port: int = 8000, + kv_request_id: str = "req-1", +) -> dict: + """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" + return { + "remote_kv_source": { + "kv_request_id": kv_request_id, + "remote_host": remote_host, + "remote_port": remote_port, + }, + } + + +def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: + """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries kv_request_id only.""" - return {"decode": {"kv_request_id": kv_request_id}} + return {"remote_decoder": {"kv_request_id": kv_request_id}} def _req_context(kv_params: dict | None = None) -> ReqContext: - return ReqContext(req_id="test", kv_transfer_params=kv_params) + ctx = ReqContext(req_id="test", kv_transfer_params=kv_params) + # Mirror on_new_request: parse the P2P routing state once and cache it, + # so lookup/submit_*/on_request_finished can read it back via get_state. + _annotate_req_context(ctx) + return ctx def _job_metadata( @@ -82,38 +104,78 @@ def _make_manager() -> P2PSecondaryTierManager: """Create a manager with stubbed __init__.""" mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager) mgr._local_id = "127.0.0.1:7777" + mgr._hash_seed = "0" mgr._finished_jobs = [] mgr._failed_req_ids = set() mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._failed_serve_ctxs = [] return mgr +def _init_offloading_spec() -> SimpleNamespace: + """Minimal offloading_spec for driving the real __init__.""" + return SimpleNamespace( + config=SimpleNamespace(parallel=SimpleNamespace(data_parallel_index=0)), + blocks_per_chunk=1, + ) + + +# --------------------------------------------------------------------------- +# Tests for __init__ PYTHONHASHSEED assertion +# --------------------------------------------------------------------------- + + +class TestInitHashSeedAssertion: + def test_missing_pythonhashseed_raises(self, monkeypatch): + """P2P instance refuses to start when PYTHONHASHSEED is unset.""" + monkeypatch.delenv("PYTHONHASHSEED", raising=False) + with pytest.raises(ValueError, match="PYTHONHASHSEED"): + P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + + def test_pythonhashseed_set_succeeds(self, monkeypatch): + """With PYTHONHASHSEED set, __init__ records it for the handshake.""" + monkeypatch.setenv("PYTHONHASHSEED", "12345") + monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) + monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) + monkeypatch.setattr( + manager_module.FileMapper, + "from_offloading_spec", + lambda **k: SimpleNamespace(get_run_config=lambda: {}), + ) + mgr = P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + assert mgr._hash_seed == "12345" + + # --------------------------------------------------------------------------- -# Tests for _remote_id_from_params +# Tests for _peer_id_from_params # --------------------------------------------------------------------------- -class TestRemoteIdFromParams: +class TestPeerIdFromParams: def test_valid_params(self): - result = P2PSecondaryTierManager._remote_id_from_params( + result = manager_module._peer_id_from_params( {"remote_host": "10.0.0.1", "remote_port": 8000} ) assert result == "10.0.0.1:8000" def test_missing_host(self): - result = P2PSecondaryTierManager._remote_id_from_params({"remote_port": 8000}) + result = manager_module._peer_id_from_params({"remote_port": 8000}) assert result is None def test_missing_port(self): - result = P2PSecondaryTierManager._remote_id_from_params( - {"remote_host": "10.0.0.1"} - ) + result = manager_module._peer_id_from_params({"remote_host": "10.0.0.1"}) assert result is None def test_empty_dict(self): - result = P2PSecondaryTierManager._remote_id_from_params({}) + result = manager_module._peer_id_from_params({}) assert result is None @@ -130,35 +192,105 @@ def test_lookup_returns_miss_without_kv_params(self): def test_lookup_returns_miss_without_required_fields(self): mgr = _make_manager() - ctx = _req_context(kv_params={"prefill": {"remote_host": "x"}}) + ctx = _req_context(kv_params={"remote_prefiller": {"remote_host": "x"}}) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_valid_request(self): mgr = _make_manager() - ctx = _req_context(kv_params=_prefill_kv_params()) + ctx = _req_context(kv_params=_remote_prefiller_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_for_failed_request(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_different_request_id(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-2")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-2")) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_without_prefill_key(self): - """No ``prefill`` sub-dict means the request was not routed for + """No ``remote_prefiller`` sub-dict means the request was not routed for remote prefill — local prefill should run instead, so lookup() - returns MISS even when a stale ``decode`` block is present.""" + returns MISS even when a stale ``remote_decoder`` block is present.""" mgr = _make_manager() - ctx = _req_context(kv_params=_decode_kv_params()) + ctx = _req_context(kv_params=_remote_decoder_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.MISS +# --------------------------------------------------------------------------- +# Tests for serve_external_requests +# --------------------------------------------------------------------------- + + +class _RecordingParent: + """Minimal ParentManager stub recording on_request_finished calls.""" + + def __init__(self) -> None: + self.finished: list[str] = [] + + def on_new_request(self, ctx): + from vllm.v1.kv_offload.base import RequestOffloadingContext + + return RequestOffloadingContext() + + def lookup(self, key, ctx): + return LookupResult.MISS + + def create_store_job(self, keys, ctx): + raise AssertionError("unreachable") + + def on_request_finished(self, ctx) -> None: + self.finished.append(ctx.req_id) + + +class _RecordingSession: + """Fake P2PSession that records the parent it was served with.""" + + def __init__(self) -> None: + self.served_with: list[object] = [] + + def serve_external_requests(self, parent) -> None: + self.served_with.append(parent) + + +class TestServeExternalRequests: + def test_flushes_failed_serve_ctxs_then_serves_each_session(self): + """serve_external_requests releases the failed serves left by + reaped sessions via parent.on_request_finished (clearing the + queue), then delegates to every live session with the same parent.""" + mgr = _make_manager() + ctx = ReqContext(req_id="p2p:peer:req-1:lu1") + mgr._failed_serve_ctxs = [ctx] + sess_a = _RecordingSession() + sess_b = _RecordingSession() + mgr._sessions = {"a": sess_a, "b": sess_b} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + # Failed serve released and queue cleared. + assert parent.finished == ["p2p:peer:req-1:lu1"] + assert mgr._failed_serve_ctxs == [] + # Every live session served with the same parent handle. + assert sess_a.served_with == [parent] + assert sess_b.served_with == [parent] + + def test_no_failed_serves_still_serves_sessions(self): + mgr = _make_manager() + sess = _RecordingSession() + mgr._sessions = {"a": sess} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + assert parent.finished == [] + assert sess.served_with == [parent] + + # --------------------------------------------------------------------------- # Tests for submit_store # --------------------------------------------------------------------------- @@ -166,16 +298,16 @@ def test_lookup_returns_miss_without_prefill_key(self): class TestSubmitStore: def test_no_decode_succeeds_immediately(self): - """Without a ``decode`` block, job succeeds immediately.""" + """Without a ``remote_decoder`` block, job succeeds immediately.""" mgr = _make_manager() job = _job_metadata(job_id=1, kv_params={}) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] def test_missing_kv_request_id_fails(self): - """Missing kv_request_id inside ``decode`` fails the job.""" + """Missing kv_request_id inside ``remote_decoder`` fails the job.""" mgr = _make_manager() - params: dict = {"decode": {}} + params: dict = {"remote_decoder": {}} job = _job_metadata(job_id=1, kv_params=params) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] @@ -189,7 +321,7 @@ def test_no_binding_yet_parks_in_unbound_stores(self): job_id=1, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -220,7 +352,7 @@ def test_routes_to_bound_session(self): job_id=7, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -233,9 +365,9 @@ def test_routes_to_bound_session(self): def test_extra_top_level_keys_are_ignored(self): """Producer-side kv_transfer_params should not pre-create a session even when a stale caller still passes a top-level - ``remote_host``/``remote_port`` next to ``decode``.""" + ``remote_host``/``remote_port`` next to ``remote_decoder``.""" mgr = _make_manager() - params = _decode_kv_params() + params = _remote_decoder_kv_params() params["remote_host"] = "stale" params["remote_port"] = 12345 job = _job_metadata(job_id=1, kv_params=params) @@ -262,7 +394,7 @@ def test_empty_keys_succeeds_immediately(self): """Empty key list succeeds immediately.""" mgr = _make_manager() job = _job_metadata( - job_id=1, keys=[], block_ids=[], kv_params=_prefill_kv_params() + job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params() ) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] @@ -270,7 +402,7 @@ def test_empty_keys_succeeds_immediately(self): def test_no_session_fails(self): """No session for peer fails and marks request failed.""" mgr = _make_manager() - job = _job_metadata(job_id=1, kv_params=_prefill_kv_params()) + job = _job_metadata(job_id=1, kv_params=_remote_prefiller_kv_params()) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] assert "req-1" in mgr._failed_req_ids @@ -286,7 +418,7 @@ def test_happy_path_with_active_session(self): job_id=42, keys=[b"k1", b"k2"], block_ids=[5, 6], - kv_params=_prefill_kv_params(kv_request_id="req-42"), + kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"), ) mgr.submit_load(job) @@ -294,6 +426,20 @@ def test_happy_path_with_active_session(self): assert mgr._finished_jobs == [] assert "req-42" not in mgr._failed_req_ids + def test_missing_consumer_flag_fails(self): + """Peer fields present but neither do_remote_prefill nor + do_p2p_fetch is set — submit_load fails the job rather than + emit a stray FetchMsg.""" + mgr = _make_manager() + params = { + "remote_host": "10.0.0.1", + "remote_port": 8000, + "kv_request_id": "req-1", + } + job = _job_metadata(job_id=1, kv_params=params) + mgr.submit_load(job) + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] + # --------------------------------------------------------------------------- # Tests for on_request_finished @@ -308,7 +454,7 @@ def _make_with_failed(self) -> P2PSecondaryTierManager: def test_prunes_failed_req_ids(self): mgr = self._make_with_failed() - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" not in mgr._failed_req_ids @@ -325,14 +471,26 @@ def test_no_kv_request_id_does_nothing(self): assert "req-1" in mgr._failed_req_ids def test_decoder_side_calls_session_finish_request(self): - """Decoder-side finish (``prefill`` set) still routes via peer_id + """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id because the consumer addresses the producer it loaded from. The session's finish_request cancels the client-role load.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) mgr._sessions[peer_id] = session - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) + mgr.on_request_finished(ctx) + assert session.finishes == ["req-1"] + + def test_p2p_consumer_side_calls_session_finish_request(self): + """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id + so the session drops any pending lookups (cancel_lookups) and + cancels any inbound load.""" + mgr = _make_manager() + peer_id = "10.0.0.1:8000" + session = _FakeSession(peer_id=peer_id) + mgr._sessions[peer_id] = session + ctx = _req_context(kv_params=_remote_kv_source_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert session.finishes == ["req-1"] @@ -342,7 +500,7 @@ def test_prefiller_bound_id_routes_via_kv_to_session(self): mgr = _make_manager() bound = _FakeSession(peer_id="some-peer:1", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert bound.finishes == ["req-1"] assert "req-1" not in mgr._kv_to_session @@ -360,7 +518,7 @@ def test_prefiller_unbound_id_leaves_batches_parked(self): _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]), _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]), ] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" in mgr._unbound_stores assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11] @@ -378,11 +536,19 @@ class _FakeServerHalf: def __init__(self) -> None: self._inflight: dict[int, object] = {} + @property + def has_inflight_transfers(self) -> bool: + return bool(self._inflight) + class _FakeClientHalf: def __init__(self) -> None: self._inbound: dict[int, object] = {} + @property + def has_active_loads(self) -> bool: + return bool(self._inbound) + class _FakeSession: """Fake bidirectional session that returns canned poll() results.""" @@ -395,8 +561,10 @@ def __init__( loads: list[LoadResult] | None = None, stores: list[StoreResult] | None = None, new_fetch_ids: list[str] | None = None, - close_loads: list[tuple[int, str]] | None = None, + close_jobs: list[int] | None = None, + close_req_ids: list[str] | None = None, close_stores: list[int] | None = None, + close_failed_serves: list[ReqContext] | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -405,18 +573,24 @@ def __init__( self._loads = loads or [] self._stores = stores or [] self._new_fetch_ids = new_fetch_ids or [] - self._close_loads = close_loads or [] + self._close_jobs = close_jobs or [] + self._close_req_ids = close_req_ids or [] self._close_stores = close_stores or [] + self._close_failed_serves = close_failed_serves or [] self.requests: list[tuple[int, str]] = [] self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] self.finishes: list[str] = [] # Mirror P2PSession._server._inflight (transfer_id → handle) and - # P2PSession._client._inbound for the shutdown-drain and drain_jobs - # paths. Tests populate _server._inflight when needed. + # P2PSession._client.has_active_loads for the shutdown-drain and + # drain_jobs paths. Tests populate _server._inflight when needed. self._server = _FakeServerHalf() self._client = _FakeClientHalf() + @property + def has_pending_work(self) -> bool: + return self._client.has_active_loads or self._server.has_inflight_transfers + def poll(self): result = SessionPollResult( loads=self._loads, @@ -442,7 +616,12 @@ def finish_request(self, kv_request_id): self.finishes.append(kv_request_id) def close(self): - return self._close_loads, self._close_stores + return SessionCloseResult( + failed_jobs=self._close_jobs, + failed_req_ids=self._close_req_ids, + failed_stores=self._close_stores, + failed_serves=self._close_failed_serves, + ) class TestGetFinished: @@ -484,7 +663,8 @@ def remove_remote_peer(self, pid): peer_id="dead:1234", alive=False, connected=True, - close_loads=[(20, "req-load")], + close_jobs=[20], + close_req_ids=["req-load"], close_stores=[10, 11], ) mgr._sessions["dead:1234"] = dead # type: ignore[assignment] @@ -498,6 +678,29 @@ def remove_remote_peer(self, pid): assert "dead:1234" not in mgr._sessions assert "req-load" in mgr._failed_req_ids + def test_reap_fails_probes(self): + """A reaped session's in-flight lookups land in _failed_req_ids so + the consumer's lookup() returns MISS instead of RETRY forever.""" + + class FakeData: + def remove_remote_peer(self, pid): + pass + + mgr = self._make() + mgr._data = FakeData() # type: ignore[assignment] + dead = _FakeSession( + peer_id="dead:1234", + alive=False, + connected=True, + close_req_ids=["req-probe-1", "req-probe-2"], + ) + mgr._sessions["dead:1234"] = dead # type: ignore[assignment] + + list(mgr.get_finished_jobs()) + assert "dead:1234" not in mgr._sessions + assert "req-probe-1" in mgr._failed_req_ids + assert "req-probe-2" in mgr._failed_req_ids + def test_unbound_store_kept_within_timeout(self): """Recently-parked unbound stores stay across a poll.""" mgr = self._make() @@ -537,7 +740,7 @@ def test_submit_store_parks_unbound_batch(self): submitted_at stamp so the unbound-store sweep can age it out.""" mgr = _make_manager() job = _job_metadata( - job_id=1, kv_params=_decode_kv_params(kv_request_id="req-1") + job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") ) before = time.monotonic() mgr.submit_store(job) @@ -881,21 +1084,21 @@ def test_both_loads_succeed(self): b_loads_kv = "req-BtoA-load" # B loads, A serves a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": a_loads_kv, "remote_host": "B", "remote_port": 2, }, } b_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": b_loads_kv, "remote_host": "A", "remote_port": 1, }, } - a_prefiller_params = {"decode": {"kv_request_id": b_loads_kv}} - b_prefiller_params = {"decode": {"kv_request_id": a_loads_kv}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": b_loads_kv}} + b_prefiller_params = {"remote_decoder": {"kv_request_id": a_loads_kv}} # 1. Both sides open client-role sessions toward the peer. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1065,7 +1268,8 @@ def test_orchestrates_accept_poll_and_reap(self): peer_id=peer_dead, alive=False, connected=True, - close_loads=[(33, "req-33")], + close_jobs=[33], + close_req_ids=["req-33"], close_stores=[44], ) mgr._sessions[peer_dead] = dead # type: ignore[assignment] @@ -1309,13 +1513,13 @@ def test_dead_connection_with_pending_work_surfaces_failures(self): mgr_a, mgr_b = _build_paired_managers() a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": "req-load", "remote_host": "B", "remote_port": 2, }, } - a_prefiller_params = {"decode": {"kv_request_id": "req-store"}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": "req-store"}} # Open the outbound session A->B and submit one load + one store. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1394,6 +1598,7 @@ def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager: identity (``host:port``) stays decoupled from the NIXL agent name (a uuid). """ + monkeypatch.setenv("PYTHONHASHSEED", "0") monkeypatch.setattr( manager_module, "FileMapper", @@ -1407,16 +1612,17 @@ def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager: monkeypatch.setattr( manager_module, "NixlTransport", - lambda agent_name, *a, **k: calls.update(nixl_name=agent_name) - or SimpleNamespace(), + lambda agent_name, *a, **k: ( + calls.update(nixl_name=agent_name) or SimpleNamespace() + ), ) monkeypatch.setattr( manager_module, "ZmqTransport", - lambda local_id, host, port, *a, **k: calls.update( - zmq_id=local_id, zmq_host=host, zmq_port=port - ) - or SimpleNamespace(), + lambda local_id, host, port, *a, **k: ( + calls.update(zmq_id=local_id, zmq_host=host, zmq_port=port) + or SimpleNamespace() + ), ) spec = SimpleNamespace( blocks_per_chunk=1, diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index d1c6ad0cda7b..20e4bad21863 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -13,9 +13,18 @@ from __future__ import annotations import time +from collections.abc import Sequence +import numpy as np import pytest +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) +from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, P2PSession, @@ -33,11 +42,14 @@ ConnectMsg, DisconnectMsg, FetchMsg, + LookupMsg, + LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( _CANCEL_DRAIN_TIMEOUT_S, _InflightXfer, + _OutboundRequestState, ) from vllm.v1.kv_offload.tiering.p2p.session.session import ( _MAX_CONSECUTIVE_DISPATCH_ERRORS, @@ -48,6 +60,11 @@ # --------------------------------------------------------------------------- +# Shared PYTHONHASHSEED used by the session under test and the fake peer's +# ConnectMsg so the handshake succeeds unless a test overrides one side. +_DEFAULT_HASH_SEED = "0" + + class FakeDataTransport: """Minimal fake DataTransport for testing sessions.""" @@ -146,12 +163,16 @@ def __init__(self, peer_id: str = "peer:8000") -> None: self._inbox: list[dict] = [] self._sent: list[dict] = [] self._closed = False + # When True, send() raises to simulate a broken/dead connection. + self.fail_send = False @property def alive(self) -> bool: return not self._closed def send(self, msg: dict) -> None: + if self.fail_send: + raise ConnectionError("simulated dead connection") self._sent.append(msg) def recv(self) -> list[dict]: @@ -173,6 +194,7 @@ def _peer_connect_msg( peer_id: str = "peer:8000", block_len: int = 4096, fingerprint: str | None = None, + hash_seed: str = _DEFAULT_HASH_SEED, ) -> dict: """Build a ConnectMsg as if the peer sent it.""" msg = { @@ -182,17 +204,82 @@ def _peer_connect_msg( ConnectMsg.BASE_ADDR: 0x2000, ConnectMsg.NUM_BLOCKS: 16, ConnectMsg.BLOCK_LEN: block_len, + ConnectMsg.HASH_SEED: hash_seed, } if fingerprint is not None: msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint return msg +class FakeParent: + """Configurable :class:`ParentManager` for server-role tests. + + ``stored`` is the dict of ready blocks (key → primary block_id). + ``pending`` and ``retry`` script the first lookup() result for those + keys; subsequent lookups behave normally (a key that promised + HIT_PENDING / RETRY can later be promoted to HIT by adding it to + ``stored`` and removing it from ``pending``/``retry``). ``calls`` + captures every parent invocation in order for assertions. + + Injected per-step via ``session.serve_external_requests(parent)`` — + not held by the session, matching how ``TieringOffloadingManager`` + hands the tier a handle valid only for that call. + """ + + def __init__( + self, + stored: dict[OffloadKey, int] | None = None, + pending: set[OffloadKey] | None = None, + retry: set[OffloadKey] | None = None, + ) -> None: + self.stored: dict[OffloadKey, int] = dict(stored or {}) + self.pending: set[OffloadKey] = set(pending or ()) + self.retry: set[OffloadKey] = set(retry or ()) + self._next_job_id: int = 1000 + self.calls: list[tuple] = [] + + def on_new_request(self, ctx: ReqContext) -> RequestOffloadingContext: + self.calls.append(("on_new_request", ctx.req_id)) + return RequestOffloadingContext() + + def lookup(self, key: OffloadKey, ctx: ReqContext) -> LookupResult: + self.calls.append(("lookup", key, ctx.req_id)) + if key in self.pending: + return LookupResult.HIT_PENDING + if key in self.retry: + return LookupResult.RETRY + if key in self.stored: + return LookupResult.HIT + return LookupResult.MISS + + def create_store_job( + self, + keys: Sequence[OffloadKey], + ctx: ReqContext, + ) -> JobMetadata: + keys_list = list(keys) + self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) + block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) + job_id = self._next_job_id + self._next_job_id += 1 + return JobMetadata( + job_id=job_id, + keys=keys_list, + block_ids=block_ids, + is_promotion=False, + req_context=ctx, + ) + + def on_request_finished(self, ctx: ReqContext) -> None: + self.calls.append(("on_request_finished", ctx.req_id)) + + def _make_session( conn: FakeConnection | None = None, transport: FakeDataTransport | None = None, peer_id: str = "peer:8000", local_id: str = "local:9000", + local_hash_seed: str = _DEFAULT_HASH_SEED, ) -> tuple[P2PSession, FakeConnection, FakeDataTransport]: if conn is None: conn = FakeConnection(peer_id=peer_id) @@ -203,11 +290,17 @@ def _make_session( local_id=local_id, transport=transport, # type: ignore[arg-type] local_block_len=transport.block_len, + local_hash_seed=local_hash_seed, conn=conn, # type: ignore[arg-type] ) return session, conn, transport +def _serve(session: P2PSession, parent: FakeParent) -> None: + """Resolve enqueued inbound lookups, as the manager does each step.""" + session.serve_external_requests(parent) # type: ignore[arg-type] + + def _activate( session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000" ) -> None: @@ -217,6 +310,60 @@ def _activate( session.poll() +# --- Accessors for the server role's consolidated per-request state. +# ServerRole keeps one _ServerRequestState per kv_request_id; entries are +# garbage-collected once fully idle, so "no outbound"/"no abort" reads as +# either a missing entry or a None field. These helpers paper over that. + + +def _client_load(session: P2PSession, kv_request_id: str): + """The single in-flight load of a kv_request_id (loads are per-round).""" + loads = session._client._requests[kv_request_id].loads + assert len(loads) == 1 + return next(iter(loads.values())) + + +def _srv_outbound(session: P2PSession, kv_request_id: str): + """Serve-side round for a kv_request_id, or None (idle / GC'd). + + Rounds are keyed by wire round_seq; surfaces the demanded round when + a fetch has bound one, else any parked supply round. + """ + st = session._server._requests.get(kv_request_id) + if st is None or not st.outbound: + return None + for rnd in st.outbound.values(): + if rnd.demand_received: + return rnd + return next(iter(st.outbound.values())) + + +def _srv_lookups(session: P2PSession) -> list: + """Every parked inbound _ActiveLookup across all requests.""" + return [ + lu for st in session._server._requests.values() for lu in st.lookups.values() + ] + + +def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: + """Pending-abort start time for a kv_request_id, or None.""" + for (kv, _), started in session._server._pending_aborts.items(): + if kv == kv_request_id: + return started + return None + + +def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: + """Inflight-transfer count tracked for a kv_request_id (0 if idle).""" + st = session._server._requests.get(kv_request_id) + return len(st.inflight_tids) if st is not None else 0 + + +def _srv_total_inflight(session: P2PSession) -> int: + """Sum of per-request inflight counts across all requests.""" + return sum(len(st.inflight_tids) for st in session._server._requests.values()) + + # --------------------------------------------------------------------------- # Connect / handshake # --------------------------------------------------------------------------- @@ -300,6 +447,29 @@ def test_missing_fingerprint_allowed(self): session.poll() assert "peer:8000" in transport._remote_peers + def test_hash_seed_mismatch_marks_dead(self): + """Mismatched PYTHONHASHSEED rejects peer and marks connection dead.""" + session, conn, transport = _make_session(local_hash_seed="0") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) # mismatch + session.poll() + assert "peer:8000" not in transport._remote_peers + assert not session.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_match_succeeds(self): + """Matching PYTHONHASHSEED registers the peer and acks.""" + session, conn, transport = _make_session(local_hash_seed="12345") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) + session.poll() + assert "peer:8000" in transport._remote_peers + assert session.alive + assert any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_advertised_in_connect_msg(self): + """Session advertises its own PYTHONHASHSEED in the ConnectMsg.""" + _, conn, _ = _make_session(local_hash_seed="777") + assert conn._sent[0][ConnectMsg.HASH_SEED] == "777" + # --------------------------------------------------------------------------- # Client-role flows @@ -316,7 +486,7 @@ def test_request_blocks_sends_fetch(self): lookup = conn._sent[-1] assert lookup[TYPE_KEY] == FetchMsg.TYPE assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1" - assert lookup[FetchMsg.BLOCK_HASHES] == [b"k1", b"k2"] + assert lookup[FetchMsg.KEYS] == [b"k1", b"k2"] assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1] def test_transfer_done_success(self): @@ -328,6 +498,7 @@ def test_transfer_done_success(self): conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -344,6 +515,7 @@ def test_transfer_done_failure(self): conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: False, } @@ -362,13 +534,47 @@ def test_finish_request_sends_abort(self): assert abort[TYPE_KEY] == AbortFetchMsg.TYPE assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1" + def test_active_loads_work_list_tracks_in_flight(self): + """collect_results / has_active_loads use the _active_loads work-list, + armed when a fetch is issued and discarded exactly when its load + clears — a probe-only request never enters it, and completion empties + it while the entry may briefly linger for GC.""" + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # A probe-only request has no in-flight load: not in _active_loads. + session.register_lookup("req-probe", b"hp") + assert client._active_loads == set() + assert client.has_active_loads is False + + # Issuing a fetch arms the work-list. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] + ) + assert client._active_loads == {"req-1"} + assert client.has_active_loads is True + + # Completion clears the load and discards it from the work-list. + conn.enqueue( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, + TransferDoneMsg.KV_REQUEST_ID: "req-1", + TransferDoneMsg.SUCCESS: True, + } + ) + session.poll() + assert client._active_loads == set() + assert client.has_active_loads is False + def test_load_timeout_sends_abort(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-1"].submitted_at = time.monotonic() - 60.0 + _client_load(session, "req-1").submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE @@ -376,7 +582,7 @@ def test_load_timeout_sends_abort(self): def test_load_abort_ack_timeout_surfaces_failure(self): """After load timeout sends AbortFetch, if no AbortAck arrives within _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed - from _inbound — the engine cannot wait forever on a peer that won't + from _requests — the engine cannot wait forever on a peer that won't ack. """ session, conn, _ = _make_session() @@ -385,7 +591,7 @@ def test_load_abort_ack_timeout_surfaces_failure(self): job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. - session._client._inbound["req-7"].submitted_at = ( + _client_load(session, "req-7").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -395,43 +601,693 @@ def test_load_abort_ack_timeout_surfaces_failure(self): and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) - assert session._client._inbound["req-7"].aborted_at is not None + assert _client_load(session, "req-7").aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. - session._client._inbound["req-7"].aborted_at = ( + _client_load(session, "req-7").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)] - assert "req-7" not in session._client._inbound + assert "req-7" not in session._client._requests def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from - _inbound — covers the on_abort_ack arrival path.""" + _requests — covers the on_abort_ack arrival path.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-8"].submitted_at = ( + _client_load(session, "req-8").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() - assert session._client._inbound["req-8"].aborted_at is not None + assert _client_load(session, "req-8").aborted_at is not None # Peer acks the abort. conn.enqueue( { TYPE_KEY: AbortAckMsg.TYPE, + AbortAckMsg.ROUND_SEQ: 0, AbortAckMsg.KV_REQUEST_ID: "req-8", } ) loads = session.poll().loads assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)] - assert "req-8" not in session._client._inbound + assert "req-8" not in session._client._requests + + +# --------------------------------------------------------------------------- +# Symmetric-P2P lookup flow (do_p2p_fetch) +# --------------------------------------------------------------------------- + + +class TestLookupFlow: + """Consumer-side state machine for do_p2p_fetch lookups.""" + + def test_aggregate_flush_resolve_round_trip(self): + """register_lookup → flush sends one LookupMsg → response + resolves entries → register_lookup returns the cached bool on + every call, and repeat probes never re-issue a LookupMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Aggregate two keys for the same kv_request_id; both return None. + assert session.register_lookup("req-1", b"hA") is None + assert session.register_lookup("req-1", b"hB") is None + + # Flush sends one LookupMsg with both keys. + sent_before = len(conn._sent) + session.flush_pending_lookups() + new = conn._sent[sent_before:] + assert len(new) == 1 + msg = new[0] + assert msg[TYPE_KEY] == LookupMsg.TYPE + assert msg[LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(msg[LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Idempotent re-flush sends nothing — the entries are now in-flight. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert conn._sent[sent_before:] == [] + + # While in-flight, register_lookup keeps returning None. + assert session.register_lookup("req-1", b"hA") is None + + # Peer answers: hA hit, hB miss. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA", b"hB"], + LookupRespMsg.HITS: [True, False], + } + ) + session.poll() + + # register_lookup returns the resolved bool. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + # The entry is cached, not popped: repeat probes keep returning the + # same result and never re-queue the key, so a flush sends nothing. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + def test_request_blocks_clears_probe_cache(self): + """A resolved HIT probe is popped when its fetch is issued, so a + re-scheduled request re-probes instead of trusting the stale True + (the served block is unpinned and may have been evicted).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Probe hA, flush, and let the peer resolve it to a HIT. + assert session.register_lookup("req-1", b"hA") is None + session.flush_pending_lookups() + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + assert session.register_lookup("req-1", b"hA") is True + + # Fetch consumes the probe. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[0] + ) + assert b"hA" not in session._client._requests["req-1"].probes + + # Re-scheduled probe of the same key is treated as brand-new: it + # returns None and re-queues, so a flush emits a fresh LookupMsg. + assert session.register_lookup("req-1", b"hA") is None + sent_before = len(conn._sent) + session.flush_pending_lookups() + fresh = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(fresh) == 1 + assert fresh[0][LookupMsg.KEYS] == [b"hA"] + + def test_flush_uses_work_list_not_full_scan(self): + """flush drains a work-list rather than scanning every live request. + + A request with no newly-registered keys is not revisited: after a + flush the work-list is empty, an idle re-flush sends nothing, and a + subsequent register re-arms exactly the one affected id — even while + an unrelated request stays live in ``_requests``. + """ + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # Two requests register keys; both are queued for flush. + session.register_lookup("req-A", b"hA") + session.register_lookup("req-B", b"hB") + assert client._flush_pending == {"req-A", "req-B"} + + # Flush drains the work-list even though both requests stay live. + session.flush_pending_lookups() + assert client._flush_pending == set() + assert set(client._requests) == {"req-A", "req-B"} + + # An idle re-flush visits nothing and sends no LookupMsg. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + # A new register re-arms only that id. + session.register_lookup("req-A", b"hA2") + assert client._flush_pending == {"req-A"} + + def test_separate_lookup_msg_per_kv_request_id(self): + """Hashes for different kv_request_ids flush as separate LookupMsgs.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-A", b"h1") + session.register_lookup("req-B", b"h2") + session.register_lookup("req-A", b"h3") + + sent_before = len(conn._sent) + session.flush_pending_lookups() + sent = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(sent) == 2 + by_req = {m[LookupMsg.KV_REQUEST_ID]: m[LookupMsg.KEYS] for m in sent} + assert sorted(by_req["req-A"]) == [b"h1", b"h3"] + assert by_req["req-B"] == [b"h2"] + + def test_multiple_lookup_msgs_across_steps(self): + """A request's block set may be discovered across scheduler steps: + each step that registers new keys flushes its own LookupMsg for + the same kv_request_id, carrying only the newly-probed keys.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Step 1: probe hA, hB. + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + sent_before = len(conn._sent) + session.flush_pending_lookups() + first = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(first) == 1 + assert first[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(first[0][LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Step 2: a new key is discovered for the same request. The + # in-flight keys from step 1 are not re-sent; a second LookupMsg + # goes out carrying only the newly-probed key. + assert session.register_lookup("req-1", b"hA") is None # in-flight no-op + session.register_lookup("req-1", b"hC") + sent_before = len(conn._sent) + session.flush_pending_lookups() + second = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(second) == 1 + assert second[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert second[0][LookupMsg.KEYS] == [b"hC"] + + def test_split_response_resolves_across_messages(self): + """Producer may answer one LookupMsg's keys across multiple + LookupRespMsgs — pairs are self-describing so each lands.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + + # Two responses, each carrying one of the two keys. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hB"], + LookupRespMsg.HITS: [False], + } + ) + session.poll() + + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + + def test_finish_request_cancels_pending_lookups(self): + """finish_request drops every pending lookup for the kv_request_id.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.register_lookup("req-2", b"hC") + session.finish_request("req-1") + + # req-1 entries gone, req-2 untouched. + assert "req-1" not in session._client._requests + assert b"hC" in session._client._requests["req-2"].probes + + def test_finish_after_flushed_lookup_sends_empty_fetch(self): + """LookupMsg flushed but no FetchMsg sent (all-miss case) → + finish_request emits an empty FetchMsg so the peer can drop its + lookup state and call parent.on_request_finished.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert len(fetches) == 1 + assert fetches[0][FetchMsg.KV_REQUEST_ID] == "req-1" + assert fetches[0][FetchMsg.KEYS] == [] + assert fetches[0][FetchMsg.BLOCK_INDEXES] == [] + + def test_finish_without_flushed_lookup_sends_no_fetch(self): + """No LookupMsg was ever sent → finish_request must not emit an + empty FetchMsg (the peer has no state to release).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Register but never flush. + session.register_lookup("req-1", b"hA") + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_finish_after_real_fetch_sends_no_second_fetch(self): + """A real FetchMsg was already sent for the id → finish_request + must not emit a second (empty) FetchMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + # Resolve the probe to a HIT before fetching, as the manager only + # loads confirmed hits (an unresolved probe yields RETRY). + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[7] + ) + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_server_lookup_deferred_until_serve_then_all_misses(self): + """``poll()`` only enqueues an inbound LookupMsg — no response is + sent until ``serve_external_requests``. With an all-miss parent + the aggregated LookupRespMsg carries the same keys and + ``hits=[False, ...]``.""" + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, + LookupMsg.KV_REQUEST_ID: "req-1", + LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], + } + ) + session.poll() + + # Dispatch alone must not answer — the parent handle is only valid + # during serve_external_requests. + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] == [] + + _serve(session, FakeParent()) + + resps = [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] + assert len(resps) == 1 + resp = resps[0] + assert resp[LookupRespMsg.KV_REQUEST_ID] == "req-1" + assert resp[LookupRespMsg.KEYS] == [b"hX", b"hY", b"hZ"] + assert resp[LookupRespMsg.HITS] == [False, False, False] + + +# --------------------------------------------------------------------------- +# Server-side handling of inbound LookupMsg (ParentManager-driven) +# +# poll() only enqueues the LookupMsg; serve_external_requests(parent) +# resolves it. Tests follow the poll() → _serve() pattern. +# --------------------------------------------------------------------------- + + +def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, + LookupMsg.KV_REQUEST_ID: kv_request_id, + LookupMsg.KEYS: list(keys), + } + ) + + +def _lookup_resps(conn: FakeConnection, since: int = 0) -> list[dict]: + return [m for m in conn._sent[since:] if m[TYPE_KEY] == LookupRespMsg.TYPE] + + +class TestServerLookupHandling: + def test_immediate_hits_create_one_store_job(self): + """All-HIT batch: one create_store_job call with all keys, one + LookupRespMsg with hits=[True]*N, on_request_finished fires at the + end of serve, and `available` is populated for the eventual fetch.""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB", b"hC"] + assert resps[0][LookupRespMsg.HITS] == [True, True, True] + + kinds = [c[0] for c in cb.calls] + assert kinds.count("create_store_job") == 1 + cs = next(c for c in cb.calls if c[0] == "create_store_job") + assert cs[1] == (b"hA", b"hB", b"hC") + assert cb.calls[-1][0] == "on_request_finished" + + # Hits are pinned in outbound state for the upcoming FetchMsg match. + assert set(_srv_outbound(session, "req-1").available) == { + b"hA", + b"hB", + b"hC", + } + + def test_all_misses_no_store_job_finish_fires(self): + """All-MISS batch: no create_store_job call; one LookupRespMsg + with hits=[False]*N; on_request_finished fires at end of serve.""" + cb = FakeParent() + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.HITS] == [False, False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert cb.calls[-1][0] == "on_request_finished" + + def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): + """HIT/MISS resolutions do not go out on first sight when any + key is still HIT_PENDING / RETRY. The lookup parks until every + key has settled (or the deadline fires), then one + LookupRespMsg carries all keys in wire order.""" + cb = FakeParent( + stored={b"hA": 1}, + pending={b"hB"}, + retry={b"hD"}, + ) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC", b"hD"]) + session.poll() + _serve(session, cb) + + # No LookupRespMsg yet — hB and hD are still pending. + assert _lookup_resps(conn, sent_before) == [] + # HIT is still pinned immediately so the eventual FetchMsg matches. + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 1 + assert cs_calls[0][1] == (b"hA",) + # Lookup is parked; on_request_finished not yet called. + assert all(c[0] != "on_request_finished" for c in cb.calls) + assert len(_srv_lookups(session)) == 1 + + def test_pending_resolves_then_aggregate_response_fires(self): + """A HIT_PENDING key that becomes HIT on a later poll releases + the deferred aggregate response: one LookupRespMsg carrying + both keys in wire order, and one create_store_job call per + HIT (the second HIT is pinned when it resolves, not when the + response goes out).""" + cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + # No response yet — hB still pending. + assert _lookup_resps(conn, sent_before) == [] + + # Promote hB. + cb.pending.discard(b"hB") + cb.stored[b"hB"] = 2 + + # Drive resolver via a second serve_external_requests. + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB"] + assert resps[0][LookupRespMsg.HITS] == [True, True] + + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 2 + assert cs_calls[0][1] == (b"hA",) + assert cs_calls[1][1] == (b"hB",) + # on_request_finished fires once after the aggregate resolve. + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + assert b"hA" in _srv_outbound(session, "req-1").available + assert b"hB" in _srv_outbound(session, "req-1").available + + def test_pending_timeout_replies_miss_no_store_job(self): + """A HIT_PENDING key that stays pending past the batch + ``deadline`` is force-MISS and never pinned; the deferred + aggregate response fires with hits=[False].""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + # Initial serve: nothing immediate, lookup parked, no LookupRespMsg. + assert _lookup_resps(conn, sent_before) == [] + + # Forge the deadline into the past to trigger the timeout branch. + lookup = _srv_lookups(session)[0] + lookup.deadline = time.monotonic() - 0.1 + + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA"] + assert resps[0][LookupRespMsg.HITS] == [False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + + def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): + """Two LookupMsgs for the same kv_request_id get distinct ctxs + and two on_request_finished calls (one per batch).""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-1", [b"hB"]) + session.poll() + _serve(session, cb) + + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 2 + # Distinct synthetic req_ids + assert finish_calls[0][1] != finish_calls[1][1] + # Both namespaced under the same kv_request_id + assert ":req-1:" in finish_calls[0][1] + assert ":req-1:" in finish_calls[1][1] + + def test_close_returns_open_batch_ctxs_as_failed_serves(self): + """Tearing the session down with a parked batch returns the + synthetic ctx as a failed serve (no parent handle at teardown) so + the manager can release the TieringManager's state on its next + serve.""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 1 + assert all(c[0] != "on_request_finished" for c in cb.calls) + + result = session.close() + + assert len(result.failed_serves) == 1 + assert ":req-1:" in result.failed_serves[0].req_id + # close() itself must not call the parent. + assert all(c[0] != "on_request_finished" for c in cb.calls) + + def test_wire_finish_drops_pending_batches_for_kv_request_id(self): + """``ServerRole.finish(kv_request_id)`` drops every parked batch + whose kv_request_id matches and queues its ctx for the next + serve's on_request_finished.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + session._server.finish("req-1") + + # req-1 batch dropped from parked lookups; its ctx queued for release. + remaining_kv_request_ids = {b.kv_request_id for b in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): + """A peer FetchMsg terminates the lookup phase for its id: parked + lookups with matching kv_request_id are dropped and their + ctx queued for on_request_finished; other kv_request_ids untouched.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + # Empty FetchMsg: peer signals "lookup phase done" without asking + # for any blocks (the all-miss case). + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [], + FetchMsg.BLOCK_INDEXES: [], + } + ) + session.poll() + + remaining_kv_request_ids = {lu.kv_request_id for lu in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + # Dispatch queues the ctx but does not call the parent yet. + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_lookup_then_fetch_round_trip_emits_store_result(self): + """End-to-end: lookup pins primary slots → fetch matches them → + NIXL transfer completes → StoreResult surfaces with the + create_store_job's job_id (the engine releases the pin).""" + cb = FakeParent(stored={b"hA": 7, b"hB": 8}) + session, conn, transport = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + cs = next(c for c in cb.calls if c[0] == "create_store_job") + # FakeParent issues monotonic job_ids starting at 1000. + expected_job_id = 1000 + + # Consumer issues FetchMsg on the resolved hits. + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"hA", b"hB"], + FetchMsg.BLOCK_INDEXES: [20, 21], + } + ) + session.poll() + + # NIXL write_blocks called with our pinned local block_ids. + assert len(transport._transfers) == 1 + _, (_peer, local, remote) = next(iter(transport._transfers.items())) + assert local == [7, 8] + assert remote == [20, 21] + + # Drive the transport completion. + transport._poll_done.append(0) + result = session.poll() + + store_results = [s for s in result.stores if s.success] + assert any(s.job_id == expected_job_id for s in store_results) + # Sanity: kv mention in synthetic ctx. + assert cs[2].startswith("p2p:") # --------------------------------------------------------------------------- @@ -448,8 +1304,9 @@ def test_store_then_fetch_matches(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -466,8 +1323,9 @@ def test_fetch_then_store_matches(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -484,8 +1342,9 @@ def test_transfer_completion_emits_store_result_and_done(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -502,17 +1361,18 @@ def test_abort_fetch_replies_with_ack(self): conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is - deferred and the abort is parked in _pending_aborts.""" + deferred and the abort is parked (abort_started_at set).""" session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends @@ -520,20 +1380,26 @@ def test_abort_fetch_defers_ack_when_cancel_pending(self): tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # First attempt happens inside _on_abort_fetch; the per-tick # drain runs again at the end of poll() — both are wait-mode. assert all(mode == "wait" for _, mode in transport._cancel_calls) @@ -547,18 +1413,24 @@ def test_abort_fetch_acks_after_drain(self): tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backend finishes draining: transport.poll() will return tid as # DONE, and the next cancel(mode='wait') call sees it's gone. @@ -569,7 +1441,7 @@ def test_abort_fetch_acks_after_drain(self): ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight def test_abort_fetch_force_cancels_after_timeout(self): @@ -580,20 +1452,26 @@ def test_abort_fetch_force_cancels_after_timeout(self): tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. - session._server._pending_aborts["req-1"] = ( + session._server._pending_aborts[("req-1", 0)] = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the @@ -604,7 +1482,7 @@ def test_abort_fetch_force_cancels_after_timeout(self): ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight assert ([tid], "immediate") in transport._cancel_calls @@ -616,29 +1494,36 @@ def test_abort_fetch_idempotent_while_draining(self): tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() - first_started_at = session._server._pending_aborts["req-1"] + first_started_at = _srv_abort_started(session, "req-1") # Second AbortFetchMsg for the same kv_request_id while still # draining must not reset the deadline. conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() - assert session._server._pending_aborts["req-1"] == first_started_at + assert _srv_abort_started(session, "req-1") == first_started_at # Now let the drain succeed and confirm exactly one ack ever. transport._cancel_still_inflight.discard(tid) @@ -668,8 +1553,9 @@ def test_store_timeout_then_late_completion_no_duplicate(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -700,8 +1586,9 @@ def test_store_timeout_then_late_failure_no_duplicate(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -742,13 +1629,14 @@ def test_no_inflight_unmatched_demand_sends_failure(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None session.finish_request("req-1") @@ -756,7 +1644,7 @@ def test_no_inflight_unmatched_demand_sends_failure(self): assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the @@ -767,8 +1655,9 @@ def test_with_inflight_defers_then_fires_on_last_transfer(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -780,8 +1669,8 @@ def test_with_inflight_defers_then_fires_on_last_transfer(self): before = len(conn._sent) session.finish_request("req-1") assert len(conn._sent) == before - assert "req-1" in session._server._outbound - assert session._server._outbound["req-1"].finishing + assert _srv_outbound(session, "req-1") is not None + assert _srv_outbound(session, "req-1").finishing # Last inflight settles -> early-fail fires. tid = next(iter(transport._transfers)) @@ -792,7 +1681,7 @@ def test_with_inflight_defers_then_fires_on_last_transfer(self): assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: @@ -802,8 +1691,9 @@ def test_full_demand_satisfied_still_sends_success(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -831,13 +1721,14 @@ def test_prefiller_first_finish_before_fetch(self): session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # finish_request first — no demand received yet -> defer. session.finish_request("req-1") - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Fetch arrives now: demand fully satisfied by available. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -857,8 +1748,9 @@ def test_prefiller_first_finish_before_fetch(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-2", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -870,7 +1762,7 @@ def test_prefiller_first_finish_before_fetch(self): session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-2" not in session._server._outbound + assert _srv_outbound(session, "req-2") is None def test_unknown_request_is_noop(self): session, conn, _ = _make_session() @@ -890,8 +1782,9 @@ def test_finish_request_no_inflight_emits_store_failure(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"demand"], + FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -900,7 +1793,7 @@ def test_finish_request_no_inflight_emits_store_failure(self): # matches demand. Without the shortcut, job 42 sits in _store_jobs # for _STORE_TIMEOUT_S. session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42) - assert session._server._outbound["req-1"].pending_job_ids == {42} + assert _srv_outbound(session, "req-1").pending_job_ids == {42} session.finish_request("req-1") @@ -908,7 +1801,7 @@ def test_finish_request_no_inflight_emits_store_failure(self): msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Local store job surfaces on the next poll, success=False. stores = session.poll().stores @@ -924,8 +1817,9 @@ def test_finish_request_remaining_zero_emits_success_via_inflight(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -933,7 +1827,7 @@ def test_finish_request_remaining_zero_emits_success_via_inflight(self): session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7) # finish_request races with the inflight transfer. session.finish_request("req-1") - assert "req-1" in session._server._outbound # deferred + assert _srv_outbound(session, "req-1") is not None # deferred # Last inflight completes -> _finalize_outbound(success=True) fires. tid = next(iter(transport._transfers)) @@ -944,7 +1838,7 @@ def test_finish_request_remaining_zero_emits_success_via_inflight(self): assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is True assert StoreResult(job_id=7, success=True) in stores - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_write_blocks_failure_finalizes_with_failure(self): """write_blocks returning None must not leave the request hanging. @@ -961,8 +1855,9 @@ def test_write_blocks_failure_finalizes_with_failure(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -973,7 +1868,7 @@ def test_write_blocks_failure_finalizes_with_failure(self): session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) # Outbound was finalized immediately (no other inflight). - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Peer notified with success=False. msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" @@ -995,15 +1890,16 @@ def test_partial_match_completes_in_two_rounds(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2", b"k3"], + FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], } ) session.poll() # Demand registered, no matches yet. assert session._server._inflight == {} - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 3 assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"} @@ -1024,7 +1920,7 @@ def test_partial_match_completes_in_two_rounds(self): assert session._server._inflight == {} assert outbound.remaining == 2 # Not yet finalized — still 2 blocks demanded. - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Round 2: k2 and k3 arrive together. session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200) @@ -1038,7 +1934,7 @@ def test_partial_match_completes_in_two_rounds(self): stores = session.poll().stores assert StoreResult(job_id=200, success=True) in stores # _finalize_outbound fired — request gone, peer notified with success. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is True @@ -1056,8 +1952,9 @@ def test_write_blocks_failure_finalizes_after_last_inflight_completes(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -1067,7 +1964,7 @@ def test_write_blocks_failure_finalizes_after_last_inflight_completes(self): session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100) assert len(session._server._inflight) == 1 tid_1 = next(iter(session._server._inflight)) - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 2 # decrement happens on completion assert outbound.finishing is False @@ -1078,7 +1975,7 @@ def test_write_blocks_failure_finalizes_after_last_inflight_completes(self): assert list(session._server._inflight.keys()) == [tid_1] # Marked finishing, but NOT finalized yet (transfer_1 still inflight). assert outbound.finishing is True - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE] assert done_msgs == [] @@ -1090,7 +1987,7 @@ def test_write_blocks_failure_finalizes_after_last_inflight_completes(self): stores_first = session.poll().stores assert StoreResult(job_id=100, success=True) in stores_first # Outbound state cleaned up; peer notified with success=False. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is False @@ -1122,8 +2019,9 @@ def test_session_handles_both_roles_concurrently(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-srv", - FetchMsg.BLOCK_HASHES: [b"served"], + FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], } ) @@ -1152,6 +2050,7 @@ def test_session_handles_both_roles_concurrently(self): conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-cli", TransferDoneMsg.SUCCESS: True, } @@ -1178,6 +2077,7 @@ def test_pending_session_buffers_stored_blocks(self): local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1197,6 +2097,7 @@ def test_attach_connection_sends_connect(self): local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) conn = FakeConnection(peer_id="peer:8000") @@ -1218,13 +2119,16 @@ def test_pending_close_returns_pending_stores(self): local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2) - failed_loads, failed_stores = session.close() - assert failed_loads == [] - assert set(failed_stores) == {1, 2} + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == [] + assert set(result.failed_stores) == {1, 2} + assert result.failed_serves == [] # --------------------------------------------------------------------------- @@ -1246,9 +2150,50 @@ def test_close_returns_pending_loads_and_stores(self): session.request_blocks(1, "req-1", [b"k"], [0]) session.request_blocks(2, "req-2", [b"k"], [0]) session.add_stored_blocks("req-srv", [b"k"], [0], job_id=10) - failed_loads, failed_stores = session.close() - assert set(failed_loads) == {(1, "req-1"), (2, "req-2")} - assert set(failed_stores) == {10} + result = session.close() + assert set(result.failed_jobs) == {1, 2} + assert set(result.failed_req_ids) == {"req-1", "req-2"} + assert set(result.failed_stores) == {10} + assert result.failed_serves == [] + + def test_send_failure_marks_connection_dead(self): + """A raising send must mark the connection dead, not silently drop + the message — otherwise the session lingers alive, is never reaped, + and in-flight lookups/loads toward the dead peer hang forever.""" + session, conn, _ = _make_session() + _activate(session, conn) + assert session.alive + + conn.fail_send = True + # request_blocks flushes a FetchMsg synchronously via _do_send. + session.request_blocks(1, "req-1", [b"k"], [0]) + + assert not session.alive + + def test_close_surfaces_inflight_lookups(self): + """close() reports kv_request_ids whose symmetric-P2P probe is still + unresolved; resolved probes are not reported (their answer is in).""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-hit", b"hA") + session.register_lookup("req-inflight", b"hB") + session.flush_pending_lookups() + + # Only req-hit is answered; req-inflight stays in flight. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-hit", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == ["req-inflight"] # --------------------------------------------------------------------------- @@ -1293,8 +2238,9 @@ def test_fetch_mismatched_lengths(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1341,8 +2287,9 @@ def test_value_error_disconnects_on_first_occurrence(self): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1371,8 +2318,9 @@ def _boom(*args, **kwargs): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1394,8 +2342,9 @@ def _boom(*args, **kwargs): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1421,8 +2370,9 @@ def _boom(*args, **kwargs): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1438,10 +2388,10 @@ def _boom(*args, **kwargs): class TestInflightPerReqInvariant: - """`_inflight_per_req` is the O(1) replacement for the previous - O(N) scan in `_has_inflight_for`. These tests check that every - mutation site keeps the counter in sync with `_inflight` and that - the lookup is correct under high fan-out. + """Per-request `inflight_tids` is the O(1) replacement for the + previous O(N) scan in `_has_inflight_for`. These tests check that + every mutation site keeps the set in sync with `_inflight` and + that the lookup is correct under high fan-out. """ def test_invariant_holds_through_lifecycle(self): @@ -1452,25 +2402,23 @@ def test_invariant_holds_through_lifecycle(self): _activate(session, conn) def _invariant_holds() -> bool: - counted = sum(session._server._inflight_per_req.values()) - return counted == len(session._server._inflight) and all( - v > 0 for v in session._server._inflight_per_req.values() - ) + return _srv_total_inflight(session) == len(session._server._inflight) assert _invariant_holds() # Two requests, two blocks each, all dispatched in one batch. session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10) session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11) - for kv_id, hashes, indexes in ( + for kv_id, keys, indexes in ( ("req-A", [b"a1", b"a2"], [100, 101]), ("req-B", [b"b1", b"b2"], [102, 103]), ): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: kv_id, - FetchMsg.BLOCK_HASHES: hashes, + FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, } ) @@ -1493,7 +2441,7 @@ def _invariant_holds() -> bool: assert _invariant_holds() assert not session._server._has_inflight_for("req-A") - assert "req-A" not in session._server._inflight_per_req # entry was removed + assert _srv_inflight_count(session, "req-A") == 0 # entry drained assert session._server._has_inflight_for("req-B") # Complete req-B; counter must drain to empty. @@ -1508,7 +2456,7 @@ def _invariant_holds() -> bool: assert _invariant_holds() assert session._server._inflight == {} - assert session._server._inflight_per_req == {} + assert _srv_total_inflight(session) == 0 def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must @@ -1522,11 +2470,14 @@ def test_has_inflight_for_correct_with_many_requests(self): tid = kv_id_idx * 10 + j session._server._inflight_add( tid, - _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}), + _InflightXfer( + kv_request_id=kv_id, + block_count=1, + job_ids={tid}, + round=_OutboundRequestState(inflight=1), + ), ) - assert sum(session._server._inflight_per_req.values()) == len( - session._server._inflight - ) + assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") assert session._server._has_inflight_for("req-99") assert not session._server._has_inflight_for("req-missing") @@ -1539,7 +2490,7 @@ def test_has_inflight_for_correct_with_many_requests(self): ] for tid in tids_50: session._server._inflight_pop(tid) - assert "req-50" not in session._server._inflight_per_req + assert _srv_inflight_count(session, "req-50") == 0 assert not session._server._has_inflight_for("req-50") # Other ids unaffected. assert session._server._has_inflight_for("req-49") @@ -1559,6 +2510,7 @@ def _valid_msg(self) -> dict: ConnectMsg.BASE_ADDR: 0x1000, ConnectMsg.NUM_BLOCKS: 8, ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.HASH_SEED: "0", } def test_valid_message_passes(self): @@ -1600,13 +2552,26 @@ def test_block_len_zero(self): with pytest.raises(ValueError, match="block_len"): ConnectMsg.validate(msg) + def test_missing_hash_seed(self): + msg = self._valid_msg() + del msg[ConnectMsg.HASH_SEED] + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + + def test_hash_seed_wrong_type(self): + msg = self._valid_msg() + msg[ConnectMsg.HASH_SEED] = 12345 # int, not str + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], } @@ -1630,6 +2595,7 @@ class TestTransferDoneMsgValidation: def test_valid_message_passes(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -1638,6 +2604,7 @@ def test_valid_message_passes(self): def test_success_wrong_type(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: 1, } diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index 101d4291cfd7..f814a04887c4 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -219,6 +219,7 @@ def test_dead_connection_removed_on_poll(self): # Pruning is synchronous within poll(). transport_a.poll() assert len(transport_a._connections) == 0 + assert new_conns[0]._sockets.dealer.closed finally: transport_a.close() transport_b.close() @@ -228,3 +229,110 @@ def test_close_is_idempotent(self): transport, _ = _make_transport() transport.close() transport.close() # should not raise + + +class TestZmqReconnect: + """Reconnecting to a peer whose connection died (real ZMQ sockets). + + A session marks its connection dead while handling messages, which happens + after the transport's own sweep has run for that tick — so a dead + connection stays registered until the next poll(). Reconnecting in that + window must succeed, and the retired connection must release its sockets. + Real sockets are required: a mock reports every attribute as closed. + """ + + def test_close_after_mark_dead_releases_sockets(self): + """close() releases sockets even when mark_dead() ran first. + + mark_dead() must not set the flag close() guards on, or every peer + disconnect leaks a DEALER and a monitor socket. + """ + transport, _ = _make_transport() + try: + conn = transport.connect(f"127.0.0.1:{_free_port()}") + dealer, monitor = conn._sockets.dealer, conn._sockets.monitor + + conn.mark_dead() + assert not conn.alive + assert not dealer.closed + + conn.close() + assert dealer.closed + assert monitor.closed + finally: + transport.close() + + def test_connect_retires_dead_connection(self): + """connect() replaces a registered-but-dead connection.""" + transport, _ = _make_transport() + try: + # The peer port is never bound — only the peer id matters here. + peer_id = f"127.0.0.1:{_free_port()}" + dead = transport.connect(peer_id) + dead.mark_dead() + + conn = transport.connect(peer_id) + + assert conn is not dead + assert conn.alive + assert transport._connections[peer_id] is conn + assert dead._sockets.dealer.closed + finally: + transport.close() + + def test_repeated_reconnect_to_same_peer(self): + """A flapping peer stays reconnectable. + + Monitor endpoints are inproc addresses that libzmq releases + asynchronously, so deriving one from peer_id alone makes each + reconnect race the previous teardown and fail with EADDRINUSE. + """ + transport, _ = _make_transport() + try: + peer_id = f"127.0.0.1:{_free_port()}" + for _ in range(10): + conn = transport.connect(peer_id) + conn.mark_dead() + transport.poll() + assert conn._sockets.dealer.closed + + assert not transport._connections + finally: + transport.close() + + def test_inbound_message_survives_dead_registration(self): + """A reconnecting peer's first message is not dropped. + + poll() must retire connections killed by their session before routing + traffic, otherwise the message is enqueued into the dead connection and + discarded when it is swept — and a session announces itself only once. + Covers only that between-polls window: a peer dying while poll() runs, + or dying silently until the heartbeat expires, is out of scope. + """ + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + + try: + conn_b = transport_b.connect(f"127.0.0.1:{port_a}") + conn_b.send({"type": "connect", "seq": 1}) + + inbound = _wait_for_inbound(transport_a)[0] + _wait_for_messages(transport_a, inbound, 1) + + inbound.mark_dead() + conn_b.send({"type": "connect", "seq": 2}) + + # Wait until the frame is readable on the ROUTER, so the message is + # known to have arrived rather than merely being slow. + poller = zmq.Poller() + poller.register(transport_a._router, zmq.POLLIN) + assert poller.poll(2000), "message never reached the ROUTER" + + new_conns = _wait_for_inbound(transport_a) + assert len(new_conns) == 1 + assert new_conns[0] is not inbound + msgs = _wait_for_messages(transport_a, new_conns[0], 1) + assert msgs == [{"type": "connect", "seq": 2}] + finally: + transport_a.close() + transport_b.close() diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index dcc92a4fa8bf..2959ac1aa03d 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -18,10 +18,10 @@ import pytest import torch -from vllm.distributed.kv_events import MEDIUM_FS from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -46,13 +46,24 @@ # Helpers # --------------------------------------------------------------------------- +_NUM_BLOCKS = 8 _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096. _DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") -def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: +def _make_offloading_spec( + enable_kv_cache_events: bool = False, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> MagicMock: """Mock spec with an explicit global KV events flag.""" + if world_size is None: + world_size = tp_size spec = MagicMock() spec.config = OffloadingConfig( groups=(), @@ -63,15 +74,16 @@ def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: model=OffloadingModelConfig(name="test-model", dtype="float32"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) spec.blocks_per_chunk = 1 spec.kv_events_config = OffloadingKVEventsConfig( @@ -162,7 +174,7 @@ def _page_aligned_rand_tensor( @pytest.fixture def fs_tier(tmp_path): - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) tier = FileSystemTierManager( offloading_spec=_MOCK_OFFLOADING_SPEC, @@ -178,7 +190,7 @@ def fs_tier(tmp_path): @pytest.fixture def fs_tier_with_events(tmp_path): - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) tier = FileSystemTierManager( offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), @@ -347,38 +359,85 @@ def test_shutdown_discards_pending_tasks(fs_tier): assert all(not t.is_alive() for t in tier._pool._threads) -def test_store_load_data_integrity(fs_tier): - """Data written by store must be exactly recovered by load.""" +@pytest.mark.parametrize("batch_size", [0, 1, 2, 5]) +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): + """Data written by store must be exactly recovered by load, for batches + of any size -- including the empty batch.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + tier, tensor = fs_tier # Populate tensor with random data - tensor[:] = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) - - # Store first 2 blocks - num_store = 2 - expected = tensor[:num_store].clone() + tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) - store_ids = list(range(num_store)) - keys = [key(i) for i in range(num_store)] + keys = [key(i) for i in range(batch_size)] + store_block_ids = list(range(batch_size)) + load_block_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS)) + expected = tensor[:batch_size].clone() - tier.submit_store(make_job(1, keys, store_ids)) - results = drain(tier) - assert all(r.success for r in results) + tier.submit_store(make_job(1, keys, store_block_ids)) + store_results = drain(tier) + assert len(store_results) == 1 + assert store_results[0].success + assert all(os.path.exists(tier.file_mapper.get_file_name(k)) for k in keys) - # Overwrite source blocks to prove data is read from disk - tensor[:num_store] = 0.0 + # reset tensor to prove data is read from disk + tensor[:] = 0.0 - # Load into last 2 blocks - load_ids = [2, 3] - tier.submit_load(make_job(2, keys, load_ids, is_promotion=True)) - results = drain(tier) - assert all(r.success for r in results) + # Load into a range disjoint by index from the store ids, to also + # exercise loading a block into a different id than it was stored from. + tier.submit_load(make_job(2, keys, load_block_ids, is_promotion=True)) + load_results = drain(tier) + assert len(load_results) == 1 + assert load_results[0].success - for i, bid in enumerate(load_ids): + for i, bid in enumerate(load_block_ids): assert torch.allclose(tensor[bid], expected[i]), ( f"Block {bid} data mismatch after store+load" ) +def test_store_load_roundtrip_without_o_direct(tmp_path, monkeypatch): + """Buffered fallback must round-trip data when O_DIRECT is unsupported. + + Simulates filesystems (e.g. overlayfs, some NFS) that reject O_DIRECT by + forcing the capability probe to report it unavailable. + """ + monkeypatch.setattr( + "vllm.v1.kv_offload.tiering.fs.manager.probe_o_direct", + lambda _dir: False, + ) + tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + tier = FileSystemTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=memoryview(tensor.numpy()), + tier_type="fs", + root_dir=str(tmp_path), + n_read_threads=4, + n_write_threads=4, + ) + try: + assert tier._use_o_direct is False + + keys = [key(0), key(1)] + expected = tensor[:2].clone() + tier.submit_store(make_job(1, keys, [0, 1])) + assert all(r.success for r in drain(tier)) + + tensor[:2] = 0.0 + tier.submit_load(make_job(2, keys, [2, 3], is_promotion=True)) + assert all(r.success for r in drain(tier)) + + for i, bid in enumerate([2, 3]): + assert torch.allclose(tensor[bid], expected[i]) + finally: + tier.shutdown() + + def test_wait_idle_blocks_until_tasks_complete(): """wait_idle must not return while a task is still in flight.""" pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1) @@ -462,6 +521,30 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): assert results == [LookupResult.HIT, LookupResult.MISS] +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): + """Smoke test: a block id beyond the primary tensor's block count must + fail the job, for both the C extension and the Python fallback.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, tensor = fs_tier + out_of_bounds_bid = tensor.shape[0] # one past the last valid block + + tier.submit_store(make_job(1, [key(1)], [out_of_bounds_bid])) + store_results = drain(tier) + assert len(store_results) == 1 + assert not store_results[0].success + + tier.submit_load(make_job(2, [key(1)], [out_of_bounds_bid], is_promotion=True)) + load_results = drain(tier) + assert len(load_results) == 1 + assert not load_results[0].success + + # --------------------------------------------------------------------------- # KV events # --------------------------------------------------------------------------- @@ -477,8 +560,7 @@ def test_successful_store_emits_stored_event(fs_tier_with_events): events = list(tier.take_events()) assert len(events) == 1 assert events[0].keys == keys - # Literal medium pins the wire contract, not just the constant choice. - assert events[0].medium == "FS" + assert events[0].medium == Medium.STORAGE assert events[0].locality is Locality.LOCAL assert not events[0].removed # take_events drains the buffer. @@ -535,14 +617,14 @@ def test_mixed_job_results_emit_event_only_for_successful_job( tier = fs_tier_with_events failing_path = tier.file_mapper.get_file_name(key(1)) - original_store_block = mgr_mod.store_block + original_batch_store_block = mgr_mod.batch_store_block - def flaky_store_block(dest_path, *args, **kwargs): - if dest_path == failing_path: + def flaky_batch_store_block(paths, *args, **kwargs): + if failing_path in paths: raise OSError("injected store failure") - return original_store_block(dest_path, *args, **kwargs) + return original_batch_store_block(paths, *args, **kwargs) - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) tier.submit_store(make_job(1, [key(1)], [0])) tier.submit_store(make_job(2, [key(2)], [1])) @@ -563,14 +645,14 @@ def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch) tier = fs_tier_with_events failing_path = tier.file_mapper.get_file_name(key(2)) - original_store_block = mgr_mod.store_block + original_batch_store_block = mgr_mod.batch_store_block - def flaky_store_block(dest_path, *args, **kwargs): - if dest_path == failing_path: + def flaky_batch_store_block(paths, *args, **kwargs): + if failing_path in paths: raise OSError("injected store failure") - return original_store_block(dest_path, *args, **kwargs) + return original_batch_store_block(paths, *args, **kwargs) - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) results = drain(tier) @@ -648,9 +730,54 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): events.extend(manager.take_events()) time.sleep(0.01) - fs_events = [e for e in events if e.medium == MEDIUM_FS] + fs_events = [e for e in events if e.medium == Medium.STORAGE] assert len(fs_events) == 1 assert set(fs_events[0].keys) == set(keys) assert not fs_events[0].removed finally: tier.shutdown() + + +def test_fs_tier_cross_tp_round_trip(tmp_path): + """TP=2 replicated writer and TP=4 reader share namespace and bytes.""" + root = str(tmp_path) + writer_tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + expected = writer_tensor[0].clone() + writer = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=2, world_size=2, rank=0, replicated_layout=True + ), + primary_kv_view=memoryview(writer_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + writer.submit_store(make_job(1, [key(7)], [0])) + assert all(r.success for r in drain(writer)) + writer_base = writer.file_mapper.base_path + writer_path = writer.file_mapper.get_file_name(key(7)) + finally: + writer.shutdown() + + reader_tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + reader = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + primary_kv_view=memoryview(reader_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + assert reader.file_mapper.base_path == writer_base + assert reader.file_mapper.get_file_name(key(7)) == writer_path + assert lookup_and_wait(reader, [key(7)]) == [LookupResult.HIT] + reader.submit_load(make_job(2, [key(7)], [1], is_promotion=True)) + assert all(r.success for r in drain(reader)) + assert torch.allclose(reader_tensor[1], expected) + finally: + reader.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 7500df8b4b94..661438dce633 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -21,6 +21,7 @@ from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingKVEventsConfig, OffloadKey, ReqContext, @@ -34,6 +35,10 @@ OffloadingParallelConfig, ) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -42,7 +47,17 @@ # --------------------------------------------------------------------------- -def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: +def _make_offloading_config( + enable_kv_cache_events: bool, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> OffloadingConfig: + if world_size is None: + world_size = tp_size return OffloadingConfig( groups=(), worker_kv_bytes_per_block=0, @@ -52,15 +67,16 @@ def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: model=OffloadingModelConfig(name="test/model", dtype="float16"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) @@ -208,12 +224,14 @@ def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: def _make_tier( num_blocks: int = 4, offloading_spec: SimpleNamespace = _OFFLOADING_SPEC, + primary_kv_view: memoryview | None = None, **tier_kwargs, ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: """Create a tier backed by a fresh MockNixlAgent.""" mock_agent = MockNixlAgent() - tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) - view = memoryview(tensor.numpy()) + if primary_kv_view is None: + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) with ( patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"), patch( @@ -223,7 +241,7 @@ def _make_tier( ): tier = ObjectStoreSecondaryTierManager( offloading_spec=offloading_spec, - primary_kv_view=view, + primary_kv_view=primary_kv_view, tier_type="obj", store_config=_STORE_CONFIG, prefix=_RUN_PREFIX, @@ -437,6 +455,105 @@ def register_once_fail(*a, **k): assert not by_id[1].success assert by_id[2].success + def test_release_xfer_failure_retries_without_losing_result(self, monkeypatch): + tier, agent = _make_tier(num_blocks=4) + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + + tier.submit_store(make_job(1, [key(1)], [0])) + + # The transfer handle could not be released safely, so the job must + # remain tracked and must not be finalized yet. + assert list(tier.get_finished_jobs()) == [] + assert 1 in tier._transfers + + # Cleanup is retried without polling again or changing the failure + # verdict. The completion is then returned exactly once. + results = list(tier.get_finished_jobs()) + assert len(results) == 1 + assert results[0].job_id == 1 + assert not results[0].success + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + @pytest.mark.parametrize( + "cleanup_method", ["release_dlist_handle", "deregister_memory"] + ) + def test_post_transfer_cleanup_failure_does_not_lose_result( + self, monkeypatch, cleanup_method + ): + tier, agent = _make_tier(num_blocks=4) + monkeypatch.setattr( + agent, + cleanup_method, + MagicMock(side_effect=RuntimeError("cleanup failed")), + ) + + tier.submit_store(make_job(1, [key(1)], [0])) + results = list(tier.get_finished_jobs()) + + assert len(results) == 1 + assert results[0].job_id == 1 + assert results[0].success + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch): + num_blocks = 4 + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) + mmap_region = MagicMock() + mmap_region.create_kv_memoryview.return_value = primary_kv_view + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=num_blocks, mmap_region=mmap_region + ) + obj_tier, agent = _make_tier( + num_blocks=num_blocks, primary_kv_view=primary_kv_view + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, secondary_tiers=[obj_tier] + ) + + keys = [key(1)] + primary_result = primary_tier.prepare_store(keys, _CTX) + assert primary_result is not None + primary_tier.complete_store(keys, _CTX, success=True) + job = manager.create_store_job(keys, _CTX) + obj_tier.submit_store(job) + + block = primary_tier._policy.get(keys[0]) + assert block is not None + assert block.ref_cnt == 1 + assert len(manager._transfer_jobs) == 1 + + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + schedule_context = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + + manager.on_schedule_end(schedule_context) + + assert len(manager._transfer_jobs) == 1 + assert block.ref_cnt == 1 + assert len(obj_tier._transfers) == 1 + assert manager.has_pending_work() + + manager.on_schedule_end(schedule_context) + + assert manager._transfer_jobs == {} + assert block.ref_cnt == 0 + assert obj_tier._transfers == {} + assert not manager.has_pending_work() + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + class TestMockObjTierShutdown: def test_shutdown_clears_in_flight_transfers(self): @@ -473,8 +590,7 @@ def test_successful_store_emits_stored_event(self): events = list(self.tier.take_events()) assert len(events) == 1 assert events[0].keys == keys - # Literal medium pins the wire contract, not just the constant choice. - assert events[0].medium == "OBJ" + assert events[0].medium == Medium.STORAGE assert events[0].locality is Locality.REMOTE assert not events[0].removed # take_events drains the buffer. @@ -617,3 +733,37 @@ def test_ca_bundle_included_when_set(self): params = cfg.to_nixl_params() assert params["ca_bundle"] == "/path/to/ca.pem" assert "access_key" not in params + + +def test_obj_tier_replicated_layout_collapses_mapper_identity(): + """TP=2 and TP=4 replicated configs share the obj FileMapper namespace.""" + tp2_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=2, world_size=2, rank=1, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp4_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp2_tier, _ = _make_tier(offloading_spec=tp2_spec) + tp4_tier, _ = _make_tier(offloading_spec=tp4_spec) + try: + assert tp2_tier._file_mapper.base_path == tp4_tier._file_mapper.base_path + assert tp2_tier._file_mapper.rank == 0 + assert tp4_tier._file_mapper.rank == 0 + assert tp2_tier._file_mapper.get_run_config() == ( + tp4_tier._file_mapper.get_run_config() + ) + finally: + tp2_tier.shutdown() + tp4_tier.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index cf546b304cd8..b19a270efb46 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -20,8 +20,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + _parse_tier_filter, +) from vllm.v1.kv_offload.base import ( + Locality, LookupResult, + Medium, OffloadingCounterMetadata, OffloadingEvent, OffloadKey, @@ -29,6 +34,8 @@ ReqContext, RequestOffloadingContext, ScheduleEndContext, + TierFilter, + TierMatcher, make_offload_key, ) from vllm.v1.kv_offload.tiering.base import ( @@ -245,9 +252,9 @@ def _start_request(self, req_context: ReqContext = _CTX): self.manager.on_new_request(req_context) def test_take_events_aggregates_tier_owned_events(self, manager_setup): - primary_event = OffloadingEvent(to_keys([1]), "CPU", removed=False) - secondary_event1 = OffloadingEvent(to_keys([2]), "tier-1", removed=False) - secondary_event2 = OffloadingEvent(to_keys([3]), "tier-2", removed=True) + primary_event = OffloadingEvent(to_keys([1]), Medium.CPU, removed=False) + secondary_event1 = OffloadingEvent(to_keys([2]), Medium.STORAGE, removed=False) + secondary_event2 = OffloadingEvent(to_keys([3]), Medium.STORAGE, removed=True) self.primary_tier.take_events = MagicMock(return_value=[primary_event]) self.secondary_tier1.take_events = MagicMock(return_value=[secondary_event1]) @@ -973,6 +980,56 @@ def test_reset_cache_drains_all_tiers(self, manager_setup): self.secondary_tier2.drain_jobs.assert_called_once() assert self.manager._transfer_jobs == {} + @pytest.mark.parametrize( + "load_tier_filter", + [ + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)), + TierFilter(matchers=()), + ], + ids=["non_matching_medium", "empty_no_load"], + ) + def test_tier_filter_skips_filtered_secondary( + self, manager_setup, load_tier_filter + ): + """Filter excluding secondary medium returns MISS from secondaries + even when they hold the block; primary is unaffected.""" + blocks = to_keys(range(2)) + # Put one block in primary, one only in secondary + self._start_request() + self.manager.prepare_store(blocks[:1], _CTX) + self.manager.complete_store(blocks[:1], _CTX, success=True) + self.secondary_tier1.blocks[blocks[1]] = True + + # Secondaries have medium=CPU, so load_tier_filter skips them. + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) + + ctx = ReqContext(req_id="r1", load_tier_filter=load_tier_filter) + assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT + assert self.manager.lookup(blocks[1], ctx) is LookupResult.MISS + self.secondary_tier1.lookup.assert_not_called() + + @pytest.mark.parametrize( + "load_tier_filter", + [ + TierFilter.ALL, + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)), + TierFilter(matchers=(TierMatcher(),)), + ], + ids=["all", "explicit_cpu", "unconstrained_matcher"], + ) + def test_tier_filter_allows_matching_secondary( + self, manager_setup, load_tier_filter + ): + """Filter that matches the secondary's medium allows lookup.""" + blocks = to_keys(range(1)) + self.secondary_tier1.blocks[blocks[0]] = True + + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) + + ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter) + assert self.manager.lookup(blocks[0], ctx) is LookupResult.RETRY + self.secondary_tier1.lookup.assert_called() + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" @@ -999,5 +1056,81 @@ def test_works_without_secondary_tiers(self): assert count_hits(manager, blocks) == 3 +@pytest.mark.parametrize( + "raw,expected", + [ + ( + [{"medium": "storage"}], + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)), + ), + ( + [{"medium": "CPU"}], + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)), + ), + ( + [{}], + TierFilter(matchers=(TierMatcher(),)), + ), + ( + [{"medium": "storage", "locality": "local"}], + TierFilter( + matchers=(TierMatcher(medium=Medium.STORAGE, locality=Locality.LOCAL),) + ), + ), + ( + [{"medium": "cpu"}, {"medium": "storage"}], + TierFilter( + matchers=( + TierMatcher(medium=Medium.CPU), + TierMatcher(medium=Medium.STORAGE), + ) + ), + ), + ( + [], + TierFilter(matchers=()), + ), + ], + ids=[ + "medium_storage", + "medium_cpu_uppercase", + "unconstrained", + "with_locality", + "multiple_matchers", + "empty_list_deny_all", + ], +) +def test_parse_tier_filter_valid(raw, expected): + assert _parse_tier_filter(raw) == expected + + +@pytest.mark.parametrize( + "raw", + [ + "not a list", + [{"medium": "unknown"}], + [{"locality": "nowhere"}], + ], + ids=["non_list", "invalid_medium", "invalid_locality"], +) +def test_parse_tier_filter_invalid_returns_all(raw): + assert _parse_tier_filter(raw) is TierFilter.ALL + + +def test_parse_tier_filter_skips_bad_entries(): + result = _parse_tier_filter( + [ + {"medium": "storage"}, + "not a dict", + {"medium": "bogus"}, + {"medium": "cpu"}, + ] + ) + assert result.matchers == ( + TierMatcher(medium=Medium.STORAGE), + TierMatcher(medium=Medium.CPU), + ) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index fc8ce50c05fa..f57ea285eb7b 100644 --- a/tests/v1/logits_processors/utils.py +++ b/tests/v1/logits_processors/utils.py @@ -11,6 +11,7 @@ from tests.utils import requires_spawn_multiprocessing from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.v1.sample.logits_processor import ( @@ -61,7 +62,7 @@ def validate_params(cls, params: SamplingParams): "target_token" ) if target_token is not None and not isinstance(target_token, int): - raise ValueError( + raise VLLMValidationError( f"target_token value {target_token} {type(target_token)} is not int" ) diff --git a/tests/v1/sample/test_head_dtype.py b/tests/v1/sample/test_head_dtype.py index 7531cdfc4501..ce07919577ea 100644 --- a/tests/v1/sample/test_head_dtype.py +++ b/tests/v1/sample/test_head_dtype.py @@ -14,6 +14,7 @@ from vllm import LLM, SamplingParams from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, UnquantizedEmbeddingMethod, ) @@ -28,6 +29,7 @@ def __init__( self.weight = weight self.quant_method = object() if quantized else UnquantizedEmbeddingMethod() self.shard_indices = shard_indices + self.tp_size = 1 def _build_processor(vocab_size: int) -> LogitsProcessor: @@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config): lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None) +def test_replicated_lm_head_skips_tp_communication_and_preserves_processing( + default_vllm_config, +): + from unittest import mock + + vocab_size, hidden_size = 12, 8 + soft_cap, scale = 2.0, 0.5 + lp = LogitsProcessor( + vocab_size, + soft_cap=soft_cap, + scale=scale, + ) + lp.head_dtype = torch.float32 + + hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16) + weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16) + world_size_getter = ( + "vllm.model_executor.layers.vocab_parallel_embedding." + "get_tensor_model_parallel_world_size" + ) + with mock.patch(world_size_getter, return_value=2): + lm_head = ParallelLMHead( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + disable_tp=True, + ) + lm_head.weight_loader(lm_head.weight, weight) + assert lm_head.tp_size == 1 + + with mock.patch.object(lp, "_gather_logits") as gather_mock: + logits = lp(lm_head, hidden_states) + + gather_mock.assert_not_called() + + expected = torch.nn.functional.linear(hidden_states.float(), weight.float()) + expected = torch.tanh(expected / soft_cap) * soft_cap * scale + torch.testing.assert_close(logits, expected) + + all_gather_path = ( + "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather" + ) + with mock.patch(all_gather_path) as all_gather: + top = lp.get_top_tokens(lm_head, hidden_states) + + all_gather.assert_not_called() + assert torch.equal(top, expected.argmax(dim=-1)) + + def test_get_top_tokens_honors_head_dtype(default_vllm_config): # The spec-decode local-argmax path (get_top_tokens) must run the lm_head # in head_dtype too, not just _get_logits. import types - from unittest import mock vocab_size, hidden_size = 64, 16 lp = _build_processor(vocab_size) @@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config): ), ) - with mock.patch( - "vllm.model_executor.layers.logits_processor." - "get_tensor_model_parallel_world_size", - return_value=1, - ): - top = lp.get_top_tokens(lm_head, hidden_states, None) - + top = lp.get_top_tokens(lm_head, hidden_states, None) expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax( dim=-1 ) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index aa17d2a1004c..d643b0b6fdec 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -405,7 +405,7 @@ def test_max_logprobs(): runner.generate(["Hello world"], sampling_params=vllm_sampling_params) bad_sampling_params = SamplingParams(logprobs=2) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): runner.generate(["Hello world"], sampling_params=bad_sampling_params) diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py index 56b93ea1e017..d385e96b7a21 100644 --- a/tests/v1/sample/test_sampling_params_e2e.py +++ b/tests/v1/sample/test_sampling_params_e2e.py @@ -4,6 +4,7 @@ import pytest from vllm import LLM, SamplingParams +from vllm.exceptions import VLLMValidationError MODEL = "hmellor/tiny-random-LlamaForCausalLM" PROMPT = "Hello my name is Robert and I" @@ -161,15 +162,15 @@ def test_allowed_token_ids(llm): assert output[0].outputs[0].token_ids[-1] == token_id # Reject empty allowed_token_ids. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[])) # Reject negative token id. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[-1])) # Reject out of vocabulary. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[10000000])) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 439ce0ea3409..55bbd1ec0c6d 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -5,13 +5,11 @@ from torch import Generator from tests.utils import large_gpu_mark -from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import set_random_seed from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p_pytorch, - flashinfer_sample, random_sample, ) from vllm.v1.sample.sampler import Sampler @@ -67,6 +65,10 @@ def test_sampler_threads_fp64_gumbel_to_topk_topp_sampler(): assert sampler.topk_topp_sampler.use_fp64_gumbel +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm aiter sampler test only runs on ROCm", +) def test_rocm_aiter_sampler_defers_import_when_generators_force_native( monkeypatch: pytest.MonkeyPatch, ): @@ -1045,39 +1047,3 @@ def _chi2_check(self, empirical, expected, chisquare_fn, *, label): f"{label}: distribution differs from theoretical: " f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}" ) - - -@pytest.mark.skipif( - not FLASHINFER_TOPK_TOPP_SUPPORTED, - reason="FlashInfer top-k/top-p sampler is not available on this platform.", -) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)]) -def test_flashinfer_sample_padded_vocab( - dtype: torch.dtype, k: int | None, p: float | None -): - """flashinfer_sample must accept the logits the sampler actually hands it. - - compute_logits slices the padding off the vocab, so for a vocab that isn't a - multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model - dtype, while FlashInfer requires contiguous fp32. - """ - torch.set_default_device(DEVICE_TYPE) - batch_size = 8 - org_vocab_size = 50272 - padded_vocab_size = pad_vocab_size(org_vocab_size) - assert padded_vocab_size != org_vocab_size - - logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[ - ..., :org_vocab_size - ] - # A single row stays contiguous despite the padded stride, hence batch_size > 1. - assert not logits.is_contiguous() - - token_ids = flashinfer_sample( - logits, - torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None, - torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None, - ) - assert token_ids.shape == (batch_size,) - assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size)) diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index 39386f3fd638..991ad6891856 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/tests/v1/shutdown/test_delete.py @@ -132,4 +132,13 @@ def test_llm_delete_inprocess( wait_for_gpu_memory_to_clear( devices=[0], threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES, + # Activate the helper's ROCm idle-runtime floor. VllmRunner already + # performed the stable-memory wait during context exit. + threshold_ratio=0.01 if current_platform.is_rocm() else None, + # Fail in the child before the outer pytest timeout expires. + timeout_s=( + SHUTDOWN_TEST_TIMEOUT_SEC // 2 + if current_platform.is_rocm() + else SHUTDOWN_TEST_TIMEOUT_SEC + ), ) diff --git a/tests/v1/shutdown/test_processor_error.py b/tests/v1/shutdown/test_processor_error.py index 0aba775e4308..665d4570976f 100644 --- a/tests/v1/shutdown/test_processor_error.py +++ b/tests/v1/shutdown/test_processor_error.py @@ -9,7 +9,7 @@ from tests.v1.shutdown.utils import SHUTDOWN_TEST_TIMEOUT_SEC from vllm import SamplingParams from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.inputs import TokensPrompt +from vllm.inputs import ExplicitEncoderDecoderPrompt from vllm.sampling_params import RequestOutputKind from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.exceptions import EngineGenerateError @@ -29,9 +29,14 @@ async def test_async_llm_processor_error(model: str) -> None: async_llm = AsyncLLM.from_engine_args(engine_args) async def generate(request_id: str): - # [] is not allowed and will raise a ValueError in Processor. + # An encoder/decoder prompt is rejected by the Processor for a + # decoder only model. generator = async_llm.generate( - TokensPrompt([]), request_id=request_id, sampling_params=SamplingParams() + ExplicitEncoderDecoderPrompt( + encoder_prompt="Hello my name is", decoder_prompt=None + ), + request_id=request_id, + sampling_params=SamplingParams(), ) try: async for _ in generator: diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py index 310b3b6db86d..02e2a0bdbec2 100644 --- a/tests/v1/spec_decode/test_dflash_causality.py +++ b/tests/v1/spec_decode/test_dflash_causality.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Config-only resolution of DFlash draft attention causality. +"""Config-only DFlash behavior. ``dflash_has_any_non_causal`` decides pre-build whether the draft needs a non-causal-capable backend, so its branch table (explicit override, SWA-derived @@ -13,8 +13,12 @@ from vllm.model_executor.models.qwen3_dflash import ( _dflash_layer_causal, + _get_dflash_fc_input_size, dflash_has_any_non_causal, ) +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) def _config(num_hidden_layers, layer_types=None, causal_override=None): @@ -53,3 +57,35 @@ def test_dflash_layer_causal_is_per_layer(): config = _config(2, layer_types=["sliding_attention", "full_attention"]) assert _dflash_layer_causal(config, 0) is True assert _dflash_layer_causal(config, 1) is False + + +def _vllm_config(**draft_config): + config = SimpleNamespace(**draft_config) + return SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=config) + ) + ) + + +def test_dflash_fc_uses_aux_layer_count(): + vllm_config = _vllm_config( + num_hidden_layers=5, + hidden_size=4096, + target_hidden_size=None, + target_layer_ids=[1, 17, 32], + ) + + assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096 + + +@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"]) +def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name): + layer_ids = [1, 17, 32] + vllm_config = _vllm_config( + **{config_name: {"layer_ids": layer_ids}}, + ) + + assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple( + layer_ids + ) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py new file mode 100644 index 000000000000..4ef99badaa39 --- /dev/null +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the EAGLE speculator's draft attention metadata builder. + +These tests guard the regression where ``_build_draft_attn_metadata`` did +not populate ``seq_lens_cpu_upper_bound`` on the per-step +``CommonAttentionMetadata``. Several downstream attention backends and +helpers (``split_decodes_prefills_and_extends``, the MLA indexer, +flex-attention, cross-attention) assert this field is non-None, so +omitting it caused crashes at the start of draft decode for certain +backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode): + + AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator +from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator + + +def _make_fake_speculator( + *, + max_num_reqs: int = 8, + max_num_tokens: int = 16, + max_model_len: int = 1024, + draft_max_seq_len: int = 1024, +) -> SimpleNamespace: + """Build a fake EagleSpeculator with just the attributes used by + ``_build_draft_attn_metadata``. We deliberately avoid constructing a + real ``EagleSpeculator`` because that requires a full ``VllmConfig`` + and a draft model. + """ + fake_input_buffers = SimpleNamespace( + query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32), + seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), + ) + fake_block_tables = SimpleNamespace( + input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)], + slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64), + ) + return SimpleNamespace( + arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"), + block_tables=fake_block_tables, + input_buffers=fake_input_buffers, + attn_groups=[], + kv_cache_config=SimpleNamespace(kv_cache_groups=[]), + max_model_len=max_model_len, + draft_max_seq_len=draft_max_seq_len, + ) + + +def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step): + captured: dict[str, object] = {} + + def fake_build_attn_metadata(**kwargs): + captured.update(kwargs) + return {} + + with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata): + EagleSpeculator._build_draft_attn_metadata( + fake, # type: ignore[arg-type] + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_tokens_padded=num_tokens_padded, + seq_lens_cpu_upper_bound=base, + step=step, + ) + return captured + + +def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): + """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None + ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus + the current draft-step offset. Padded entries are zeroed (matching the + main model runner's convention).""" + fake = _make_fake_speculator() + base = torch.tensor([100, 200, 300, 0], dtype=torch.int32) + + captured = _run_build( + fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2 + ) + + bound = captured["seq_lens_cpu_upper_bound"] + assert isinstance(bound, torch.Tensor), ( + "seq_lens_cpu_upper_bound must be a tensor, not None" + ) + assert bound.shape == (4,), ( + f"expected shape (num_reqs_padded=4,), got {bound.shape}" + ) + assert bound.device.type == "cpu" + assert bound.dtype == torch.int32 + # base[:num_reqs] + step, padded tail zeroed. + assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32)) + + +def test_build_draft_attn_metadata_handles_zero_unpadded_reqs(): + """Edge case: when ``num_reqs == 0`` the upper-bound tensor must + still be a valid all-zero tensor of length ``num_reqs_padded``.""" + fake = _make_fake_speculator() + base = torch.zeros(2, dtype=torch.int32) + + captured = _run_build( + fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1 + ) + + bound = captured["seq_lens_cpu_upper_bound"] + assert isinstance(bound, torch.Tensor) + assert bound.shape == (2,) + assert torch.equal(bound, torch.zeros(2, dtype=torch.int32)) + + +def test_build_draft_attn_metadata_clamps_to_max_model_len(): + """The per-request upper bound (target bound + step) is clamped to the + model length so it never exceeds the allocated KV range.""" + fake = _make_fake_speculator(max_model_len=1024) + base = torch.tensor([1023, 500], dtype=torch.int32) + + captured = _run_build( + fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3 + ) + + bound = captured["seq_lens_cpu_upper_bound"] + # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected. + assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32)) diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py new file mode 100644 index 000000000000..5510c36e60d2 --- /dev/null +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for SpecDecodeBaseProposer.initialize_attn_backend. + +Block tables are stored at kernel-block granularity, so the proposer's +``block_size`` (used for slot-mapping math) must be the kernel block size, +not the KV cache manager's block size — the two differ when manager blocks +are split for the attention kernel. The value must also be deterministic: +``_draft_attn_layer_names`` is a set, whose iteration order varies across +processes, so anything derived from iteration order must not leak into +``block_size``. +""" + +from types import SimpleNamespace + +import pytest + +import vllm.v1.spec_decode.llm_base_proposer as llm_base_proposer +from vllm.v1.spec_decode.eagle import EagleProposer + +SCHEDULER_BLOCK_SIZE = 256 +KERNEL_BLOCK_SIZE = 64 + + +class _FakeAttentionGroup: + def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id): + self.backend = backend + self.layer_names = list(layer_names) + self.kv_cache_spec = kv_cache_spec + self.kv_cache_group_id = kv_cache_group_id + self.kernel_block_size = None + + def create_metadata_builders(self, vllm_config, device, kernel_block_size=None): + self.kernel_block_size = kernel_block_size + + def get_metadata_builder(self): + return SimpleNamespace(kv_cache_spec=self.kv_cache_spec) + + +def _make_proposer( + monkeypatch: pytest.MonkeyPatch, layer_names: set[str] +) -> EagleProposer: + fake_layers = {} + for name in layer_names: + backend = SimpleNamespace(full_cls_name=lambda: "FakeBackend") + fake_layers[name] = SimpleNamespace( + get_attn_backend=lambda backend=backend: backend + ) + monkeypatch.setattr( + llm_base_proposer, "get_layers_from_vllm_config", lambda *a, **k: fake_layers + ) + monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup) + + proposer = EagleProposer.__new__(EagleProposer) + proposer.vllm_config = None + proposer.device = None + proposer._draft_attn_layer_names = set(layer_names) + proposer.kv_cache_gid = -1 + proposer.draft_attn_groups = [] + proposer.block_size = -1 + return proposer + + +def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace: + spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE) + group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec) + return SimpleNamespace(kv_cache_groups=[group]) + + +def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): + """The proposer's slot-mapping math runs against the kernel-granularity + block table, so block_size must come from kernel_block_sizes.""" + layer_names = {"draft.0.self_attn.attn"} + proposer = _make_proposer(monkeypatch, layer_names) + + proposer.initialize_attn_backend( + _make_kv_cache_config(layer_names), + kernel_block_sizes=[KERNEL_BLOCK_SIZE], + ) + + assert proposer.block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size != SCHEDULER_BLOCK_SIZE + # The metadata builder keeps receiving the kernel block size as well. + assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE + + +def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch): + layer_names = {"draft.0.self_attn.attn"} + proposer = _make_proposer(monkeypatch, layer_names) + + proposer.initialize_attn_backend( + _make_kv_cache_config(layer_names), kernel_block_sizes=None + ) + + assert proposer.block_size == SCHEDULER_BLOCK_SIZE + + +def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch): + """_draft_attn_layer_names is a set; the attention groups built from it + must not depend on its (process-random) iteration order.""" + layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"} + expected_order = sorted(layer_names) + + for insertion_order in (expected_order, expected_order[::-1]): + proposer = _make_proposer(monkeypatch, set(insertion_order)) + proposer.initialize_attn_backend( + _make_kv_cache_config(set(insertion_order)), + kernel_block_sizes=[KERNEL_BLOCK_SIZE], + ) + assert len(proposer.draft_attn_groups) == 1 + assert proposer.draft_attn_groups[0].layer_names == expected_order + assert proposer.block_size == KERNEL_BLOCK_SIZE diff --git a/tests/v1/spec_decode/test_llm_base_proposer_sampling.py b/tests/v1/spec_decode/test_llm_base_proposer_sampling.py index 9c7ec760ebb1..cb390b992f3a 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer_sampling.py +++ b/tests/v1/spec_decode/test_llm_base_proposer_sampling.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest import torch from vllm.platforms import current_platform @@ -8,6 +11,7 @@ from vllm.v1.sample.logits_processor import LogitsProcessors from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.llm_base_proposer import ( + SpecDecodeBaseProposer, compute_probs_and_sample_next_token, ) @@ -68,3 +72,21 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race(): assert torch.equal(actual_ids, expected_ids) assert torch.allclose(actual_probs, probs) + + +@pytest.mark.parametrize( + ("architecture", "expected"), + [ + ("DeepSeekMTPModel", True), + ("KimiK3MTPModel", True), + ("MiniMaxM3ForCausalLM", False), + ], +) +def test_mtp_model_returns_tuple(architecture: str, expected: bool): + proposer = object.__new__(SpecDecodeBaseProposer) + proposer.method = "mtp" + proposer.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=[architecture]) + ) + + assert proposer.model_returns_tuple() is expected diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 77c041d84a94..81f842419375 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -5,7 +5,7 @@ import pytest from tests.utils import get_attn_backend_list_based_on_platform -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.platforms import current_platform from vllm.sampling_params import StructuredOutputsParams @@ -18,10 +18,12 @@ @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) -def test_ngram_max_len(num_speculative_tokens: int): - llm = LLM( - model="facebook/opt-125m", +def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): + with vllm_runner( + "facebook/opt-125m", + trust_remote_code=False, max_model_len=100, + enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ "method": "ngram", @@ -29,21 +31,26 @@ def test_ngram_max_len(num_speculative_tokens: int): "prompt_lookup_min": 3, "num_speculative_tokens": num_speculative_tokens, }, - ) - sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) - llm.generate(_PROMPTS, sampling_params) + ) as runner: + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + runner.llm.generate(_PROMPTS, sampling_params) @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) def test_eagle_max_len( - monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str + monkeypatch: pytest.MonkeyPatch, + num_speculative_tokens: int, + attn_backend: str, + vllm_runner, ): if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") - llm = LLM( - model="meta-llama/Meta-Llama-3-8B-Instruct", + with vllm_runner( + "meta-llama/Meta-Llama-3-8B-Instruct", + trust_remote_code=False, + enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ "method": "eagle", @@ -53,31 +60,34 @@ def test_eagle_max_len( }, max_model_len=200, attention_config={"backend": attn_backend}, - ) - sampling_params = SamplingParams(max_tokens=200, ignore_eos=True) - outputs = llm.generate(_PROMPTS, sampling_params) - for o in outputs: - assert o.outputs[0].finish_reason == "length", ( - "This test is only meaningful if the output is truncated due to max length" - ) + ) as runner: + sampling_params = SamplingParams(max_tokens=200, ignore_eos=True) + outputs = runner.llm.generate(_PROMPTS, sampling_params) + for o in outputs: + assert o.outputs[0].finish_reason == "length", ( + "This test is only meaningful if the output is truncated " + "due to max length" + ) - sampling_params = SamplingParams( - max_tokens=200, - structured_outputs=StructuredOutputsParams(regex="^" + "a b c d e " * 15 + "$"), - ) - output = llm.generate(_PROMPTS, sampling_params) - for o in output: - assert o.prompt_token_ids is not None - assert ( - len(o.prompt_token_ids) - < 80 - < len(o.prompt_token_ids) + len(o.outputs[0].token_ids) - <= 200 - ), ( - "This test is only meaningful if the output " - "is longer than the eagle max length" + sampling_params = SamplingParams( + max_tokens=200, + structured_outputs=StructuredOutputsParams( + regex="^" + "a b c d e " * 15 + "$" + ), ) - assert o.outputs[0].text == "a b c d e " * 15 + output = runner.llm.generate(_PROMPTS, sampling_params) + for o in output: + assert o.prompt_token_ids is not None + assert ( + len(o.prompt_token_ids) + < 80 + < len(o.prompt_token_ids) + len(o.outputs[0].token_ids) + <= 200 + ), ( + "This test is only meaningful if the output " + "is longer than the eagle max length" + ) + assert o.outputs[0].text == "a b c d e " * 15 @pytest.mark.parametrize("spec_max_model_len", [80, 150]) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index bf9bea80bf72..982582ffaee8 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -412,3 +412,65 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int) f"Block verification mean accepted length {mean_block:.4f} is worse " f"than standard {mean_standard:.4f}." ) + + +@pytest.mark.parametrize("has_draft_logits", [True, False]) +def test_chunked_requests_match_full_batch(has_draft_logits: bool): + torch.manual_seed(7) + device = "cuda" + num_reqs = 5 + num_speculative_steps = 3 + vocab_size = 257 + + target_logits = torch.randn(vocab_size, device=device) + draft_logits = torch.randn(vocab_size, device=device) + inputs = _build_rejection_sample_inputs( + target_logits, + draft_logits, + num_speculative_steps, + temperature=0.6, + num_trials=num_reqs, + ) + padded_target_logits = torch.empty( + inputs["target_logits"].shape[0], vocab_size + 3, device=device + ) + padded_target_logits[:, :vocab_size].copy_(inputs["target_logits"]) + inputs["target_logits"] = padded_target_logits[:, :vocab_size] + assert inputs["target_logits"].stride(-1) == 1 + assert not inputs["target_logits"].is_contiguous() + if not has_draft_logits: + inputs["draft_logits"] = None + + sampled, num_sampled = rejection_sample( + **inputs, num_speculative_steps=num_speculative_steps + ) + + sampled_chunks = [] + num_sampled_chunks = [] + for start, end in ((0, 2), (2, 5)): + lo = start * (num_speculative_steps + 1) + hi = end * (num_speculative_steps + 1) + chunk_inputs = dict(inputs) + for name in ( + "target_logits", + "draft_sampled", + "pos", + "expanded_idx_mapping", + "expanded_local_pos", + ): + chunk_inputs[name] = inputs[name][lo:hi] + chunk_inputs["cu_num_logits"] = inputs["cu_num_logits"][start : end + 1] - lo + chunk_inputs["idx_mapping"] = inputs["idx_mapping"][start:end] + + chunk_sampled, chunk_num_sampled = rejection_sample( + **chunk_inputs, num_speculative_steps=num_speculative_steps + ) + sampled_chunks.append(chunk_sampled) + num_sampled_chunks.append(chunk_num_sampled) + + chunked_sampled = torch.cat(sampled_chunks) + chunked_num_sampled = torch.cat(num_sampled_chunks) + assert torch.equal(chunked_num_sampled, num_sampled) + steps = torch.arange(num_speculative_steps + 1, device=device) + valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1) + assert torch.equal(chunked_sampled[valid], sampled[valid]) diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index fd619610b767..a8f7221da81b 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -16,7 +16,8 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner -pytestmark = pytest.mark.cpu_test +# Not cpu_test: InputBatch allocates pinned (UVA) memory, which requires a +# CUDA device even though the batch tensors live on the CPU. @pytest.fixture @@ -28,6 +29,7 @@ def mock_model_runner_with_input_batch(): runner.requests = {} runner.max_num_reqs = 10 runner.max_model_len = 1024 + runner.late_interaction_runner = Mock() # Create a real InputBatch for e2e testing runner.input_batch = InputBatch( diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index 8fde0f117ca2..489284d6677b 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py @@ -16,7 +16,8 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.states import RequestState -pytestmark = pytest.mark.cpu_test +# Not cpu_test: RequestState allocates pinned (UVA) memory, which requires a +# CUDA device even though the request state itself lives on the CPU. @pytest.fixture @@ -31,16 +32,16 @@ def mock_model_runner_with_req_states(): num_speculative_steps=0, vocab_size=32000, device=torch.device("cpu"), - model_dtype=torch.float32, - cache_draft_logits=False, ) runner.encoder_cache = None runner.model_state = Mock() runner.block_tables = Mock() runner.lora_state = Mock() + runner.pp_handler = None runner.sampler = None runner.prompt_logprobs_worker = None runner.is_last_pp_rank = False + runner.pooling_runner = None # Mock staged writes — they use Triton kernels that require GPU runner.req_states.apply_staged_writes = Mock() diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 7d680895b836..822b4e49da34 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -53,6 +53,8 @@ def create_scheduler() -> Scheduler: vllm_config.model_config = MagicMock() vllm_config.model_config.skip_tokenizer_init = True vllm_config.model_config.is_multimodal_model = False + vllm_config.model_config.is_encoder_decoder = False + vllm_config.model_config.is_diffusion = False vllm_config.model_config.max_model_len = 1024 vllm_config.model_config.enable_return_routed_experts = False vllm_config.cache_config = MagicMock() @@ -496,7 +498,9 @@ def test_streaming_e2e_lifecycle(self): eco_cycle2 = eco_dict_cycle2[session.client_index].outputs[0] assert eco_cycle2.finish_reason == FinishReason.STOP assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ - assert session in scheduler.waiting + # Sessions paused for streaming input are blocked-waiting, so they + # live in the skipped_waiting queue rather than the main waiting queue. + assert session in scheduler.skipped_waiting assert session._all_token_ids == [1, 2, 3, 10, STOP_TOKEN] # CRITICAL ASSERTION: Cached prompt_token_ids STILL must not have changed diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index 861e919c102a..ad5f1d5d7951 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -73,6 +73,7 @@ def mock_request_with_structured_output(self): request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8] request.num_computed_tokens = 5 request.num_output_placeholders = 0 + request.request_id = "mock_req" return request @pytest.fixture @@ -208,53 +209,144 @@ def test_should_advance_reasoning_just_ended( mock_request_with_structured_output ) - # Should set reasoning_ended to True but return False for this step + # The scheduler trims the reasoning prefix before advancing the grammar. assert ( mock_request_with_structured_output.structured_output_request.reasoning_ended is True ) - assert result is False + assert result is True - def test_should_advance_reasoning_just_ended_with_spec_decode_structural_tag( + def test_should_advance_reasoning_already_ended( self, manager_with_reasoner, mock_request_with_structured_output, ): - """When reasoning ends this step, advance immediately for structural - tags with speculative decoding.""" + """Test should_advance when reasoning has already ended.""" + # Set reasoning as already ended + ( + mock_request_with_structured_output.structured_output_request + ).reasoning_ended = True + + result = manager_with_reasoner.should_advance( + mock_request_with_structured_output + ) + + # Should return True since reasoning has ended + assert result is True + + def test_should_advance_uses_new_token_ids_when_provided( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Regression for #43388: when caller passes new_token_ids, the + reasoner sees the exact multi-token delta rather than the + placeholder-derived window. + """ structured_req = mock_request_with_structured_output.structured_output_request structured_req.reasoning_ended = False - structured_req.structured_output_key = ( - StructuredOutputOptions.STRUCTURAL_TAG, - "{}", - ) + + end_token_id = 248069 + reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = True + # Detection mirrors the real Qwen3 parser: end token in the delta. + reasoner.is_reasoning_end_streaming = Mock( + side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids) + ) structured_req.reasoner = reasoner - manager_with_reasoner.vllm_config.speculative_config = Mock() + # Scenario from #43388: async + spec decode K=4, 4 tokens accepted + # but only 1 placeholder remains (some drafts were rejected). + # The placeholder math would yield delta=[271] and miss . + # Passing new_token_ids must override that. + new_token_ids = [9, 198, end_token_id, 271] + mock_request_with_structured_output.all_token_ids = [ + 1, + 2, + 3, + 4, + 5, + ] + new_token_ids + mock_request_with_structured_output.num_computed_tokens = 9 + mock_request_with_structured_output.num_output_placeholders = 1 result = manager_with_reasoner.should_advance( - mock_request_with_structured_output + mock_request_with_structured_output, + new_token_ids=new_token_ids, ) + # First call to is_reasoning_end_streaming was with the full + # new_token_ids (not the truncated placeholder window). + first_call = reasoner.is_reasoning_end_streaming.call_args_list[0] + _, called_delta = first_call.args + assert list(called_delta) == new_token_ids + assert structured_req.reasoning_ended is True assert result is True - def test_should_advance_reasoning_already_ended( + def test_should_advance_without_new_token_ids_falls_back( self, manager_with_reasoner, mock_request_with_structured_output, ): - """Test should_advance when reasoning has already ended.""" - # Set reasoning as already ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = True + """Backward compat: callers that don't pass new_token_ids keep + the original placeholder-derived delta window. + """ + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + reasoner = MockReasoner(tokenizer=Mock()) + reasoner.is_reasoning_end_streaming.return_value = False + structured_req.reasoner = reasoner + + mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] + mock_request_with_structured_output.num_computed_tokens = 5 + mock_request_with_structured_output.num_output_placeholders = 2 result = manager_with_reasoner.should_advance( mock_request_with_structured_output ) - # Should return True since reasoning has ended + # placeholder window: start = 5 - 2 = 3, delta = [4, 5] + _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0] + assert list(called_delta) == [4, 5] + assert result is False + + def test_should_advance_trims_reasoning_prefix_for_json( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """JSON uses the common trim-then-advance path at the boundary.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + structured_req.structured_output_key = ( + StructuredOutputOptions.JSON_OBJECT, + "{}", + ) + + marker = 248069 + + class MarkerReasoner: + def __init__(self, *_, **__): + pass + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return marker in list(delta_ids) + + structured_req.reasoner = MarkerReasoner() + + new_token_ids = [9, 198, marker, 271, 5005] + mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids + + result = manager_with_reasoner.should_advance( + mock_request_with_structured_output, + new_token_ids=new_token_ids, + ) + + structured_req.grammar.accept_tokens.assert_not_called() + assert structured_req.reasoning_ended is True assert result is True + assert structured_req.reasoning_end_token_index == 5 + assert manager_with_reasoner.trim_reasoning_for_advance( + mock_request_with_structured_output, new_token_ids + ) == [271, 5005] diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 31ce961ff616..7ea60cc6609b 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,6 +5,7 @@ import pytest from vllm.config import StructuredOutputsConfig +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams, StructuredOutputsParams pytestmark = pytest.mark.cpu_test @@ -32,7 +33,7 @@ def test_structured_outputs_rejected_for_diffusion_models(): params = SamplingParams( structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) ) - with pytest.raises(ValueError, match="not yet supported for diffusion"): + with pytest.raises(VLLMValidationError, match="not yet supported for diffusion"): params._validate_structured_outputs( _StubModelConfig(is_diffusion=True), StructuredOutputsConfig(), @@ -63,7 +64,7 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match): rejected at request validation (-> 400) instead of reaching and crashing the engine.""" params = SamplingParams(structured_outputs=structured_outputs) - with pytest.raises(ValueError, match=match): + with pytest.raises(VLLMValidationError, match=match): params._validate_structured_outputs( _StubModelConfig(is_diffusion=False), StructuredOutputsConfig(), diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 89d551e344cf..4696eefa29fe 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -2,7 +2,32 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest import TestCase -from vllm.v1.outputs import LogprobsLists +import torch + +from vllm.v1.outputs import LogprobsLists, LogprobsTensors + + +def test_logprobs_tensors_cat(): + first = LogprobsTensors( + torch.tensor([[1, 2]]), + torch.tensor([[0.1, 0.2]]), + torch.tensor([1]), + ) + second = LogprobsTensors( + torch.tensor([[3, 4]]), + torch.tensor([[0.3, 0.4]]), + torch.tensor([2]), + ) + + result = LogprobsTensors.cat([first, second], [0, 1, 2]) + + assert result.logprob_token_ids.tolist() == [[1, 2], [3, 4]] + assert result.logprobs.tolist() == ( + first.logprobs.tolist() + second.logprobs.tolist() + ) + assert result.selected_token_ranks.tolist() == [1, 2] + assert result.cu_num_generated_tokens == [0, 1, 2] + assert LogprobsTensors.cat([first]) is first class TestLogprobsLists(TestCase): diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 4ed8724e60fb..9f7761c22b5d 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc(): assert torch.allclose(decoded.prompt_embeds, original_tensor), ( f"Value mismatch for sender {sender_idx} msg {msg_idx}" ) + + +def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): + """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends + it: many requests, each with per-token tensors small enough that pyzmq + copies their frames, while the accumulated payload frame is large enough + that pyzmq sends it zero-copy.""" + from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs + from vllm.v1.outputs import LogprobsTensors + + outputs = [] + for req in range(num_reqs): + num_tokens = num_prompt_tokens + req % 4 + outputs.append( + EngineCoreOutput( + request_id=f"req-{req:08d}", + new_token_ids=[req], + new_prompt_logprobs_tensors=LogprobsTensors( + logprob_token_ids=torch.arange( + num_tokens * 2, dtype=torch.int64 + ).view(num_tokens, 2), + logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32), + selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32), + ), + ) + ) + return EngineCoreOutputs(outputs=outputs) + + +def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages(): + """The engine core recycles the msgpack payload buffer across messages + (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished + sending that buffer, otherwise a newer payload is delivered alongside the + older message's zero-copy tensor frames. + + `Socket.send_multipart(track=True)` cannot be used to detect this: it + returns a tracker for the last frame only, and pyzmq copies frames below + `zmq.COPY_THRESHOLD` and reports them as already-sent. + """ + import zmq + + from vllm.v1.engine import EngineCoreOutputs + from vllm.v1.engine.core import EngineCoreProc + + num_msgs = 100 + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(EngineCoreOutputs) + # Enough requests that the payload frame is zero-copied rather than copied + # by pyzmq, which is what makes early reuse observable. + messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)] + assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD + + reuse_buffers: list[bytearray] = [] + pending: list[tuple[zmq.MessageTracker, bytearray]] = [] + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-payload-reuse") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-payload-reuse") + + for outputs in messages: + while pending and pending[0][0].done: + reuse_buffers.append(pending.pop(0)[1]) + buffer = reuse_buffers.pop() if reuse_buffers else bytearray() + buffers = encoder.encode_into(outputs, buffer) + tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers) + if tracker.done: + reuse_buffers.append(buffer) + else: + pending.append((tracker, buffer)) + + for i, sent in enumerate(messages): + received = decoder.decode(pull.recv_multipart(copy=False)) + assert len(received.outputs) == len(sent.outputs), f"message {i}" + for expected, actual in zip(sent.outputs, received.outputs): + sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids + got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids + assert actual.request_id == expected.request_id, f"message {i}" + assert torch.equal(got_ids, sent_ids), ( + f"message {i} request {actual.request_id}: corrupted " + f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}" + ) + push.close(linger=0) + pull.close(linger=0) + + +def test_zero_copy_frames_survive_without_caller_side_references(): + """Callers don't need to retain the encoded object until zmq has sent it: + for a zero-copy frame, zmq holds its own reference to the backing buffer. + + The engine core clients rely on this when sending requests that carry + tensors (e.g. prompt embeds) without tracking the messages. + + What makes that safe is that `tensor_data()` hands zmq a memoryview which + transitively references the source tensor, so refcounting - not timing - + keeps the memory from being freed and reused underneath zmq. + """ + import gc + + import zmq + + from vllm.v1.utils import tensor_data + + num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD + expected = torch.arange(num_elems, dtype=torch.int64) + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(RequestWithTensor) + + # The buffer handed to zmq must keep the tensor's storage alive by itself. + holder = tensor_data(expected).obj + while getattr(holder, "base", None) is not None: + holder = holder.base + assert isinstance(holder, torch.Tensor) + assert holder.data_ptr() == expected.data_ptr() + + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-zero-copy-lifetime") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-zero-copy-lifetime") + + request = RequestWithTensor(prompt_embeds=expected.clone(), data="req") + buffers = encoder.encode(request) + assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD + push.send_multipart(buffers, copy=False) + + # Drop every reference the sender holds, then churn the allocator. + del request, buffers + gc.collect() + torch.arange(num_elems * 4, dtype=torch.int64) + + decoded = decoder.decode(pull.recv_multipart(copy=False)) + assert decoded.prompt_embeds is not None + assert torch.equal(decoded.prompt_embeds, expected) + push.close(linger=0) + pull.close(linger=0) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 31acd475adec..365e365a574c 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -130,3 +130,66 @@ def test_block_tables_apply_staged_writes_single_group(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) + + +def test_v1_block_table_move_row_clears_vacated_row(): + """condense() moves the last row into a freed slot; the vacated row must + not keep stale block ids. Padded dummy-run batches dereference stale rows + as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode + padding) and write state in place there — corrupting the blocks' new + owner once they are reallocated, e.g. to an in-flight NIXL load.""" + from vllm.v1.worker.block_table import BlockTable + + block_table = BlockTable( + block_size=16, + max_num_reqs=4, + max_num_blocks_per_req=8, + max_num_batched_tokens=64, + pin_memory=False, + device=torch.device("cuda"), + kernel_block_size=16, + cp_kv_cache_interleave_size=1, + ) + block_table.add_row([7, 8, 9], row_idx=0) + block_table.add_row([4, 5], row_idx=1) + + block_table.move_row(1, 0) + + assert block_table.block_table.np[0, :2].tolist() == [4, 5] + assert block_table.num_blocks_per_row[0] == 2 + # The vacated source row routes to the reserved null block. + assert block_table.num_blocks_per_row[1] == 0 + assert (block_table.block_table.np[1] == 0).all() + + +def test_get_dummy_block_tables_returns_zeroed_rows(): + """Dummy runs bypass the gather, so the persistent input_block_tables + hold the previous real step's rows. Mamba/GDN metadata routes in-place + state writes through block_table[:, 0] (dummy slot mappings are + PAD-filled, state indices are not), so stale rows would direct dummy + state writes at freed — possibly reallocated — blocks. + get_dummy_block_tables must hand out zeroed (null block) rows while + preserving the persistent storage address for CUDA graphs.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=4, + max_num_batched_tokens=64, + max_num_blocks_per_group=[8], + device=device, + kernel_block_sizes=[16], + ) + # Simulate a real step: stage a request's blocks and gather them into + # the persistent input block tables. + block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) + torch.accelerator.synchronize() + assert block_tables.input_block_tables[0][0, 0].item() == 1 + + dummy = block_tables.get_dummy_block_tables(num_reqs=1) + torch.accelerator.synchronize() + assert (dummy[0] == 0).all() + # CUDA graph invariant: same persistent tensor, not a fresh allocation. + assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py new file mode 100644 index 000000000000..f517abb2b14d --- /dev/null +++ b/tests/v1/worker/test_gpu_input_batch_v2.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch).""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers + +DEVICE = current_platform.device_type + + +@pytest.mark.parametrize( + "num_reqs,num_tokens", + [ + (256, 496), # remainder 240: previously gave the last request 241 tokens + (128, 512), # no remainder + (3, 8), + (1, 7), + ], +) +def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int): + """No dummy request may exceed ceil(num_tokens / num_reqs) tokens. + + Dumping the remainder on a single request can produce a dummy request with + seq_len > max_model_len, which the block tables cannot back; attention + kernels running on the dummy batch during cudagraph capture then read + block-table entries out of bounds (https://github.com/vllm-project/vllm/pull/49364 + CI failure). + """ + buffers = InputBuffers( + max_num_reqs=num_reqs, max_num_tokens=num_tokens, device=torch.device(DEVICE) + ) + batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers) + + max_per_req = -(-num_tokens // num_reqs) + assert batch.num_scheduled_tokens.sum() == num_tokens + assert batch.num_scheduled_tokens.max() == max_per_req + assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs + # Requests with an extra token are placed at the end of the batch. + assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all() + + # seq_len == query_len for the dummy prefill-shaped batch, on GPU and CPU. + query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1] + assert (query_lens == batch.num_scheduled_tokens).all() + assert torch.equal( + batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens).to(DEVICE) + ) + assert batch.query_start_loc_np[-1] == num_tokens + assert torch.equal( + batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np) + ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 57931669c548..282c02295817 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -46,6 +46,11 @@ from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.worker.block_table import ( + MultiGroupBlockTable, + SlotMappingMode, + get_block_table_width, +) from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras @@ -831,7 +836,7 @@ def test_update_config(model_runner): model_runner.update_config({"load_config": {"load_format": "dummy"}}) assert model_runner.load_config.load_format == "dummy" # Raise error on non-existing config - with pytest.raises(AssertionError): + with pytest.raises(ValueError, match="do_not_exist_config"): model_runner.update_config({"do_not_exist_config": "dummy"}) @@ -1380,6 +1385,29 @@ def test_hybrid_block_table_initialization(): ) +def test_get_block_table_width_aligns_to_128_tokens(): + assert get_block_table_width(1875, 64) == 1876 + + +def test_get_block_table_width_splits_virtual_blocks(): + assert get_block_table_width(235, 256, 64) == 940 + + +def test_mamba_state_table_width_is_not_aligned(): + block_tables = MultiGroupBlockTable( + max_num_reqs=1, + max_num_batched_tokens=1, + pin_memory=False, + device=torch.device("cpu"), + block_sizes=[39664], + kernel_block_sizes=[39664], + max_num_blocks=[1], + slot_mapping_modes=[SlotMappingMode.NONE], + ) + + assert block_tables[0].max_num_blocks_per_req == 1 + + def test_input_batch_with_kernel_block_sizes(): """Test InputBatch initialization with kernel_block_sizes parameter.""" max_num_reqs = 10 @@ -1669,3 +1697,59 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): with pytest.raises(ValueError, match="max_num_seqs"): runner.initialize_kv_cache(kv_cache_config) + + +class TestInitFp8KvScalesHybridModels: + """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries. + + Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors + rather than a single tensor. init_fp8_kv_scales must iterate both forms. + """ + + @staticmethod + def _make_runner_stub(kv_caches): + runner = Mock(spec=GPUModelRunner) + runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3") + runner.kv_caches = kv_caches + runner.compilation_config = SimpleNamespace(static_forward_context={}) + runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__( + runner, GPUModelRunner + ) + return runner + + def test_zeroes_both_tensor_and_list_entries(self): + single_tensor = torch.ones(4, 8) + list_tensors = [torch.ones(2, 4), torch.ones(3, 6)] + + runner = self._make_runner_stub([single_tensor, list_tensors]) + runner.init_fp8_kv_scales() + + assert (single_tensor == 0).all() + assert all((t == 0).all() for t in list_tensors) + + def test_skips_none_entries(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([None, tensor, None]) + runner.init_fp8_kv_scales() + + assert (tensor == 0).all() + + def test_noop_when_kv_cache_not_quantized(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([tensor]) + runner.cache_config.cache_dtype = "auto" + runner.init_fp8_kv_scales() + + assert (tensor == 1).all() + + def test_mixed_none_tensor_and_list(self): + t1 = torch.ones(2, 2) + t2 = torch.ones(3, 3) + list_entry = [torch.ones(1, 1), torch.ones(1, 1)] + + runner = self._make_runner_stub([None, t1, list_entry, None, t2]) + runner.init_fp8_kv_scales() + + assert (t1 == 0).all() + assert (t2 == 0).all() + assert all((t == 0).all() for t in list_entry) diff --git a/tests/v1/worker/test_gpu_model_runner_mm_gather.py b/tests/v1/worker/test_gpu_model_runner_mm_gather.py index 586acd504638..a86e77c8f52f 100644 --- a/tests/v1/worker/test_gpu_model_runner_mm_gather.py +++ b/tests/v1/worker/test_gpu_model_runner_mm_gather.py @@ -46,6 +46,7 @@ def _gather(features, cached, *, num_scheduled, shift, num_computed=0): input_batch=SimpleNamespace(req_ids=["req0"]), requests={"req0": req_state}, encoder_cache=encoder_cache, + _get_encoder_output_from_cache=lambda mm_hash: encoder_cache.get(mm_hash), is_multimodal_pruning_enabled=False, uses_mrope=False, ) diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py new file mode 100644 index 000000000000..22cb71e67b8d --- /dev/null +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import MethodType, SimpleNamespace +from typing import get_args + +import numpy as np +import pytest +import torch + +from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( + RejectionSampler, + _iter_request_chunks, +) + + +def test_iter_request_chunks_preserves_request_boundaries(): + cu_num_logits = np.array([0, 3, 4, 11, 13], dtype=np.int32) + + assert list(_iter_request_chunks(cu_num_logits, max_chunk_logits=5)) == [ + (0, 2), + (2, 3), + (3, 4), + ] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) +def test_chunked_scores_match_full_batch(logprobs_mode: str): + device = torch.device("cuda") + cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32) + num_logits_per_req = np.diff(cu_num_logits_np) + idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32) + input_batch = SimpleNamespace( + num_reqs=4, + cu_num_logits_np=cu_num_logits_np, + cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device), + idx_mapping_np=idx_mapping_np, + idx_mapping=torch.from_numpy(idx_mapping_np).to(device), + expanded_idx_mapping=torch.from_numpy( + np.repeat(idx_mapping_np, num_logits_per_req) + ).to(device), + expanded_local_pos=torch.from_numpy( + np.concatenate( + [np.arange(count, dtype=np.int32) for count in num_logits_per_req] + ) + ).to(device), + ) + rejection_sampler = object.__new__(RejectionSampler) + rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode) + rejection_sampler.num_speculative_steps = 3 + + def fake_verify( + self, + logits, + _draft_logits, + _draft_sampled, + _pos, + cu_num_logits, + idx_mapping, + *_mappings, + ): + num_sampled = torch.diff(cu_num_logits).to(torch.int32) + sampled = ( + idx_mapping.to(torch.int64).unsqueeze(1) + torch.arange(4, device=device) + ) % logits.shape[1] + return logits.float() + 1, sampled, num_sampled + + rejection_sampler._verify = MethodType(fake_verify, rejection_sampler) + logits = torch.arange(170, dtype=torch.float32, device=device).view(10, 17) + + sampled, num_sampled, chunked_logprobs = rejection_sampler._verify_in_chunks( + logits, + input_batch, + draft_logits=None, + draft_sampled=torch.arange(10, device=device), + pos=torch.arange(10, device=device), + max_chunk_logits=5, + max_num_logprobs=2, + ) + score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits + full_logprobs = rejection_sampler._get_logprobs_tensors( + sampled, + num_sampled, + score_logits, + input_batch.cu_num_logits, + input_batch.cu_num_logits_np, + max_num_logprobs=2, + ) + + assert sampled[:, 0].tolist() == idx_mapping_np.tolist() + assert num_sampled.tolist() == num_logits_per_req.tolist() + assert chunked_logprobs is not None + assert full_logprobs is not None + assert torch.equal( + chunked_logprobs.logprob_token_ids, + full_logprobs.logprob_token_ids, + ) + assert torch.equal(chunked_logprobs.logprobs, full_logprobs.logprobs) + assert torch.equal( + chunked_logprobs.selected_token_ranks, + full_logprobs.selected_token_ranks, + ) + assert ( + chunked_logprobs.cu_num_generated_tokens + == full_logprobs.cu_num_generated_tokens + ) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index cdaa644b62ec..43232ca6be40 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -6,122 +6,13 @@ import pytest -import vllm.v1.worker.gpu_worker as gpu_worker_module -from vllm.multimodal.video import ( - PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, - PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, - PYNVVIDEOCODEC_VIDEO_BACKEND, -) from vllm.utils.mem_constants import GiB_bytes from vllm.v1.worker import startup_plan -from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.startup_plan import ( maybe_apply_startup_plan, maybe_save_startup_plan, ) - -def _worker_with_mm_config( - mm_config: SimpleNamespace, - *, - api_process_count: int = 1, -) -> Worker: - worker = object.__new__(Worker) - worker.model_config = SimpleNamespace(multimodal_config=mm_config) - worker.parallel_config = SimpleNamespace(_api_process_count=api_process_count) - return worker - - -def _mm_config( - *, - mm_ipc_gpu_memory_gb: float = 0, - video_backend: str | None = None, -) -> SimpleNamespace: - video_kwargs = {} if video_backend is None else {"video_backend": video_backend} - return SimpleNamespace( - mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, - media_io_kwargs={"video": video_kwargs} if video_kwargs else {}, - ) - - -def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int: - return api_process_count * ( - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS - + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES - ) - - -@pytest.mark.parametrize("video_backend", [None, "opencv"]) -def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( - monkeypatch: pytest.MonkeyPatch, - video_backend: str | None, -): - monkeypatch.setattr( - gpu_worker_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - "opencv", - ) - worker = _worker_with_mm_config( - _mm_config(mm_ipc_gpu_memory_gb=0.25, video_backend=video_backend) - ) - - assert worker._reserve_mm_ipc_gpu_memory(GiB_bytes) == int(0.75 * GiB_bytes) - - -def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr( - gpu_worker_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - "opencv", - ) - worker = _worker_with_mm_config( - _mm_config( - mm_ipc_gpu_memory_gb=0.25, - video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, - ) - ) - available_bytes = 4 * GiB_bytes - - assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( - available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget() - ) - - -def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr( - gpu_worker_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - PYNVVIDEOCODEC_VIDEO_BACKEND, - ) - worker = _worker_with_mm_config(_mm_config()) - available_bytes = 4 * GiB_bytes - - assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( - available_bytes - _pynvvideocodec_decoder_budget() - ) - - -def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr( - gpu_worker_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - PYNVVIDEOCODEC_VIDEO_BACKEND, - ) - worker = _worker_with_mm_config(_mm_config(), api_process_count=3) - available_bytes = 8 * GiB_bytes - - assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( - available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) - ) - - # Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and # saved by Worker.determine_available_memory / compile_or_warm_up_model. diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index aeb727d9ce32..6a97d64c6be0 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -9,6 +9,7 @@ import pytest +from vllm.config import VllmConfig, get_current_vllm_config from vllm.v1.worker.gpu_worker import Worker @@ -21,29 +22,55 @@ def __init__(self, raise_on_update: bool = False): self.finished = False self.reset_count = 0 self.update_calls: list[dict] = [] + self.seen_configs: list[VllmConfig] = [] + + def _record_config(self) -> None: + self.seen_configs.append(get_current_vllm_config()) def start_weight_update(self) -> None: + self._record_config() self.started = True def update_weights(self, update_info: dict) -> None: + self._record_config() self.update_calls.append(update_info) if self.raise_on_update: raise ValueError("boom") def finish_weight_update(self) -> None: + self._record_config() self.finished = True def reset_weight_update_target(self) -> None: self.reset_count += 1 +class _RecordingModelRunner: + def __init__(self) -> None: + self.seen_config: VllmConfig | None = None + + def reload_weights(self) -> None: + self.seen_config = get_current_vllm_config() + + def _make_worker(engine: _RecordingEngine | None) -> Worker: worker = object.__new__(Worker) + worker.vllm_config = VllmConfig() worker.weight_transfer_engine = engine worker._weight_update_active = False return worker +def test_reload_weights_sets_current_config(): + worker = _make_worker(None) + model_runner = _RecordingModelRunner() + worker.model_runner = model_runner # type: ignore[assignment] + + Worker.reload_weights(worker) + + assert model_runner.seen_config is worker.vllm_config + + def test_start_update_finish_delegates_to_engine(): engine = _RecordingEngine() worker = _make_worker(engine) @@ -60,6 +87,7 @@ def test_start_update_finish_delegates_to_engine(): assert engine.finished is True assert engine.reset_count == 1 assert worker._weight_update_active is False + assert engine.seen_configs == [worker.vllm_config] * 3 def test_double_start_raises(): diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 8f15229912d9..17aa1bf38d47 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.worker.utils import KVBlockZeroer +from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -14,25 +14,23 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight(): page_size_el = 4 storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device) - # Build the minimal zeroer state directly so the test can focus on ID-buffer - # lifetime without constructing model attention groups. + # Build the minimal zeroer state directly so the test can focus on the + # in-flight copy behavior without constructing model attention groups. zeroer = KVBlockZeroer.__new__(KVBlockZeroer) zeroer.device = device - zeroer.pin_memory = True - zeroer.max_concurrency = 2 - zeroer._id_cap = 8 - zeroer._allocate_id_buffers() zeroer._meta = ( torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), - page_size_el, - page_size_el, - 1, + torch.tensor([page_size_el], dtype=torch.int64, device=device), + page_size_el // page_size_el, # max_chunks = 1 + page_size_el, # blk_size + 1, # n_segs ) stream = torch.cuda.Stream() with torch.cuda.stream(stream): # Keep the first nonblocking H2D copy pending while the host submits the - # second call. A single shared pinned source would be overwritten here. + # second call. Each call must stage from its own pinned source so the + # first copy is not corrupted before it runs. torch.cuda._sleep(10_000_000) zeroer.zero_block_ids([1]) zeroer.zero_block_ids([2]) @@ -42,3 +40,113 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight(): assert torch.all(storage[1] == 0) assert torch.all(storage[2] == 0) assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_non_uniform_page_sizes(): + """Two segments with different page sizes (e.g. MLA + DSA indexer).""" + device = torch.device("cuda") + num_blocks = 4 + page_size_a = 10496 # int32 elements + page_size_b = 2112 + + storage_a = torch.ones((num_blocks, page_size_a), dtype=torch.int32, device=device) + storage_b = torch.ones((num_blocks, page_size_b), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + + seg_page_sizes = [page_size_a, page_size_b] + max_ps = max(seg_page_sizes) + + def largest_power_of_2_divisor(n): + return n & -n + + blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024) + + zeroer._meta = ( + torch.tensor( + [storage_a.data_ptr(), storage_b.data_ptr()], + dtype=torch.uint64, + device=device, + ), + torch.tensor(seg_page_sizes, dtype=torch.int64, device=device), + max_ps // blk_size, + blk_size, + 2, + ) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + zeroer.zero_block_ids([1, 2]) + stream.synchronize() + + for storage in (storage_a, storage_b): + assert torch.all(storage[0] == 1) + assert torch.all(storage[1] == 0) + assert torch.all(storage[2] == 0) + assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_compiles_every_n_blocks_specialization(): + """After warmup, no launch should trigger a first-request JIT compile. + + ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must + cover every block count. + """ + device = torch.device("cuda") + num_blocks = 64 + page_size_el = 4 + storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, # max_chunks + page_size_el, # blk_size + 1, # n_segs + ) + + def compiled_variants() -> set: + return { + key + for caches in _zero_kv_blocks_kernel.device_caches.values() + for key in caches[0] + } + + zeroer.warmup(num_blocks) + torch.accelerator.synchronize() + warmed = compiled_variants() + assert warmed + + for n_blocks in (1, 2, 3, 16, 32): + zeroer.zero_block_ids(list(range(n_blocks))) + torch.accelerator.synchronize() + + assert compiled_variants() == warmed + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_respects_available_block_count(): + """An empty KV cache must not be warmed with out-of-range block IDs.""" + device = torch.device("cuda") + page_size_el = 4 + storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, + page_size_el, + 1, + ) + + zeroer.warmup(0) + torch.accelerator.synchronize() + + assert torch.all(storage == 1) diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index b4eabe18ca1d..04bc3ef27680 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -11,6 +11,38 @@ Step 2 is necessary for multi-node deployment. All scripts accept a positional argument as workspace path for staging the build, defaulting to `$(pwd)/ep_kernels_workspace`. +## NCCL version requirement (CUDA 13+) + +DeepEPv2 uses the NCCL GIN (GPU-Initiated Networking) backend, which requires +NCCL >= 2.30.4 at both compile time and runtime. PyTorch 2.11 pins +`nvidia-nccl-cu13==2.28.9` as a transitive dependency, so you need to +override it. + +**With uv** (recommended): + +```bash +# Create an override file +echo "nvidia-nccl-cu13>=2.30.4" > /tmp/nccl-override.txt +export UV_OVERRIDE=/tmp/nccl-override.txt + +# All subsequent uv pip install commands will respect the override +uv pip install vllm +``` + +**With pip**: + +```bash +pip install vllm +pip install "nvidia-nccl-cu13>=2.30.4" --no-deps +``` + +The override / reinstall must happen before building DeepEP (for GIN device +headers) and must remain in place at runtime. You can verify with: + +```bash +python -c "from vllm.utils.import_utils import has_deep_ep_v2; print(has_deep_ep_v2())" +``` + ## Usage ```bash diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 94beef897d08..5f5a597bacab 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -170,6 +170,48 @@ do_build() { sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py" fi + # DeepEPv2 requires Linux 5.6+ at runtime for pidfd_getfd (pidfd_open was + # added in Linux 5.3), but manylinux headers predate both definitions. + # DeepEP is built as a separate wheel for the vLLM container image and is + # not included in the vLLM wheel, so this does not change its manylinux ABI. + if [[ "$name" == "DeepEP" ]] && \ + ! grep -q "vLLM manylinux syscall compatibility" \ + csrc/kernels/backend/symmetric.hpp; then + sed -i '1i\ +// vLLM manylinux syscall compatibility\ +#if defined(__x86_64__) || defined(__aarch64__)\ +#ifndef SYS_pidfd_open\ +#ifdef __NR_pidfd_open\ +#define SYS_pidfd_open __NR_pidfd_open\ +#else\ +#define SYS_pidfd_open 434\ +#endif\ +#endif\ +#ifndef SYS_pidfd_getfd\ +#ifdef __NR_pidfd_getfd\ +#define SYS_pidfd_getfd __NR_pidfd_getfd\ +#else\ +#define SYS_pidfd_getfd 438\ +#endif\ +#endif\ +#endif' csrc/kernels/backend/symmetric.hpp + fi + + if [[ "$name" == "DeepEP" ]]; then + # DeepEP links against the CUDA driver API in driverless build images. + local cuda_driver_stub + local cuda_driver_stub_dir + cuda_driver_stub=$( + find -H "$CUDA_HOME" -path "*/stubs/libcuda.so" -print -quit + ) + if [[ -z "$cuda_driver_stub" ]]; then + echo "CUDA driver stub not found under $CUDA_HOME" >&2 + exit 1 + fi + cuda_driver_stub_dir=$(dirname "$cuda_driver_stub") + export LIBRARY_PATH="${cuda_driver_stub_dir}${LIBRARY_PATH:+:$LIBRARY_PATH}" + fi + if [ "$MODE" = "install" ]; then echo "Installing $name into environment" eval "$extra_env" uv pip install --no-build-isolation -vvv . diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 52f73e9050f7..3f80ab567302 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,9 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This is currently targeting nv-dev branch due to sm120 support -DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a" +DEEPGEMM_GIT_REF="e21c821f39a2056d68067a466c64ddc942200106" WHEEL_DIR="" # Parse command line arguments diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index a2fc173f0357..52a95ce1d8d9 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -6,6 +6,13 @@ import regex as re +# Hub entry points that must go through the vLLM-tagged repo_utils helpers. +_HF_NAMES = ( + r"HfApi|HfFileSystem|hf_hub_download|snapshot_download" + r"|list_repo_files|file_exists|try_to_load_from_cache" + r"|list_repo_refs|repo_exists" +) + @dataclass class ForbiddenImport: @@ -13,6 +20,7 @@ class ForbiddenImport: tip: str allowed_pattern: re.Pattern = re.compile(r"^$") # matches nothing by default allowed_files: set[str] = field(default_factory=set) + allowed_dirs: set[str] = field(default_factory=set) CHECK_IMPORTS = { @@ -40,6 +48,7 @@ class ForbiddenImport: "vllm/distributed/device_communicators/shm_object_storage.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", + "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", "vllm/utils/hashing.py", "tests/multimodal/media/test_base.py", @@ -83,6 +92,23 @@ class ForbiddenImport: ), allowed_files={"vllm/triton_utils/importing.py"}, ), + "huggingface_hub repo API": ForbiddenImport( + # Catch `from huggingface_hub import `, including parenthesized, + # multi-line imports. + pattern=( + r"^\s*from\s+huggingface_hub\s+import\s*\([^)]*\b(?:" + _HF_NAMES + r")\b" + r"|" + r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b" + ), + tip=( + "Use the shared, vLLM-tagged helpers from " + "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), " + "list_repo_files, file_exists) instead of calling " + "huggingface_hub directly." + ), + allowed_files={"vllm/transformers_utils/repo_utils.py"}, + allowed_dirs={"examples/"}, + ), } @@ -95,11 +121,18 @@ def check_file(path: str) -> int: # Skip files that are allowed for this import if path in forbidden_import.allowed_files: continue + # Skip directories that are allowed for this import + if any(path.startswith(prefix) for prefix in forbidden_import.allowed_dirs): + continue # Search for forbidden imports for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE): # Check if it's allowed if forbidden_import.allowed_pattern.match(match.group()): continue + # Skip matches inside a comment + line_start = content.rfind("\n", 0, match.start()) + 1 + if "#" in content[line_start : match.start()]: + continue # Calculate line number from match position line_num = content[: match.start() + 1].count("\n") + 1 print( @@ -119,7 +152,10 @@ def main(): def test_regex(): - test_cases = [ + def matches(rule: str, content: str) -> bool: + return bool(re.search(CHECK_IMPORTS[rule].pattern, content, re.MULTILINE)) + + pickle_cases = [ # Should match ("import pickle", True), ("import cloudpickle", True), @@ -139,11 +175,47 @@ def test_regex(): ("print('import pickle')", False), ("import pickleas as asdf", False), ] - for i, (line, should_match) in enumerate(test_cases): - result = bool(CHECK_IMPORTS["pickle/cloudpickle"].pattern.match(line)) + for i, (content, should_match) in enumerate(pickle_cases): + result = matches("pickle/cloudpickle", content) assert result == should_match, ( - f"Test case {i} failed: '{line}' (expected {should_match}, got {result})" + f"pickle case {i} failed: {content!r} " + f"(expected {should_match}, got {result})" ) + + hf_cases = [ + # Should match + ("from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import hf_hub_download", True), + ("from huggingface_hub import HfApi", True), + ("from huggingface_hub import HfFileSystem", True), + ("from huggingface_hub import list_repo_files", True), + ("from huggingface_hub import try_to_load_from_cache", True), + (" from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True), + ("from huggingface_hub import (snapshot_download)", True), + # Parenthesized multi-line import must not bypass the hook + ("from huggingface_hub import (\n snapshot_download,\n)", True), + ( + "from huggingface_hub import (\n PyTorchModelHubMixin,\n HfApi,\n)", + True, + ), + # Should not match + ("import huggingface_hub", False), + ("import huggingface_hub as hf", False), + ("from huggingface_hub import PyTorchModelHubMixin", False), + ("from huggingface_hub.constants import HF_HUB_CACHE", False), + ("from huggingface_hub.utils import EntryNotFoundError", False), + ("from vllm.transformers_utils.repo_utils import hf_api", False), + ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False), + ("# from huggingface_hub import snapshot_download", False), + ] + for i, (content, should_match) in enumerate(hf_cases): + result = matches("huggingface_hub repo API", content) + assert result == should_match, ( + f"huggingface_hub case {i} failed: {content!r} " + f"(expected {should_match}, got {result})" + ) + print("All regex tests passed.") diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index c52e0e7ce630..585585e3520f 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -21,34 +21,55 @@ import regex as re +# Paths verified clean under follow_imports="silent". Matched before +# SEPARATE_GROUPS, so these files join the default group and are checked at the +# stricter setting even while a parent directory remains in SEPARATE_GROUPS. +# +# Fixing a directory means moving it from SEPARATE_GROUPS to here. Without that +# move the fixes are not enforced because "tests" claims every file below it. +SILENT_GROUPS = [ + "tests/compile/correctness_e2e", + "tests/compile/fullgraph", + "tests/compile/fusions_e2e", + "tests/config", + "tests/entrypoints/generate", + "tests/entrypoints/tool_parsers", + "tests/entrypoints/weight_transfer", + "tests/kernels/core", + "tests/kernels/mamba", + "tests/models/language", + "tests/models/quantization", + "tests/plugins/bge_m3_sparse_plugin", + "tests/plugins/prithvi_io_processor_plugin", + "tests/plugins/vllm_add_dummy_platform", + "tests/plugins/vllm_add_dummy_stat_logger", + "tests/plugins_tests/gguf", + "tests/plugins_tests/lora_resolvers", + "tests/spec_decode", + "tests/transformers_utils", + "tests/v1/distributed", + "tests/v1/shutdown", +] + # After fixing errors resulting from changing follow_imports -# from "skip" to "silent", remove its directory from SEPARATE_GROUPS. +# from "skip" to "silent", move its directory to SILENT_GROUPS. SEPARATE_GROUPS = [ "tests", "tests/benchmarks", - "tests/compile/correctness_e2e", - "tests/config", "tests/compile", - "tests/compile/fullgraph", - "tests/compile/fusions_e2e", "tests/compile/passes", "tests/distributed", "tests/entrypoints/anthropic", - "tests/entrypoints/generate", "tests/entrypoints/llm", "tests/entrypoints/multimodal", "tests/entrypoints/openai", "tests/entrypoints/pooling", "tests/entrypoints/serve", "tests/entrypoints/speech_to_text", - "tests/entrypoints/tool_parsers", "tests/entrypoints/unit_tests", - "tests/entrypoints/weight_transfer", "tests/kernels", "tests/kernels/attention", - "tests/kernels/core", "tests/kernels/helion", - "tests/kernels/mamba", "tests/kernels/moe", "tests/kernels/quantization", "tests/lora", @@ -57,34 +78,23 @@ "tests/model_executor/model_loader", "tests/models", "tests/models/test_initialization.py", - "tests/models/language", "tests/models/multimodal", - "tests/models/quantization", "tests/multimodal", "tests/parser", - "tests/plugins_tests/gguf", - "tests/plugins_tests/lora_resolvers", - "tests/plugins/bge_m3_sparse_plugin", - "tests/plugins/prithvi_io_processor_plugin", - "tests/plugins/vllm_add_dummy_platform", - "tests/plugins/vllm_add_dummy_stat_logger", "tests/plugins_tests", "tests/quantization", "tests/reasoning", "tests/renderers", "tests/samplers", - "tests/spec_decode", "tests/tokenizers_", "tests/tool_parsers", "tests/tool_use", - "tests/transformers_utils", "tests/utils_", "tests/v1", "tests/v1/attention", "tests/v1/core", "tests/v1/cudagraph", "tests/v1/determinism", - "tests/v1/distributed", "tests/v1/e2e", "tests/v1/ec_connector", "tests/v1/engine", @@ -94,7 +104,6 @@ "tests/v1/logits_processors", "tests/v1/metrics", "tests/v1/sample", - "tests/v1/shutdown", "tests/v1/simple_kv_offload", "tests/v1/spec_decode", "tests/v1/streaming_input", @@ -142,14 +151,22 @@ def group_files(changed_files: list[str]) -> dict[str, list[str]]: A dictionary mapping file group names to lists of changed files. """ exclude_pattern = re.compile(f"^{'|'.join(EXCLUDE)}.*") - file_groups = {"": []} + silent_pattern = re.compile(f"^({'|'.join(SILENT_GROUPS)}).*") + file_groups: dict[str, list[str]] = {"": []} file_groups.update({k: [] for k in SEPARATE_GROUPS}) + # Longest path first so a sub-directory is not shadowed by its parent + separate_groups = sorted(SEPARATE_GROUPS, key=len, reverse=True) for changed_file in changed_files: # Skip files which should be ignored completely if exclude_pattern.match(changed_file): continue + # Already-fixed paths go in the default group, which runs at the + # stricter follow_imports setting from pyproject.toml + if silent_pattern.match(changed_file): + file_groups[""].append(changed_file) + continue # Group files by mypy call - for directory in SEPARATE_GROUPS: + for directory in separate_groups: if re.match(f"^{directory}.*", changed_file): file_groups[directory].append(changed_file) break diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 8ab06d446bf4..30bfcb80b017 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -52,7 +52,7 @@ def is_aiter_found() -> bool: def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. - Checks: platform (ROCm), device arch (gfx9), and library existence. + Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence. Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled(). This function determines if aiter CAN be used, not if it SHOULD be used. @@ -66,9 +66,9 @@ def is_aiter_found_and_supported() -> bool: VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery. """ if current_platform.is_rocm() and IS_AITER_FOUND: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return on_mi3xx() + return get_cdna_version() > 2 return False @@ -133,6 +133,8 @@ def _rocm_aiter_fused_moe_impl( bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe @@ -143,6 +145,12 @@ def _rocm_aiter_fused_moe_impl( extra_kwargs: dict = {} if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): extra_kwargs["gate_mode"] = gate_mode + if ( + getattr(ActivationType, "Situv2", None) is not None + and activation == ActivationType.Situv2 + ): + extra_kwargs["beta"] = beta + extra_kwargs["linear_beta"] = linear_beta return fused_moe( hidden_states, @@ -193,6 +201,8 @@ def _rocm_aiter_fused_moe_fake( bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: if output_dtype is not None: return torch.empty_like(hidden_states, dtype=output_dtype) @@ -409,37 +419,6 @@ def _rocm_aiter_fused_topk_fake( # Cache whether aiter supports FP8 MLA parameters _AITER_MLA_SUPPORTS_FP8: bool | None = None -_AITER_HAS_FUSED_QK_RMSNORM: bool | None = None - - -def check_aiter_fused_qk_rmsnorm() -> bool: - """Check if aiter provides fused_qk_rmsnorm. - - Supports both the new private name ``_fused_qk_rmsnorm`` - (AITER >= PR #2958) and the old public name ``fused_qk_rmsnorm`` - (AITER >= PR #2442). - - TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once - AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207). - """ - global _AITER_HAS_FUSED_QK_RMSNORM - if _AITER_HAS_FUSED_QK_RMSNORM is None: - try: - from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401 - _fused_qk_rmsnorm, - ) - - _AITER_HAS_FUSED_QK_RMSNORM = True - except (ImportError, ModuleNotFoundError, AttributeError): - try: - from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401 - fused_qk_rmsnorm, - ) - - _AITER_HAS_FUSED_QK_RMSNORM = True - except (ImportError, ModuleNotFoundError, AttributeError): - _AITER_HAS_FUSED_QK_RMSNORM = False - return _AITER_HAS_FUSED_QK_RMSNORM def _check_aiter_mla_fp8_support() -> bool: @@ -708,6 +687,32 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake( return Y +def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + from aiter import gemm_a8w8_blockscale_bpreshuffle + + # B preshuffled (shuffle_weight (16,16)); As column-major group scale. + return gemm_a8w8_blockscale_bpreshuffle(A, B, As, Bs, dtype=output_dtype) + + +def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + m = A.shape[0] + n = B.shape[0] + Y = torch.empty(m, n, dtype=output_dtype, device=A.device) + return Y + + def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl( x: torch.Tensor, residual: torch.Tensor, @@ -795,7 +800,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] - token_num = input_.shape[0] + token_num = input_.numel() // hidden_dim if input_.dtype in (torch.bfloat16, torch.float16): pack_size = 16 // input_.element_size() hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 @@ -855,7 +860,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] - token_num = input_.shape[0] + token_num = input_.numel() // hidden_dim if input_.dtype in (torch.bfloat16, torch.float16): pack_size = 16 // input_.element_size() hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 @@ -1166,17 +1171,21 @@ def _rocm_aiter_fused_rms_gated_fp8_group_quant_fake( def _rocm_aiter_group_fp8_quant_impl( x: torch.Tensor, group_size: int, + transpose_scale: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size" from aiter import QuantType, get_hip_quant aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128) - return aiter_per1x128_quant(x.contiguous(), quant_dtype=FP8_DTYPE) + return aiter_per1x128_quant( + x.contiguous(), quant_dtype=FP8_DTYPE, transpose_scale=transpose_scale + ) def _rocm_aiter_group_fp8_quant_fake( x: torch.Tensor, group_size: int, + transpose_scale: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: M, N = x.shape x_fp8 = torch.empty((M, N), dtype=FP8_DTYPE, device=x.device) @@ -1267,43 +1276,17 @@ def _fused_mla_dual_rms_norm_impl( x1_epsilon: float, x2_epsilon: float, ) -> tuple[torch.Tensor, torch.Tensor]: - try: - import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops - except (ImportError, ModuleNotFoundError, AttributeError) as exc: - raise ImportError( - "fused_qk_rmsnorm requires AITer >= PR #2442. " - "Please upgrade aiter or disable the " - "fuse_mla_dual_rms_norm pass." - ) from exc - - if hasattr(aiter_ops, "_fused_qk_rmsnorm"): - return aiter_ops._fused_qk_rmsnorm( - q_out=None, - q=x1, - q_weight=x1_weight, - q_eps=x1_epsilon, - k_out=None, - k=x2, - k_weight=x2_weight, - k_eps=x2_epsilon, - ) - - # TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once - # AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207). - if hasattr(aiter_ops, "fused_qk_rmsnorm"): - return aiter_ops.fused_qk_rmsnorm( - q=x1, - q_weight=x1_weight, - q_eps=x1_epsilon, - k=x2, - k_weight=x2_weight, - k_eps=x2_epsilon, - ) - - raise ImportError( - "fused_qk_rmsnorm requires AITer >= PR #2442. " - "Please upgrade aiter or disable the " - "fuse_mla_dual_rms_norm pass." + from aiter.ops.fused_qk_norm_rope_cache_quant import _fused_qk_rmsnorm + + return _fused_qk_rmsnorm( + q_out=None, + q=x1, + q_weight=x1_weight, + q_eps=x1_epsilon, + k_out=None, + k=x2, + k_weight=x2_weight, + k_eps=x2_epsilon, ) @@ -1461,6 +1444,65 @@ def _triton_rotary_embedding_fake( return +def _rocm_aiter_fp8_attn_impl( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_descale: torch.Tensor, + k_descale: torch.Tensor, + v_descale: torch.Tensor, + batch_size: int, + output_dtype: torch.dtype, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, +) -> torch.Tensor: + """Run AITER FP8 attention for fixed or packed inputs.""" + from aiter import flash_attn_varlen_fp8_pertensor_func + + q_len = q.size(1) + if cu_seqlens is None: + cu_seqlens = torch.arange( + 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device + ) + max_seqlen_value = q_len if max_seqlen is None else max_seqlen.item() + + q, k, v = (x.flatten(0, 1) for x in (q, k, v)) + output = flash_attn_varlen_fp8_pertensor_func( + q, + k, + v, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen_value, + max_seqlen_k=max_seqlen_value, + causal=False, + softmax_scale=scale, + ) + return output.to(output_dtype).reshape(batch_size, q_len, *output.shape[1:]) + + +def _rocm_aiter_fp8_attn_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_descale: torch.Tensor, + k_descale: torch.Tensor, + v_descale: torch.Tensor, + batch_size: int, + output_dtype: torch.dtype, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, +) -> torch.Tensor: + return torch.empty( + (*q.shape[:-1], v.shape[-1]), device=q.device, dtype=output_dtype + ) + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -1616,6 +1658,7 @@ def get_aiter_activation_type(activation_str: str): "silu": ActivationType.Silu, "gelu": ActivationType.Gelu, "swiglu": ActivationType.Swiglu, + "situ": getattr(ActivationType, "Situv2", None), } return mapping.get(name) @@ -1755,16 +1798,21 @@ def is_fp8bmm_enabled(cls) -> bool: @classmethod @if_aiter_supported def is_fp4bmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import get_cdna_version - return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() + # TODO GFX1250: Enable for cdna 4+ when aiter supports batched_gemm_a16wfp4 on gfx1250 + return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4 @classmethod @if_aiter_supported def is_linear_hipbmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED + return ( + cls.is_linear_enabled() + and (get_cdna_version() > 2) + and cls._LINEAR_HIPBMM_ENABLED + ) @classmethod @if_aiter_supported @@ -1942,6 +1990,12 @@ def register_ops_once() -> None: fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake, ) + direct_register_custom_op( + op_name="rocm_aiter_gemm_a8w8_blockscale_bpreshuffle", + op_func=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl, + fake_impl=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake, + ) + direct_register_custom_op( op_name="rocm_aiter_rmsnorm_fused_dynamic_quant", op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl, @@ -2016,6 +2070,14 @@ def register_ops_once() -> None: dispatch_key=current_platform.dispatch_key, ) + direct_register_custom_op( + op_name="aiter_fp8_attn_wrapper", + op_func=_rocm_aiter_fp8_attn_impl, + mutates_args=[], + fake_impl=_rocm_aiter_fp8_attn_fake, + dispatch_key=current_platform.dispatch_key, + ) + direct_register_custom_op( op_name="rocm_aiter_gemm_a8wfp4", op_func=_rocm_aiter_gemm_a8wfp4_impl, @@ -2188,6 +2250,18 @@ def gemm_a8w8_blockscale( A, B, As, Bs, output_dtype ) + @staticmethod + def gemm_a8w8_blockscale_bpreshuffle( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype = torch.float16, + ) -> torch.Tensor: + return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale_bpreshuffle( + A, B, As, Bs, output_dtype + ) + @staticmethod def fused_moe( hidden_states: torch.Tensor, @@ -2212,6 +2286,8 @@ def fused_moe( bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: return torch.ops.vllm.rocm_aiter_fused_moe( hidden_states, @@ -2236,6 +2312,8 @@ def fused_moe( bias2, moe_sorting_dispatch_policy, swiglu_limit, + beta, + linear_beta, ) @staticmethod @@ -2677,9 +2755,12 @@ def triton_fp8_bmm( def group_fp8_quant( input_2d: torch.Tensor, group_size: int = 128, + transpose_scale: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: assert group_size == 128, "Group size must be 128" - return torch.ops.vllm.rocm_aiter_group_fp8_quant(input_2d, group_size) + return torch.ops.vllm.rocm_aiter_group_fp8_quant( + input_2d, group_size, transpose_scale + ) @staticmethod def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool: @@ -2887,6 +2968,34 @@ def flash_attn_varlen_func( sink_ptr=sink_ptr, ) + @staticmethod + def fp8_attn_wrapper( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_descale: torch.Tensor, + k_descale: torch.Tensor, + v_descale: torch.Tensor, + batch_size: int, + output_dtype: torch.dtype, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, + ) -> torch.Tensor: + return torch.ops.vllm.aiter_fp8_attn_wrapper( + q, + k, + v, + q_descale, + k_descale, + v_descale, + batch_size, + output_dtype, + scale, + cu_seqlens, + max_seqlen, + ) + @staticmethod def pa_fwd_asm( Q: torch.Tensor, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 599cac0ed6f2..8c1cc105e2de 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1022,7 +1022,7 @@ def cutlass_fp4_moe_mm( An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs the gemms for each combination based on the specified problem sizes. - This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward. + This is used as the MoE gemm during NVFP4 Quantized MoERunner forward. - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized input and expert weights. - a_/b_scales: The blockscales in FP8-E4M3 precision @@ -2070,6 +2070,93 @@ def selective_scan_fwd( ) +def causal_conv1d_update_cpu_vec( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = 0, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_update_cpu_vec( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + query_start_loc, + pad_slot_id, + ) + + +def selective_state_update_cpu( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + dt_bias: torch.Tensor | None, + dt_softplus: bool, + state_batch_indices: torch.Tensor | None, + dst_state_batch_indices: torch.Tensor | None, + null_block_id: int, + out: torch.Tensor, + num_accepted_tokens: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +): + torch.ops._C.selective_state_update_cpu( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + ) + + +def mamba_chunk_scan_fwd_cpu( + out: torch.Tensor, + final_states: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + cu_seqlens: torch.Tensor, +) -> None: + """Prefill SSM scan kernel. out and final_states are written in-place.""" + torch.ops._C.mamba_chunk_scan_fwd_cpu( + out, + final_states, + x, + dt, + A, + B, + C, + D, + z, + cu_seqlens, + ) + + # ROCm skinny gemms def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor: return torch.ops._rocm_C.LLMM1(a, b, rows_per_block) @@ -2302,6 +2389,7 @@ def topk_softmax( gating_output: torch.Tensor, renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_softmax( topk_weights, @@ -2310,6 +2398,7 @@ def topk_softmax( gating_output, renormalize, e_score_correction_bias, + is_padding, ) @@ -2321,6 +2410,7 @@ def topk_sigmoid( renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_sigmoid( topk_weights, @@ -2330,6 +2420,7 @@ def topk_sigmoid( renormalize, e_score_correction_bias, routed_scaling_factor, + is_padding, ) @@ -2343,6 +2434,7 @@ def topk_hash_softplus_sqrt( e_score_correction_bias: torch.Tensor | None = None, input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_softplus_sqrt( topk_weights, @@ -2354,6 +2446,7 @@ def topk_hash_softplus_sqrt( e_score_correction_bias, input_tokens, hash_indices_table, + is_padding, ) @@ -2562,6 +2655,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert( index_q_out: torch.Tensor | None = None, kv_cache_dtype: str = "auto", skip_index_branch: bool = False, + q_fp8_out: torch.Tensor | None = None, + q_fp8_scale: float = 1.0, ) -> None: """Fused MiniMax-M3 attention pre-processing (in-place). @@ -2584,6 +2679,9 @@ def fused_minimax_m3_qknorm_rope_kv_insert( callers skip a separate ``.contiguous()`` copy before the SM100 sparse attention's flat TMA descriptor. + If ``q_fp8_out`` is given, the same normalized q is also written in FP8 + E4M3 using ``q_fp8_scale`` as its dequantization scale. + When ``skip_index_branch`` is true, sparse rows still keep their packed ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v branches and main KV cache. This is used by MiniMax-M3 index-topk reuse @@ -2611,9 +2709,57 @@ def fused_minimax_m3_qknorm_rope_kv_insert( index_q_out, kv_cache_dtype, skip_index_branch, + q_fp8_out, + q_fp8_scale, ) +def fused_kda_decode( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_indices: torch.Tensor, + state: torch.Tensor, + out: torch.Tensor | None = None, + lower_bound: float | None = None, + output_gate: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-5, +) -> torch.Tensor: + if out is None: + out = torch.empty( + 1, + x.shape[0], + raw_g.shape[2], + raw_g.shape[3], + dtype=x.dtype, + device=x.device, + ) + torch.ops._C.fused_kda_decode( + x, + weight, + bias, + conv_state, + raw_g, + raw_beta, + A_log, + dt_bias, + state_indices, + state, + out, + lower_bound, + output_gate, + norm_weight, + norm_eps, + ) + return out + + def concat_and_cache_mla( kv_c: torch.Tensor, k_pe: torch.Tensor, @@ -2627,6 +2773,53 @@ def concat_and_cache_mla( ) +def concat_and_cache_mla_grouped( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache_ptrs: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int, + block_stride: int, + entry_stride: int, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla_grouped( + kv_c, + k_pe, + kv_cache_ptrs, + slot_mapping, + block_size, + block_stride, + entry_stride, + ) + + +def kimi_k3_attn_res( + prefix: torch.Tensor, + delta: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor, + num_blocks: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + output = torch.empty_like(prefix) + torch.ops._C.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + num_blocks, + eps, + output_norm_eps, + ) + return output + + def concat_and_cache_mla_rope_fused( positions: torch.Tensor, q_pe: torch.Tensor, @@ -2894,6 +3087,63 @@ def all_reduce( torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) +def custom_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_all_gather( + fa, inp, out, reg_buffer, reg_buffer_sz_bytes + ) + + +def mnnvl_lamport_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + multicast_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_all_gather( + fa, + inp, + out, + local_buffer, + multicast_buffer, + epoch_buffer, + stage_sz_bytes, + ) + + +def custom_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_reduce_scatter( + fa, inp, out, reg_buffer, reg_buffer_sz_bytes + ) + + +def mnnvl_lamport_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter( + fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes + ) + + def dispose(fa: int) -> None: torch.ops._C_custom_ar.dispose(fa) @@ -2998,18 +3248,18 @@ def dsv3_fused_a_gemm( output: torch.Tensor, mat_a: torch.Tensor, mat_b: torch.Tensor, + enable_pdl: bool = False, ) -> None: - """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens). - Computes output = mat_a @ mat_b.T where: - mat_a: [num_tokens, 7168] row-major bf16 (hidden states) - mat_b: [7168, 2112] column-major bf16 (weight transposed) - output: [num_tokens, 2112] row-major bf16 + Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and + DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major; + ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits + programmatic dependent launch for callers that have validated it. - Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes. - Requires SM 9.0+ (Hopper). + Requires SM 9.0+. """ - torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b) + torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl) if hasattr(torch.ops._C, "weight_packed_linear"): @@ -3248,6 +3498,7 @@ def chunk_gated_delta_rule_cpu( cu_seqlens: torch.Tensor, head_first: bool, use_qk_l2norm_in_kernel: bool, + initial_state_indices: torch.Tensor, eps: float = 1e-5, ) -> tuple[torch.Tensor, torch.Tensor]: return torch.ops._C.chunk_gated_delta_rule_cpu( @@ -3261,6 +3512,7 @@ def chunk_gated_delta_rule_cpu( cu_seqlens, head_first, use_qk_l2norm_in_kernel, + initial_state_indices, eps, ) @@ -3386,6 +3638,7 @@ def causal_conv1d_update_cpu( silu_activation: bool, conv_state_indices: torch.Tensor | None, is_vnni: bool, + num_accepted_tokens: torch.Tensor | None = None, ) -> torch.Tensor: return torch.ops._C.causal_conv1d_update_cpu( x, @@ -3393,7 +3646,7 @@ def causal_conv1d_update_cpu( weight, bias, silu_activation, - None, + num_accepted_tokens, conv_state_indices, -1, is_vnni, @@ -3544,6 +3797,7 @@ def cpu_attn_get_scheduler_metadata( isa: str, enable_kv_split: bool, dynamic_causal: torch.Tensor | None = None, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: scheduler_metadata = torch.ops._C.get_scheduler_metadata( num_reqs, @@ -3558,6 +3812,7 @@ def cpu_attn_get_scheduler_metadata( isa, enable_kv_split, dynamic_causal, + kv_cache_dtype, ) return scheduler_metadata @@ -3668,6 +3923,15 @@ def cpu_prepack_moe_weight( return output +def cpu_prepack_moe_weight_int8( + weight: torch.Tensor, + isa: str, +) -> torch.Tensor: + output = torch.empty_like(weight) + torch.ops._C.prepack_moe_weight_int8(weight, output, isa) + return output + + def cpu_fused_moe( input: torch.Tensor, w13: torch.Tensor, @@ -3697,6 +3961,39 @@ def cpu_fused_moe( return output +def cpu_fused_moe_int8( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + act: str, + isa: str, + skip_weighted: bool = False, +) -> torch.Tensor: + output = torch.empty_like(input) + torch.ops._C.cpu_fused_moe_int8( + output, + input, + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + skip_weighted, + act, + isa, + ) + return output + + if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"): @register_fake("_qutlass_C::matmul_mxf4_bf16_tn") @@ -3776,10 +4073,10 @@ def fusedQuantizeMx( raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'") -if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"): +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNvAbsMax"): - @register_fake("_qutlass_C::fusedQuantizeNv") - def _fake_fused_quantize_nv( + @register_fake("_qutlass_C::fusedQuantizeNvAbsMax") + def _fake_fused_quantize_nv_absmax( a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, @@ -3805,7 +4102,40 @@ def fusedQuantizeNv( padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device ) - return torch.ops._qutlass_C.fusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale) + safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale) + return xh_e2m1, xh_e4m3 + + +@torch.library.custom_op( + "vllm::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3") +) +def safeFusedQuantizeNv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, +) -> None: + """ + Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place + rather than returning them, to prevent torch 2.12+ errors that outputs of custom + operators may not alias any inputs to the custom operator. + """ + torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) + return + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"): + + @register_fake("vllm::safeFusedQuantizeNv") + def _fake_fused_quantize_nv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, + ) -> None: + return def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 8875ed49f6e7..4cdba1e745ba 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -219,6 +219,63 @@ def _xpu_ops_deepseek_scaling_rope_fake( return query, key +def _xpu_fp8_bmm_impl( + a: torch.Tensor, + b: torch.Tensor, + out_dtype: torch.dtype, + a_scale: torch.Tensor, + b_scale: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + """XPU FP8 batched GEMM implementation for ``torch.ops.vllm.xpu_fp8_bmm``. + + Computes batched matrix multiplication over the leading group dimension: + ``[G, M, K] @ [G, K, N] -> [G, M, N]``. + + Args: + a: FP8 activation tensor with shape ``[G, M, K]``. + Does not need to be contiguous. + b: FP8 weight tensor with shape ``[G, K, N]``. + Does not need to be contiguous. + out_dtype: Output dtype accepted by the kernel (typically + ``torch.bfloat16`` for the DeepSeek-V4 O-proj path). + a_scale: Activation scale tensor for ``a``. + In current DeepSeek-V4 XPU usage it is block-scaled with shape + ``[G, M, K/bs]`` (``bs`` is the quant block size, e.g. 128). + Must be contiguous. + b_scale: Weight scale tensor for ``b``. + In current DeepSeek-V4 XPU usage it is block-scaled with shape + ``[G, K/bs, N/bs]`` (``bs`` is the quant block size, e.g. 128). + Must be contiguous. + bias: Optional bias tensor. Pass ``None`` when no bias is required. + + Returns: + Output tensor with shape ``[G, M, N]`` and dtype ``out_dtype``. + + Notes: + This implementation centralizes access to + ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while + ``a`` and ``b`` may be non-contiguous views. + """ + return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias) + + +def _xpu_fp8_bmm_fake( + a: torch.Tensor, + b: torch.Tensor, + out_dtype: torch.dtype, + a_scale: torch.Tensor, + b_scale: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + # [G, M, K] @ [G, K, N] => [G, M, N] + return torch.empty( + (a.shape[0], a.shape[1], b.shape[2]), + dtype=out_dtype, + device=a.device, + ) + + def _xpu_fp8_mqa_logits_impl( q: torch.Tensor, k_quant: torch.Tensor, @@ -318,6 +375,115 @@ def _topk_topp_sample_fake( return +def _xpu_deepseek_fused_indexer_q_rope_fp8_impl( + index_q: torch.Tensor, + positions: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_fp8: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + """Fused RoPE + FP8 quant of the DeepSeek-V4 sparse-indexer Q (XPU). + Writes ``index_q_fp8`` and ``index_weights_out`` in place (no return). + Requires head_dim=128, rope_dim=64, and num_heads divisible by 2. + + Args: + index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous. + positions: (T,) int64 absolute token positions. + index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache. + index_weights: (T, H) bfloat16 raw indexer weights. + index_weights_softmax_scale: scalar softmax scale. + index_weights_head_scale: scalar per-head scale. + index_q_fp8: (T, H, 128) fp8 (e4m3) output, preallocated. [written] + index_weights_out: (T, H) float32 output, preallocated; + = index_weights * q_scale * softmax_scale * head_scale (the + per-(token, head) q_scale is folded in here). [written] + """ + torch.ops._xpu_C.deepseek_fused_indexer_q_rope_fp8( + index_q, + positions, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_fp8, + index_weights_out, + ) + + +def _xpu_deepseek_fused_indexer_q_rope_fp8_fake( + index_q: torch.Tensor, + positions: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_fp8: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + return + + +def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl( + index_q: torch.Tensor, + positions: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_packed: torch.Tensor, + index_q_scale: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + """Fused RoPE + MXFP4 quant of the DeepSeek-V4 sparse-indexer Q (XPU). + Writes ``index_q_packed``, ``index_q_scale`` and ``index_weights_out`` in + place (no return). Requires head_dim=128, rope_dim=64, and num_heads + divisible by 4. + + Args: + index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous. + positions: (T,) int64 absolute token positions. + index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache. + index_weights: (T, H) bfloat16 raw indexer weights. + index_weights_softmax_scale: scalar softmax scale. + index_weights_head_scale: scalar per-head scale. + index_q_packed: (T, H, 64) uint8 packed E2M1 nibbles (2 per byte), + preallocated. [written] + index_q_scale: (T, H, 4) uint8 ue8m0 per-block (32-elem) scales, + preallocated. [written] + index_weights_out: (T, H) float32 output, preallocated; + = index_weights * softmax_scale * head_scale (no q_scale folded; + per-block scales live in index_q_scale). [written] + """ + torch.ops._xpu_C.deepseek_fused_indexer_q_rope_mxfp4( + index_q, + positions, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_packed, + index_q_scale, + index_weights_out, + ) + + +def _xpu_deepseek_fused_indexer_q_rope_mxfp4_fake( + index_q: torch.Tensor, + positions: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_packed: torch.Tensor, + index_q_scale: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + return + + def _xpu_mxfp8_quantize_impl( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1053,6 +1219,12 @@ def register_ops_once() -> None: fake_impl=_xpu_mxfp4_quantize_fake, ) + direct_register_custom_op( + op_name="xpu_fp8_bmm", + op_func=_xpu_fp8_bmm_impl, + fake_impl=_xpu_fp8_bmm_fake, + ) + direct_register_custom_op( op_name="xpu_fp8_mqa_logits", op_func=_xpu_fp8_mqa_logits_impl, @@ -1078,6 +1250,24 @@ def register_ops_once() -> None: fake_impl=_topk_topp_sample_fake, ) + direct_register_custom_op( + op_name="xpu_deepseek_fused_indexer_q_rope_fp8", + op_func=_xpu_deepseek_fused_indexer_q_rope_fp8_impl, + mutates_args=["index_q_fp8", "index_weights_out"], + fake_impl=_xpu_deepseek_fused_indexer_q_rope_fp8_fake, + ) + + direct_register_custom_op( + op_name="xpu_deepseek_fused_indexer_q_rope_mxfp4", + op_func=_xpu_deepseek_fused_indexer_q_rope_mxfp4_impl, + mutates_args=[ + "index_q_packed", + "index_q_scale", + "index_weights_out", + ], + fake_impl=_xpu_deepseek_fused_indexer_q_rope_mxfp4_fake, + ) + _OPS_REGISTERED = True diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 83bda31241e7..de98ab135a66 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -2091,7 +2091,7 @@ def _parse_range_ratio(value: str) -> RangeRatio: return json.loads(value) -def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: +def get_samples(args, tokenizer: TokenizerLike | None) -> list[SampleRequest]: if not hasattr(args, "request_id_prefix"): args.request_id_prefix = "" @@ -2148,6 +2148,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: ) elif args.dataset_name == "sonnet": + assert tokenizer is not None, ( + "Tokenizer must be initialized for the 'sonnet' dataset." + ) sonnet_dataset = SonnetDataset( dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle ) @@ -2317,6 +2320,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: "Multi-modal content is only supported on 'openai-chat' and " "'openai-audio' backends." ) + assert tokenizer is not None, ( + "Tokenizer must be initialized for the 'hf' dataset." + ) input_requests = dataset_class( dataset_path=args.dataset_path, dataset_subset=args.hf_subset, @@ -2338,6 +2344,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: ) elif args.dataset_name == "timed_trace": + assert tokenizer is not None, ( + "Tokenizer must be initialized for the 'timed_trace' dataset." + ) dataloader = TimedTrace(**vars(args)) input_requests = dataloader.sample( num_requests=args.num_prompts, @@ -2347,6 +2356,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: else: # For datasets that follow a similar structure, use a mapping. + assert tokenizer is not None, ( + f"Tokenizer must be initialized for the '{args.dataset_name}' dataset." + ) dataset_mapping = { "spec_bench": lambda: SpecBench( dataset_path=args.dataset_path, @@ -2532,7 +2544,7 @@ def load_data(self) -> None: def sample( self, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, num_requests: int, request_id_prefix: str = "", no_oversample: bool = False, @@ -2809,7 +2821,7 @@ def _process_image_files( def sample( self, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, num_requests: int, request_id_prefix: str = "", no_oversample: bool = False, @@ -2907,7 +2919,7 @@ class CustomAudioDataset(CustomDataset): def sample( self, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, num_requests: int, request_id_prefix: str = "", no_oversample: bool = False, @@ -2929,7 +2941,9 @@ def sample( prompt = item.get("prompt", "") if tokenizer is None: prompt_len = 1 - new_output_len = output_len if output_len not in (None, -1) else 256 + new_output_len = ( + output_len if (output_len is not None and output_len != -1) else 256 + ) mm_content = None else: use_chat_template = ( @@ -2979,6 +2993,7 @@ def sample( ) new_output_len = int(item["output_tokens"]) else: + assert output_len is not None new_output_len = output_len sampled_requests.append( SampleRequest( @@ -3037,7 +3052,7 @@ def load_data(self) -> None: def sample( self, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, num_requests: int, request_id_prefix: str = "", no_oversample: bool = False, diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 76993a7d6b7f..e4ab5a583a04 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -557,7 +557,7 @@ def calculate_metrics( input_requests: list[SampleRequest], outputs: list[RequestFuncOutput], dur_s: float, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, selected_percentiles: list[float], goodput_config_dict: dict[str, float], ) -> tuple[BenchmarkMetrics, list[int]]: @@ -772,7 +772,7 @@ async def benchmark( base_url: str, model_id: str, model_name: str, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, input_requests: list[SampleRequest], logprobs: int | None, request_rate: float, @@ -795,6 +795,7 @@ async def benchmark( ready_check_timeout_sec: int = 600, ssl_context: ssl.SSLContext | bool | None = None, self_timed: bool = False, + probe_request_rate: float = 0.0, ): try: request_func = ASYNC_REQUEST_FUNCS[endpoint_type] @@ -971,6 +972,30 @@ async def limited_request_func(request_func_input, session, pbar): request_func_input=request_func_input, session=session, pbar=pbar ) + probe_outputs: list[RequestFuncOutput] = [] + probe_stop = asyncio.Event() + + async def probe_loop(): + probe_input = replace( + test_input, + prompt="Hi", + prompt_len=1, + output_len=1, + multi_modal_content=None, + chat_messages=None, + ) + interval = 1 / probe_request_rate + while not probe_stop.is_set(): + probe_outputs.append( + await request_func(request_func_input=probe_input, session=session) + ) + await asyncio.sleep(interval) + + probe_task: asyncio.Task | None = None + if probe_request_rate > 0: + print(f"Probe request rate: {probe_request_rate} req/s") + probe_task = asyncio.create_task(probe_loop()) + benchmark_start_time = time.perf_counter() tasks: list[asyncio.Task] = [] @@ -1042,6 +1067,10 @@ async def limited_request_func(request_func_input, session, pbar): ) outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks) + if probe_task is not None: + probe_stop.set() + await probe_task + if pbar is not None: pbar.close() @@ -1189,6 +1218,44 @@ async def limited_request_func(request_func_input, session, pbar): ) ) + probe_stats: dict[str, Any] | None = None + if probe_task is not None: + probe_lats = [o.latency for o in probe_outputs if o.success] + if probe_lats: + probe_stats = { + "probe_completed": len(probe_lats), + "probe_failed": len(probe_outputs) - len(probe_lats), + "probe_median_e2el_ms": float(np.median(probe_lats)) * 1000, + "probe_p99_e2el_ms": float(np.percentile(probe_lats, 99)) * 1000, + "probe_max_e2el_ms": float(max(probe_lats)) * 1000, + } + print("{s:{c}^{n}}".format(s="Probe Requests", n=50, c="-")) + print( + "{:<40} {:<10}".format( + "Probe requests completed:", probe_stats["probe_completed"] + ) + ) + print( + "{:<40} {:<10}".format( + "Probe requests failed:", probe_stats["probe_failed"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Median probe E2EL (ms):", probe_stats["probe_median_e2el_ms"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "P99 probe E2EL (ms):", probe_stats["probe_p99_e2el_ms"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Max probe E2EL (ms):", probe_stats["probe_max_e2el_ms"] + ) + ) + result: dict[str, Any] if isinstance(metrics, BenchmarkMetrics): result = { @@ -1223,6 +1290,9 @@ async def limited_request_func(request_func_input, session, pbar): "errors": [output.error for output in outputs], } + if probe_stats is not None: + result.update(probe_stats) + if rps_change_events: result["rps_change_events"] = rps_change_events @@ -1605,6 +1675,15 @@ def add_cli_args(parser: FlexibleArgumentParser): "bursty requests. A higher burstiness value (burstiness > 1) " "results in a more uniform arrival of requests.", ) + parser.add_argument( + "--probe-request-rate", + type=float, + default=0.0, + help="If positive, send single-token text-only probe requests at " + "this rate (req/s) alongside the main workload, bypassing " + "--max-concurrency, and report their latency separately. Useful " + "for measuring how the main workload stalls unrelated requests.", + ) parser.add_argument( "--disable-tqdm", action="store_true", @@ -2036,7 +2115,6 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: args.self_timed = False # Load the dataset. - assert tokenizer is not None, "Tokenizer must be initialized before loading dataset" input_requests = get_samples(args, tokenizer) if args.dataset_name in ("random", "prefix_repetition"): @@ -2122,6 +2200,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: ready_check_timeout_sec=args.ready_check_timeout_sec, ssl_context=ssl_context, self_timed=args.self_timed, + probe_request_rate=args.probe_request_rate, ) # Save config and results to json diff --git a/vllm/compilation/passes/fusion/add_rms_fusion.py b/vllm/compilation/passes/fusion/add_rms_fusion.py new file mode 100644 index 000000000000..b61c2b36adba --- /dev/null +++ b/vllm/compilation/passes/fusion/add_rms_fusion.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product + +import torch + +import vllm.ir.ops +from vllm.config import VllmConfig + +from ..vllm_inductor_pass import ( + VllmFusionPatternMatcherPass, + VllmPatternReplacement, +) + + +class AddRMSNormPattern(VllmPatternReplacement): + def __init__(self, epsilon: float, residual_first: bool) -> None: + self.epsilon = epsilon + self.residual_first = residual_first + + @property + def pattern(self): + def _pattern( + branch: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + residual_out = ( + residual + branch if self.residual_first else branch + residual + ) + rms = vllm.ir.ops.rms_norm(residual_out, weight, self.epsilon) + return rms, residual_out + + return _pattern + + @property + def replacement(self): + def _replacement( + branch: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return vllm.ir.ops.fused_add_rms_norm( + branch, residual, weight, self.epsilon + ) + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty_bf16(5, 16), # branch + self.empty_bf16(5, 16), # residual + self.empty_bf16(16), # weight + ] + + +class RMSNormReshapePattern(VllmPatternReplacement): + """Move a prefix-flatten before RMSNorm.""" + + def __init__(self, epsilon: float) -> None: + self.epsilon = epsilon + + @property + def pattern(self): + def _pattern( + input: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon) + return rms.reshape(-1, rms.shape[-1]) + + return _pattern + + @property + def replacement(self): + def _replacement( + input: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + input = input.reshape(-1, input.shape[-1]) + return vllm.ir.ops.rms_norm(input, weight, self.epsilon) + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty_bf16(1, 5, 16), # input + self.empty_bf16(16), # weight + ] + + +class FusedAddRMSNormReshapePattern(VllmPatternReplacement): + """Move a prefix-flatten before fused add RMSNorm.""" + + def __init__(self, epsilon: float) -> None: + self.epsilon = epsilon + + @property + def pattern(self): + def _pattern( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + rms, residual_out = vllm.ir.ops.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) + return rms.reshape(-1, rms.shape[-1]), residual_out + + return _pattern + + @property + def replacement(self): + def _replacement( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + original_shape = input.shape + hidden_size = input.shape[-1] + rms, residual_out = vllm.ir.ops.fused_add_rms_norm( + input.reshape(-1, hidden_size), + residual.reshape(-1, hidden_size), + weight, + self.epsilon, + ) + return rms, residual_out.reshape(original_shape) + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty_bf16(1, 5, 16), # input + self.empty_bf16(1, 5, 16), # residual + self.empty_bf16(16), # weight + ] + + +class AddRMSNormFusionPass(VllmFusionPatternMatcherPass): + """Fuse residual Add and RMSNorm emitted separately by the Transformers backend.""" + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "add_rmsnorm_fusion_pass") + + for epsilon, residual_first in product([1e-5, 1e-6], [True, False]): + self.register(AddRMSNormPattern(epsilon, residual_first)) + + self.dump_patterns(config, self.pm_pass) + + +class RMSNormReshapeFusionPass(VllmFusionPatternMatcherPass): + """Move the Transformers backend's post-RMSNorm flatten before the norm + for downstream 2D fusions. + """ + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "rmsnorm_reshape_fusion_pass") + + for epsilon in [1e-5, 1e-6]: + self.register(FusedAddRMSNormReshapePattern(epsilon)) + self.register(RMSNormReshapePattern(epsilon)) + + self.dump_patterns(config, self.pm_pass) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 4716e7012f3f..cdfa83002ada 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -119,6 +119,11 @@ def _norm_input_weight_dtype_match(match: pm.Match) -> bool: 8: 2, # 2MB 16: 64, # 64MB (mnnvl multi-node) }, + 107: { + 2: 64, # 64MB + 4: 64, # 64MB + 8: 2, # 2MB + }, } # Max size of the input tensor per world size per device capability @@ -140,6 +145,11 @@ def _norm_input_weight_dtype_match(match: pm.Match) -> bool: 4: 4, # 4MB 8: 2, # 2MB }, + 107: { + 2: 32, # 32MB + 4: 4, # 4MB + 8: 2, # 2MB + }, } MiB = 1024 * 1024 @@ -225,7 +235,7 @@ def call_trtllm_fused_allreduce_norm( max_token_num=max_token_num, hidden_dim=hidden_size, dtype=allreduce_in.dtype, - group=get_tp_group().device_group, + group=get_tp_group().cpu_group, ) assert workspace is not None, ( "Flashinfer allreduce workspace must be initialized when using flashinfer" @@ -996,7 +1006,7 @@ def __init__(self, config: VllmConfig) -> None: ) return self.hidden_dim = config.model_config.get_hidden_size() - self.group = get_tp_group().device_group + self.group = get_tp_group().cpu_group rank = get_tensor_model_parallel_rank() if flashinfer_comm is None: logger.warning( @@ -1416,8 +1426,7 @@ class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent with the sibling patterns above), which traces both ``QuantFP8.forward_hip`` and ``forward_native`` paths and so matches whichever op the call site - lowers to (``vllm.triton_per_token_group_quant_fp8`` or - ``vllm.rocm_aiter_group_fp8_quant``). + lowers to (``vllm.rocm_aiter_group_fp8_quant``). """ def __init__( diff --git a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py index fa641897bfad..2364281761f7 100644 --- a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py @@ -74,7 +74,7 @@ def fused_rope_and_unified_kv_cache_update_impl( layer_slot_mapping, ) - return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype) + return query.new_empty(0) def fused_rope_and_unified_kv_cache_update_fake( diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 67a3e4cbae5c..acc90044983b 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -7,7 +7,7 @@ from torch import fx as fx from vllm import envs -from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass from vllm.config import VllmConfig, set_current_vllm_config from vllm.logger import init_logger @@ -30,19 +30,23 @@ ) if current_platform.is_cuda_alike() or current_platform.is_xpu(): + from .fusion.add_rms_fusion import ( + AddRMSNormFusionPass, + RMSNormReshapeFusionPass, + ) + from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.sequence_parallelism import SequenceParallelismPass + from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass - from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass from .fusion.rms_quant_fusion import RMSNormQuantFusionPass from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass from .utility.scatter_split_replace import ScatterSplitReplacementPass - from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass @@ -138,6 +142,16 @@ def __call__(self, graph: fx.Graph) -> None: def configure(self, config: VllmConfig) -> None: self.pass_config = config.compilation_config.pass_config + model_config = config.model_config + enable_transformers_norm_canonicalization = ( + ( + self.pass_config.fuse_act_padding + or self.pass_config.fuse_allreduce_rms + or self.pass_config.fuse_norm_quant + ) + and model_config is not None + and model_config.using_transformers_backend() + ) # Set the current vllm config to allow tracing CustomOp instances with set_current_vllm_config(config, check_compile=False): @@ -149,6 +163,9 @@ def configure(self, config: VllmConfig) -> None: if self.pass_config.fuse_gemm_comms: self.passes += [AsyncTPPass(config)] + if enable_transformers_norm_canonicalization: + self.passes += [AddRMSNormFusionPass(config)] + if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled(): # Run the more specific RMSNorm+router-pad fusion before # AR+RMS, since both consume fused_add_rms_norm. @@ -160,6 +177,11 @@ def configure(self, config: VllmConfig) -> None: else: self.passes += [AllReduceFusionPass(config)] + if enable_transformers_norm_canonicalization: + # Let AR+RMS match before moving output reshapes ahead of the + # remaining RMSNorms, exposing them to RMS+Quant fusion. + self.passes += [RMSNormReshapeFusionPass(config)] + if self.pass_config.fuse_norm_quant: if rocm_aiter_ops.is_enabled(): self.passes += [ @@ -177,11 +199,7 @@ def configure(self, config: VllmConfig) -> None: self.passes += [ScatterSplitReplacementPass(config)] self.passes += [QkNormRopeKvCacheFusionPass(config)] - if ( - self.pass_config.fuse_mla_dual_rms_norm - and rocm_aiter_ops.is_enabled() - and check_aiter_fused_qk_rmsnorm() - ): + if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled(): self.passes += [MLADualRMSNormFusionPass(config)] if self.pass_config.fuse_rope_kvcache: diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 82ab1842fe9a..24a8b5a31713 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -11,7 +11,9 @@ ) from vllm.config.device import DeviceConfig from vllm.config.diffusion import DiffusionConfig +from vllm.config.ec_manager_config import EncoderCacheManagerConfig from vllm.config.ec_transfer import ECTransferConfig +from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.config.kernel import KernelConfig from vllm.config.kv_events import KVEventsConfig from vllm.config.kv_transfer import KVTransferConfig @@ -75,6 +77,8 @@ "DeviceConfig", # From vllm.config.diffusion "DiffusionConfig", + # From vllm.config.ec_manager_config + "EncoderCacheManagerConfig", # From vllm.config.ec_transfer "ECTransferConfig", # From vllm.config.kernel @@ -121,6 +125,8 @@ "StructuredOutputsConfig", # From vllm.config.profiler "ProfilerConfig", + # From vllm.config.fault_tolerance + "FaultToleranceConfig", # From vllm.config.utils "ConfigType", "SupportsMetricsInfo", diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 78bc1f9537df..994be05f54ea 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -11,6 +11,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"] +MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] @config @@ -20,6 +21,9 @@ class AttentionConfig: backend: AttentionBackendEnum | None = None """Attention backend to use. Use "auto" or None for automatic selection.""" + minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton" + """Sparse decode kernel used by the MiniMax M3 MSA backend.""" + backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict) """Per-KV-cache-group attention backend overrides, keyed by `KVCacheSpecKind` (e.g. `{"mla_attention": "FLASHINFER_MLA", @@ -89,13 +93,26 @@ class AttentionConfig: flex_attn_q_block_size: int | None = None """Logical Q block size for the flex attention block mask. Must be a power of 2 and divisible by flex_attn_block_m. - If None, uses the default (16 on PyTorch >= 2.9, 128 otherwise).""" + If None, uses 16 for paged KV attention on PyTorch >= 2.9, and 128 + for encoder-only attention or older PyTorch versions.""" flex_attn_kv_block_size: int | None = None """Logical KV block size for the flex attention block mask. Must be a power of 2 and divisible by flex_attn_block_n. - If None, uses the default (kv_cache_block_size on PyTorch >= 2.9, - 128 otherwise).""" + If None, uses the KV cache block size for paged KV attention on + PyTorch >= 2.9, and 128 for encoder-only attention or older PyTorch + versions.""" + + def __post_init__(self) -> None: + msa_aliases: dict[AttentionBackendEnum, MiniMaxM3MSADecodeBackend] = { + AttentionBackendEnum.CUTLASS_MSA: "cutlass", + AttentionBackendEnum.TRITON_MSA: "triton", + } + if self.backend in msa_aliases: + self.minimax_m3_msa_decode_backend = msa_aliases[self.backend] + # The alias selects only MiniMax's sparse decode kernel. Dense + # layers still use the platform's normal automatic backend. + self.backend = None def compute_hash(self) -> str: """ diff --git a/vllm/config/cache.py b/vllm/config/cache.py index a628e7d7cdd0..d267e9af3584 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -137,14 +137,25 @@ class CacheConfig: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers. + """The cache strategy for Mamba layers: + - "none": set when prefix caching is disabled. - "all": cache the mamba state of all tokens at position i * block_size. This is - the default behavior (for models that support it) when prefix caching is - enabled. + the default behavior (for models that support it) when prefix caching is enabled. - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. + when the token is at position i * block_size. """ + replayssm_buffer_len: int = Field(default=16, gt=0) + """ReplaySSM history buffer length B: with use_replayssm, standard decode + caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint + state to HBM every B steps. Default 16.""" + use_replayssm: bool = False + """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip + the per-step full-state store, writing the checkpoint back only on flush. + Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton + mamba backend; standard (non-speculative) decode only. In align mode flushes + are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, + but this is not required.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 810e97c4cadc..4eb66613a81b 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -768,11 +768,9 @@ class CompilationConfig: "vllm::mamba_mixer", "vllm::short_conv", "vllm::linear_attention", - "vllm::plamo2_mamba_mixer", "vllm::qwen_gdn_attention_core", "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", - "vllm::kda_attention", "vllm::sparse_attn_indexer", "vllm::rocm_aiter_sparse_attn_indexer", "vllm::deepseek_v4_attention", @@ -905,10 +903,6 @@ def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: return handler(value) def __post_init__(self) -> None: - count_none = self.custom_ops.count("none") - count_all = self.custom_ops.count("all") - assert count_none + count_all <= 1, "Can only specify 'none' or 'all'" - # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2: # 1. A bug in PyTorch, fixed in 2.7: # https://github.com/pytorch/pytorch/issues/147924 @@ -1001,13 +995,28 @@ def __post_init__(self) -> None: ) for op in self.custom_ops: - if op[0] not in {"+", "-"} and op not in {"all", "none"}: + if op not in {"all", "none"} and (len(op) < 2 or op[0] not in {"+", "-"}): raise ValueError( f"Invalid syntax '{op}' for custom op, " "must be 'all', 'none', '+op' or '-op' " "(where 'op' is the registered op name)" ) + base_modes = [op for op in self.custom_ops if op in {"all", "none"}] + if len(base_modes) > 1: + raise ValueError( + "custom_ops can contain only one base mode: 'all' or 'none'" + ) + + enabled_ops = {op[1:] for op in self.custom_ops if op.startswith("+")} + disabled_ops = {op[1:] for op in self.custom_ops if op.startswith("-")} + conflicting_ops = sorted(enabled_ops & disabled_ops) + if conflicting_ops: + raise ValueError( + "custom_ops cannot both enable and disable the same operation(s): " + f"{', '.join(conflicting_ops)}. Remove either the '+' or '-' directive" + ) + # Currently only eager and inductor backend are supported. # for piecewise compilation. Custom backends are not supported for # piecewise compilation. Update when more backends are supported. @@ -1344,10 +1353,16 @@ def custom_op_log_check(self): ) def is_custom_op_enabled(self, op: str) -> bool: - if "all" in self.custom_ops: + count_all = self.custom_ops.count("all") + count_none = self.custom_ops.count("none") + if count_all + count_none != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) + + if count_all: return f"-{op}" not in self.custom_ops - assert "none" in self.custom_ops return f"+{op}" in self.custom_ops def resolve_cudagraph_mode_and_sizes( diff --git a/vllm/config/ec_manager_config.py b/vllm/config/ec_manager_config.py new file mode 100644 index 000000000000..9e7cda64e633 --- /dev/null +++ b/vllm/config/ec_manager_config.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC + +from vllm.config.utils import config +from vllm.utils.import_utils import resolve_obj_by_qualname + + +@config +class EncoderCacheManagerConfig: + encoder_cache_manager_cls: str | None = None + """The qualified class name of the custom encoder cache manager. + """ + + def get_encoder_cache_manager_obj(self): + cls_path = self.encoder_cache_manager_cls + if cls_path is None: + return None + return resolve_obj_by_qualname(cls_path) + + +class EncoderCacheManagerMetadata(ABC): # noqa: B024 + """ + Abstract Metadata used to communicate between the + Scheduler EncoderCacheManager and Worker EncoderCacheManager. + """ + + pass diff --git a/vllm/config/fault_tolerance.py b/vllm/config/fault_tolerance.py new file mode 100644 index 000000000000..7ed095d204f6 --- /dev/null +++ b/vllm/config/fault_tolerance.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from vllm.config.utils import config + + +@config +class FaultToleranceConfig: + """Configuration for fault tolerance.""" + + engine_recovery_timeout_sec: int = 120 + """Timeout (in seconds) to wait for error handling instructions + before raising an exception. If the EngineCore encounters an + error, it waits up to this many seconds for vLLM to receive + instructions on how to handle the error and then recover from the fault. + If vLLM does not recover during this time, the original error is raised. + """ diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index f13562f59c0e..59195f307b80 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -122,6 +122,7 @@ def with_default( MoEBackend = Literal[ "auto", "triton", + "batched_triton", "deep_gemm", "deep_gemm_mega_moe", "cutlass", @@ -175,9 +176,15 @@ class KernelConfig: enable_flashinfer_autotune: bool = None # type: ignore[assignment] """If True, run FlashInfer autotuning during kernel warmup.""" + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 enable_cutedsl_warmup: bool = True """If True, run CuTeDSL compile warmup during kernel warmup.""" + enable_jit_warmup: bool = True + """If True, run JIT compile warmup during kernel warmup.""" + enable_bf16x3_router_gemm: bool = False """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM.""" @@ -186,6 +193,8 @@ class KernelConfig: - "auto": Automatically select the best backend based on model and hardware - "triton": Use Triton-based fused MoE kernels + - "batched_triton": Use batched Triton experts (moe_mmk) on the batched + activation format ([E_local, max_num_tokens, K]) - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only) - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels - "cutlass": Use vLLM CUTLASS kernels @@ -250,6 +259,7 @@ def compute_hash(self) -> str: """ ignored_factors = { "enable_cutedsl_warmup", + "enable_jit_warmup", "enable_flashinfer_autotune", "ir_op_priority", # handled separately below } @@ -260,6 +270,7 @@ def compute_hash(self) -> str: @field_validator( "enable_flashinfer_autotune", "enable_cutedsl_warmup", + "enable_jit_warmup", mode="wrap", ) @classmethod diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index 996478c36760..e8988aab1940 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum, EnumMeta -from typing import Any +from typing import Any, Literal, get_args from pydantic import field_validator @@ -27,6 +27,10 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): TRITON = "triton" FLASHINFER = "flashinfer" + CPU = "cpu" + + +MambaSSUAlgorithm = Literal["auto", "simple", "vertical", "horizontal"] @config @@ -45,6 +49,12 @@ class MambaConfig: generation. 0 uses the Triton default. Higher values improve randomness quality at the cost of compute.""" + ssu_algorithm: MambaSSUAlgorithm | None = None + """Selective state update algorithm to use with the FlashInfer backend. + None defaults to FlashInfer's "auto" algorithm. Forced algorithms must + be supported by FlashInfer for the active GPU, state dtype, and decoding + mode.""" + @field_validator("backend", mode="before") @classmethod def validate_backend_before(cls, value: Any) -> Any: @@ -53,7 +63,25 @@ def validate_backend_before(cls, value: Any) -> Any: return MambaBackendEnum[value.upper()] return value + def validate_ssu_algorithm(self) -> None: + if self.ssu_algorithm is None: + return + valid_algorithms = get_args(MambaSSUAlgorithm) + if self.ssu_algorithm not in valid_algorithms: + valid = ", ".join(valid_algorithms) + raise ValueError( + f"Unknown Mamba SSU algorithm: '{self.ssu_algorithm}'. " + f"Valid options are: {valid}" + ) + if self.backend != MambaBackendEnum.FLASHINFER: + raise ValueError( + "Mamba SSU algorithm selection is only supported with the " + "FlashInfer backend. Please set `--mamba-backend flashinfer`, " + "or omit `--mamba-ssu-algorithm`." + ) + def __post_init__(self): + self.validate_ssu_algorithm() if self.enable_stochastic_rounding: from vllm.platforms import current_platform diff --git a/vllm/config/model.py b/vllm/config/model.py index 6b032ae76210..91eb3c9bce17 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -17,6 +17,7 @@ from vllm.config.multimodal import ( MMCacheType, MMEncoderTPMode, + MMHasherAlgorithm, MMTensorIPC, MultiModalConfig, ) @@ -84,12 +85,24 @@ ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] TokenizerMode = Literal[ - "auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "inkling" + "auto", + "hf", + "slow", + "mistral", + "deepseek_v32", + "deepseek_v4", + "inkling", + "kimi_k3", + "cohere", ] ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal[ "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs" ] +PROCESSED_LOGPROBS_MODES: tuple[LogprobsMode, ...] = ( + "processed_logits", + "processed_logprobs", +) HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig] ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"] LayerBlockType = Literal["attention", "linear_attention", "mamba"] @@ -137,6 +150,11 @@ class ModelConfig: - "mistral" will always use the tokenizer from `mistral_common`. - "deepseek_v32" will always use the tokenizer from `deepseek_v32`. - "deepseek_v4" will always use the tokenizer from `deepseek_v4`. + - "kimi_k3" will always use the "hf" tokenizer but render chat prompts + with Kimi K3's Python XTML encoding instead of a Jinja template. + - "cohere" uses the standard HF tokenizer but renders the chat template + via the `cohere_melody` library (cmd3 / cmd4 templates) instead of + Jinja, and surfaces grounded-citation metadata on responses. - Other custom values can be supported via plugins. To swap the Rust BPE backend that powers HF fast tokenizers for the @@ -361,6 +379,7 @@ class ModelConfig: mm_processor_kwargs: InitVar[dict[str, Any] | None] = None mm_processor_cache_gb: InitVar[float | None] = None mm_processor_cache_type: InitVar[MMCacheType | None] = None + mm_hasher_algorithm: InitVar[MMHasherAlgorithm | None] = None mm_shm_cache_max_object_size_mb: InitVar[int | None] = None mm_encoder_only: InitVar[bool | None] = None mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None @@ -372,6 +391,7 @@ class ModelConfig: interleave_mm_strings: InitVar[bool | None] = None skip_mm_profiling: InitVar[bool | None] = None video_pruning_rate: InitVar[float | None] = None + video_pruning_method: InitVar[str | None] = None mm_tensor_ipc: InitVar[MMTensorIPC] = None mm_ipc_gpu_memory_gb: InitVar[float | None] = None @@ -489,6 +509,7 @@ def __post_init__( mm_processor_kwargs: dict[str, Any] | None, mm_processor_cache_gb: float | None, mm_processor_cache_type: MMCacheType | None, + mm_hasher_algorithm: MMHasherAlgorithm | None, mm_shm_cache_max_object_size_mb: int | None, mm_encoder_only: bool | None, mm_encoder_tp_mode: MMEncoderTPMode | None, @@ -500,6 +521,7 @@ def __post_init__( interleave_mm_strings: bool | None, skip_mm_profiling: bool | None, video_pruning_rate: float | None, + video_pruning_method: str | None, mm_tensor_ipc: MMTensorIPC, mm_ipc_gpu_memory_gb: float | None, ) -> None: @@ -575,9 +597,9 @@ def __post_init__( self.hf_text_config, "attention_chunk_size", None ) self.encoder_config = self._get_encoder_config() - self.hf_image_processor_config = get_hf_image_processor_config( - self.model, hf_token=self.hf_token, revision=self.revision - ) + # Image-processor metadata is only consumed by multimodal models. + # Probing it for text-only models causes avoidable Hub requests. + self.hf_image_processor_config: dict[str, Any] = {} architectures = self.architectures registry = self.registry @@ -628,6 +650,8 @@ def __post_init__( self.tokenizer_mode = "terratorch" elif arch == "MoonshotKimiaForCausalLM": self.tokenizer_mode = "kimi_audio" + elif arch == "KimiK3ForConditionalGeneration": + self.tokenizer_mode = "kimi_k3" elif arch == "DeepseekV32ForCausalLM": self.tokenizer_mode = "deepseek_v32" elif arch == "DeepseekV4ForCausalLM": @@ -699,6 +723,9 @@ def __post_init__( # Init multimodal config if needed if self._model_info.supports_multimodal: + self.hf_image_processor_config = get_hf_image_processor_config( + self.model, hf_token=self.hf_token, revision=self.revision + ) if ( mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data @@ -717,6 +744,7 @@ def __post_init__( mm_processor_kwargs=mm_processor_kwargs, mm_processor_cache_gb=mm_processor_cache_gb, mm_processor_cache_type=mm_processor_cache_type, + mm_hasher_algorithm=mm_hasher_algorithm, mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, mm_encoder_only=mm_encoder_only, mm_encoder_tp_mode=mm_encoder_tp_mode, @@ -728,6 +756,7 @@ def __post_init__( interleave_mm_strings=interleave_mm_strings, skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, + video_pruning_method=video_pruning_method, mm_tensor_ipc=mm_tensor_ipc, mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, ) @@ -738,19 +767,37 @@ def __post_init__( self.multimodal_config = MultiModalConfig(**mm_config_kwargs) # type: ignore[arg-type] + pruning_spec = self.multimodal_config.get_video_pruning_spec() + supported_pruning = self._model_info.supported_video_pruning_methods + if ( + pruning_spec is not None + and supported_pruning + and pruning_spec[0] not in supported_pruning + ): + raise ValueError( + f"Video pruning method '{pruning_spec[0]}' is not " + f"supported by {self._model_info.architecture} " + f"(supported methods: {supported_pruning})." + ) + if ( self.renderer_num_workers > 1 and self.multimodal_config.mm_processor_cache_gb > 0 + and self.runner_type == "pooling" ): raise ValueError( "Cannot use --renderer-num-workers > 1 with the " - "multimodal processor cache enabled. The cache is " - "not thread-safe and does not support concurrent " - "renderer workers. Please set " + "multimodal processor cache enabled for pooling models. " + "Pooling preprocessing runs on the renderer workers, and " + "the cache is not thread-safe. Please set " "--renderer-num-workers 1 (the default), or " "disable the cache with --mm-processor-cache-gb 0." ) + # Rebuild after multimodal_config exists so text-only mm_prefix + # clearing is applied (and cached for later with_hf_config calls). + self.model_arch_config = self.get_model_arch_config() + if self.disable_sliding_window: # Set after get_and_verify_max_len to ensure that max_model_len # can be correctly capped to sliding window size @@ -763,6 +810,42 @@ def __post_init__( self._verify_cuda_graph() self._verify_bnb_config() + def _supports_multimodal_for_mm_prefix(self) -> bool: + """Whether multimodal inputs can still appear for this deployment. + + This runs more than once per config: once early in ``__post_init__`` + (before ``multimodal_config`` exists), again after it is created, and + then for every ``get_model_arch_config`` regeneration -- notably + ``with_hf_config``, which deep-copies this ``ModelConfig`` and swaps + ``hf_config`` for a text-only submodule (e.g. ``Gemma4ForCausalLM``). + + The result is cached for correctness, not just to save work: on the + ``with_hf_config`` copy the submodule architecture has no registered + multimodal processor, so re-querying the registry would raise and be + treated as text-only, wrongly clearing ``is_mm_prefix_lm`` even when a + vision modality is still enabled (e.g. ``image=0`` but video allowed). + The deep-copied cache preserves the top-level decision instead. + """ + cached = getattr(self, "_supports_multimodal_inputs_cached", None) + if cached is not None: + return cached + + if self.multimodal_config is None: + # Early call before multimodal init — do not clear mm_prefix yet. + return True + + from vllm.multimodal import MULTIMODAL_REGISTRY + + supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(self) + self._supports_multimodal_inputs_cached = supports_mm + if not supports_mm: + logger.info_once( + "Disabled mm_prefix attention mode because multimodal inputs " + "are configuration-disabled. Attention backends without " + "mm_prefix support may now be selected." + ) + return supports_mm + def get_model_arch_config( self, ) -> ModelArchitectureConfig: @@ -770,7 +853,9 @@ def get_model_arch_config( self.hf_config.model_type, ModelArchConfigConvertorBase ) convertor = convertor_cls(self.hf_config, self.hf_text_config) - return convertor.convert() + return convertor.convert( + supports_multimodal=self._supports_multimodal_for_mm_prefix() + ) @field_validator("tokenizer", "max_model_len", mode="wrap") @classmethod @@ -1249,27 +1334,33 @@ def verify_with_parallel_config( decode_context_parallel_size = parallel_config.decode_context_parallel_size if decode_context_parallel_size > 1 and not self.use_mla: total_num_kv_heads = self.get_total_num_kv_heads() - assert tensor_parallel_size > total_num_kv_heads, ( - f"tensor parallel size {tensor_parallel_size} must be greater " - f"than total num kv heads {total_num_kv_heads} when enable " - f"decode context parallel for GQA/MQA" - ) + if tensor_parallel_size <= total_num_kv_heads: + raise ValueError( + "Decode context parallelism for GQA/MQA requires " + f"`--tensor-parallel-size` ({tensor_parallel_size}) to be " + "greater than the model's total number of KV heads " + f"({total_num_kv_heads}). Increase `--tensor-parallel-size` " + "or set `--decode-context-parallel-size 1`." + ) max_dcp_size = tensor_parallel_size // total_num_kv_heads - assert decode_context_parallel_size <= max_dcp_size, ( - f"decode context parallel size must less than or equal to " - f"(tensor parallel size {tensor_parallel_size} // total " - f"num kv heads {total_num_kv_heads}) = {max_dcp_size}, " - f"but got {decode_context_parallel_size}" - ) + if decode_context_parallel_size > max_dcp_size: + raise ValueError( + "`--decode-context-parallel-size` " + f"({decode_context_parallel_size}) exceeds the maximum " + f"supported value ({max_dcp_size}) for " + f"`--tensor-parallel-size` ({tensor_parallel_size}) and " + f"{total_num_kv_heads} model KV heads." + ) num_q_per_kv = total_num_attention_heads // total_num_kv_heads - assert num_q_per_kv % decode_context_parallel_size == 0, ( - f"Total number of q per kv attn heads ({num_q_per_kv})" - " must be divisible by dcp world size when enable " - "decode context parallel for GQA " - f"({parallel_config.decode_context_parallel_size})." - ) + if num_q_per_kv % decode_context_parallel_size != 0: + raise ValueError( + "The model's number of query heads per KV head " + f"({num_q_per_kv}) must be divisible by " + "`--decode-context-parallel-size` " + f"({decode_context_parallel_size}) for GQA/MQA." + ) # torch_shm uses a single IPC queue to rank 0; DP>1 is # incompatible because API servers can't know which @@ -1441,7 +1532,7 @@ def get_mamba_chunk_size(self) -> int: """ Returns the mamba chunk size if it exists """ - # used by e.g. Bamba, FalconH1, Granite, PLaMo2 + # used by e.g. Bamba, FalconH1, Granite chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None) if chunk_size is None: # used by e.g. Mamba2, NemotronH, Zamba @@ -1692,6 +1783,10 @@ def has_inner_state(self): def supports_mamba_prefix_caching(self) -> bool: return self._model_info.supports_mamba_prefix_caching + @property + def supports_replayssm(self) -> bool: + return self._model_info.supports_replayssm + @property def use_mla(self) -> bool: return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE @@ -2011,7 +2106,6 @@ def str_dtype_to_torch_dtype(type: str): "gemma2": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3_text": "Numerical instability. Please use bfloat16 or float32 instead.", - "plamo2": "Numerical instability. Please use bfloat16 or float32 instead.", "glm4": "Numerical instability. Please use bfloat16 or float32 instead.", } diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 150d58cf1f74..6c85fc7c9319 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -3,12 +3,13 @@ from collections.abc import Mapping from pathlib import Path -from typing import Any, Literal, TypeAlias, TypedDict, final +from typing import Any, Literal, TypeAlias, TypedDict, cast, final from pydantic import ConfigDict, Field, field_validator, model_validator from pydantic.dataclasses import dataclass -from vllm.config.utils import config +import vllm.envs as envs +from vllm.config.utils import config, get_from_deprecated_env_if_set from vllm.utils.hashing import safe_hash from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -60,7 +61,21 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False): MMEncoderTPMode = Literal["weights", "data"] MMCacheType = Literal["shm", "lru"] +VideoPruningMethod = Literal["evs", "vidcom2"] MMTensorIPC = Literal["direct_rpc", "torch_shm"] +MMHasherAlgorithm = Literal["blake3", "sha256", "sha512"] + + +def _get_mm_hasher_algorithm() -> MMHasherAlgorithm: + env_value = get_from_deprecated_env_if_set( + "VLLM_MM_HASHER_ALGORITHM", + "v0.27", + "mm_hasher_algorithm", + ) + env_value = "blake3" if env_value is None else env_value + return cast(MMHasherAlgorithm, env_value.lower()) + + MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions] """ A dictionary containing an entry for each modality type of dummy data. @@ -132,6 +147,11 @@ class MultiModalConfig: mm_processor_cache_type: MMCacheType = "lru" """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`, use shared memory FIFO cache. If `lru`, use mirrored LRU cache.""" + mm_hasher_algorithm: MMHasherAlgorithm = Field( + default_factory=_get_mm_hasher_algorithm + ) + """Hash algorithm to use for multi-modal input caching. Use `"sha256"` or + `"sha512"` for FIPS-compliant deployments.""" mm_shm_cache_max_object_size_mb: int = Field(default=128, ge=0) """Size limit (in MiB) for each object stored in the multi-modal processor shared memory cache. Only effective when `mm_processor_cache_type` is @@ -188,9 +208,14 @@ class MultiModalConfig: estimating the peak memory usage of the activation of multimodal encoder and embedding cache.""" video_pruning_rate: float | None = Field(default=None, ge=0.0, lt=1.0) - """Sets pruning rate for video pruning via Efficient Video Sampling. - Value sits in range [0;1) and determines fraction of media tokens - from each video to be pruned. + """Fraction of video tokens to prune from each video. Value sits in range + [0;1); pruning is enabled when it is greater than 0. The pruning algorithm + is selected by `video_pruning_method`. + """ + video_pruning_method: VideoPruningMethod = "evs" + """Video token pruning algorithm applied when `video_pruning_rate` > 0: + - "evs": Efficient Video Sampling. + - "vidcom2": Video Compression Commander. """ mm_tensor_ipc: MMTensorIPC = "direct_rpc" """IPC (inter-process communication) method for multimodal tensors. @@ -344,5 +369,26 @@ def merge_mm_processor_kwargs( kwargs = self.mm_processor_kwargs or {} return kwargs | dict(inference_kwargs) + def use_gpu_video_backend(self) -> bool: + """Return whether the configured video loader or codec uses the GPU.""" + from vllm.multimodal.video import VIDEO_LOADER_REGISTRY + + video_kwargs = self.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_REGISTRY.backend_requires_gpu(video_loader_backend) or ( + codec_backend is not None + and VIDEO_LOADER_REGISTRY.backend_requires_gpu(codec_backend) + ) + def is_multimodal_pruning_enabled(self): - return self.video_pruning_rate is not None and self.video_pruning_rate > 0 + return self.get_video_pruning_spec() is not None + + def get_video_pruning_spec(self) -> tuple[VideoPruningMethod, float] | None: + """Return `(method, rate)` when video pruning is enabled, else None. + `rate` is the fraction of video tokens to prune.""" + if self.video_pruning_rate is not None and self.video_pruning_rate > 0: + return (self.video_pruning_method, float(self.video_pruning_rate)) + return None diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 53688c05d92d..2e908ead1762 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -13,6 +13,7 @@ from typing_extensions import Self import vllm.envs as envs +from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.config.utils import config from vllm.logger import init_logger from vllm.platforms import current_platform @@ -22,6 +23,7 @@ from ray.runtime_env import RuntimeEnv from ray.util.placement_group import PlacementGroup + from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.v1.executor import Executor else: RuntimeEnv = Any @@ -92,7 +94,7 @@ class EPLBConfig: Backend for EPLB expert weight communication: - "torch_nccl": Use torch.distributed on the device process group - "torch_gloo": Use torch.distributed gloo with CPU staging - - "nixl": Use NIXL/ RIXL with staged send/recv buffers + - "nixl": Use NIXL with staged send/recv buffers - "pynccl": Use PyNccl send/recv - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo") """ @@ -393,6 +395,16 @@ class is dynamically inherited by the worker class. This is used to inject should only be set by API server scale-out. """ + enable_fault_tolerance: bool = False + """Enable fault tolerance for detailed error recovery, + such as scaling down fault DPEngineCore. + """ + + fault_tolerance_config: FaultToleranceConfig = Field( + default_factory=FaultToleranceConfig + ) + """The configurations for fault tolerance.""" + @field_validator("disable_nccl_for_dp_synchronization", mode="wrap") @classmethod def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: @@ -445,6 +457,13 @@ def _validate_parallel_config(self) -> Self: f"but found: {self._api_process_rank}" ) + if self.enable_fault_tolerance and self._api_process_count > 1: + raise ValueError( + "Fault tolerance requires a single API server process " + f"(--api-server-count=1), but got {self._api_process_count}. " + "The FT system assumes one AsyncMPClient manages all engines." + ) + if self.all2all_backend in ["pplx", "naive"]: logger.warning( "The '%s' all2all backend has been removed. " @@ -667,6 +686,14 @@ def use_sequence_parallel_moe(self) -> bool: and self.data_parallel_size > 1 ) + @property + def use_all2all(self) -> bool: + return ( + self.data_parallel_size > 1 + or self.use_sequence_parallel_moe + or (self.enable_expert_parallel and self.prefill_context_parallel_size > 1) + ) + @property def use_batched_dp_moe(self) -> bool: return ( @@ -767,6 +794,7 @@ def compute_hash(self): "data_parallel_master_ip", "data_parallel_master_port", "_data_parallel_master_port_list", + "_coord_store_port", "data_parallel_rpc_port", "rank", "master_addr", @@ -1000,11 +1028,6 @@ def _verify_args(self) -> Self: "Disabled the custom all-reduce kernel because it is not " "supported on current platform." ) - if self.nnodes > 1: - self.disable_custom_all_reduce = True - logger.debug( - "Disabled the custom all-reduce since we are running on multi-node." - ) if self.ray_workers_use_nsight and not self.use_ray: raise ValueError( "Unable to use nsight profiling unless workers run with Ray." diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 370fbaf2731f..4f2a40426cdd 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -88,7 +88,7 @@ class QuantizationConfigArgs: """Spec applied to ``LinearBase`` layers.""" moe: QuantSpec | None = None - """Spec applied to ``FusedMoE`` layers.""" + """Spec applied to ``FusedMoEFactory`` layers.""" ignore: list[str] = Field(default_factory=list) """Layers to skip quantization for.""" diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 3773a47ec94f..1b82b2c51283 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -67,16 +67,6 @@ class SchedulerConfig: In real usage, this should be set in `EngineArgs.create_engine_config`. """ - max_num_partial_prefills: int = Field(default=1, ge=1) - """For chunked prefill, the maximum number of sequences that can be - partially prefilled concurrently.""" - - max_long_partial_prefills: int = Field(default=1, ge=1) - """For chunked prefill, the maximum number of prompts longer than - long_prefill_token_threshold that will be prefilled concurrently. Setting - this less than max_num_partial_prefills will allow shorter prompts to jump - the queue in front of longer prompts in some cases, improving latency.""" - long_prefill_token_threshold: int = Field(default=0, ge=0) """For chunked prefill, a request is considered long if the prompt is longer than this number of tokens. 0 disables the cap (default).""" @@ -254,19 +244,6 @@ def __post_init__(self, max_model_len: int, is_encoder_decoder: bool) -> None: self.max_num_batched_tokens, ) - if self.max_num_partial_prefills > 1: - if self.long_prefill_token_threshold == 0: - self.long_prefill_token_threshold = int(max_model_len * 0.04) - - logger.info( - "Concurrent partial prefills enabled with " - "max_num_partial_prefills=%d, max_long_partial_prefills=%d, " - "long_prefill_token_threshold=%d", - self.max_num_partial_prefills, - self.max_long_partial_prefills, - self.long_prefill_token_threshold, - ) - self.verify_max_model_len(max_model_len) def verify_max_model_len(self, max_model_len: int) -> Self: @@ -298,24 +275,11 @@ def verify_max_model_len(self, max_model_len: int) -> Self: self.max_num_seqs * max_model_len, ) - if self.max_num_partial_prefills > 1: - if not self.enable_chunked_prefill: - raise ValueError( - "Chunked prefill must be enabled to set " - "max_num_partial_prefills > 1." - ) - - if self.long_prefill_token_threshold > max_model_len: - raise ValueError( - "long_prefill_token_threshold " - f"({self.long_prefill_token_threshold}) cannot be greater " - f"than the max_model_len ({max_model_len})." - ) - - if self.max_long_partial_prefills > self.max_num_partial_prefills: + if self.long_prefill_token_threshold > max_model_len: raise ValueError( - f"{self.max_long_partial_prefills=} must be less than or equal to " - f"{self.max_num_partial_prefills=}." + "long_prefill_token_threshold " + f"({self.long_prefill_token_threshold}) cannot be greater " + f"than the max_model_len ({max_model_len})." ) return self diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 8ba55a2ec964..07bdcac68b84 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -51,6 +51,7 @@ "minimax_m3_mtp", "bailing_hybrid_mtp", "mtp", + "kimi_k3_mtp", "pangu_ultra_moe_mtp", "step3p5_mtp", "hy_v3_mtp", @@ -355,6 +356,16 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: {"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]} ) + if hf_config.model_type == "kimi_k3": + # Kimi-K3 keeps the text-model fields (incl. the MTP layer count) + # nested under ``text_config`` (a KimiLinearConfig). + text_config = getattr(hf_config, "text_config", hf_config) + n_predict = getattr(text_config, "num_nextn_predict_layers", None) + hf_config.model_type = "kimi_k3_mtp" + hf_config.update( + {"n_predict": n_predict, "architectures": ["KimiK3MTPModel"]} + ) + if hf_config.architectures[0] == "MiMoForCausalLM": hf_config.model_type = "mimo_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", None) @@ -910,6 +921,15 @@ def __post_init__(self): f"Unsupported speculative method: '{self.method}'" ) + if self.method in ("eagle", "eagle3"): + # EAGLE drafts share the target's positional space; a + # draft checkpoint with a smaller max_position_embeddings + # than the target under-sizes its rotary cache (#48894). + SpeculativeConfig._maybe_override_draft_max_position_embeddings( + self.draft_model_config.hf_config, + self.target_model_config.max_model_len, + ) + # Replace hf_config for EAGLE draft_model if self.method in ("eagle", "eagle3", "dflash"): from vllm.transformers_utils.configs.eagle import EAGLEConfig @@ -934,6 +954,7 @@ def __post_init__(self): if self.method == "dspark" and ( "Qwen3DSparkModel" not in self.draft_model_config.architectures and "Gemma4DSparkModel" not in self.draft_model_config.architectures + and "K3DSparkModel" not in self.draft_model_config.architectures ): # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config # and its weights ship in the target checkpoint. @@ -963,6 +984,16 @@ def __post_init__(self): if self.method in ("dflash", "dspark"): self.parallel_drafting = True + if ( + self.method == "dspark" + and "K3DSparkModel" in self.draft_model_config.architectures + and self.target_parallel_config.decode_context_parallel_size > 1 + ): + raise ValueError( + "MLA DSpark does not currently support decode context " + "parallelism; set decode_context_parallel_size=1." + ) + if self.num_speculative_tokens is not None and hasattr( self.draft_model_config.hf_config, "num_lookahead_tokens" ): @@ -1130,6 +1161,40 @@ def _maybe_override_draft_max_model_len( ) return result + @staticmethod + def _maybe_override_draft_max_position_embeddings( + draft_hf_config: PretrainedConfig, + target_max_model_len: int, + ) -> None: + """Raise an EAGLE draft's max_position_embeddings up to the target's. + + The proposer feeds the draft positions up to the target's + max_model_len, while max_position_embeddings sizes the draft's + rotary cos_sin_cache. A smaller checkpoint value (e.g. 2048 for + yuhuili/EAGLE3-LLaMA3.1-Instruct-8B) makes that cache gather go + out of bounds (#48894). + + Args: + draft_hf_config: The draft model's HF config, mutated in place. + target_max_model_len: The target model's max_model_len. + """ + draft_max_position_embeddings = getattr( + draft_hf_config, "max_position_embeddings", None + ) + if ( + draft_max_position_embeddings is None + or draft_max_position_embeddings >= target_max_model_len + ): + return + logger.info( + "Overriding draft model max_position_embeddings from %d to the " + "target model's max_model_len (%d); EAGLE drafts share the " + "target's positional space.", + draft_max_position_embeddings, + target_max_model_len, + ) + draft_hf_config.max_position_embeddings = target_max_model_len + @staticmethod def _verify_and_get_draft_tp( target_parallel_config: ParallelConfig, @@ -1345,6 +1410,14 @@ def uses_extract_hidden_states(self) -> bool: def use_ngram_gpu(self) -> bool: return self.method == "ngram_gpu" + def use_multi_module_mtp(self) -> bool: + if self.method != "mtp" or self.draft_model_config is None: + return False + num_mtp_layers = getattr( + self.draft_model_config.hf_config, "num_nextn_predict_layers", 1 + ) + return min(num_mtp_layers, self.num_speculative_tokens) > 1 + def __repr__(self) -> str: method = self.method model = ( diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 3df0f7210f7b..6ab346d1a030 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -13,7 +13,7 @@ from collections.abc import Callable, Mapping, Sequence, Set from dataclasses import MISSING, field, fields, is_dataclass from itertools import pairwise -from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, get_type_hints, overload import torch from pydantic import ConfigDict @@ -227,22 +227,38 @@ class SupportsMetricsInfo(Protocol): def metrics_info(self) -> dict[str, str]: ... -def update_config(config: ConfigT, overrides: dict[str, Any]) -> ConfigT: - processed_overrides = {} +def update_config(config: ConfigT, overrides: Mapping[str, Any]) -> ConfigT: + return _update_config(config, overrides, type(config).__name__) + + +def _update_config( + config: ConfigT, overrides: Mapping[str, Any], config_path: str +) -> ConfigT: + processed_overrides: dict[str, Any] = {} + field_types = get_type_hints(type(config)) for field_name, value in overrides.items(): - assert hasattr(config, field_name), ( - f"{type(config)} has no field `{field_name}`" - ) + field_path = f"{config_path}.{field_name}" + if not hasattr(config, field_name): + raise ValueError(f"{field_path} is not a valid config field") + current_value = getattr(config, field_name) - if is_dataclass(current_value) and not is_dataclass(value): - assert isinstance(value, dict), ( - f"Overrides to {type(config)}.{field_name} must be a dict" - f" or {type(current_value)}, but got {type(value)}" - ) - value = update_config( - current_value, # type: ignore[type-var] - value, - ) + if is_dataclass(current_value): + expected_type = field_types[field_name] + if isinstance(value, Mapping): + value = _update_config( + current_value, # type: ignore[type-var] + value, + field_path, + ) + elif not isinstance(value, expected_type): + expected_type_name = getattr( + expected_type, "__name__", str(expected_type) + ) + raise ValueError( + f"Override for {field_path} must be a mapping or " + f"{expected_type_name}, got {type(value).__name__}" + ) + processed_overrides[field_name] = value return replace(config, **processed_overrides) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4b4a97f41b6b..d6e55798cae8 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -32,13 +32,14 @@ from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig from .diffusion import DiffusionConfig +from .ec_manager_config import EncoderCacheManagerConfig from .ec_transfer import ECTransferConfig from .kernel import KernelConfig from .kv_events import KVEventsConfig from .kv_transfer import KVTransferConfig from .load import LoadConfig from .lora import LoRAConfig -from .mamba import MambaConfig +from .mamba import MambaBackendEnum, MambaConfig from .model import ModelConfig from .observability import ObservabilityConfig from .offload import OffloadConfig @@ -71,11 +72,34 @@ "GraniteMoeForCausalLM", "InklingForCausalLM", "InklingForConditionalGeneration", + "KimiK3ForConditionalGeneration", "LongcatFlashNgramForCausalLM", "Qwen2MoeForCausalLM", } ) +# Architectures that default to V1 on ROCm: the V2 runner faults during the +# profile run. VLLM_USE_V2_MODEL_RUNNER=1 still forces V2. +# TODO: fix V2 enablement +ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "KimiK3ForConditionalGeneration", + } +) + + +@lru_cache +def default_v2_model_runner_architectures() -> frozenset[str]: + """Architectures defaulting to the V2 model runner on this platform.""" + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + return ( + DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES + - ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES + ) + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -187,10 +211,10 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool: def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: - """Enable MLA dual RMS norm fusion when AITer has fused_qk_rmsnorm.""" - from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops + """Enable MLA dual RMS norm fusion on ROCm with AITER.""" + from vllm._aiter_ops import rocm_aiter_ops - return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm() + return rocm_aiter_ops.is_enabled() def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: @@ -367,6 +391,10 @@ class VllmConfig: """The configurations for event publishing.""" ec_transfer_config: ECTransferConfig | None = None """The configurations for distributed EC cache transfer.""" + ec_manager_config: EncoderCacheManagerConfig = Field( + default_factory=EncoderCacheManagerConfig + ) + """The configurations for custom encoder cache manager.""" reasoning_config: ReasoningConfig | None = None """The configurations for reasoning model.""" # some opaque config, only used to provide additional information @@ -552,6 +580,10 @@ def use_v2_model_runner(self) -> bool: if use_v2_model_runner is not None: return use_v2_model_runner + # PCP runtime support is implemented only by the V2 model runner. + if self.parallel_config.prefill_context_parallel_size > 1: + return True + # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4 # is not otherwise a default-V2 architecture, so force V2 for it. If V2 # is unsupported for the rest of the config, _validate_v2_model_runner @@ -610,16 +642,20 @@ def _is_default_v2_model_runner_model(self) -> bool: if model_config.runner_type != "generate": return False - if getattr(model_config, "is_hybrid", False): + architectures = getattr(model_config, "architectures", []) + default_architectures = default_v2_model_runner_architectures() + is_default_v2_architecture = any( + arch in default_architectures for arch in architectures + ) + + if getattr(model_config, "is_hybrid", False) and ( + not is_default_v2_architecture + ): return False if getattr(model_config, "is_attention_free", False): return False - architectures = getattr(model_config, "architectures", []) - return ( - any(arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures) - or not model_config.is_moe - ) + return is_default_v2_architecture or not model_config.is_moe @property def needs_dp_coordinator(self) -> bool: @@ -1106,11 +1142,6 @@ def __post_init__(self): else: self.scheduler_config.async_scheduling = True - logger.info_once( - "Asynchronous scheduling is %s.", - "enabled" if self.scheduler_config.async_scheduling else "disabled", - ) - if self.parallel_config.disable_nccl_for_dp_synchronization is None: if self.scheduler_config.async_scheduling: if self.parallel_config.data_parallel_size > 1 and ( @@ -1160,7 +1191,7 @@ def __post_init__(self): ) if self.model_config is not None and self.model_config.enforce_eager: - logger.warning( + logger.warning_once( "Enforce eager set, disabling torch.compile and CUDAGraphs. " "This is equivalent to setting -cc.mode=none -cc.cudagraph_mode=none" ) @@ -1168,7 +1199,7 @@ def __post_init__(self): self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE if os.environ.get("TORCH_COMPILE_DISABLE") == "1": - logger.warning( + logger.warning_once( "TORCH_COMPILE_DISABLE is set, disabling torch.compile. " "This is equivalent to setting -cc.mode=none" ) @@ -1187,6 +1218,9 @@ def __post_init__(self): "DeepSeekV4MTPModel", "InklingForCausalLM", "InklingForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "KimiK3MTPModel", + "KimiLinearForCausalLM", "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", ) @@ -1210,7 +1244,7 @@ def __post_init__(self): self.compilation_config.mode is not None and self.compilation_config.mode != CompilationMode.VLLM_COMPILE ): - logger.warning( + logger.warning_once( "Inductor compilation was disabled by user settings, " "optimizations settings that are only active during " "inductor compilation will be ignored." @@ -1278,7 +1312,7 @@ def has_blocked_weights(): and self.compilation_config.mode != CompilationMode.VLLM_COMPILE and not envs.VLLM_USE_BREAKABLE_CUDAGRAPH ): - logger.info( + logger.info_once( "Cudagraph mode %s is not compatible with compilation mode %s." "Overriding to NONE.", self.compilation_config.cudagraph_mode, @@ -1292,7 +1326,7 @@ def has_blocked_weights(): pass_config.enable_sp = True if pass_config.enable_sp: if self.parallel_config.tensor_parallel_size == 1: - logger.warning("Sequence Parallelism requires TP>1, disabling") + logger.warning_once("Sequence Parallelism requires TP>1, disabling") pass_config.enable_sp = False pass_config.fuse_gemm_comms = False else: @@ -1310,7 +1344,7 @@ def has_blocked_weights(): ) if pass_config.sp_min_token_num is None: - logger.warning( + logger.warning_once( "Model hidden_size too small for the SP " "threshold heuristic, disabling. To force SP, " "set pass_config.sp_min_token_num manually." @@ -1389,7 +1423,7 @@ def has_blocked_weights(): # disable cudagraph when enforce eager execution if self.model_config is not None and self.model_config.enforce_eager: - logger.info("Cudagraph is disabled under eager mode") + logger.info_once("Cudagraph is disabled under eager mode") self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE # override related settings when enforce eager self.compilation_config.max_cudagraph_capture_size = 0 @@ -1424,7 +1458,7 @@ def has_blocked_weights(): and self.model_config.architecture == "WhisperForConditionalGeneration" and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn" ): - logger.warning( + logger.warning_once( "Whisper is known to have issues with " "forked workers. If startup is hanging, " "try setting 'VLLM_WORKER_MULTIPROC_METHOD' " @@ -1436,7 +1470,7 @@ def has_blocked_weights(): and self.kv_events_config.enable_kv_cache_events and not self.cache_config.enable_prefix_caching ): - logger.warning( + logger.warning_once( "KV cache events are on, but prefix caching is not enabled. " "Use --enable-prefix-caching to enable." ) @@ -1445,7 +1479,7 @@ def has_blocked_weights(): and self.kv_events_config.publisher != "null" and not self.kv_events_config.enable_kv_cache_events ): - logger.warning( + logger.warning_once( "KV cache events are disabled, " "but the scheduler is configured to publish them. " "Modify KVEventsConfig.enable_kv_cache_events " @@ -1455,6 +1489,11 @@ def has_blocked_weights(): if self.use_v2_model_runner: self._validate_v2_model_runner() + elif self.parallel_config.prefill_context_parallel_size > 1: + raise ValueError( + "Prefill context parallelism requires Model Runner V2. " + "Remove VLLM_USE_V2_MODEL_RUNNER=0." + ) # Re-compute compile ranges after platform-specific config updates # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled) @@ -1478,7 +1517,7 @@ def has_blocked_weights(): # the pass will operate on higher-level IR to avoid the issue. # TODO: https://github.com/vllm-project/vllm/issues/27894 if self.compilation_config.mode != CompilationMode.VLLM_COMPILE: - logger.warning( + logger.warning_once( "Sequence parallelism is enabled, but running in wrong " "vllm compile mode: %s.", self.compilation_config.mode, @@ -2263,14 +2302,6 @@ def validate_block_size(self) -> None: # Mamba cache align-mode constraints if self.cache_config.mamba_cache_mode == "align": - assert block_size <= self.scheduler_config.max_num_batched_tokens, ( - "In Mamba cache align mode, block_size " - f"({block_size}) must be <= " - "max_num_batched_tokens " - f"({self.scheduler_config.max_num_batched_tokens})." - ) - if self.scheduler_config.long_prefill_token_threshold > 0: - assert self.scheduler_config.long_prefill_token_threshold >= block_size assert not self.scheduler_config.disable_chunked_mm_input, ( "Chunked MM input is required because we need the flexibility " "to schedule a multiple of block_size tokens even if they are " @@ -2303,6 +2334,37 @@ def validate_mamba_block_size(self) -> "VllmConfig": ) return self + @model_validator(mode="after") + def validate_mamba_cached_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm: + return self + # ReplaySSM adds a 3-tensor ring to the mamba state; only models that + # opt in (supports_replayssm) build a consistent shape on both the layer + # and config paths. Reject others so the mamba page size cannot desync. + if self.model_config is not None and not self.model_config.supports_replayssm: + raise ValueError( + "--use-replayssm is only supported for Nemotron-H models " + f"(got architecture {self.model_config.architecture!r})" + ) + if self.cache_config.mamba_cache_mode == "all": + raise ValueError( + "--use-replayssm supports prefix caching only in align mode; " + "pass --mamba-cache-mode align" + ) + if self.num_speculative_tokens > 0: + raise ValueError("--use-replayssm does not support speculative decoding") + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError("--use-replayssm requires --mamba-backend triton") + if ( + self.kv_transfer_config is not None + and self.kv_transfer_config.is_kv_transfer_instance + ): + raise ValueError( + "--use-replayssm is incompatible with KV connectors " + "(P/D disaggregation, KV cache offload)" + ) + return self + _current_vllm_config: VllmConfig | None = None _current_prefix: str | None = None diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 7c4fedd34a3b..2cb9805bae39 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -291,6 +291,9 @@ def wake_up(self, tags: list[str] | None = None) -> None: back to GPU memory. If None, all memory allocation will be loaded back to GPU memory. """ + gc.collect() + torch.accelerator.empty_cache() + for ptr, data in self.pointer_to_data.items(): if tags is None or data.tag in tags: handle = data.handle diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 33841de306e3..5abe7568a292 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -8,6 +8,7 @@ import torch.distributed as dist import vllm.envs as envs +from vllm.config import get_current_vllm_config from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group from vllm.distributed.utils import StatelessProcessGroup from vllm.forward_context import get_forward_context @@ -278,7 +279,9 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) - self.support_fault_tolerance = False # TODO: set to True when FT is supported. + self.support_fault_tolerance = ( + get_current_vllm_config().parallel_config.enable_fault_tolerance + ) def _make_all2all_kwargs( self, @@ -360,6 +363,16 @@ def query_fault(self) -> torch.Tensor: has_fault = (current != DeepEPLLAll2AllManager._last_mask).any() return has_fault + def clean_buffers(self) -> None: + buf = DeepEPLLAll2AllManager._buffer + if buf is None: + return + buf.get_local_buffer_tensor(dtype=torch.int8, use_rdma_buffer=True).zero_() + torch.accelerator.synchronize() + buf.low_latency_clean_mask_buffer() + torch.accelerator.synchronize() + DeepEPLLAll2AllManager._last_mask = None + @dataclass class _NixlEPBufferState: @@ -565,6 +578,18 @@ def query_fault(self) -> torch.Tensor: has_fault = (current != last).any() return has_fault + def clean_buffers(self) -> None: + if NixlEPAll2AllManager._buffer is None: + return + state = NixlEPAll2AllManager._buffer + state.buffer.get_local_buffer_tensor( + dtype=torch.int8, use_rdma_buffer=True + ).zero_() + torch.accelerator.synchronize() + state.buffer.clean_mask_buffer() + torch.accelerator.synchronize() + NixlEPAll2AllManager._last_mask = None + class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): """ @@ -860,6 +885,20 @@ def cleanup(self): self.mapping = None self.initialized = False + def checkpoint_prepare(self) -> None: + if self.initialized: + assert self.moe_alltoall is not None + self.moe_alltoall.checkpoint_prepare() + + def checkpoint_restore(self) -> None: + if self.initialized: + assert self.moe_alltoall is not None + from vllm.distributed.device_communicators.mnnvl_compat import ( + CustomCommunicator, + ) + + self.moe_alltoall.checkpoint_restore(CustomCommunicator(self.cpu_group)) + class MoriAll2AllManager(All2AllManagerBase): def __init__(self, cpu_group, all2all_backend: str): @@ -978,6 +1017,7 @@ def __init__(self, cpu_group, tcp_store_group=None, device_group=None): self._device_group = device_group self.handle_cache = Cache() self._num_sms: int | None = None + self._gin_checked = False def _make_all2all_kwargs( self, @@ -1000,11 +1040,37 @@ def _make_all2all_kwargs( explicitly_destroy=True, ) + def _check_gin_support(self, group) -> None: + from vllm.utils.nccl import query_nccl_gin_type + + # ProcessGroupNCCL creates communicators lazily. Initialize this exact + # group before querying so a null comm pointer is not mistaken for + # missing GIN support. + probe = torch.zeros(1, device="cuda") + torch.distributed.all_reduce(probe, group=group) + + gin_type = query_nccl_gin_type(group) + if gin_type is None: + raise RuntimeError( + "DeepEPv2 communicator properties query failed; " + "networking capability could not be determined." + ) + if gin_type == 0: + raise RuntimeError( + "DeepEPv2 requires NCCL GIN (GPU-Initiated Networking). " + "This usually means IBGDA-capable InfiniBand NICs or drivers " + "are not available. See tools/ep_kernels/README.md for " + "requirements." + ) + def get_handle(self, kwargs): import deep_ep # type: ignore[import-not-found] num_experts = kwargs.pop("num_experts", 256) buffer_kwargs = self._make_all2all_kwargs(**kwargs) + if not self._gin_checked: + self._check_gin_support(buffer_kwargs["group"]) + self._gin_checked = True logger.debug("DeepEP v2 all2all args %s", buffer_kwargs) handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create( buffer_kwargs, deep_ep.ElasticBuffer diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index 423cc5373451..86b709ed41c7 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -47,6 +47,12 @@ 6: 8 * MiB, # 8 MB 8: 4 * MiB, # 4 MB }, + "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds + 2: 4 * MiB, # 4 MB + 4: 4 * MiB, # 4 MB + 6: 8 * MiB, # 8 MB + 8: 4 * MiB, # 4 MB + }, } SYMM_MEM_ALL_REDUCE_MAX_SIZES = { @@ -68,6 +74,12 @@ 6: 32 * MiB, # 32 MB 8: 64 * MiB, # 64 MB }, + "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds + 2: 4 * MiB, # 4 MB + 4: 32 * MiB, # 32 MB + 6: 32 * MiB, # 32 MB + 8: 64 * MiB, # 64 MB + }, } # NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks. diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 70f1fb5d62c5..73fd1331f5ce 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -7,8 +7,11 @@ import torch.distributed as dist from torch.distributed import ProcessGroup +from vllm.logger import init_logger from vllm.utils import is_moe_layer +logger = init_logger(__name__) + class Cache: def __init__(self): @@ -105,10 +108,30 @@ def dispatch( raise NotImplementedError def query_active_mask(self) -> torch.Tensor: + """Return the all2all liveness mask for the EP ranks. + + Returns: + An int32 device tensor where 0 marks a live rank and 1 marks a + masked (dead/unreachable) rank. + """ raise NotImplementedError def query_fault(self) -> torch.Tensor: - """Returns has_fault scalar.""" + """Return a scalar bool tensor, True if a new fault appeared. + + Compares the current mask against the baseline recorded at the last + recovery point. + """ + raise NotImplementedError + + def clean_buffers(self) -> None: + """Reset this rank's RDMA buffers and all2all mask state (rank-local). + + Post-fault cleanup: a dispatch/combine that hit a dead peer or timed + out can leave partially-written or stale tokens in the RDMA receive + buffer, so it is zeroed to stop the next forward from reading that + contaminated data. + """ raise NotImplementedError def set_num_sms(self, num_sms: int): @@ -117,6 +140,18 @@ def set_num_sms(self, num_sms: int): def max_sms_used(self) -> int | None: return None # None means it could use the whole GPU + def checkpoint_prepare(self) -> None: + logger.warning_once( + "%s.checkpoint_prepare is not implemented; skipping.", + type(self).__name__, + ) + + def checkpoint_restore(self) -> None: + logger.warning_once( + "%s.checkpoint_restore is not implemented; skipping.", + type(self).__name__, + ) + def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False): raise NotImplementedError @@ -140,6 +175,7 @@ def __init__( unique_name: str = "", global_ranks: list[int] | None = None, global_world_size: int | None = None, + use_all2all: bool = False, ): self.device = device or torch.device("cpu") self.cpu_group = cpu_group @@ -169,26 +205,15 @@ def __init__( self.global_world_size = dist.get_world_size() self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) - use_ep = False all2all_backend = None from vllm.config import get_current_vllm_config_or_none config = get_current_vllm_config_or_none() if config is not None: - # initialize the all2all manager for DP or sequence-parallel EP. - parallel_config = config.parallel_config - use_ep = ( - parallel_config.data_parallel_size > 1 - or parallel_config.use_sequence_parallel_moe - or ( - parallel_config.enable_expert_parallel - and parallel_config.prefill_context_parallel_size > 1 - ) - ) - all2all_backend = parallel_config.all2all_backend + all2all_backend = config.parallel_config.all2all_backend self.is_ep_communicator = unique_name.split(":")[0] == "ep" - self.use_all2all = self.is_ep_communicator and use_ep + self.use_all2all = self.is_ep_communicator and use_all2all self.all2all_backend = all2all_backend self.all2all_manager: All2AllManagerBase | None = None @@ -196,6 +221,12 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: dist.all_reduce(input_, group=self.device_group) return input_ + def checkpoint_prepare(self) -> None: + """Prepare reclaimable communicator state for checkpoint (default: no-op).""" + + def checkpoint_restore(self) -> None: + """Restore communicator state after checkpoint (default: no-op).""" + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: if dim < 0: # Convert negative dim to positive. diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 9ec4b72f80d5..8ea12d9255a9 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -24,8 +24,11 @@ def __init__( device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.dist_module = torch.distributed if ( diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b6138194775d..06b441c5a416 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -36,6 +36,7 @@ def __init__( global_ranks: list[int] | None = None, global_world_size: int | None = None, tcp_store_group: StatelessProcessGroup | None = None, + use_all2all: bool = False, ): super().__init__( cpu_group, @@ -44,6 +45,7 @@ def __init__( unique_name, global_ranks, global_world_size, + use_all2all=use_all2all, ) if "tp" not in unique_name: # custom allreduce or torch symm mem can be used only by tp @@ -338,6 +340,18 @@ def all_reduce(self, input_): torch.distributed.all_reduce(out, group=self.device_group) return out + def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None: + ca_comm = self.ca_comm + if ca_comm is None: + return None + return ca_comm.custom_all_gather(input_.contiguous()) + + def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None: + ca_comm = self.ca_comm + if ca_comm is None: + return None + return ca_comm.custom_reduce_scatter(input_.contiguous()) + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: # Route uniform dim-0 all-gathers through NVLS symmetric memory when # enabled (mirrors reduce_scatter); otherwise fall back to the @@ -571,6 +585,22 @@ def destroy(self): self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] + def checkpoint_prepare(self) -> None: + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. + from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces + + checkpoint_prepare_fi_ar_workspaces(self.cpu_group) + if self.all2all_manager is not None: + self.all2all_manager.checkpoint_prepare() + + def checkpoint_restore(self) -> None: + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. + from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces + + checkpoint_restore_fi_ar_workspaces(self.cpu_group) + if self.all2all_manager is not None: + self.all2all_manager.checkpoint_restore() + def all_gatherv( self, input_: torch.Tensor | list[torch.Tensor], diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 95db6cc92459..86180a51a71e 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -25,6 +25,11 @@ # For CPUs custom_ar = False +try: + import torch.distributed._symmetric_memory as torch_symm_mem +except ImportError: + torch_symm_mem = None + logger = init_logger(__name__) @@ -49,7 +54,17 @@ def _can_p2p(rank: int, world_size: int) -> bool: class CustomAllreduce: - _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8] + _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16] + _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024 + _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = { + 2: 8 * 1024 * 1024, + 4: 4 * 1024 * 1024, + 6: 2 * 1024 * 1024, + 8: 2 * 1024 * 1024, + 16: 2 * 1024 * 1024, + } + _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 + _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 # max_size: max supported allreduce size def __init__( @@ -57,6 +72,10 @@ def __init__( group: ProcessGroup, device: int | str | torch.device, max_size=8192 * 1024, + max_all_gather_size=_DEFAULT_ALL_GATHER_MAX_SIZE, + max_mnnvl_all_gather_size=None, + max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE, + max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE, symm_mem_enabled=False, ) -> None: """ @@ -70,12 +89,25 @@ def __init__( are in the same node. """ self._IS_CAPTURING = False + self._ptr = 0 self.disabled = True + self.mnnvl_buffer = None + self.mnnvl_handle = None + self.mnnvl_peer_buffers: list[torch.Tensor] | None = None + self.mnnvl_multicast_ptr = 0 + self.mnnvl_buffer_size = 0 + self.mnnvl_lamport_ag_local_ptr = 0 + self.mnnvl_lamport_ag_multicast_ptr = 0 + self.mnnvl_lamport_rs_local_ptr = 0 + self.mnnvl_lamport_epochs = None + self.mnnvl_lamport_ag_epoch_ptr = 0 + self.mnnvl_lamport_rs_epoch_ptr = 0 + self.mnnvl_only = False if not custom_ar: # disable because of missing custom allreduce library # e.g. in a non-GPU environment - logger.info( + logger.info_once( "Custom allreduce is disabled because " "of missing custom allreduce library" ) @@ -87,13 +119,8 @@ def __init__( "CustomAllreduce should be attached to a non-NCCL group." ) - if not all(in_the_same_node_as(group, source_rank=0)): - # No need to initialize custom allreduce for multi-node case. - logger.warning( - "Custom allreduce is disabled because this process group" - " spans across nodes." - ) - return + same_node = all(in_the_same_node_as(group, source_rank=0)) + self.mnnvl_only = not same_node rank = dist.get_rank(group=self.group) self.rank = rank @@ -103,7 +130,7 @@ def __init__( return if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES: - logger.warning( + logger.warning_once( "Custom allreduce is disabled due to an unsupported world" " size: %d. Supported world sizes: %s. To silence this " "warning, specify disable_custom_all_reduce=True explicitly.", @@ -126,28 +153,30 @@ def __init__( and device_capability is not None ): device_capability_str = device_capability.as_version_str() - if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES: + if ( + device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES + and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str] + ): max_size = min( CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size], max_size, ) # device.index is a visible ordinal, not a logical local ID. - physical_device_id = current_platform.visible_device_id_to_physical_device_id( - device.index - ) - tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") - gather_list = [ - torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) - ] - dist.all_gather(gather_list, tensor, group=self.group) - physical_device_ids = [t.item() for t in gather_list] - - # test nvlink first, this will filter out most of the cases - # where custom allreduce is not supported - # this checks hardware and driver support for NVLink - assert current_platform.is_cuda_alike() - fully_connected = current_platform.is_fully_connected(physical_device_ids) - if world_size > 2 and not fully_connected: + fully_connected = False + if same_node: + physical_device_id = ( + current_platform.visible_device_id_to_physical_device_id(device.index) + ) + tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") + gather_list = [ + torch.tensor([0], dtype=torch.int, device="cpu") + for _ in range(world_size) + ] + dist.all_gather(gather_list, tensor, group=self.group) + physical_device_ids = [t.item() for t in gather_list] + assert current_platform.is_cuda_alike() + fully_connected = current_platform.is_fully_connected(physical_device_ids) + if same_node and world_size > 2 and not fully_connected: logger.warning( "Custom allreduce is disabled because it's not supported on" " more than two PCIe-only GPUs. To silence this warning, " @@ -158,7 +187,11 @@ def __init__( # this is expensive to compute at the first time # then we cache the result # On AMD GPU, p2p is always enabled between XGMI connected GPUs - if not current_platform.is_rocm() and not _can_p2p(rank, world_size): + if ( + same_node + and not current_platform.is_rocm() + and not _can_p2p(rank, world_size) + ): logger.warning( "Custom allreduce is disabled because your platform lacks " "GPU P2P capability or P2P test failed. To silence this " @@ -170,21 +203,40 @@ def __init__( # Buffers memory are owned by this Python class and passed to C++. # Metadata composes of two parts: metadata for synchronization and a # temporary buffer for storing intermediate allreduce results. - self.meta_ptrs = self.create_shared_buffer( - ops.meta_size() + max_size, group=group, uncached=True - ) + if same_node: + self.meta_ptrs = self.create_shared_buffer( + ops.meta_size() + max_size, group=group, uncached=True + ) + else: + meta_ptr, _ = ops.allocate_shared_buffer_and_handle(ops.meta_size()) + self.meta_ptrs = [meta_ptr] * world_size # This is a pre-registered IPC buffer. In eager mode, input tensors - # are first copied into this buffer before allreduce is performed - self.buffer_ptrs = self.create_shared_buffer(max_size, group=group) - # This is a buffer for storing the tuples of pointers pointing to - # IPC buffers from all ranks. Each registered tuple has size of - # 8*world_size bytes where world_size is at most 8. Allocating 8MB - # is enough for 131072 such tuples. The largest model I've seen only - # needs less than 10000 of registered tuples. + # are first copied into this buffer before the operation is performed + legacy_buffer_size = max(max_size, max_all_gather_size, max_reduce_scatter_size) + if same_node: + self.buffer_ptrs = self.create_shared_buffer( + legacy_buffer_size, + group=group, + ) + else: + buffer_ptr, _ = ops.allocate_shared_buffer_and_handle(legacy_buffer_size) + self.buffer_ptrs = [buffer_ptr] * world_size + # This stores tuples of pointers to IPC buffers from all ranks. + # Each registered tuple contains at most 16 addresses. + # Allocating 8MB is enough for 65536 such tuples. The largest model uses + # fewer than 10000 registered tuples. self.rank_data = torch.empty( 8 * 1024 * 1024, dtype=torch.uint8, device=self.device ) self.max_size = max_size + self.max_all_gather_size = max_all_gather_size + if max_mnnvl_all_gather_size is None: + max_mnnvl_all_gather_size = self._DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES[ + world_size + ] + self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size + self.max_reduce_scatter_size = max_reduce_scatter_size + self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size self.rank = rank self.world_size = world_size self.fully_connected = fully_connected @@ -192,6 +244,72 @@ def __init__( self.meta_ptrs, self.rank_data, rank, self.fully_connected ) ops.register_buffer(self._ptr, self.buffer_ptrs) + self._init_mnnvl_buffer( + max( + max_mnnvl_all_gather_size * world_size, + max_mnnvl_reduce_scatter_size, + ) + ) + if not same_node and not self.mnnvl_multicast_ptr: + logger.warning( + "Custom collectives are disabled because this multi-node " + "group does not support MNNVL multicast." + ) + self.close() + self.disabled = True + + def _init_mnnvl_buffer(self, stage_size: int) -> None: + if torch_symm_mem is None or not current_platform.is_cuda(): + return + try: + buffer_size = stage_size * 6 + buffer = torch_symm_mem.empty( + buffer_size, dtype=torch.uint8, device=self.device + ) + handle = torch_symm_mem.rendezvous(buffer, self.group.group_name) + if handle.multicast_ptr == 0: + return + peer_buffers = [ + handle.get_buffer( + peer, + (buffer_size,), + torch.uint8, + storage_offset=0, + ) + for peer in range(self.world_size) + ] + ptrs = [peer_buffer.data_ptr() for peer_buffer in peer_buffers] + lamport_ag_offset = 0 + lamport_rs_offset = stage_size * 3 + lamport_ag_ptrs = [ptr + lamport_ag_offset for ptr in ptrs] + lamport_rs_ptrs = [ptr + lamport_rs_offset for ptr in ptrs] + ops.register_buffer(self._ptr, lamport_ag_ptrs) + ops.register_buffer(self._ptr, lamport_rs_ptrs) + + buffer.view(torch.int32).fill_(-2147483648) + epochs = torch.zeros( + (2, 32), + dtype=torch.int32, + device=self.device, + ) + torch.accelerator.synchronize() + dist.barrier(group=self.group) + + self.mnnvl_buffer = buffer + self.mnnvl_handle = handle + self.mnnvl_peer_buffers = peer_buffers + self.mnnvl_multicast_ptr = handle.multicast_ptr + self.mnnvl_buffer_size = stage_size + self.mnnvl_lamport_ag_local_ptr = lamport_ag_ptrs[self.rank] + self.mnnvl_lamport_ag_multicast_ptr = ( + handle.multicast_ptr + lamport_ag_offset + ) + self.mnnvl_lamport_rs_local_ptr = lamport_rs_ptrs[self.rank] + self.mnnvl_lamport_epochs = epochs + self.mnnvl_lamport_ag_epoch_ptr = epochs[0].data_ptr() + self.mnnvl_lamport_rs_epoch_ptr = epochs[1].data_ptr() + except RuntimeError as error: + logger.debug("MNNVL AG/RS initialization failed: %s", error) @contextmanager def capture(self): @@ -210,7 +328,7 @@ def capture(self): def register_graph_buffers(self): handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) - logger.info("Registering %d cuda graph addresses", len(offset)) + logger.debug("Registering %d cuda graph addresses", len(offset)) # We cannot directly use `dist.all_gather_object` here # because it is incompatible with `gloo` backend under inference mode. # see https://github.com/pytorch/pytorch/issues/126032 for details. @@ -228,7 +346,7 @@ def register_graph_buffers(self): ops.register_graph_buffers(self._ptr, handles, offsets) def should_custom_ar(self, inp: torch.Tensor): - if self.disabled: + if self.disabled or self.world_size > 8: return False inp_size = inp.numel() * inp.element_size() # custom allreduce requires input byte size to be multiples of 16 @@ -279,6 +397,111 @@ def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None: # latency) compared to the performance gain of using custom kernels return self.all_reduce(input, registered=False) + def should_custom_all_gather(self, inp: torch.Tensor) -> bool: + if self.disabled or not current_platform.is_cuda(): + return False + if self.world_size == 16 and not self.mnnvl_only: + return False + inp_size = inp.nbytes + if inp.dtype not in ( + torch.float32, + torch.float16, + torch.bfloat16, + ): + return False + max_size = ( + self.max_mnnvl_all_gather_size + if self.mnnvl_multicast_ptr + else self.max_all_gather_size + ) + return ( + 0 < inp_size <= max_size + and inp_size % 16 == 0 + and is_weak_contiguous(inp) + and (self.fully_connected or bool(self.mnnvl_multicast_ptr)) + ) + + def custom_all_gather(self, inp: torch.Tensor) -> torch.Tensor | None: + if not self.should_custom_all_gather(inp): + return None + out_shape = (inp.shape[0] * self.world_size,) + inp.shape[1:] + if self.mnnvl_multicast_ptr: + logger.info_once( + "Using the MNNVL Lamport all-gather kernel.", + scope="global", + ) + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + ops.mnnvl_lamport_all_gather( + self._ptr, + inp, + out, + self.mnnvl_lamport_ag_local_ptr, + self.mnnvl_lamport_ag_multicast_ptr, + self.mnnvl_lamport_ag_epoch_ptr, + self.mnnvl_buffer_size, + ) + else: + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + ops.custom_all_gather( + self._ptr, + inp, + out, + self.buffer_ptrs[self.rank], + self.max_all_gather_size, + ) + return out + + def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool: + if self.disabled or not current_platform.is_cuda(): + return False + if self.world_size == 16 and not self.mnnvl_only: + return False + inp_size = inp.nbytes + if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return False + if inp.shape[0] % self.world_size != 0: + return False + output_size = inp_size // self.world_size + max_size = ( + self.max_mnnvl_reduce_scatter_size + if self.mnnvl_multicast_ptr + else self.max_reduce_scatter_size + ) + return ( + 0 < inp_size <= max_size + and output_size % 16 == 0 + and is_weak_contiguous(inp) + and (self.fully_connected or bool(self.mnnvl_multicast_ptr)) + ) + + def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None: + if not self.should_custom_reduce_scatter(inp): + return None + out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:] + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + if self.mnnvl_multicast_ptr: + logger.info_once( + "Using the MNNVL Lamport reduce-scatter kernel.", + scope="global", + ) + ops.mnnvl_lamport_reduce_scatter( + self._ptr, + inp, + out, + self.mnnvl_lamport_rs_local_ptr, + self.mnnvl_lamport_rs_epoch_ptr, + self.mnnvl_buffer_size, + ) + else: + ops.custom_reduce_scatter( + self._ptr, + inp, + out, + self.buffer_ptrs[self.rank], + self.max_reduce_scatter_size, + ) + return out + def close(self): if not self.disabled and self._ptr: if ops is not None: @@ -286,6 +509,10 @@ def close(self): self._ptr = 0 self.free_shared_buffer(self.meta_ptrs, rank=self.rank) self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) + self.mnnvl_peer_buffers = None + self.mnnvl_handle = None + self.mnnvl_buffer = None + self.mnnvl_lamport_epochs = None def __del__(self): self.close() diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index 881b265d2426..b7fe09081ac6 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -6,6 +6,7 @@ import os import random import threading +from typing import Any import torch import torch.distributed as dist @@ -39,6 +40,7 @@ # allreduce backend or a fallback backend when the primary workspace is not # available on the current topology. _fi_ar_quant_workspace = None +_fi_ar_workspace_groups: dict[int, ProcessGroup] = {} def _create_workspace( @@ -65,6 +67,12 @@ def _create_workspace( comm_backend=comm_backend, group=group, ) + if backend == "mnnvl" and not getattr(workspace, "mc_ptr", 0): + workspace.destroy() + logger.warning_once( + "FlashInfer MNNVL multicast is unavailable on the current topology." + ) + return None except Exception as e: if "multicast" in str(e).lower(): logger.warning_once( @@ -81,6 +89,14 @@ def _create_workspace( return None finally: random.setstate(rng_state) + workspace_id = id(workspace) + workspace_group = _fi_ar_workspace_groups.get(workspace_id) + if workspace_group is not None and workspace_group is not group: + raise RuntimeError( + "FlashInfer returned an all-reduce workspace already associated " + "with a different process group" + ) + _fi_ar_workspace_groups[workspace_id] = group logger.debug( "Initialized FlashInfer All Reduce workspace: backend=%s, " "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s", @@ -105,20 +121,20 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]: """ backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND if backend != "auto": - logger.info_once(f"Using flashinfer allreduce backend: {backend}") + logger.debug_once("Using flashinfer allreduce backend: %s", backend) return backend, False # Default to mnnvl for both single- and multi-node setups. The mnnvl # cudagraph hang that previously forced single-node to trtllm # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in - # FlashInfer (>= 0.6.12, vLLM pins 0.6.13), so mnnvl is safe here. trtllm + # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm # does not support multi-node allreduce, so mnnvl is required there anyway. # mnnvl needs NVSwitch multicast; on single-node topologies without it, # fall back to trtllm so fused allreduce stays enabled. backend = "mnnvl" allow_trtllm_fallback = get_node_count() == 1 - logger.info_once(f"Auto-selected flashinfer allreduce backend: {backend}") + logger.debug_once("Auto-selected flashinfer allreduce backend: %s", backend) return backend, allow_trtllm_fallback @@ -268,6 +284,36 @@ def destroy_fi_ar_workspace(): _fi_ar_quant_workspace.destroy() _fi_ar_workspace = _fi_ar_quant_workspace = None + _fi_ar_workspace_groups.clear() + + +def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]: + workspaces = [_fi_ar_workspace] + if _fi_ar_quant_workspace is not _fi_ar_workspace: + workspaces.append(_fi_ar_quant_workspace) + + group_workspaces = [] + for workspace in workspaces: + if workspace is None: + continue + workspace_group = _fi_ar_workspace_groups.get(id(workspace)) + if workspace_group is None: + raise RuntimeError( + "FlashInfer all-reduce workspace process group was not retained" + ) + if workspace_group is group: + group_workspaces.append(workspace) + return group_workspaces + + +def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None: + for workspace in _fi_ar_workspaces_for_group(group): + workspace.checkpoint_prepare() + + +def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None: + for workspace in _fi_ar_workspaces_for_group(group): + workspace.checkpoint_restore(TorchDistBackend(group=group)) atexit.register(destroy_fi_ar_workspace) diff --git a/vllm/distributed/device_communicators/pynccl_allocator.py b/vllm/distributed/device_communicators/pynccl_allocator.py index a5d7faceb2b9..4898c6212f59 100644 --- a/vllm/distributed/device_communicators/pynccl_allocator.py +++ b/vllm/distributed/device_communicators/pynccl_allocator.py @@ -14,7 +14,7 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.utils.nccl import find_nccl_include_paths +from vllm.utils.nccl import find_nccl_include_paths, find_nccl_library_paths logger = init_logger(__name__) @@ -74,11 +74,15 @@ def compile_nccl_allocator(): out_dir = tempfile.gettempdir() nccl_allocator_libname = "nccl_allocator" nccl_include_paths = find_nccl_include_paths() + ldflags = ["-l:libnccl.so.2"] + nccl_lib_paths = find_nccl_library_paths() + if nccl_lib_paths: + ldflags = [f"-L{p}" for p in nccl_lib_paths] + ldflags load_inline( name=nccl_allocator_libname, cpp_sources=nccl_allocator_source, with_cuda=True, - extra_ldflags=["-lnccl"], + extra_ldflags=ldflags, verbose=envs.VLLM_LOGGING_LEVEL == "DEBUG", is_python_module=False, build_directory=out_dir, diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py index 5ca8cc7c77f4..78ee81e7b5ee 100644 --- a/vllm/distributed/device_communicators/pynccl_wrapper.py +++ b/vllm/distributed/device_communicators/pynccl_wrapper.py @@ -51,6 +51,26 @@ class ncclUniqueId(ctypes.Structure): _fields_ = [("internal", ctypes.c_byte * 128)] +# NCCL 2.30+ ncclCommProperties_t. Only fields through ginType are read; +# trailing fields keep the layout aligned with NCCL's versioned structure. +class ncclCommProperties(ctypes.Structure): + _fields_ = [ + ("size", ctypes.c_size_t), + ("magic", ctypes.c_uint), + ("version", ctypes.c_uint), + ("rank", ctypes.c_int), + ("nRanks", ctypes.c_int), + ("cudaDev", ctypes.c_int), + ("nvmlDev", ctypes.c_int), + ("deviceApiSupport", ctypes.c_bool), + ("multimemSupport", ctypes.c_bool), + ("ginType", ctypes.c_int), + ("nLsaTeams", ctypes.c_int), + ("hostRmaSupport", ctypes.c_bool), + ("railedGinType", ctypes.c_int), + ] + + cudaStream_t = ctypes.c_void_p buffer_type = ctypes.c_void_p @@ -317,6 +337,12 @@ class NCCLLibrary: # ncclResult_t ncclCommWindowDeregister( # ncclComm_t comm, ncclWindow_t win); Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]), + # Query runtime properties of a specific initialized communicator. + Function( + "ncclCommQueryProperties", + ncclResult_t, + [ncclComm_t, ctypes.POINTER(ncclCommProperties)], + ), ] # class attribute to store the mapping from the path to the library @@ -375,6 +401,9 @@ def __init__(self, so_file: str | None = None): # Having an exception here on ROCm platform is # not allowed during graph capturing continue + elif func.name == "ncclCommQueryProperties": + # Optional on NCCL versions older than 2.29. + continue raise NCCLLibrary.path_to_dict_mapping[so_file] = _funcs self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 43e066c44b08..f6d2a24755fa 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copyreg import functools +import io +import os import pickle +import shutil import sys import threading import time @@ -57,6 +61,11 @@ from _typeshed import SizedBuffer VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL +# Cap on how long an idle reader parks before re-reading the authoritative SHM +# written-flag. Bounds lost-notify recovery latency to ~5s while the periodic +# wakeup stays negligible (one flag check per reader every 5s). +SHM_READER_RECHECK_INTERVAL_MS = 5000 + from_bytes_big = functools.partial(int.from_bytes, byteorder="big") @@ -213,6 +222,32 @@ def notify(self): self.local_notify_socket.send(b"\x00") +SHM_PATH = "/dev/shm" + + +def check_shm_free_space(required_bytes: int, shm_path: str = SHM_PATH) -> None: + """Raise if ``shm_path`` cannot fit a ``required_bytes`` shared segment. + + Args: + required_bytes: Size of the shared-memory segment to be created. + shm_path: Mount point backing POSIX shared memory; skipped if absent. + + Raises: + RuntimeError: If ``required_bytes`` exceeds the free space. + """ + if not os.path.isdir(shm_path): + return + free_bytes = shutil.disk_usage(shm_path).free + if required_bytes <= free_bytes: + return + mib = 1 << 20 + raise RuntimeError( + f"Insufficient space in {shm_path}: {required_bytes / mib:.0f} MiB " + f"required, {free_bytes / mib:.0f} MiB free. Increase {shm_path} " + "(e.g. --shm-size or --ipc=host)." + ) + + class ShmRingBuffer: def __init__( self, @@ -283,6 +318,7 @@ def __init__( if name is None: # we are creating a buffer self.is_creator = True + check_shm_free_space(self.total_bytes_of_buffer) self.shared_memory = shared_memory.SharedMemory( create=True, size=self.total_bytes_of_buffer ) @@ -351,6 +387,70 @@ def get_metadata(self, current_idx: int): yield buf +def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor: + """Rebuild a tensor from an out-of-band pickle buffer. + + Counterpart of `_reduce_tensor`. Note that pickle passes the original + buffer-providing object from `loads(buffers=...)` straight to this + function (no `PickleBuffer` wrapper on the receiving side), so `buf` is + a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` + if the buffer was serialized in-band. + """ + dtype = getattr(torch, dtype_str) + assert isinstance(dtype, torch.dtype) + if isinstance(buf, zmq.Frame): + # ZMQ frames own their message memory independently of any context, + # so the tensor can safely alias it with zero copies. The tensor's + # storage keeps the frame (and thus its bytes) alive via a strong + # reference for as long as the tensor is. + try: + return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape) + except ValueError: + # Empty or read-only frame buffer; fall through to the copy path. + pass + # Shared-memory ring buffer chunks are reused by the writer once all + # readers have marked them read, so we must copy out of them. bytearray + # (vs bytes) keeps the resulting tensor writable, matching normal tensor + # semantics. + raw = bytearray(buf) + if not raw: + assert 0 in shape + return torch.empty(shape, dtype=dtype) + return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape) + + +def _reduce_tensor(tensor: torch.Tensor): + """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling. + + `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle + byte stream via `torch.serialization` and never emits a `PickleBuffer`, + which defeats the out-of-band buffer handling in `MessageQueue.enqueue`. + This reducer instead exposes the tensor's memory directly, so large + tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue + without being copied into and back out of the pickled message. + """ + if ( + tensor.device.type == "cpu" + and tensor.layout == torch.strided + and not tensor.requires_grad + ): + try: + # The uint8 view exposes the raw bytes via the buffer protocol, + # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...). + # reshape(-1) first so that 0-dim tensors can be viewed as well. + raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy() + except RuntimeError: + # Exotic tensors (e.g. with the conjugate bit set) that don't + # support aliasing views; let torch handle them. + pass + else: + dtype_str = str(tensor.dtype).removeprefix("torch.") + return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str) + + # Fall back to torch's default (copying) reduction. + return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL) + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -631,25 +731,22 @@ def __init__(self, timeout: float | None, should_warn: bool) -> None: self.n_warning = 1 self.timeout = timeout - def timeout_ms(self) -> int | None: - """Returns a timeout that is: + def timeout_ms(self) -> int: + """Returns a timeout, capped at the recheck interval, that is: - min(time to deadline, time to next warning) if we're logging warnings - time to deadline, if we're not logging warnings - - None if the timeout is None and we're not logging warnings + - recheck interval if the timeout is None and we're not logging warnings - raise TimeoutError if we are past the deadline """ - warning_wait_time = self.warning_wait_time_ms + wait_ms = SHM_READER_RECHECK_INTERVAL_MS + if self.warning_wait_time_ms is not None: + wait_ms = min(wait_ms, self.warning_wait_time_ms) if self.timeout is None: - return warning_wait_time - + return wait_ms time_left_ms = int((self.deadline - time.monotonic()) * 1000) if time_left_ms <= 0: raise TimeoutError - - if warning_wait_time and warning_wait_time < time_left_ms: - return warning_wait_time - - return time_left_ms + return min(wait_ms, time_left_ms) def should_warn(self) -> bool: """Returns true if it's time to log a warning for a timeout that is not @@ -710,18 +807,18 @@ def check(): # found a block that is not read by this reader # let caller read from the buffer with self.buffer.get_data(self.current_idx) as buf: - yield buf - - # caller has read from the buffer - # set the read flag - metadata_buffer[self.local_reader_rank + 1] = 1 - # Memory fence ensures the read flag is visible to the writer. - # Without this, writer may not see our read completion and - # could wait indefinitely for all readers to finish. - memory_fence() - self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks - - self._spin_condition.record_read() + try: + yield buf + finally: + # caller has read from the buffer; set the read flag. + metadata_buffer[self.local_reader_rank + 1] = 1 + # Memory fence ensures the read flag is visible to the writer. + # Without this, writer may not see our read completion and + # could wait indefinitely for all readers to finish. + memory_fence() + next_idx = self.current_idx + 1 + self.current_idx = next_idx % self.buffer.max_chunks + self._spin_condition.record_read() break def enqueue(self, obj, timeout: float | None = None): @@ -740,9 +837,23 @@ def oob_callback(buf: PickleBuffer) -> bool: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps( - obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback - ) + # CPU tensors are routed through `_reduce_tensor` so that their + # bytes are emitted as out-of-band buffers instead of being + # copied into the pickle stream by torch's default reducer. + # Start from `copyreg.dispatch_table` to preserve globally + # registered reducers (e.g. `re.Pattern`); the per-pickler + # dispatch table would otherwise shadow them. + dispatch_table = dict(copyreg.dispatch_table) + dispatch_table[torch.Tensor] = _reduce_tensor + with io.BytesIO() as bio: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) + pickler.dispatch_table = dispatch_table + pickler.dump(obj) + all_buffers[0] = bio.getvalue() if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: @@ -798,7 +909,8 @@ def dequeue( @staticmethod def recv(socket: zmq.Socket, timeout: float | None) -> Any: - timeout_ms = None if timeout is None else int(timeout * 1000) + # Ensure non-negative timeout passed to zmq poll. + timeout_ms = None if timeout is None else max(0, int(timeout * 1000)) if not socket.poll(timeout=timeout_ms): raise TimeoutError recv, *recv_oob = socket.recv_multipart(copy=False) diff --git a/vllm/distributed/device_communicators/symm_mem.py b/vllm/distributed/device_communicators/symm_mem.py index 8c174602c3c6..26c6ec1ff47e 100644 --- a/vllm/distributed/device_communicators/symm_mem.py +++ b/vllm/distributed/device_communicators/symm_mem.py @@ -27,6 +27,7 @@ class SymmMemCommunicator: "9.0": [4, 6, 8], "10.0": [6, 8], "10.3": [6, 8], + "10.7": [6, 8], # sm_107 (Rubin): reuse 10.3 thresholds } def __init__( diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index 1b6ce9e8aae4..7ca132824eca 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -20,8 +20,11 @@ def __init__( device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.ca_comm: None = None if self.use_all2all: if self.all2all_backend in ("naive", "allgather_reducescatter"): diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 1d5f467027e7..3c20a4a1f749 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -275,3 +275,15 @@ def request_finished( get_finished(). """ return False, None + + def has_pending_push_work(self) -> bool: + """Return True if the connector has push-mode work that requires + the engine main loop to keep stepping (e.g. for EPD, + Producer has push work when Xfer is in progress - Consumer + is reading it). + This mirrors exactly the KV Connector's has_pending_push_work(). + + Connectors that don't implement push-based EC transfer should + leave this as False. + """ + return False diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index b0c3740f57ea..4ccc4d2c010f 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy import gc import weakref from collections.abc import Iterable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import replace from typing import TYPE_CHECKING @@ -43,6 +43,7 @@ from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) +from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper @@ -84,16 +85,20 @@ def batch_transfer_weights( assert len(all_params) > 0 p2p_ops = [] for param in all_params: + # PyNccl transfers flat memory and does not honor tensor strides. + transfer_param = param.contiguous() op = object.__new__(P2POp) - if is_sender: - op.op = torch.distributed.isend - op.tensor = param - else: - op.op = torch.distributed.irecv - op.tensor = param + op.op = torch.distributed.isend if is_sender else torch.distributed.irecv + op.tensor = transfer_param op.group_peer = peer_rank p2p_ops.append(op) - device_comm.batch_isend_irecv(p2p_ops) + if transfer_param is not param: + device_comm.batch_isend_irecv(p2p_ops) + p2p_ops.clear() + if not is_sender: + param.copy_(transfer_param) + if p2p_ops: + device_comm.batch_isend_irecv(p2p_ops) def broadcast_expert_mapping( @@ -145,6 +150,10 @@ def __init__(self, worker): self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} + self._async_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPAsync" + ) + self._async_future: Future[None] | None = None @property def worker(self): @@ -159,59 +168,68 @@ def execute(self, execute_method: str, *args, **kwargs): raise ValueError(f"Unknown execute method: {execute_method}") return method(*args, **kwargs) - def _set_eplb_suppressed(self, suppressed: bool) -> None: - self.worker.model_runner.eep_eplb_suppressed = suppressed - ep_group = get_standby_ep_group() or get_ep_group() - if ep_group.rank == 0: - logger.info( - "[Elastic EP] EPLB %s elastic scaling transition", - "disabled during" if suppressed else "re-enabled after", - ) + def start_async(self, execute_method: str, *args, **kwargs) -> str: + if self._async_future is not None: + raise RuntimeError("Another Elastic EP async method is active") + if args and isinstance(args[0], ReconfigureDistributedRequest): + self.reconfig_request = args[0] + dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank + done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}" + self._async_future = self._async_executor.submit( + self._run_async, execute_method, *args, **kwargs + ) + self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key)) + return done_key + + def _run_async(self, execute_method: str, *args, **kwargs) -> None: + from vllm.platforms import current_platform + + self.worker.vllm_config.enable_trace_function_call_for_thread() + assert hasattr(self.worker, "device") + current_platform.set_device(self.worker.device) + with set_current_vllm_config(self.worker.vllm_config): + self.execute(execute_method, *args, **kwargs) + + def _mark_async_done(self, done_key: str) -> None: + from vllm.distributed.utils import get_cached_tcp_store_client + + assert self.reconfig_request is not None + get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ).set(done_key, b"1") + + def clear_async(self) -> None: + future = self._async_future + if future is None: + raise RuntimeError("No Elastic EP async method is active") + if not future.done(): + raise RuntimeError("Elastic EP async method is not done") + self._async_future = None + future.result() def load_model(self) -> None: - ( - expanded_physical_to_logical, - num_logical_experts, - old_num_physical_experts, - ) = self.receive_expert_mapping() - num_physical_experts = expanded_physical_to_logical.shape[1] - self.worker.parallel_config.eplb_config.num_redundant_experts = ( - num_physical_experts - num_logical_experts - ) self.worker.load_model(load_dummy_weights=True) - self.worker.model_runner.setup_eplb_from_mapping( - expanded_physical_to_logical, old_num_physical_experts - ) - self._set_eplb_suppressed(True) def create_standby_groups( - self, reconfig_request: ReconfigureDistributedRequest + self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size - world_size = self.worker.vllm_config.parallel_config.world_size + parallel_config = self.worker.vllm_config.parallel_config + world_size = parallel_config.world_size new_world_size_across_dp = world_size * new_dp_size - updated_config = copy.copy(self.worker.vllm_config) - updated_config.parallel_config = copy.deepcopy( - self.worker.vllm_config.parallel_config + create_standby_groups( + new_dp_size=new_dp_size, + new_world_size_across_dp=new_world_size_across_dp, + master_ip=reconfig_request.new_data_parallel_master_ip, + coord_store_port=reconfig_request.coord_store_port, + use_all2all=use_all2all, + enable_eplb=parallel_config.enable_eplb, ) - updated_config.parallel_config.data_parallel_size = new_dp_size - with set_current_vllm_config(updated_config): - create_standby_groups( - new_dp_size=new_dp_size, - new_world_size_across_dp=new_world_size_across_dp, - master_ip=reconfig_request.new_data_parallel_master_ip, - coord_store_port=reconfig_request.coord_store_port, - enable_eplb=updated_config.parallel_config.enable_eplb, - ) - if new_dp_size > old_dp_size: - self._set_eplb_suppressed(True) - eplb_state = self.worker.model_runner.eplb_state - if eplb_state is not None: - eplb_state.drain_async() - elif new_dp_size < old_dp_size: - self._stage_standby_moe_quant_methods() + if new_dp_size < old_dp_size: + self.stage_standby_moe_quant_methods() def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() @@ -265,6 +283,7 @@ def broadcast_expert_mapping(self) -> None: model_config = self.worker.model_runner.model_config eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None + eplb_state.drain_async() eplb_model_state = eplb_state.model_states[model_config.compute_hash()] physical_to_logical = eplb_model_state.physical_to_logical_map num_physical_experts = physical_to_logical.shape[1] @@ -278,10 +297,6 @@ def broadcast_expert_mapping(self) -> None: src_rank=0, device=self.worker.device, ) - # New workers enter load_model after receiving the expert mapping. - # Stage replacement MoE kernels before returning to the state machine - # so existing ranks can participate in collective EP comm creation. - self._stage_standby_moe_quant_methods() def _make_eep_moe_config(self, module, dp_group, ep_group): parallel_config = self.worker.vllm_config.parallel_config @@ -300,7 +315,7 @@ def _make_eep_moe_config(self, module, dp_group, ep_group): moe_parallel_config=moe_parallel_config, ) - def _stage_standby_moe_quant_methods(self) -> None: + def stage_standby_moe_quant_methods(self) -> None: standby_dp_group = get_standby_dp_group() standby_ep_group = get_standby_ep_group() model = self.worker.model_runner.get_model() @@ -500,26 +515,6 @@ def switch_and_prepare(self) -> None: compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) - multi_block_table = self.worker.model_runner.input_batch.block_table - saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] - for bt in multi_block_table.block_tables: - saved_block_tables.append( - (bt.block_table.gpu.clone(), bt.block_table.cpu.clone()) - ) - multi_block_table.clear() - - unlock_workspace() - self.worker.compile_or_warm_up_model() - lock_workspace() - - for bt, (saved_gpu, saved_cpu) in zip( - multi_block_table.block_tables, saved_block_tables - ): - bt.block_table.gpu.copy_(saved_gpu) - bt.block_table.cpu.copy_(saved_cpu) - if new_dp_size < old_dp_size: - self._set_eplb_suppressed(False) - def _perform_eplb_reshuffle( self, rank_mapping: dict[int, int] | None = None ) -> None: @@ -553,12 +548,25 @@ def _perform_eplb_reshuffle( if get_ep_group().rank == 0: logger.info("[Elastic EP] Expert resharding completed") - def perform_eplb_reshuffle(self) -> None: + def commit_scale_up(self, is_existing_worker: bool) -> None: + if is_existing_worker: + self.broadcast_expert_mapping() + self.switch_and_prepare() + else: + mapping, _, num_valid_experts = self.receive_expert_mapping() + self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts) self._perform_eplb_reshuffle() - self._set_eplb_suppressed(False) + self.warm_and_capture() + + def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: + self.perform_scale_down_eplb_reshuffle(new_dp_size) + if removing: + self.switch_and_remove() + else: + self.switch_and_prepare() + self.warm_and_capture() def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: - self._set_eplb_suppressed(True) eplb_state = self.worker.model_runner.eplb_state if eplb_state is not None: eplb_state.drain_async() @@ -599,12 +607,17 @@ def receive_weights(self) -> None: ) model = self.worker.model_runner.get_model() + expert_weights = [ + module.get_expert_weights() + for module in model.modules() + if is_moe_layer(module) + ] batch_transfer_weights( model=model, is_sender=False, peer_rank=sender_rank, dp_group=dp_group, - expert_weights=model.expert_weights, + expert_weights=expert_weights, ) torch.accelerator.synchronize() @@ -643,14 +656,17 @@ def prepare_new_worker(self) -> None: with set_current_vllm_config(self.worker.vllm_config): prepare_communication_buffer_for_model(self.worker.model_runner.get_model()) - def rewarm_workspace(self) -> None: + def warmup_local_kernels(self) -> None: + with set_current_vllm_config(self.worker.vllm_config): + kernel_warmup(self.worker, process_local_only=True) + + def warm_and_capture(self) -> None: # Must run on every DP sibling in lockstep: _dummy_run calls # coordinate_batch_across_dp whenever data_parallel_size > 1 # (gpu_model_runner.py:3663), which deadlocks if any rank skips it. - # Save and clear block tables so profile_run/compile_or_warm_up_model - # don't write dummy slot mappings into real KV-cache blocks (mirrors - # switch_and_prepare's pattern). + # Save and clear block tables so the dummy MoE forward doesn't + # write dummy slot mappings into real KV-cache blocks. multi_block_table = self.worker.model_runner.input_batch.block_table saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] for bt in multi_block_table.block_tables: @@ -660,19 +676,16 @@ def rewarm_workspace(self) -> None: multi_block_table.clear() # _ensure_workspace_size allocates a fresh tensor on grow, leaving - # captured CUDA graphs with stale data pointers; drop graphs before - # re-warm so captures realign with the resized buffer. + # any captured CUDA graph with a stale data pointer; drop graphs + # before re-warm so captures realign with the resized buffer. self._release_cuda_graphs() unlock_workspace() - # Grow the MoE workspace at max_num_tokens. - # compile_or_warm_up_model alone only exercises cudagraph-capture - # sizes (≤64 tokens for this test) and leaves the workspace at - # ~10-14 MB; the post-all-to-all per-rank token count under real - # post-reshuffle routing needs hundreds of MB. Use _dummy_run - # directly (rather than profile_run) with skip_eplb=True so dummy - # routing doesn't pollute the just-rebalanced EPLB stats — same - # convention compile_or_warm_up_model itself uses. + # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model + # alone only exercises cudagraph-capture sizes and can leave the + # workspace too small for post-reshuffle routing. Use _dummy_run + # directly with skip_eplb=True so dummy routing doesn't pollute the + # just-rebalanced EPLB stats. runner = self.worker.model_runner runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) self.worker.compile_or_warm_up_model() diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 256efe46a4a4..f33fd90e8637 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -1,18 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import enum -import time import weakref -from datetime import timedelta -from typing import TYPE_CHECKING, Literal, TypeAlias +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, Literal, TypeAlias import torch.distributed from vllm.config import ParallelConfig from vllm.distributed import ( - sched_yield, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.utils import get_cached_tcp_store_client from vllm.logger import init_logger from vllm.v1.engine import ( EEPNotificationType, @@ -31,35 +30,29 @@ class ScaleUpExistingEngineState(enum.IntEnum): - WAIT_NEW_CORE_ENGINES_INIT = 0 - CREATE_STANDBY_GROUPS = 1 - TRANSFER_EXPERT_MAPPING = 2 - WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3 - TRANSFER_WEIGHTS = 4 - SYNC_KV_CACHE_MEMORY_SIZE = 5 - SWITCH_AND_PREPARE = 6 - EPLB_RESHUFFLE = 7 - COMPLETE = 8 + CREATE_STANDBY_GROUPS = 0 + STAGE_QUANT_METHODS = 1 + TRANSFER_WEIGHTS = 2 + SYNC_KV_CACHE_MEMORY_SIZE = 3 + COMMIT_SCALE_UP = 4 # Blocks forward passes. + COMPLETE = 5 class ScaleUpNewEngineState(enum.IntEnum): PRE_KV_INIT = 0 PREPARE = 1 - EPLB_RESHUFFLE = 2 - COMPLETE = 3 + COMPLETE = 2 class ScaleDownRemainingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - SWITCH_AND_PREPARE = 2 - COMPLETE = 3 + COMMIT_SCALE_DOWN = 1 # Blocks forward passes. + COMPLETE = 2 class ScaleDownRemovingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - COMPLETE = 2 + COMPLETE = 1 EngineState: TypeAlias = ( @@ -70,15 +63,6 @@ class ScaleDownRemovingEngineState(enum.IntEnum): ) -class _BarrierTimeoutError(RuntimeError): - """ - Exception raised for timeout - in the first stage of our two-staged - TCPStore based barrier to synchronize the - execution of all engines in the DP group. - """ - - class ElasticEPScalingState: def __init__( self, @@ -94,20 +78,24 @@ def __init__( self.engine_core_ref = weakref.ref(engine_core) self.vllm_config = vllm_config self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None - self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None self.new_parallel_config: ParallelConfig = new_parallel_config self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None self.worker_type = worker_type self.scale_type = scale_type self.reconfig_request = reconfig_request - + self.commit_requested = False + self._prepare_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPPrepare" + ) + self._prepare_future: Future[Any] | None = None + self._new_dp_sync: tuple[object, Any] | None = None self.state: EngineState if scale_type == "scale_up": self.state = ( ScaleUpNewEngineState.PRE_KV_INIT if worker_type == "new" - else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT + else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS ) else: self.state = ( @@ -130,6 +118,31 @@ def engine_core(self) -> "DPEngineCoreProc": raise RuntimeError("Engine core has been garbage collected") return engine_core + def _collective_rpc(self, *args, **kwargs): + return self.model_executor.collective_rpc(*args, **kwargs) + + def _execute_async(self, execute_method: str, *args) -> bool: + if self._prepare_future is None: + done_keys = self._collective_rpc( + "elastic_ep_execute", + args=("start_async", execute_method, *args), + ) + assert self.reconfig_request is not None + coord_store = get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ) + self._prepare_future = self._prepare_executor.submit( + coord_store.wait, done_keys + ) + if not self._prepare_future.done(): + return False + + self._prepare_future.result() + self._collective_rpc("elastic_ep_execute", args=("clear_async",)) + self._prepare_future = None + return True + def progress(self) -> bool: if self.scale_type == "scale_up": return ( @@ -149,157 +162,43 @@ def run_pre_kv_init_states(self) -> None: assert self.progress() assert self.state == ScaleUpNewEngineState.PREPARE - def _execute_tcp_store_barrier( - self, dp_store, group_rank, group_size, barrier_id, timeout=None - ): - arrival_key = f"arrival_{barrier_id}_{group_rank}" - dp_store.set(arrival_key, b"1") - - start_time = time.time() - processes_arrived: set[int] = set() - - while len(processes_arrived) < group_size: - if ( - timeout is not None - and time.time() - start_time > timeout.total_seconds() - ): - raise _BarrierTimeoutError( - f"Barrier timed out after {timeout.total_seconds()} seconds" - ) - - for i in range(group_size): - if i in processes_arrived: - continue - - key = f"arrival_{barrier_id}_{i}" - present = dp_store.check([key]) - if present: - processes_arrived.add(i) - - if len(processes_arrived) < group_size: - sched_yield() - - def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool: - """ - Execute a two-staged barrier to synchronize all engines in the DP group. - - Some DP EngineCores may receive the reconfiguration notifications - later than others, and already proceed to engine step (model forward) - in the busy loop. - In this case, EngineCores that already proceed to reconfiguration - should skip reconfiguration and execute model forward for one more - step, so in the next step, all EngineCores will be synchronized. - We use a two-staged barrier to achieve this. The first time each - EngineCore executes the barrier, if a timeout is reached before the - barrier completes, that means some EngineCores have already entered - engine step. The EngineCores that timed out will then proceed to - engine step, and will synchronize with the other EngineCores in the - next step with a barrier without timeout. - """ - dp_group = self.new_dp_group if use_new_group else self.old_dp_group - dp_store = self.new_dp_store if use_new_group else self.old_dp_store - assert dp_group is not None and dp_store is not None - - group_rank = dp_group.rank() - group_size = dp_group.size() - barrier_id = f"eep_barrier_{barrier_name}" - sync_key = f"{barrier_id}_sync" - - # TODO(yongji): figure out appropriate timeout for the barrier - timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5) - - try: - self._execute_tcp_store_barrier( - dp_store, group_rank, group_size, barrier_id, timeout=timeout - ) - torch.distributed.barrier(dp_group) - if group_rank == 0: - dp_store.delete_key(sync_key) - for i in range(group_size): - dp_store.delete_key(f"arrival_{barrier_id}_{i}") - return True - except _BarrierTimeoutError as e: - if timeout is None: - raise RuntimeError("Unexpected timeout encountered") from e - dp_store.compare_set(sync_key, "", b"1") - return False - def _progress_existing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None - - if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT: - return False + assert self.old_dp_group is not None - elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: - # NOTE(yongji): wait for all existing workers to receive the request - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=False, barrier_name="create_standby_groups" - ): + if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: + if not self._create_standby_groups(): return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._create_standby_groups() - self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING + self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS return True - elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING: - self._transfer_expert_mapping() - self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT + elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS: + if not self._execute_async("stage_standby_moe_quant_methods"): + return False + self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS return True - elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT: - return False - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + if not self._transfer_weights(): return False - if not self._staged_barrier( - use_new_group=False, barrier_name="transfer_weights" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._transfer_weights() self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE return True elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE: - self._sync_kv_cache_memory_size() - self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE - return True - - elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE: - self._switch_and_prepare() - self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE - assert self.new_dp_store is not None - self.new_dp_store.add("eep_barrier_engine_count", 1) + if not self._sync_kv_cache_memory_size(): + return False + self.state = ScaleUpExistingEngineState.COMMIT_SCALE_UP + self._mark_ready_for_switch() return True - elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE: - assert self.new_dp_group is not None and self.new_dp_store is not None - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): + elif state == ScaleUpExistingEngineState.COMMIT_SCALE_UP: + if not self.commit_requested: return False - if self.new_dp_group.rank() == 0: - self.new_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle() + self._commit_new_dp_group() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", True)) self.state = ScaleUpExistingEngineState.COMPLETE self._update_parallel_config() + self._send_reconfigure_finished() return True else: @@ -311,22 +210,17 @@ def _progress_new_engine(self) -> bool: assert self.new_dp_group is not None and self.new_dp_store is not None if state == ScaleUpNewEngineState.PRE_KV_INIT: - self.engine_core._eep_send_engine_core_notification( - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("receive_weights",) - ) + self._collective_rpc("elastic_ep_execute", args=("receive_weights",)) self.engine_core.available_gpu_memory_for_kv_cache = ( ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("prepare_new_worker",) - ) + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) self.state = ScaleUpNewEngineState.PREPARE return True elif state == ScaleUpNewEngineState.PREPARE: + self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",)) + self._mark_ready_for_switch() tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") torch.distributed.all_reduce( tensor, @@ -337,22 +231,7 @@ def _progress_new_engine(self) -> bool: self.engine_core.engines_running = bool(data[0]) self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) - self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE - self.new_dp_store.add("eep_barrier_engine_count", 1) - return True - - elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE: - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): - return False - assert self.new_dp_group.rank() > 0 - self._eplb_reshuffle() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", False)) self.state = ScaleUpNewEngineState.COMPLETE return True @@ -362,38 +241,23 @@ def _progress_new_engine(self) -> bool: def _progress_remaining_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemainingEngineState.PREPARE: - self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True + if self._create_standby_groups(): + self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + self._mark_ready_for_switch() + return True + return False - elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): + elif state == ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN: + if not self.commit_requested: return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle_before_scale_down() - self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE - # NOTE(yongji): currently, after EPLB reshuffle - # that redistributes experts to remaining workers, workers - # to be removed will immediately initiate shutdown; - # existing workers can no longer execute forward steps using - # the old setup. In the future, we may keep - # the removing workers alive a bit longer, - # e.g., to drain in-batch requests. - self._create_standby_groups() - self._switch_and_prepare() + self._commit_scale_down(removing=False) + self._commit_new_dp_group() self._update_parallel_config() self.state = ScaleDownRemainingEngineState.COMPLETE + self._send_reconfigure_finished() return True else: @@ -402,26 +266,11 @@ def _progress_remaining_engine(self) -> bool: def _progress_removing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemovingEngineState.PREPARE: - self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True - - if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): - return False assert self.old_dp_group.rank() > 0 - self._eplb_reshuffle_before_scale_down() - self._switch_and_remove() + self._commit_scale_down(removing=True) self.state = ScaleDownRemovingEngineState.COMPLETE self.engine_core._eep_send_engine_core_notification( EEPNotificationType.SHUTDOWN_COMPLETE @@ -432,22 +281,22 @@ def _progress_removing_engine(self) -> bool: assert self.state == ScaleDownRemovingEngineState.COMPLETE return True - def handle_notification(self, notification_type: EEPNotificationType): - assert self.worker_type != "new" - assert self.old_dp_store is not None - if ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY - and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS - elif ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - and self.state - == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS + def is_ready_for_switch(self) -> bool: + return self.worker_type == "existing" and ( + self.state is ScaleUpExistingEngineState.COMMIT_SCALE_UP + or self.state is ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + ) + + @property + def ready_key(self) -> str: + return f"eep_ready/{self.engine_core.dp_rank}" + + def _mark_ready_for_switch(self) -> None: + parallel_config = self.new_parallel_config + get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ).set(self.ready_key, b"1") def is_complete(self) -> bool: if self.scale_type == "scale_up": @@ -462,50 +311,78 @@ def is_complete(self) -> bool: else self.state == ScaleDownRemainingEngineState.COMPLETE ) - def _create_standby_groups(self): + def _init_new_dp_group(self) -> tuple[Any, Any]: + return self.new_parallel_config.stateless_init_dp_group(return_store=True) + + def _ensure_new_dp_group(self) -> bool: + if self.new_dp_group is not None: + return True + + if self._prepare_future is None: + self._prepare_future = self._prepare_executor.submit( + self._init_new_dp_group + ) + if not self._prepare_future.done(): + return False + + self.new_dp_group, self.new_dp_store = self._prepare_future.result() + self._prepare_future = None + return True + + def _create_standby_groups(self) -> bool: assert self.old_dp_group is not None - self.new_dp_group, self.new_dp_store = ( - self.new_parallel_config.stateless_init_dp_group(return_store=True) - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request) - ) + if not self._ensure_new_dp_group(): + return False + if not self._execute_async( + "create_standby_groups", + self.reconfig_request, + self.new_parallel_config.use_all2all, + ): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Created standby communication groups") + return True - def _transfer_weights(self): + def _transfer_weights(self) -> bool: assert self.reconfig_request is not None and self.old_dp_group is not None old_dp_size = self.old_dp_group.size() new_dp_size = self.reconfig_request.new_data_parallel_size - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size) - ) + if not self._execute_async("transfer_weights", old_dp_size, new_dp_size): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Transferred weights to new workers") + return True - def _transfer_expert_mapping(self): - assert self.old_dp_group is not None - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("broadcast_expert_mapping",) - ) - if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Broadcasted expert mapping to new workers") - - def _sync_kv_cache_memory_size(self): + def _sync_kv_cache_memory_size(self) -> bool: assert self.engine_core.available_gpu_memory_for_kv_cache > 0 assert self.new_dp_group is not None and self.old_dp_group is not None - ParallelConfig.sync_kv_cache_memory_size( - self.new_dp_group, - self.engine_core.available_gpu_memory_for_kv_cache, - ) + + if self._new_dp_sync is None: + tensor = torch.tensor( + [self.engine_core.available_gpu_memory_for_kv_cache], + dtype=torch.int64, + device="cpu", + ) + work = torch.distributed.all_reduce( + tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.new_dp_group, + async_op=True, + ) + self._new_dp_sync = (tensor, work) + return False + + _, work = self._new_dp_sync + if not work.is_completed(): + return False + work.wait() + self._new_dp_sync = None if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Synced KV cache memory size to new workers") + return True - def _switch_and_prepare(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_prepare",) - ) + def _commit_new_dp_group(self): old_dp_group = self.old_dp_group stateless_destroy_torch_distributed_process_group(old_dp_group) assert self.new_dp_group is not None @@ -529,41 +406,28 @@ def _switch_and_prepare(self): self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) if new_dp_group.rank() == 0: - self.engine_core._eep_send_engine_core_notification( - EEPNotificationType.RECONFIGURE_FINISHED - ) logger.info("[Elastic EP] Switched to new setup") - def _eplb_reshuffle(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("perform_eplb_reshuffle",) - ) - # Reshuffle changes per-rank token routing; the locked MoE workspace - # may now be too small. Rewarm covers both new and existing engines. - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("rewarm_workspace",) - ) + def _send_reconfigure_finished(self): assert self.new_dp_group is not None if self.new_dp_group.rank() == 0: - logger.info("[Elastic EP] EPLB reshuffle completed") + self.engine_core._eep_send_engine_core_notification( + EEPNotificationType.RECONFIGURE_FINISHED + ) - def _eplb_reshuffle_before_scale_down(self): + def _commit_scale_down(self, removing: bool): assert self.reconfig_request is not None and self.old_dp_group is not None - self.model_executor.collective_rpc( + self._collective_rpc( "elastic_ep_execute", args=( - "perform_scale_down_eplb_reshuffle", + "commit_scale_down", self.reconfig_request.new_data_parallel_size, + removing, ), ) if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] EPLB reshuffle completed") - def _switch_and_remove(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_remove",) - ) - def _update_parallel_config(self): assert self.reconfig_request is not None reconfig_request = self.reconfig_request diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py index 846793a955f6..1892f3e79426 100644 --- a/vllm/distributed/elastic_ep/standby_state.py +++ b/vllm/distributed/elastic_ep/standby_state.py @@ -39,6 +39,7 @@ def create_standby_groups( new_world_size_across_dp: int, master_ip: str, coord_store_port: int, + use_all2all: bool, enable_eplb: bool = True, backend: str | None = None, ) -> None: @@ -86,7 +87,7 @@ def create_standby_groups( ) standby_ep_ranks = [x.tolist() for x in standby_ep_ranks] _STANDBY_EP = _init_stateless_group( - standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store + standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all ) if enable_eplb: diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 891b57bcf18f..f9a9a8a90a81 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -38,7 +38,7 @@ def has_nixl() -> bool: - """Whether the optional NIXL / RIXL package is available.""" + """Whether the optional NIXL package is available.""" return nixl_utils.NixlWrapper is not None @@ -266,7 +266,7 @@ def __init__( assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer." nixl_wrapper_cls = nixl_utils.NixlWrapper if nixl_wrapper_cls is None: - raise RuntimeError("NIXL/ RIXL is unavailable.") + raise RuntimeError("NIXL is unavailable.") self._cpu_group = cpu_group self._world_size = cpu_group.size() diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index feacb03d28b5..d34b844bd3d2 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -46,6 +46,7 @@ from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.platforms import current_platform from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator @@ -825,23 +826,80 @@ def rearrange( eplb_model_state.physical_to_logical_map.cpu(), ) - # Update expert weights - rearrange_expert_weights_inplace( - eplb_model_state.physical_to_logical_map, - new_physical_to_logical_map, - eplb_model_state.model.expert_weights, - eplb_model_state.expert_buffer, - ep_group, - eplb_model_state.communicator, - is_profile, - rank_mapping, - ) - - if not is_profile: - _commit_eplb_maps( - eplb_model_state, - new_physical_to_logical_map=new_physical_to_logical_map, + skip_rearrange = False + if ( + current_platform.is_rocm() + and not is_profile + and rank_mapping is None + and bool((eplb_model_state.physical_to_logical_map >= 0).all()) + ): + logical_loads = global_expert_load_window.float() + ep_size = ep_group.size() + + def rank_load_imbalance( + mapping: torch.Tensor, + logical_loads: torch.Tensor = logical_loads, + ep_size: int = ep_size, + ) -> float: + mapping = mapping.to( + device=logical_loads.device, + dtype=torch.long, + ) + replica_counts = torch.zeros_like(logical_loads) + replica_counts.scatter_add_( + dim=1, + index=mapping, + src=torch.ones_like(mapping, dtype=logical_loads.dtype), + ) + loads_per_replica = torch.gather( + logical_loads / replica_counts.clamp_min(1), + dim=1, + index=mapping, + ) + loads_per_rank = loads_per_replica.reshape( + logical_loads.shape[0], ep_size, -1 + ).sum(dim=(0, 2)) + mean_load = loads_per_rank.mean() + if mean_load == 0: + return 1.0 + return (loads_per_rank.max() / mean_load).item() + + current_imbalance = rank_load_imbalance( + eplb_model_state.physical_to_logical_map ) + proposed_imbalance = rank_load_imbalance( + new_physical_to_logical_map + ) + relative_improvement = ( + current_imbalance - proposed_imbalance + ) / current_imbalance + skip_rearrange = relative_improvement < 0.05 + if skip_rearrange and is_main_rank: + logger.info( + "[EPLB] Skip rearrange: imbalance %.4f -> " + "%.4f (no material gain)", + current_imbalance, + proposed_imbalance, + ) + + if not skip_rearrange: + # Update expert weights + rearrange_expert_weights_inplace( + eplb_model_state.physical_to_logical_map, + new_physical_to_logical_map, + eplb_model_state.model.expert_weights, + eplb_model_state.expert_buffer, + ep_group, + eplb_model_state.communicator, + is_profile, + rank_mapping, + ) + + if not is_profile: + _commit_eplb_maps( + eplb_model_state, + new_physical_to_logical_map=new_physical_to_logical_map, + ) if is_main_rank: assert start_event is not None diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index be7e7363fb9d..cd222567a33c 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -44,8 +44,7 @@ class KVCacheEvent( MEDIUM_GPU = "GPU" MEDIUM_CPU = "CPU" -MEDIUM_FS = "FS" -MEDIUM_OBJ = "OBJ" +MEDIUM_STORAGE = "STORAGE" class BlockStored(KVCacheEvent): @@ -273,6 +272,10 @@ def publish(self, events: EventBatch) -> None: def shutdown(self) -> None: """Shutdown the publisher.""" + def get_publisher_config(self) -> KVEventsConfig | None: + """Return the publisher's resolved runtime configuration.""" + return None + class NullEventPublisher(EventPublisher): """No-op implementation (default when disabled).""" @@ -336,6 +339,17 @@ def __init__( self._replay_endpoint = self.offset_endpoint_port( replay_endpoint, self._dp_rank ) + assert self._endpoint is not None + self._publisher_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=self._endpoint, + replay_endpoint=self._replay_endpoint, + buffer_steps=buffer_steps, + hwm=hwm, + max_queue_size=max_queue_size, + topic=topic, + ) self._hwm = hwm self._socket_setup() @@ -352,6 +366,9 @@ def __init__( ) self._thread.start() + def get_publisher_config(self) -> KVEventsConfig: + return self._publisher_config + def publish(self, events: EventBatch) -> None: if not self._running: raise RuntimeError("Publisher is closed") diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 77f043ef5270..eff9bec8ee93 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -21,12 +21,10 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.platforms import current_platform from vllm.v1.attention.backend import AttentionBackend -from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase - from vllm.v1.kv_cache_interface import KVCacheSpec logger = init_logger(__name__) @@ -260,8 +258,12 @@ def kv_postprocess_layout_on_receive(cache, indices): This method corrects layout mismatches from direct memory copies by permuting the tensor dimensions. + 4D cache: - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]` - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]` + 5D cache: + - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]` + - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]` Implementation: - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout @@ -272,7 +274,7 @@ def kv_postprocess_layout_on_receive(cache, indices): blocks_to_update = cache.index_select(0, indices) target_shape = list(blocks_to_update.shape) target_shape[0] = -1 - inv_order = [0, 2, 1, 3] + inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3] src_shape = tuple(target_shape[i] for i in inv_order) blocks_to_update = cache.index_select(0, indices) permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order) @@ -594,32 +596,6 @@ def target_remote_ranks( abs_ratio = -tp_ratio return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)] - def get_transfer_cache_regions( - self, cache: torch.Tensor, layer_spec: "KVCacheSpec" - ) -> list[torch.Tensor] | torch.Tensor: - """Return the cache tensor(s) to register as NIXL memory regions, - also accounting for hybrid SSM models specificities. - """ - if isinstance(layer_spec, MambaSpec): - # Register the whole kv cache shared tensor, including - # SSM/Conv. - conv, ssm = cache - return [conv] - - # Check may be hacky but it's matching - # `_update_hybrid_attention_mamba_layout`. - if self.is_mamba and cache.shape[0] == 2: - # When MAMBA is present, all backends are blocks first, so - # that blocks can be shared between attention layers and mamba - # layers. Runner already adjusted strides for FlashAttn-like - # backends so its num_blocks first. - # Swap [2<>num_blocks] dims for hybrid SSM layout. - cache = cache.transpose(0, 1) - - # K and V are packed into one tensor (content dim), so each layer - # registers as a single region. - return [cache] - def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str: """One-line summary of transfer config for logging.""" info = self._engines[(remote_engine_id, remote_pp_rank)] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c0eb24729793..076d40a187c8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -181,6 +181,18 @@ def prefer_cross_layer_blocks(self) -> bool: """ return False + @property + def requires_kv_delivery(self) -> bool: + """Whether this connector hands off KV that must be reliably delivered. + + If True, a request preempted while its hand-off is still pending is + recomputed rather than allowed to finish and hand off blocks that the + preemption already freed. Defaults to the producer role, since only a + producer hands KV off when a request completes. Best-effort caches + return False, as a dropped save is just a future cache miss. + """ + return self._kv_transfer_config.is_kv_producer + def __init__( self, vllm_config: "VllmConfig", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index db64b893e93e..58195d135605 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -170,7 +170,7 @@ def inject_kv_into_layer( # Only process layers that have kv_cache # attribute (attention layers) Skip non-attention - # layers like FusedMoE/MLP etc. + # layers like FusedMoEFactory/MLP etc. kv_cache_layer = getattr(layer, "kv_cache", None) if kv_cache_layer is None: continue diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 34f99b7cfacc..dd94ffb2f9ed 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -1678,9 +1678,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): conv, _ = cache_or_caches cache_list = [conv] else: - cache_list = self.transfer_topo.get_transfer_cache_regions( - cache_or_caches, layer_spec - ) + # K and V are packed into one blocks-first tensor per layer, + # so each layer registers as a single region. + cache_list = [cache_or_caches] logger.debug( "registering layer %s with %d cache tensor(s)", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index bf6038a897aa..ca32c3345c38 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -138,9 +138,17 @@ def __init__( ) assert vllm_config.kv_transfer_config is not None assert kv_cache_config is not None, "kv_cache_config is required" - self._validate_kv_cache_config(vllm_config, kv_cache_config) - self._kv_cache_config = kv_cache_config self.kv_role = vllm_config.kv_transfer_config.kv_role + # Capacity-only: contributes its segment to the store pool but transfers + # no KV, so the KV-cache-shape invariants below cannot be reached. + self._capacity_only = self.kv_role == "kv_consumer" and not ( + vllm_config.kv_transfer_config.kv_connector_extra_config.get( + "enable_lookup", True + ) + ) + if not self._capacity_only: + self._validate_kv_cache_config(vllm_config, kv_cache_config) + self._kv_cache_config = kv_cache_config self._kv_cache_events: MooncakeStoreKVEvents | None = None self.connector_scheduler: MooncakeStoreScheduler | None = None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 92283d3a537f..de5c8463542f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -10,17 +10,16 @@ ) from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import SpecGroup from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, ) -from vllm.v1.core.single_type_kv_cache_manager import ( - SingleTypeKVCacheManager, -) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, KVCacheSpec, + MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry @@ -83,48 +82,57 @@ def __init__( self.kv_cache_groups = kv_cache_groups self.hash_block_size = hash_block_size self.lcm_block_size = scheduler_block_size + self.enable_partial_hash_hits = partial_hash_hits_enabled( + kv_cache_groups, hash_block_size + ) self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval - self.eagle_group_ids = { - i for i, g in enumerate(kv_cache_groups) if g.is_eagle_group - } - if use_eagle and not self.eagle_group_ids: - self.eagle_group_ids = set(range(len(kv_cache_groups))) self._verify_and_split_kv_cache_groups() + def align_lookup_length(self, length: int) -> int: + alignment = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) + return length // alignment * alignment + def _verify_and_split_kv_cache_groups(self) -> None: """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but dispatches via spec_manager_map (we don't allocate managers). """ - attention_groups: list[ - tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]] - ] = [] + attention_groups: list[SpecGroup] = [] for i, g in enumerate(self.kv_cache_groups): spec = _unwrap_spec(g.kv_cache_spec) manager_cls = KVCacheSpecRegistry.get_manager_class(spec) assert manager_cls is not None, ( f"No manager registered for KVCacheSpec {spec}" ) - for existing_spec, group_ids, existing_cls in attention_groups: - if existing_spec == spec: - assert manager_cls is existing_cls - group_ids.append(i) + for idx, group in enumerate(attention_groups): + if group.spec == spec: + assert manager_cls is group.manager_cls + group.group_ids.append(i) + if g.is_eagle_group and not group.use_eagle: + attention_groups[idx] = group._replace(use_eagle=True) break else: - attention_groups.append((spec, [i], manager_cls)) + attention_groups.append( + SpecGroup(spec, [i], manager_cls, g.is_eagle_group) + ) # Full attention first (matches upstream convergence ordering). - self.attention_groups = sorted( - attention_groups, - key=lambda x: not isinstance(x[0], FullAttentionSpec), - ) - self.eagle_attn_group_indices: set[int] = { - i - for i, (_, group_ids, _) in enumerate(self.attention_groups) - if any(self.kv_cache_groups[gid].is_eagle_group for gid in group_ids) + attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec)) + # Conservatively flag all groups when use_eagle is set but none is flagged. + if self.use_eagle and not any(g.use_eagle for g in attention_groups): + attention_groups = [g._replace(use_eagle=True) for g in attention_groups] + self.attention_groups = attention_groups + # Per-group eagle bits. SpecGroup carries use_eagle for the whole + # merged spec group, so the per-group store/lookup masks agree with + # the merged-group hit check, which applies the eagle drop to every + # group sharing the spec. + self.eagle_group_ids = { + gid for g in attention_groups if g.use_eagle for gid in g.group_ids } - if self.use_eagle and not self.eagle_attn_group_indices: - self.eagle_attn_group_indices = set(range(len(self.attention_groups))) def find_longest_cache_hit( self, @@ -223,9 +231,14 @@ def _reachable_masks( retention_interval: int | None, num_prompt_tokens: int | None, ) -> tuple[list[bool] | None, ...]: - assert aligned_token_len % self.lcm_block_size == 0, ( + mask_alignment = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) + assert aligned_token_len % mask_alignment == 0, ( f"aligned_token_len ({aligned_token_len}) must be a multiple of " - f"lcm_block_size ({self.lcm_block_size})" + f"{mask_alignment}" ) masks: list[list[bool] | None] = [] for g_idx, g in enumerate(self.kv_cache_groups): @@ -270,22 +283,25 @@ def _find_hit_blocks( """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but dispatches via spec_manager_map (we don't allocate managers). - When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` — + When ``apply_eagle`` is False, ignore each group's ``use_eagle`` — used by ``load_mask`` to avoid popping a second block on top of the one already removed by the lookup. """ - eagle_indices = self.eagle_attn_group_indices if apply_eagle else set() + alignment_tokens = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) if len(self.attention_groups) == 1: - spec, group_ids, manager_cls = self.attention_groups[0] - hashes = self.block_hashes_for_spec(block_hashes, spec) + spec, group_ids, manager_cls, group_eagle = self.attention_groups[0] hit_blocks, hit_length = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + block_hashes=block_hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, - drop_eagle_block=(0 in eagle_indices), - alignment_tokens=spec.block_size, + drop_eagle_block=apply_eagle and group_eagle, + alignment_tokens=alignment_tokens, ) num_groups = len(self.kv_cache_groups) blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] @@ -299,14 +315,16 @@ def _find_hit_blocks( hit_length_by_group: list[int] = [0] * num_groups is_simple_hybrid = len(self.attention_groups) == 2 and isinstance( - self.attention_groups[0][0], FullAttentionSpec + self.attention_groups[0].spec, FullAttentionSpec ) eagle_verified: set[int] = set() while True: curr_hit_length = hit_length - for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): + for idx, (spec, group_ids, manager_cls, group_eagle) in enumerate( + self.attention_groups + ): first_group_id = group_ids[0] cached = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached is not None: @@ -315,19 +333,30 @@ def _find_hit_blocks( ) continue - drop_eagle_block = idx in eagle_indices and idx not in eagle_verified + drop_eagle_block = ( + apply_eagle and group_eagle and idx not in eagle_verified + ) _max_length = curr_hit_length - if drop_eagle_block: - _max_length = min(curr_hit_length + spec.block_size, max_length) - hashes = self.block_hashes_for_spec(block_hashes, spec) + # No eagle peek margin for a recurrent (Mamba) group: its finder + # never drops a block, so a widened bound would match past the + # attention-verified hit and resume from speculative state (#43559). + if drop_eagle_block and not isinstance(spec, MambaSpec): + eagle_margin = ( + self.hash_block_size + if self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup + and spec.block_size > self.hash_block_size + else spec.block_size + ) + _max_length = min(curr_hit_length + eagle_margin, max_length) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, - alignment_tokens=self.lcm_block_size, + alignment_tokens=alignment_tokens, ) if drop_eagle_block: eagle_verified.add(idx) @@ -345,15 +374,16 @@ def _find_hit_blocks( break # Truncate full-attention hit_blocks to final converged length; - # other specs already trim themselves inside their hit logic. - spec0, group_ids0, _ = self.attention_groups[0] - if isinstance(spec0, FullAttentionSpec): - num_blocks = hit_length // spec0.block_size - for gid in group_ids0: - full_blks = hit_blocks_by_group[gid] + # other specs already trim themselves inside their hit logic. cdiv keeps + # the partial tail block when hit_length is not block-aligned. + first_group = self.attention_groups[0] + if isinstance(first_group.spec, FullAttentionSpec): + num_blocks = cdiv(hit_length, first_group.spec.block_size) + for group_id in first_group.group_ids: + full_blks = hit_blocks_by_group[group_id] assert full_blks is not None del full_blks[num_blocks:] - hit_length_by_group[gid] = hit_length + hit_length_by_group[group_id] = hit_length return ( tuple(blks if blks is not None else [] for blks in hit_blocks_by_group), @@ -365,3 +395,18 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: if isinstance(spec, UniformTypeKVCacheSpecs): return next(iter(spec.kv_cache_specs.values())) return spec + + +def partial_hash_hits_enabled( + kv_cache_groups: list[KVCacheGroupSpec], hash_block_size: int +) -> bool: + """Mirror of core's ``HybridKVCacheCoordinator.enable_partial_hash_hits`` + (its dcp == 1 clause holds: the connector rejects hybrid + DCP/PCP > 1). + Single copy on purpose — scheduler and coordinator must not disagree. + """ + return any( + isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec) + and spec.mamba_cache_mode == "align" + and spec.block_size > hash_block_size + for g in kv_cache_groups + ) 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 ef98ec0d4e43..6daa1e82ea6d 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 @@ -9,6 +9,7 @@ from dataclasses import dataclass from typing import cast +import numpy as np import torch from vllm.distributed.kv_transfer.kv_connector.v1.base import ( @@ -196,22 +197,58 @@ def set_block_len(self, block_len: list[int]): def prepare_value( self, start: int, end: int, block_ids: list[int] ) -> tuple[list[int], list[int], int]: - """Compute memory addresses and sizes for a token range. + """Compute memory addresses and sizes for a single token range. Returns: (addr_list, size_list, block_id) """ + addr_lists, size_lists, chunk_block_ids = self.prepare_values( + ((start, end),), block_ids + ) + return addr_lists[0], size_lists[0], chunk_block_ids[0] + + def prepare_values( + self, + chunks: Sequence[tuple[int, int]], + block_ids: list[int], + ) -> tuple[list[list[int]], list[list[int]], list[int]]: + """Compute memory addresses and sizes for multiple token ranges. + + Returns: + (addr_lists, size_lists, chunk_block_ids), one entry per chunk. + """ + if not chunks: + return [], [], [] + base = np.asarray(self.kv_caches_base_addr, dtype=np.int64) + length = len(self.block_len) + blen = np.asarray( + [self.block_len[i % length] for i in range(base.shape[0])], + dtype=np.int64, + ) + n = len(chunks) + starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n) + spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts + assert not (spans % self.hash_block_size).any() + bids = np.fromiter( + (block_ids[i] for i in (starts // self.block_size).tolist()), + dtype=np.int64, + count=n, + ) + addrs = base[None, :] + bids[:, None] * blen[None, :] + block_counts = (spans + self.block_size - 1) // self.block_size + sizes = blen[None, :] * block_counts[:, None] + return addrs.tolist(), sizes.tolist(), bids.tolist() + + def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]: + """Return addresses and sizes for one physical block slot.""" addr_list = [] size_list = [] - block_id = block_ids[start // self.block_size] length = len(self.block_len) for index, base_addr in enumerate(self.kv_caches_base_addr): addr = base_addr + block_id * self.block_len[index % length] - assert (end - start) % self.block_size == 0 - size = self.block_len[index % length] * cdiv(end - start, self.block_size) addr_list.append(addr) - size_list.append(size) - return addr_list, size_list, block_id + size_list.append(self.block_len[index % length]) + return addr_list, size_list def process_tokens( self, @@ -231,7 +268,8 @@ def process_tokens( rank regardless of where the processed suffix begins. Args: - token_len: Total number of tokens. + token_len: Total number of tokens. Must be hash-block aligned and + covered by ``block_hashes`` when hashes are present. block_hashes: Block hashes computed at ``hash_block_size`` granularity. When ``block_size > hash_block_size`` each group's ``block_size`` chunk is keyed by its last sub-hash via ``chunk_hashes_for_block_size``. @@ -244,11 +282,10 @@ def process_tokens( assert put_step > 0 if not block_hashes: return - chunk_hashes: Sequence[BlockHash] = chunk_hashes_for_block_size( - block_hashes, self.hash_block_size, self.block_size - ) + assert token_len % self.hash_block_size == 0 + assert token_len // self.hash_block_size <= len(block_hashes) start_chunk = max(0, cdiv(mask_num, self.block_size)) - max_chunks = min(len(chunk_hashes), cdiv(token_len, self.block_size)) + max_chunks = cdiv(token_len, self.block_size) if chunk_mask is not None: max_chunks = min(max_chunks, start_chunk + len(chunk_mask)) for chunk_id in range(start_chunk, max_chunks): @@ -256,9 +293,9 @@ def process_tokens( continue if chunk_id % put_step != put_step_rank: continue - h = chunk_hashes[chunk_id] start_idx = chunk_id * self.block_size end_idx = min(start_idx + self.block_size, token_len) + h = block_hashes[end_idx // self.hash_block_size - 1] yield start_idx, end_idx, h @@ -281,6 +318,7 @@ class RequestTracker: allocated_block_ids: tuple[list[int], ...] num_saved_tokens: int = 0 token_ids: list[int] | None = None + has_pending_offload: bool = False # Snapshot of the prefill range length at tracker creation time. # For a fresh request this is len(prompt). For a resumed-from-preemption # request it includes previously-generated tokens, which are re-prefilled. @@ -291,6 +329,7 @@ def reset(self) -> None: self.allocated_block_ids = () self.num_saved_tokens = 0 self.token_ids = None + self.has_pending_offload = False self.prefill_end_tokens = 0 def update( @@ -327,6 +366,11 @@ class ReqMeta: token_ids: list[int] | None = None num_prompt_tokens: int | None = None + # Core-provided per-mamba-group + # (group_id, cow_block_id, boundary_tokens) for this request's partial tail. + # Present only on the producer's CoW step; drives the connector's offload + # (the FA group's block is derived from block_ids and boundary_tokens). + partial_tail_offloads: list[tuple[int, int, int]] | None = None @staticmethod def from_request_tracker( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py index fc91b0aeebc6..eb6c65afa335 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py @@ -10,7 +10,8 @@ Request: [msg_type: bytes] [payload_frames...] msg_type == LOOKUP_MSG: - frame 1: token_len (u32 big-endian, 4 bytes) + frame 1: num_tokens (u32 big-endian, 4 bytes); the worker derives + the aligned lookup length frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each fixed-size block hash (0 when there are no hashes) frame 3: raw block hashes concatenated back-to-back (each hash_len diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 58dfd5e428e2..42c6f3fa99af 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -11,6 +11,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 + partial_hash_hits_enabled, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 LoadSpec, MooncakeStoreConnectorMetadata, @@ -57,12 +60,17 @@ def __init__( kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config self.load_async = kvc_extra_config.get("load_async", True) self.lookup_async = kvc_extra_config.get("lookup_async", False) + # Skips lookup CPU cost on instances that never load KV from the store. + self.enable_lookup = kvc_extra_config.get("enable_lookup", True) self.client = LookupKeyClient(vllm_config) # Align with the engine's own scheduler_block_size and hash_block_size. self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) + self.enable_partial_hash_hits = partial_hash_hits_enabled( + kv_cache_config.kv_cache_groups, self._hash_block_size + ) # Per-request state self.load_specs: dict[str, LoadSpec] = {} # to be loaded @@ -80,14 +88,20 @@ def get_num_new_matched_tokens( Returns ``(None, False)`` when an async lookup is still in flight, signaling the scheduler to retry this request on a later step. """ - # Look up against the full prefill range, not just the prompt. - token_len = request.num_tokens // self._block_size * self._block_size - if token_len < self._block_size: + if not self.enable_lookup: + return 0, False + + # Fine-grained hits may land on a hash boundary inside a block; without + # partial hits, prefixes shorter than one physical block are skipped. + align = ( + self._hash_block_size if self.enable_partial_hash_hits else self._block_size + ) + if request.num_tokens < align: return 0, False num_external_hit_tokens = self.client.lookup( request.request_id, - token_len, + request.num_tokens, request.block_hashes, non_block=self.lookup_async, ) @@ -95,14 +109,6 @@ def get_num_new_matched_tokens( # Lookup not ready yet; scheduler will retry on a later step. return None, False - if num_external_hit_tokens == request.num_tokens: - # Leave a sub-block tail uncomputed for sampling, on a block - # boundary so the recv-side load mask covers every yielded chunk. - num_external_hit_tokens = max( - 0, - (request.num_tokens - 1) // self._block_size * self._block_size, - ) - if num_external_hit_tokens < num_computed_tokens: need_to_allocate = 0 else: @@ -348,6 +354,44 @@ def build_connector_meta( if req_meta is not None: meta.add_request(req_meta) + # Flush partial-tail offloads in the step they arrive: the CoW copy is + # enqueued before the connector event records, so this step's event + # fences the cow block. Ride the request's save meta when present, else + # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal + # save; can_save=True takes the normal enqueue path). + step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None) + if step_partial_tails and not force_skip_save: + pending = dict(step_partial_tails) + for req_meta in meta.requests: + if req_meta.can_save: + groups = pending.pop(req_meta.req_id, None) + if groups: + req_meta.partial_tail_offloads = groups + tracker = self._request_trackers.get(req_meta.req_id) + if tracker is not None: + tracker.has_pending_offload = True + for req_id, groups in pending.items(): + tracker = self._request_trackers.get(req_id) + req_tuple = self._unfinished_requests.get(req_id) + if tracker is None or req_tuple is None: + # Request finished/preempted within this step; its blocks + # are going away, so the offload is conservatively dropped. + logger.debug("Dropping partial-tail offload for request %s", req_id) + continue + assert len({boundary for _, _, boundary in groups}) == 1 + tracker.has_pending_offload = True + meta.add_request( + ReqMeta( + req_id=req_id, + token_len_chunk=0, + block_ids=tracker.allocated_block_ids, + block_hashes=req_tuple[0].block_hashes, + can_save=True, + num_prompt_tokens=tracker.prefill_end_tokens, + partial_tail_offloads=groups, + ) + ) + return meta def request_finished( @@ -362,7 +406,9 @@ def request_finished( # Missing tracker can happen when the request is aborted before the # connector observes the normal finished lifecycle or is preempted # before finishing. - if tracker is None or tracker.num_saved_tokens <= 0: + if tracker is None or ( + tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload + ): return False, None total_blocks = sum(len(g) for g in block_ids) delay_free_blocks = total_blocks > 0 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 38f4cb0c3a1e..c62eaefc3543 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 @@ -61,12 +61,21 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, maybe_convert_block_hash, resolve_kv_cache_block_sizes, ) -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + UniformTypeKVCacheSpecs, +) from .metrics import MooncakeStoreConnectorStats @@ -448,7 +457,7 @@ def __init__( token_databases: list[ChunkedTokenDatabase], block_size: int, tp_rank: int, - put_step: int, + group_put_steps: Sequence[int], kv_role: str, ready_event: threading.Event, enable_kv_event: bool = False, @@ -464,7 +473,8 @@ def __init__( name="KVCacheStoreSendingThread", record_operation=record_operation, ) - self.put_step = put_step + # Only ranks with identical group bytes may stripe PUTs (e.g., MLA). + self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role self.stored_requests: defaultdict[str, int] = defaultdict(int) @@ -522,6 +532,164 @@ def _clear_store_pressure(self) -> bool: self._skip_store_requests.clear() return True + def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: + """Offload the request's sub-block partial tail (its last prompt hash + boundary) so a later request can hit the sub-block prefix. + + Covers every block from the normal save's lcm floor to the boundary: + the normal save floors to ``lcm_block_size``, so a smaller-block + group's full blocks in that gap are never persisted elsewhere, and + the consumer's lookup needs every group at every probed boundary. + Full blocks are keyed by their block-end hash, the partial boundary + block by the boundary sub-hash; the mamba "align" boundary block is + the core-provided CoW block. All keys are deduped against the store. + + Returns: + True when no put is needed or every put succeeds, False otherwise. + """ + if not self.coord.enable_partial_hash_hits or not req_meta.block_hashes: + return True + partial_tail_offloads = req_meta.partial_tail_offloads + if not partial_tail_offloads: + return True + hash_block_size = self.coord.hash_block_size + boundaries = {boundary for _, _, boundary in partial_tail_offloads} + if len(boundaries) != 1: + raise ValueError( + "Partial-tail offloads for one request must share a boundary" + ) + boundary = boundaries.pop() + if boundary == 0: + return True + if boundary // hash_block_size - 1 >= len(req_meta.block_hashes): + return True + mamba_offloads = { + group_id: block_id for group_id, block_id, _ in partial_tail_offloads + } + + keys: list[str] = [] + addrs: list[list[int]] = [] + sizes: list[list[int]] = [] + saved = self._saved_offset.get(req_meta.req_id, 0) + for g_idx, db in enumerate(self.token_databases): + group_blocks = req_meta.block_ids[g_idx] + # Distribute across ranks by the same rule as normal chunks. + put_step = self.group_put_steps[g_idx] + put_step_rank = (self.tp_rank + g_idx) % put_step + # Always include the boundary block: its sub-hash key is written + # only here, even if normal saves already advanced past it. + last_block = cdiv(boundary, db.block_size) - 1 + for block_idx in range( + min(saved // db.block_size, last_block), last_block + 1 + ): + if block_idx % put_step != put_step_rank: + continue + valid_end = min((block_idx + 1) * db.block_size, boundary) + key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1] + if ( + g_idx in mamba_offloads + and valid_end == boundary + and boundary % db.block_size != 0 + ): + block_id = mamba_offloads[g_idx] + else: + if block_idx >= len(group_blocks): + continue + block_id = group_blocks[block_idx] + if block_id == NULL_BLOCK_ID: + logger.debug( + "Skipping unavailable partial-tail source block " + "(req=%s, group=%d, block=%d)", + req_meta.req_id, + g_idx, + block_idx, + ) + continue + addr, size = db.prepare_value_for_block(block_id) + keys.append(db.key_for(key_hash)) + addrs.append(addr) + sizes.append(size) + + if not keys: + return True + exists_start = time.perf_counter() + try: + exists = self.store.batch_is_exist(keys) + except Exception as e: + self._record_operation( + "save_exists", + exists_start, + len(keys), + status="error", + num_failed_keys=len(keys), + ) + logger.error( + "Failed to check partial-tail keys for request %s: %s", + req_meta.req_id, + e, + ) + return False + self._record_operation("save_exists", exists_start, len(keys)) + missing = [i for i, e in enumerate(exists) if e != 1] + if not missing: + return True + keys = [keys[i] for i in missing] + addrs = [addrs[i] for i in missing] + sizes = [sizes[i] for i in missing] + if req_meta.current_event is not None: + # Fence the CoW block copy enqueued earlier this step. + req_meta.current_event.synchronize() + batch_bytes = _sum_batch_bytes(sizes) + put_start = time.perf_counter() + try: + res = self.store.batch_put_from_multi_buffers( + keys, addrs, sizes, self.replicate_config + ) + except Exception as e: + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="error", + num_failed_keys=len(keys), + ) + logger.error( + "Failed to put partial-tail keys for request %s: %s", + req_meta.req_id, + e, + ) + return False + + failed = [i for i, value in enumerate(res) if value < 0] + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="partial_failure" if failed else "ok", + num_failed_keys=len(failed), + ) + if failed: + failed_codes = {res[i] for i in failed} + logger.warning( + "Partial-tail put failed for request %s: %d/%d keys failed (codes=%s)", + req_meta.req_id, + len(failed), + len(keys), + failed_codes, + ) + if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes: + self._mark_request_skipped_for_pressure(req_meta.req_id) + return False + + if self._clear_store_pressure(): + logger.info( + "Mooncake CPU/disk offloading pressure cleared after a " + "successful partial-tail batch" + ) + return True + def _handle_request(self, req_meta: ReqMeta): # Cache hits are always a multiple of ``lcm_block_size`` tokens, which # is also ``store_mask``'s precondition. @@ -539,9 +707,6 @@ def _handle_request(self, req_meta: ReqMeta): # so the scheduler can release the GPU blocks it pinned for this # request (via `delay_free_blocks`) even when the store path raises. try: - if token_len == 0: - return - if self._should_skip_request(req_id): logger.debug( "Skipping Mooncake store for request %s while CPU/disk " @@ -550,6 +715,16 @@ def _handle_request(self, req_meta: ReqMeta): ) return + # Offload the sub-block partial tail (independent of the normal + # block-aligned save, which may be skipped this step). + if req_meta.partial_tail_offloads is not None and not ( + self._maybe_offload_partial_tail(req_meta) + ): + return + + if token_len == 0: + return + # Resume from where this rank left off; only the new suffix is saved. save_start = self._saved_offset.get(req_id, 0) @@ -568,13 +743,14 @@ def _handle_request(self, req_meta: ReqMeta): group_indices: list[int] = [] for g_idx, db in enumerate(self.token_databases): # Rotate the stride phase per group to balance load across ranks. - put_step_rank = (self.tp_rank + g_idx) % self.put_step + put_step = self.group_put_steps[g_idx] + put_step_rank = (self.tp_rank + g_idx) % put_step for start, end, block_hash in db.process_tokens( token_len, req_meta.block_hashes, mask_num=save_start, chunk_mask=store_masks[g_idx], - put_step=self.put_step, + put_step=put_step, put_step_rank=put_step_rank, ): starts.append(start) @@ -634,6 +810,21 @@ def _handle_request(self, req_meta: ReqMeta): addrs: list[list[int]] = [] sizes: list[list[int]] = [] stored_events: list[BlockStored] = [] + chunks_per_group: list[list[tuple[int, int]]] = [ + [] for _ in self.token_databases + ] + for start, end, g_idx in zip(starts, ends, group_indices, strict=True): + chunks_per_group[g_idx].append((start, end)) + for g_idx, chunks in enumerate(chunks_per_group): + if not chunks: + continue + db = self.token_databases[g_idx] + group_addrs, group_sizes, _ = db.prepare_values( + chunks, block_ids_per_group[g_idx] + ) + addrs.extend(group_addrs) + sizes.extend(group_sizes) + # parent_block_hash chains live within a group, not across. if self.enable_kv_event: prev_key_per_group: dict[int, Any] = {} @@ -645,10 +836,6 @@ def _handle_request(self, req_meta: ReqMeta): zip(starts, ends, group_indices, strict=True) ): db = self.token_databases[g_idx] - addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx]) - addrs.append(addr) - sizes.append(size) - if self.enable_kv_event: token_ids = ( req_meta.token_ids[s:e] @@ -805,19 +992,21 @@ def _handle_request(self, req_meta: ReqMeta): block_id_list: list[int] = [] for g_idx, db in enumerate(self.token_databases): mask = load_mask_per_group[g_idx] + chunks: list[tuple[int, int]] = [] for start, end, block_hash in db.process_tokens( token_len, req_meta.block_hashes, mask_num ): chunk_idx = start // db.block_size if chunk_idx >= len(mask) or not mask[chunk_idx]: continue - addr, size, block_id = db.prepare_value( - start, end, req_meta.block_ids[g_idx] - ) key_list.append(db.key_for(block_hash)) - addr_list.append(addr) - size_list.append(size) - block_id_list.append(block_id) + chunks.append((start, end)) + g_addrs, g_sizes, g_block_ids = db.prepare_values( + chunks, req_meta.block_ids[g_idx] + ) + addr_list.extend(g_addrs) + size_list.extend(g_sizes) + block_id_list.extend(g_block_ids) # Rotate aligned lists by tp_rank for load balancing. rotation = self.tp_rank % len(key_list) @@ -983,50 +1172,19 @@ def __init__( self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get( "load_async", True ) + # Mirrors MooncakeStoreConnector._capacity_only. + self._capacity_only = self.kv_role == "kv_consumer" and not ( + vllm_config.kv_transfer_config.kv_connector_extra_config.get( + "enable_lookup", True + ) + ) self.cache_config = vllm_config.cache_config self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) self.num_layers = model_config.get_num_layers(parallel_config) - self.use_mla = False - if ( - hasattr(model_config, "use_mla") - and isinstance(model_config.use_mla, bool) - and model_config.use_mla - ): - self.use_mla = True - - if self.use_mla: - self.num_kv_head = 1 - else: - self.num_kv_head = model_config.get_total_num_kv_heads() - - if self.num_kv_head < self.tp_size and self.dcp_size <= 1: - # Dedup: TP ranks holding the same KV heads stripe PUTs across - # one shared key namespace. DCP splits the TP group, so with - # DCP>1 those ranks have different `@dcpN` namespaces and - # striping would leave keys unwritten (OBJECT_NOT_FOUND on - # GET). PCP is outer to TP (pcp_rank is constant within a TP - # group), so it needs no guard. - self.put_step = self.tp_size // self.num_kv_head - self.head_or_tp_rank = self.tp_rank // self.put_step - else: - self.head_or_tp_rank = self.tp_rank - self.put_step = 1 - - self.metadata = KeyMetadata( - model_name=model_config.model.rstrip("/").split("/")[-1], - tp_rank=self.head_or_tp_rank, - pcp_rank=self.pcp_rank, - dcp_rank=self.dcp_rank, - pp_rank=self.pp_rank, - cache_prefix=str( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "cache_prefix", "" - ) - ), - ) + self.num_kv_head = model_config.get_total_num_kv_heads() # Initialize MooncakeDistributedStore with its own TransferEngine store_config = MooncakeStoreConfig.load_from_config() @@ -1114,6 +1272,17 @@ def __init__( self.kv_connector_stats = MooncakeStoreConnectorStats() self._kv_cache_config = kv_cache_config + self.token_dbs: list[ChunkedTokenDatabase] = [] + + # a capacity-only instance does not need below utils + if self._capacity_only: + logger.info( + "Mooncake store in capacity-only mode: segment mounted " + "(global_segment_size=%d), KV transfer disabled.", + store_config.global_segment_size, + ) + return + # Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to # self.block_size (= scheduler_block_size) so the coordinator's # ``block_size % hash_block_size == 0`` invariant holds. @@ -1143,10 +1312,32 @@ def __init__( retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL, ) # One ChunkedTokenDatabase per group; addresses populated in - # register_kv_caches once the kv-cache layout is known. - self.token_dbs: list[ChunkedTokenDatabase] = [ + # register_kv_caches once the kv-cache layout is known. Each group's + # key namespace is its TP shard id: ranks holding identical bytes + # (MLA / shared GQA KV heads) share a namespace, TP-sharded Mamba + # state gets one namespace per rank. + metadata = KeyMetadata( + model_name=model_config.model.rstrip("/").split("/")[-1], + tp_rank=self.tp_rank, + pcp_rank=self.pcp_rank, + dcp_rank=self.dcp_rank, + pp_rank=self.pp_rank, + cache_prefix=str( + vllm_config.kv_transfer_config.kv_connector_extra_config.get( + "cache_prefix", "" + ) + ), + ) + self._group_tp_replication_factors: tuple[int, ...] = ( + self._compute_group_tp_replication_factors() + ) + self.token_dbs = [ ChunkedTokenDatabase( - dataclasses.replace(self.metadata, group_id=g_idx), + dataclasses.replace( + metadata, + group_id=g_idx, + tp_rank=self.tp_rank // self._group_tp_replication_factors[g_idx], + ), g.kv_cache_spec.block_size, hash_block_size=self.hash_block_size, ) @@ -1154,27 +1345,50 @@ def __init__( ] self._init_lookup_key_prefixes() - def _init_lookup_key_prefixes(self) -> None: - """Prepare per-group key prefixes across parallel rank namespaces.""" - # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: if self.dcp_size > 1: - # DCP reuses the TP workers and splits each TP group into - # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. - # Store/load paths do not apply KV-head dedup under DCP - rank_namespaces = tuple( - (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) - for pcp_rank in range(self.pcp_size) - for tp_rank in range(self.tp_size) - for pp_rank in range(self.pp_size) - ) - else: - # Without DCP, TP ranks that share a KV head write identical KV, so - # lookup only needs one TP namespace per unique KV head. - tp_count = min(self.tp_size, self.num_kv_head) - rank_namespaces = tuple( - (tp_rank, pcp_rank, 0, pp_rank) + return 1 + inner_specs = ( + tuple(spec.kv_cache_specs.values()) + if isinstance(spec, UniformTypeKVCacheSpecs) + else (spec,) + ) + # Any rank-specific state makes the whole packed value rank-specific. + if any(isinstance(inner, MambaSpec) for inner in inner_specs): + return 1 + # A pure MLA packed value is replicated on every TP rank. + if all( + isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec)) + for inner in inner_specs + ): + return self.tp_size + return max(1, self.tp_size // self.num_kv_head) + + def _compute_group_tp_replication_factors(self) -> tuple[int, ...]: + """Return the number of byte-identical TP replicas per cache group. + + DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses + ``tp_size // num_kv_head``. + """ + return tuple( + self._spec_tp_replication_factor(group.kv_cache_spec) + for group in self._kv_cache_groups + ) + + def _init_lookup_key_prefixes(self) -> None: + def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]: + if self.dcp_size > 1: + # DCP is a TP subdivision: dcp_rank == tp_rank % dcp_size. + return tuple( + (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) + for pcp_rank in range(self.pcp_size) + for tp_rank in range(self.tp_size) + for pp_rank in range(self.pp_size) + ) + return tuple( + (shard_rank, pcp_rank, 0, pp_rank) for pcp_rank in range(self.pcp_size) - for tp_rank in range(tp_count) + for shard_rank in range(self.tp_size // factor) for pp_rank in range(self.pp_size) ) @@ -1187,11 +1401,12 @@ def _init_lookup_key_prefixes(self) -> None: dcp_rank=dcp_rank, pp_rank=pp_rank, ) - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces( + self._group_tp_replication_factors[g_idx] + ) ) - for db in self.token_dbs + for g_idx, db in enumerate(self.token_dbs) ) - self._lookup_expected_per_key = len(rank_namespaces) def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: """Register a cross-layers KV cache tensor. @@ -1207,6 +1422,8 @@ def register_kv_caches( kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], ) -> None: """Register KV cache tensors and start transfer threads.""" + if self._capacity_only: + return if not kv_caches: logger.warning("No KV caches to offload.") return @@ -1284,7 +1501,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: self.token_dbs, self.block_size, self.tp_rank, - self.put_step, + self._group_tp_replication_factors, self.kv_role, ready_event_sending, self.enable_kv_events, @@ -1343,6 +1560,9 @@ def get_finished( compute is launched on the compute stream) for better compute-I/O overlap. """ + if self._capacity_only: + return set(), set() + # Issue async loads for request in meta.requests: load_spec = request.load_spec @@ -1353,7 +1573,7 @@ def get_finished( self.recv_request_queue.put(request) assert self.load_async, "load_async must be True for better performance." - # Issue stores with CUDA event synchronization + # Issue stores with CUDA event synchronization. if self.kv_role in ["kv_producer", "kv_both"]: current_event = None for request in meta.requests: @@ -1454,11 +1674,17 @@ def _get_and_clear_finished_sending( return finished_sending - def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: + def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int: """Check how many prefix tokens exist in the store. - Checks across all rank-specific key namespaces that may be loaded. + Checks across all rank-specific key namespaces that may be loaded. A + hit covering all ``num_tokens`` is re-derived below the request end so + the last token is recomputed for sampling. """ + if self._capacity_only: + return 0 + + token_len = self.coord.align_lookup_length(num_tokens) if not block_hashes or token_len <= 0: return 0 @@ -1466,21 +1692,32 @@ def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: # candidate_meta stores the (group, hash_bytes) for key slice. candidate_keys: list[str] = [] candidate_meta: list[tuple[int, bytes]] = [] - lookup_masks = self.coord.lookup_mask(token_len) + fine_grained = self.coord.enable_partial_hash_hits + lookup_masks = None if fine_grained else self.coord.lookup_mask(token_len) for g_idx, db in enumerate(self.token_dbs): spec_block_size = db.block_size - lookup_mask = lookup_masks[g_idx] key_prefixes = self._lookup_key_prefixes[g_idx] - group_hashes = self.coord.block_hashes_for_spec( - block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec - ) - max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size)) - mask_limit = ( - max_chunks if lookup_mask is None else min(max_chunks, len(lookup_mask)) - ) - for chunk_id in range(mask_limit): - if lookup_mask is not None and not lookup_mask[chunk_id]: - continue + if fine_grained: + max_units = min(len(block_hashes), token_len // self.hash_block_size) + unit_ids: range | list[int] = range(max_units) + group_hashes: Sequence[BlockHash] = block_hashes + else: + lookup_mask = lookup_masks[g_idx] # type: ignore[index] + group_hashes = self.coord.block_hashes_for_spec( + block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec + ) + max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size)) + mask_limit = ( + max_chunks + if lookup_mask is None + else min(max_chunks, len(lookup_mask)) + ) + unit_ids = [ + chunk_id + for chunk_id in range(mask_limit) + if lookup_mask is None or lookup_mask[chunk_id] + ] + for chunk_id in unit_ids: h = group_hashes[chunk_id] hash_hex = h.hex() for key_prefix in key_prefixes: @@ -1511,22 +1748,35 @@ def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: logger.error("Remote connection failed in lookup: %s", e) return 0 - # A (group, hash) is "present" only when every TP*PP rank has it. - ranks_per_candidate = self._lookup_expected_per_key - exists_set = { - (g_idx, hash_bytes) - for i, (g_idx, hash_bytes) in enumerate(candidate_meta) - if all( - res[i * ranks_per_candidate + j] == 1 - for j in range(ranks_per_candidate) - ) - } - + # A (group, hash) is "present" only when every namespace that will be + # loaded has it (per-group count: sharded groups need every rank's + # shard, replicated groups one namespace per unique KV head). + exists_set = set() + pos = 0 + for g_idx, hash_bytes in candidate_meta: + count = len(self._lookup_key_prefixes[g_idx]) + if all(res[pos + j] == 1 for j in range(count)): + exists_set.add((g_idx, hash_bytes)) + pos += count + + cached_block_pool = ExternalCachedBlockPool( + self.hash_block_size, + exists_set, + ) _masks, hit_length = self.coord.find_longest_cache_hit( block_hashes, token_len, - ExternalCachedBlockPool(self.hash_block_size, exists_set), + cached_block_pool, ) + if hit_length >= num_tokens: + usable_length = self.coord.align_lookup_length(num_tokens - 1) + if usable_length <= 0: + return 0 + _masks, hit_length = self.coord.find_longest_cache_hit( + block_hashes, + usable_length, + cached_block_pool, + ) return hit_length def get_kv_events(self) -> list[BlockStored]: @@ -1592,11 +1842,11 @@ def process_request(): msg_type = bytes(all_frames[0]) if msg_type == LOOKUP_MSG: - token_len = int.from_bytes(all_frames[1], byteorder="big") + num_tokens = int.from_bytes(all_frames[1], byteorder="big") hash_len = int.from_bytes(all_frames[2], byteorder="big") blob = all_frames[3].buffer block_hashes = BlobBlockHashes(blob, hash_len) - result = self.store_worker.lookup(token_len, block_hashes) + result = self.store_worker.lookup(num_tokens, block_hashes) self.socket.send(result.to_bytes(4, "big")) elif msg_type == RESET_MSG: @@ -1659,11 +1909,11 @@ def __init__(self, vllm_config: VllmConfig): ) self.futures: dict[str, Future[int]] = {} - def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: + def _lookup(self, num_tokens: int, block_hashes: list[BlockHash]) -> int: hash_len = len(block_hashes[0]) if block_hashes else 0 all_frames = ( LOOKUP_MSG, - token_len.to_bytes(4, byteorder="big"), + num_tokens.to_bytes(4, byteorder="big"), hash_len.to_bytes(2, byteorder="big"), b"".join(block_hashes), ) @@ -1674,7 +1924,7 @@ def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: def lookup( self, req_id: str, - token_len: int, + num_tokens: int, block_hashes: list[BlockHash], non_block: bool = False, ) -> int | None: @@ -1682,7 +1932,7 @@ def lookup( so the caller retries on a later step.""" future = self.futures.get(req_id) if future is None: - future = self.executor.submit(self._lookup, token_len, list(block_hashes)) + future = self.executor.submit(self._lookup, num_tokens, list(block_hashes)) self.futures[req_id] = future if non_block and not future.done(): return None 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 15585123e5c9..811bedb08415 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 @@ -422,6 +422,10 @@ class ReqMeta: remote_engine_id: str tp_size: int remote_dp_size: int + # Prefill DP rank that owns this request's KV (forwarded by the proxy). The + # read must target this rank's memory registration; the default 0 preserves + # the symmetric single-DP behaviour. + remote_dp_rank: int = 0 class MoRIIOConnectorMetadata(KVConnectorMetadata): @@ -473,8 +477,17 @@ def add_new_req( remote_port=int(remote_handshake_port), remote_handshake_port=int(remote_handshake_port), remote_notify_port=int(remote_notify_port), - tp_size=kv_transfer_params.get("tp_size", 1), + # Remote peer TP degree (used as remote_tp_size downstream). The + # proxy advertises it under "remote_tp_size"; #46332 read "tp_size" + # which is absent on WRITE producer requests -> defaulted to 1 -> + # rank collapse. Read the right key; 0 == unknown (== homogeneous). + tp_size=int( + kv_transfer_params.get("remote_tp_size") + or kv_transfer_params.get("tp_size") + or 0 + ), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), + remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), ) if write_mode: self.reqs_to_save[request_id] = _req 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 ea41adf437ce..2a4c028e15b0 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 @@ -1014,6 +1014,8 @@ def __init__( self._handshake_futures: dict[EngineId, Future[set[str]]] = {} # Protects _handshake_futures and _remote_agents. self._handshake_lock = threading.RLock() + # Remote engines already covered by the eager pre-forward handshake. + self._eager_handshaked_engines: set[EngineId] = set() self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config @@ -1261,8 +1263,16 @@ def _moriio_handshake( remote_tp_size: int, expected_engine_id: str, remote_dp_rank: int = 0, + remote_tp_rank: int | None = None, ) -> set[str]: - """Do a MoRIIO handshake with a remote instance.""" + """Do a MoRIIO handshake with a remote instance. + + remote_tp_rank: explicit remote TP index to dial. Flexible-read callers + pass the chosen prefill TP rank so the handshake, the (dp, tp) session + key and the notify port all address the SAME rank. None falls back to the + local-rank mapping _remote_tp_rank -- byte-identical for callers not yet + TP-aware. + """ start_time = time.perf_counter() @@ -1270,9 +1280,12 @@ def _moriio_handshake( # a hack to keep us moving. We will switch when moving to etcd # or where we have a single ZMQ socket in the scheduler. - port_offset = get_port_offset( - remote_dp_rank, self._remote_tp_rank(remote_tp_size) + dial_tp_rank = ( + self._remote_tp_rank(remote_tp_size) + if remote_tp_rank is None + else int(remote_tp_rank) ) + port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size) path = make_zmq_path("tcp", host, port + port_offset) logger.debug("handshake Querying metadata on path: %s", path) @@ -1338,6 +1351,9 @@ def _moriio_handshake( return {remote_agent_name} def _remote_tp_rank(self, remote_tp_size: int) -> int: + # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0). + if remote_tp_size == 0: + remote_tp_size = self.world_size return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size) def _background_moriio_handshake( @@ -1729,6 +1745,149 @@ def save_kv_layer( def get_engine_name_with_dp(self, engine_name, dp_rank): return f"{engine_name}_dp{dp_rank}" + def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank): + # Per-(dp, tp) session key. The flexible mirror read keys sessions per + # (dp, tp) so one decode worker can hold a session to EACH prefill TP + # rank and spread reads across them; other configs keep the DP-only key. + return f"{engine_name}_dp{dp_rank}_tp{tp_rank}" + + def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None: + """Handshake EVERY remote prefill DP rank BEFORE the decode forward pass, + identically across all local TP workers. + + Why this exists (the deadlock it prevents): with heterogeneous DP prefill + a decode TP worker reads KV from whichever prefill DP rank owns the + request, so across requests every worker must reach several prefill DP + ranks. The decode forward issues per-layer TP collectives (e.g. an + all-gather) that all local TP workers must enter together. If the + handshakes are left to fire lazily on the read path, the workers diverge: + a worker whose target rank is already cached races ahead into the forward + collective while a peer is still blocked in a handshake recv(). The first + worker then waits inside the collective for the stuck peer -> 600s NCCL + timeout / hang. This was observed directly with mixed TP<->DP configs. + + Fix: complete ALL prefill-DP-rank handshakes for every referenced remote + engine HERE, before any read enters the forward, so no worker is still + handshaking once its peers reach a collective. Fires ONCE per remote + engine (first contact), gated by _eager_handshaked_engines. The engine + set comes from scheduler-built metadata (identical on every TP worker), + so all workers run the same handshakes in the same order and reach the + all-reduce barrier below together. + + Failure handling: handshake exceptions are caught, never raised before + the collective (raising early would hang the peers still waiting for it). + Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL + raise the same error AFTER the collective, so the step fails fast and + uniformly in ~seconds instead of one rank hanging the forward for 600s. + """ + import torch.distributed as dist + + # Distinct remote engines referenced this step, in metadata (== + # scheduler) order so every TP worker iterates engines identically. + engines: dict[str, ReqMeta] = {} + for _req_id, meta in metadata.reqs_to_recv.items(): + remote_engine_id = ( + str(meta.remote_host) + ":" + str(meta.remote_handshake_port) + ) + engines.setdefault(remote_engine_id, meta) + + for remote_engine_id, meta in engines.items(): + if remote_engine_id in self._eager_handshaked_engines: + continue + + remote_dp_size = int(meta.remote_dp_size) + port = int(meta.remote_handshake_port) + tp_size = int(meta.tp_size) + + # Flexible mirror (TP prefill + MLA, world_size==1 decode): the read + # round-robins over prefill TP ranks, so pre-warm a session to EVERY + # (dp, tp) rank. Other configs pre-warm per DP rank (tp resolved by + # the fixed local-rank mapping) -- byte-identical to before. The + # mirror's decode is DP+EP, whose forward all-to-all is the collective + # that the eager barrier keeps everyone in step for. + flexible = ( + self.world_size == 1 + and self.use_mla + and remote_dp_size == 1 + and tp_size > 1 + ) + # (engine_id, dp_rank, tp_rank_or_None); tp_rank is None on the legacy + # path so _moriio_handshake falls back to its _remote_tp_rank mapping. + targets: list[tuple[Any, int, int | None]] + if flexible: + targets = [ + (self.get_engine_name_with_dp_tp(remote_engine_id, dp, tp), dp, tp) + for dp in range(remote_dp_size) + for tp in range(max(1, tp_size)) + ] + else: + targets = [ + (self.get_engine_name_with_dp(remote_engine_id, dp), dp, None) + for dp in range(remote_dp_size) + ] + + # Submit handshakes for every not-yet-known target UNDER the lock; do + # NOT hold it across the join or the collective (a stalled recv must + # not block another thread's lock acquisition). Gate on BOTH + # _remote_agents AND layer metadata: a rank with an agent entry but no + # layer metadata is half-handshaked and would KeyError at read time. + futures: list[tuple[str, Future[set[str]]]] = [] + with self._handshake_lock: + for eid, cur_dp_rank, cur_tp_rank in targets: + if ( + eid in self._remote_agents + and eid in self.layer_name_to_remote_kv_cache_metadata + ): + continue + fut = self._handshake_initiation_executor.submit( + self._moriio_handshake, + meta.remote_host, + port, + tp_size, + eid, + cur_dp_rank, + cur_tp_rank, + ) + futures.append((eid, fut)) + + # Join outside the lock. Bounded handshake errors are recorded here + # and reported after the all-reduce. + all_ok = True + results: dict[str, set[str]] = {} + for eid, fut in futures: + try: + results[eid] = fut.result() + except Exception: + logger.exception("Eager MoRIIO handshake failed for %s", eid) + all_ok = False + + with self._handshake_lock: + for eid, agents in results.items(): + self._remote_agents[eid] = agents + + logger.info( + "Eager MoRIIO handshake: engine=%s dp_size=%d new_ranks=%d " + "ok=%s tp_rank=%d", + remote_engine_id, + remote_dp_size, + len(futures), + all_ok, + self.tp_rank, + ) + # CPU all-reduce = TP-uniform success vote AND lockstep barrier: it + # blocks until every TP worker arrives, gives them the same verdict, + # and stays off the model compute stream. + vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32) + dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN) + if int(vote.item()) == 0: + raise HandshakeError( + f"Eager MoRIIO handshake failed for {remote_engine_id} on " + "at least one TP rank; failing this step fast to avoid a " + "TP collective hang" + ) + + self._eager_handshaked_engines.add(remote_engine_id) + def start_load_kv(self, metadata: MoRIIOConnectorMetadata): """ Start loading by triggering non-blocking moriio_xfer. @@ -1750,6 +1909,12 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): if self.mode == MoRIIOMode.WRITE: return + # Handshake every referenced remote prefill rank up front, before any + # read enters the forward pass. A lazy per-rank handshake on the read + # path lets TP workers diverge into a forward collective while a peer is + # still blocked handshaking -> NCCL hang (see below). + self._eager_handshake_all_dp_ranks(metadata) + wait_handshake_readd_req = False remote_engine_id = None @@ -1758,8 +1923,15 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): str(meta.remote_host) + ":" + str(meta.remote_handshake_port) ) meta.remote_engine_id = remote_engine_id + # The eager handshake above already covered every referenced engine + # (and keys the mirror per (dp, tp), which the DP-only dp0 probe below + # would miss). Only fall back to the lazy background handshake for an + # engine it did not cover. dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0) - if dp0_remote_engine_id not in self._remote_agents: + if ( + remote_engine_id not in self._eager_handshaked_engines + and dp0_remote_engine_id not in self._remote_agents + ): # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: if remote_engine_id not in self._remote_agents: @@ -1809,12 +1981,53 @@ def wait_for_save(self, metadata: MoRIIOConnectorMetadata): self.save_kv_layer(metadata, layer_name, kv_layer, None) self._writer.seal_pending_transfers() + def _next_flex_tp_rank(self, remote_tp_size: int) -> int: + """Deterministic round-robin over prefill tp0..N-1 for the flexible read. + + Round-robin (not random): exactly uniform and testable, with the same + prefill-NIC balancing. Seeded from this decode rank's dp_rank so + concurrent decode DP ranks are phase-staggered -- at a given read index + distinct decode ranks target distinct prefill TP ranks. + """ + rr = getattr(self, "_flex_tp_rr", None) + if rr is None: + rr = int(getattr(self, "dp_rank", 0) or 0) + self._flex_tp_rr = rr + 1 + return rr % remote_tp_size + + def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]: + """Resolve (chosen_tp, flexible) for reading this request's KV. + + Flexible mirror (decode world_size==1 + MLA + pure-TP prefill): MLA + replicates the latent KV across the prefill TP ranks, so any is a valid + source; round-robin across them to spread RDMA/NIC load. Otherwise the + source TP rank is fixed by the local-rank mapping (_remote_tp_rank) -- + forward DP8EP->TP8 -> tp0; symmetric TP -> tp_rank -- byte-identical to + prior behaviour. chosen_tp is the single value threaded into the (dp, tp) + session key, the handshake dial and the notify port, so all three address + the SAME prefill rank (drift -> read one rank but notify another -> the + read rank's prefill buffer is never freed). + """ + remote_tp_size = int(meta.tp_size) + flexible = ( + self.world_size == 1 + and self.use_mla + and int(meta.remote_dp_size) == 1 + and remote_tp_size > 1 + ) + if flexible: + chosen_tp = self._next_flex_tp_rank(remote_tp_size) + else: + chosen_tp = self._remote_tp_rank(remote_tp_size) + return chosen_tp, flexible + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): logger.debug( "Remote agent %s available, calling _read_blocks for req %s", meta.remote_engine_id, req_id, ) + chosen_tp, flexible = self._resolve_read_source(meta) self._read_blocks( request_id=req_id, transfer_id=meta.transfer_id, @@ -1824,6 +2037,9 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_host=meta.remote_host, remote_notify_port=meta.remote_notify_port, remote_tp_size=meta.tp_size, + remote_dp_rank=meta.remote_dp_rank, + chosen_tp=chosen_tp, + flexible=flexible, ) def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer): @@ -1932,7 +2148,9 @@ def _compute_block_transfer_offsets( validate_moriio_heterogeneous_tp_kv_heads( local_tp_size=self.world_size, remote_tp_size=( - remote_tp_size if remote_tp_size is not None else self.world_size + remote_tp_size + if remote_tp_size and remote_tp_size > 0 + else self.world_size ), total_num_kv_heads=self.model_config.get_total_num_kv_heads(), is_mla=self._is_mla_cache_layer(layer_name), @@ -1976,12 +2194,37 @@ def _read_blocks( remote_host: str, remote_notify_port: int, remote_tp_size: int, + remote_dp_rank: int = 0, + chosen_tp: int | None = None, + flexible: bool = False, ) -> None: if self.mode == MoRIIOMode.WRITE: return - dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) - sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) + # Read from the prefill rank that actually computed this request's KV + # (forwarded by the proxy). Hardcoding DP0 reads from a different rank's + # memory registration; per-rank num_blocks differ, so high block ids can + # overrun the wrong rank's region. + # + # eff_tp = the remote TP rank this read targets. The flexible mirror + # reads from a round-robin-chosen prefill TP rank and keys the session + # per (dp, tp); other configs use the fixed local-rank mapping (eff_tp == + # _remote_tp_rank), byte-identical to before. This key MUST match the one + # the eager handshake stored the session under. + eff_tp = ( + int(chosen_tp) + if chosen_tp is not None + else self._remote_tp_rank(remote_tp_size) + ) + if flexible: + remote_dp_engine_id = self.get_engine_name_with_dp_tp( + dst_engine_id, int(remote_dp_rank), eff_tp + ) + else: + remote_dp_engine_id = self.get_engine_name_with_dp( + dst_engine_id, int(remote_dp_rank) + ) + sessions, remote_moriio_meta = self._get_built_session(remote_dp_engine_id) # SQ-full backpressure deadline, shared across this request's layers. _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout @@ -2030,6 +2273,13 @@ def _read_blocks( self._recving_transfers[request_id].append(transfer_status) self._recving_transfers_callback_addr[request_id] = ( remote_host, - str(remote_notify_port + self._remote_tp_rank(remote_tp_size)), + str( + remote_notify_port + + get_port_offset( + int(remote_dp_rank), + eff_tp, + remote_tp_size, + ) + ), transfer_id, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index f3e89f6623c4..59dad186c610 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -209,6 +209,10 @@ def prefer_cross_layer_blocks(self) -> bool: return False return all(c.prefer_cross_layer_blocks for c in self._connectors) + @property + def requires_kv_delivery(self) -> bool: + return any(c.requires_kv_delivery for c in self._connectors) + @classmethod def _get_connector_classes_and_configs( cls, vllm_config: "VllmConfig" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 7a3629c7229f..afa16279ff1c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -135,6 +135,20 @@ def __init__( for n_tokens, block_size in sw_sizes_tokens ] + # Trailing scratch slots that mamba managers co-allocate per request + # for speculative decoding; None for non-SSM groups. + self._ssm_spec_blocks = [ + g.kv_cache_spec.num_speculative_blocks + if isinstance(g.kv_cache_spec, MambaSpec) + else None + for g in kv_cache_config.kv_cache_groups + ] + # Only "all" mode keeps a state per block position; the other modes + # keep a single running state in the last non-speculative slot. + self._ssm_state_slots_are_positional = ( + vllm_config.cache_config.mamba_cache_mode == "all" + ) + # Threshold to decide whether to compute kv cache locally # or pull from a remote node: minimum number of remote # tokens to amortize the xfer latencies @@ -218,13 +232,25 @@ def _stop_heartbeat(self, req_id: ReqId) -> None: # Clean up empty engines so we don't leak a key when remote dies. del self._heartbeat_by_engine[engine_id] - def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: - """ - Clip the number of blocks to the sliding window size for each kv cache group - that employs SWA. - This is necessary because the KV Cache manager initially allocates blocks for - the entire sequence length, and successively cleans up blocks that are outside - the window prior to the `request_finished_all_groups` hook. + def get_exchange_clipped_blocks( + self, block_ids: BlockIds, clip_ssm: bool = True + ) -> BlockIds: + """Clip a request's block lists down to the transferable blocks. + + Sliding-window groups keep only the in-window tail: the KV cache + manager allocates blocks for the entire sequence length and cleans up + out-of-window blocks only prior to the `request_finished_all_groups` + hook. + + SSM groups keep only their state-bearing slots: the trailing + speculative scratch slots always go, and in single-state cache modes + so does everything before the running state (null placeholders and + the previous step's superseded state). "all" mode keeps its remaining + slots, which the worker pairs position-wise. + + Use this at every block-id exchange point. Pass ``clip_ssm=False`` + for per-step partial lists (host-buffer save), where the SSM strip + does not apply. """ if len(block_ids) == 0 or not self._is_hma_required: # No blocks to clip eg Full prefix cache hit or not a hybrid model. @@ -235,15 +261,22 @@ def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: assert len(block_ids) == len(self.blocks_per_sw), ( "Number of KV cache groups must match" ) - # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged - return tuple( - [ - blocks[-self.blocks_per_sw[i] :] - if self.blocks_per_sw[i] > 0 - else blocks - for i, blocks in enumerate(block_ids) - ] - ) + clipped = [] + for i, blocks in enumerate(block_ids): + if n_sw := self.blocks_per_sw[i]: + blocks = blocks[-n_sw:] + elif ( + clip_ssm + and blocks + and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None + ): + if n_spec := min(n_spec_blocks, len(blocks) - 1): + blocks = blocks[:-n_spec] + if not self._ssm_state_slots_are_positional: + # Never empty: downstream reads that as a full prefix hit. + blocks = blocks[-1:] + clipped.append(blocks) + return tuple(clipped) def set_xfer_handshake_metadata( self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] @@ -380,7 +413,9 @@ def _build_save_meta( req = req_to_save assert req.kv_transfer_params is not None - clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups) + clipped_block_id_groups = self.get_exchange_clipped_blocks( + new_block_id_groups, clip_ssm=False + ) meta.add_new_req_to_save( request_id=req_id, local_block_ids=clipped_block_id_groups, @@ -418,6 +453,10 @@ def build_connector_meta( self._build_save_meta(meta, scheduler_output) meta.reqs_to_send = self._reqs_need_send + # Clock reference for reqs_to_send: deadlines above are in this + # process's perf_counter domain; workers (possibly on other nodes, + # where perf_counter has a different epoch) rebase against this. + meta.scheduler_clock = time.perf_counter() meta.reqs_in_batch = self._reqs_in_batch meta.reqs_not_processed = self._reqs_not_processed diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 834c5fccc7ac..13f45dfcc4ac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -12,6 +12,7 @@ from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor +from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -67,6 +68,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -96,7 +98,6 @@ def _compute_desc_ids( physical_blocks_per_logical: int, ) -> np.ndarray: """Compute NIXL descriptor IDs for given block IDs.""" - num_fa_regions = self.num_regions num_ssm_regions = 0 if self._has_mamba: assert self._conv_decomp is not None @@ -108,7 +109,7 @@ def _compute_desc_ids( num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - num_fa_descs = num_fa_regions * num_blocks + num_fa_descs = self.num_regions * num_blocks # All-attention fast path: single vectorized broadcast. if num_ssm_regions == 0: @@ -119,18 +120,20 @@ def _compute_desc_ids( # always differ (different areas). Therefore we can just flatten the # block_ids and compute the descs ids for all groups at once. block_arr = np.concatenate(block_ids)[None, :] - region_ids = np.arange(num_fa_regions)[:, None] + region_ids = np.arange(self.num_regions)[:, None] return (region_ids * num_blocks + block_arr).flatten() # Compute desc ids per group using the right stride: FA descs have - # num_blocks entries per region (kernel granularity), SSM descs have - # logical_blocks entries per region (no kernel splitting). - logical_blocks = num_blocks // physical_blocks_per_logical + # num_blocks entries per region (kernel granularity, expanded by + # block_size_ratio for heterogeneous block sizes), SSM descs have + # logical_blocks entries per region (no kernel splitting, and never + # ratio-expanded since state blocks are indivisible). + logical_blocks = dst_num_blocks // physical_blocks_per_logical all_descs: list[np.ndarray] = [] for i, group in enumerate(block_ids): group_arr = np.asarray(group) if _is_attention_spec(self._group_spec_types[i]): - fa_region_ids = np.arange(num_fa_regions)[:, None] + fa_region_ids = np.arange(self.num_regions)[:, None] all_descs.append( (fa_region_ids * num_blocks + group_arr[None, :]).flatten() ) @@ -162,6 +165,7 @@ def _build_local_splits_from_plan( plan: TPMapping, src_blocks_data: np.ndarray, num_fa_descs: int, + block_size_ratio: int = 1, ) -> Iterator[list[tuple[int, int, int]]]: """Build split handle data for P_TP > D_TP scenario. @@ -188,6 +192,11 @@ def _build_local_splits_from_plan( # Per-FA-descriptor replicate flag, in _build_fa_local emission order. fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + + assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), ( + "Head-sharded attention reads with P_TP > D_TP and heterogeneous " + "block sizes are not supported" + ) src_blocks_list = src_blocks_data.tolist() for p_idx, p_rank in enumerate(plan.all_source_ranks): @@ -210,7 +219,7 @@ def _build_local_splits_from_plan( def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: """Per-FA-descriptor replicate flag, in _build_fa_local emission order - (region-major; K then optional V per region). Length ``num_fa_descs``. + (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``. """ assert self.transfer_topo is not None n_regions = len(self.block_len_per_layer) @@ -443,9 +452,12 @@ def __init__( # nixl_prepped_dlist_handle. self.src_xfer_handles_by_block_size: dict[int, int] = {} + # Local descriptor arrays per remote block size (block_size_ratio>1), + # kept for building per-tp-ratio splits at the same granularity. + self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {} # Populated dynamically during handshake based on remote configuration. - # Keep track of regions at different tp_ratio values. tp_ratio->handles - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} + # Per-source split handles, keyed by (tp_ratio, remote_block_size). + self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) @@ -701,8 +713,9 @@ def _nixl_handshake( ) setup_agent_time = time.perf_counter() logger.debug( - "NIXL handshake: add agent took: %s", + "NIXL handshake: add agent took: %s (notif_agents_only=%s)", setup_agent_time - got_metadata_time, + notif_agents_only, ) remote_ranks = (remote_pp_rank, remote_rank) remote_rank_to_agent_name[remote_ranks] = remote_agent_name @@ -792,7 +805,7 @@ def _log_failure( failure_type: str, req_id: str | None, msg: str = "", - error: Exception | None = None, + error: BaseException | None = None, meta: ReqMeta | None = None, **extra_context, ): @@ -1080,23 +1093,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): caches_data = [] # With hybrid allocator, layers can share a kv cache tensor - seen_base_addresses = [] - - # Note(tms): I modified this from the original region setup code. - # K and V are now in different regions. Advantage is that we can - # elegantly support MLA and any cases where the K and V tensors - # are non-contiguous (it's not locally guaranteed that they will be) - # Disadvantage is that the encoded NixlAgentMetadata is now larger - # (roughly 8KB vs 5KB). - # Conversely for FlashInfer, K and V are registered in the same region - # to better exploit the memory layout (ie num_blocks is the first dim). + seen_base_addresses: list[int] = [] + + # K and V are packed into the content dim, so each attention layer is a + # single NIXL region whose block transfers as one unit. Mamba layers instead + # register separate conv/ssm sub-regions (see `_build_mamba_local`). tensor_size_bytes = None - for layer_name, cache_or_caches in xfer_buffers.items(): - # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to - # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. - # However, physical page_size may differ when kernel requires a specific - # block size. This leads to SSM and FA layers having different num_blocks. + for layer_name, cache in xfer_buffers.items(): + # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when + # kernel requires a specific block size. This leads to SSM and FA layers + # having different num_blocks. # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. layer_spec = self._layer_specs.get(layer_name) if layer_spec is None: @@ -1107,11 +1114,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) continue if isinstance(layer_spec, UniformTypeKVCacheSpecs): - # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs + # DSA Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs layer_spec = layer_spec.kv_cache_specs[layer_name] - cache_list = self.transfer_topo.get_transfer_cache_regions( - cache_or_caches, layer_spec - ) # `layer_spec.page_size_bytes` only accounts for logical page_size, that is # the page_size assuming constant `self._logical_num_blocks`. physical_page_size = ( @@ -1120,8 +1124,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): else layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) - # For when registering multiple tensors eg K/V in separate regions. - physical_page_size = physical_page_size // len(cache_list) if self.transfer_topo._cross_layers_blocks: # When cross-layers blocks are used, multiply by number of layers physical_page_size = physical_page_size * len( @@ -1136,66 +1138,66 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # [`num_blocks` * `page_size`] curr_tensor_size_bytes = num_blocks * physical_page_size - # TODO (NickLucche) we could eventually unify how we handle FA/FI regions, - # registering a single tensor for both K/V and splitting logically like FI. - for cache in cache_list: - base_addr = cache.data_ptr() - if base_addr in seen_base_addresses: - # NOTE (NickLucche) HMA employs memory pooling to share tensors - # across groups. This results in skipping all tensors but the ones - # pointed to by group0. Also, generally we will have more blocks - # per tensor but fewer regions. - logger.debug("Skipping %s because it's already seen", layer_name) - continue - logger.debug( - "Registering layer %s with cache shape: %s", layer_name, cache.shape + base_addr = cache.data_ptr() + is_mla_region = isinstance( + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) + ) + if base_addr in seen_base_addresses: + # NOTE (NickLucche) HMA employs memory pooling to share tensors + # across groups. This results in skipping all tensors but the ones + # pointed to by group0. Also, generally we will have more blocks + # per tensor but fewer regions. + # A shared tensor may back both SSM and attention layers (e.g. + # KDA+MLA in KimiLinear); the region's FA view is MLA whichever + # layer registered it first. + idx = seen_base_addresses.index(base_addr) + self._region_is_mla[idx] |= is_mla_region + logger.debug("Skipping %s because it's already seen", layer_name) + continue + logger.debug( + "Registering layer %s with cache shape: %s", layer_name, cache.shape + ) + seen_base_addresses.append(base_addr) + # Only record non-Mamba page sizes. + if isinstance(layer_spec, MambaSpec): + self.block_len_per_layer.append( + physical_page_size // self._physical_blocks_per_logical_kv_block ) - seen_base_addresses.append(base_addr) - # Only record non-Mamba page sizes. - if isinstance(layer_spec, MambaSpec): - self.block_len_per_layer.append( - physical_page_size // self._physical_blocks_per_logical_kv_block - ) - else: - self.block_len_per_layer.append(physical_page_size) - is_mla_region = isinstance( - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) + else: + self.block_len_per_layer.append(physical_page_size) + self._region_is_mla.append(is_mla_region) + + if not is_mla_region: + if tensor_size_bytes is None: + tensor_size_bytes = curr_tensor_size_bytes + assert tensor_size_bytes == curr_tensor_size_bytes, ( + "All non-MLA kv cache tensors must have the same size" ) - self._region_is_mla.append(is_mla_region) - - if not is_mla_region: - if tensor_size_bytes is None: - tensor_size_bytes = curr_tensor_size_bytes - assert tensor_size_bytes == curr_tensor_size_bytes, ( - "All non-MLA kv cache tensors must have the same size" - ) - - # When there's a mismatch between kbs<>bs, we rely on HMA to ensure - # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs. - if ( - self._physical_blocks_per_logical_kv_block == 1 - and cache.shape[0] != num_blocks - ): - raise AssertionError( - "All kv cache tensors must have the same number of " - f"blocks; layer={layer_name}, " - f"expected_num_blocks={num_blocks}, " - f"cache_shape={tuple(cache.shape)}, " - f"cache_stride={tuple(cache.stride())}, " - f"layer_spec={type(layer_spec).__name__}, " - f"backend={self.backend_name}, " - "all_backends=" - f"{[backend.get_name() for backend in self.attn_backends]}, " - f"kv_cache_layout={self.kv_cache_layout}" - ) - # Need to make sure the device ID is non-negative for NIXL, - # Torch uses -1 to indicate CPU tensors. - self.device_id = max(cache.get_device(), 0) - caches_data.append( - (base_addr, curr_tensor_size_bytes, self.device_id, "") + # When there's a mismatch between kbs<>bs, we rely on HMA to ensure + # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs. + if ( + self._physical_blocks_per_logical_kv_block == 1 + and cache.shape[0] != num_blocks + ): + raise AssertionError( + "All kv cache tensors must have the same number of " + f"blocks; layer={layer_name}, " + f"expected_num_blocks={num_blocks}, " + f"cache_shape={tuple(cache.shape)}, " + f"cache_stride={tuple(cache.stride())}, " + f"layer_spec={type(layer_spec).__name__}, " + f"backend={self.backend_name}, " + "all_backends=" + f"{[backend.get_name() for backend in self.attn_backends]}, " + f"kv_cache_layout={self.kv_cache_layout}" ) + # Need to make sure the device ID is non-negative for NIXL, + # Torch uses -1 to indicate CPU tensors. + self.device_id = max(cache.get_device(), 0) + caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, "")) + logger.debug( "Different block lengths collected: %s", set(self.block_len_per_layer) ) @@ -1274,22 +1276,44 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): agent_metadata_bytes=encoder.encode(agent_metadata), ) - def _build_mamba_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> np.ndarray: + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for - local mamba blocks with DS conv layout, as an Nx3 uint64 array.""" - assert block_size_ratio == 1, ( - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " - f"Got block_size_ratio={block_size_ratio}." - ) + local mamba blocks with DS conv layout, as an Nx3 uint64 array. + + A Mamba block interleaves conv and SSM state, which crucially differ in + size, so the two are indexed as separate sub-regions. Attention blocks + instead pack K and V into the content dim and transfer as a single unit. + Reference diagram: + KVCacheTensor (Shared) + / \\ + / \\ + / \\ + Attention (FlashInfer) View Mamba View + | | + | | + +-------------------+ +-------------------+ + | KVCacheTensor | | KVCacheTensor | + | | | | + |<----- page ------>| |<----- page ------->| + | size | | size | + | Key 0 | Val 0 | |Conv 0 | SSM 0 | + | Key 1 | Val 1 | |Conv 1 | SSM 1 | + | ... | ... | | ... | ... | + | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | + | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | + +-------------------+ +--------------------+ + |1st_split-2nd_split| |1st_split-2nd_split | + + Mamba state blocks are indivisible (not token-extent data), so the + descriptors always use the local page geometry regardless of any + attention block-size ratio; their desc ids are likewise never + ratio-expanded (see _compute_desc_ids). + """ assert base_addresses, "Local KV cache base addresses must not be empty." assert self._conv_decomp is not None conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size - num_blocks = self._logical_num_blocks * block_size_ratio + num_blocks = self._logical_num_blocks physical_per_logical = self._physical_blocks_per_logical_kv_block device_id = self.device_id block_arange = np.arange(num_blocks, dtype=np.uint64) @@ -1298,9 +1322,7 @@ def _build_mamba_local( for i, base_addr in enumerate(base_addresses): # Jump one page_size, but ssm page_size may be bigger when kernel # locks block size to a specific value (physical_per_logical scale). - page_stride = ( - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical - ) + page_stride = self.block_len_per_layer[i] * physical_per_logical blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) @@ -1372,15 +1394,11 @@ def _build_fa_local( block_arange = np.arange(num_blocks, dtype=np.uint64) parts: list[np.ndarray] = [] for i, base_addr in enumerate(base_addresses): - kv_block_len = ( - self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) - // block_size_ratio - ) - page_stride = self.block_len_per_layer[i] // block_size_ratio - addrs = base_addr + block_arange * page_stride - parts.append(self._stack_descs(addrs, kv_block_len, device_id)) + # K/V are packed into the content dim, so the whole block transfers + # as one unit: desc length equals the block stride. + block_len = self.block_len_per_layer[i] // block_size_ratio + addrs = base_addr + block_arange * block_len + parts.append(self._stack_descs(addrs, block_len, device_id)) return np.concatenate(parts) def _build_fa_remote( @@ -1407,9 +1425,7 @@ def _build_fa_remote( for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): replicated = self._is_region_replicated(i) # Read our whole local region size from remote.. - local_block_len = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) + local_block_len = self.block_len_per_layer[i] remote_kv_block_len = local_block_len // block_size_ratio if block_size_ratio > 1: # ..using remote kv_block_len as transfer unit @@ -1456,7 +1472,7 @@ def register_local_xfer_handler( self.device_id, ) if self._has_mamba: - assert self.num_descs == len(blocks_data) + assert self.num_descs * block_size_ratio == len(blocks_data) # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split # is unnecessary — a single conv desc per block suffices. Consider # adding a fast path that falls back to the standard 2-region @@ -1464,7 +1480,7 @@ def register_local_xfer_handler( # remote has been seen. Currently we always register 4 regions # because local descs are created before knowing the remote TP. logger.debug("Registering local Mamba descriptors (4 regions/layer)") - mamba = self._build_mamba_local(local_base_addresses, block_size_ratio) + mamba = self._build_mamba_local(local_base_addresses) blocks_data = np.concatenate([blocks_data, mamba]) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1531,8 +1547,7 @@ def add_remote_agent( ) return self._remote_agents[engine_id][(0, remote_tp_rank)] - # Compare physical regions, not self.num_regions (doubled by - # FlashInfer's virtual K/V split). + # Number of physical regions registered locally (one per layer/tensor). num_local_regions = len(self.block_len_per_layer) if ( self.pp_size > 1 @@ -1605,27 +1620,44 @@ def add_remote_agent( plan = self.tp_mappings[engine_id] + ### (Optional) Register a local handler at the remote engine's block + ### granularity (remote/prefill blocks smaller than local). + remote_block_size = nixl_agent_meta.block_size + src_blocks_data = self.src_blocks_data + if block_size_ratio > 1: + if remote_block_size not in self.src_xfer_handles_by_block_size: + handle, blocks_data = self.register_local_xfer_handler( + remote_block_size + ) + self.src_xfer_handles_by_block_size[remote_block_size] = handle + self.src_blocks_data_by_block_size[remote_block_size] = blocks_data + src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size] + ### (Optional) Register local agent memory regions. MLA is not split. + split_key = (tp_ratio, remote_block_size) if ( tp_ratio < 0 - and not self.use_mla - and tp_ratio not in self.src_xfer_handles_by_tp_ratio + and (not self.use_mla or len(plan.all_source_ranks) > 1) + and split_key not in self.src_xfer_handles_by_tp_ratio ): # Remote tp_size > local tp_size: read from multiple remote ranks. - # Logically "split" own regions into |tp_ratio| chunks. Mind that - # we only do this once per remote tp_size (replica-friendly). - self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] + # Logically "split" own regions into per-source chunks. Hybrid + # MLA+SSM also needs this path: MLA is replicated and read once, + # while the SSM state is sharded across every remote TP rank. + # We only do this once per remote (tp_size, block_size). + self.src_xfer_handles_by_tp_ratio[split_key] = [] for handle_data in self._build_local_splits_from_plan( plan, - self.src_blocks_data, - self.num_descs, + src_blocks_data, + self.num_descs * block_size_ratio, + block_size_ratio, ): descs = self.nixl_wrapper.get_xfer_descs( handle_data, self.nixl_memory_type ) handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + self.src_xfer_handles_by_tp_ratio[split_key].append(handle) ### Register remote agent memory regions # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With @@ -1661,13 +1693,6 @@ def add_remote_agent( self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) ) - if block_size_ratio > 1: - # when prefill with smaller block_size, we need to init a - # new handler with same block_len to match - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] - ) - return remote_agent_name def _validate_remote_agent_handshake( @@ -1712,9 +1737,13 @@ def _validate_remote_agent_handshake( "Disable prefix caching with --no-enable-prefix-caching." ) - if self._is_hma_required: - assert block_size_ratio == 1, ( - "HMA does not support different remote block size yet" + if block_size_ratio != 1: + # Heterogeneous block sizes transfer at remote-block granularity; + # the untransferred tail of the last local attention block is + # zeroed in the receive post-process, and mamba state pages + # transfer 1:1 (never sub-split). + assert not self.use_host_buffer, ( + "Heterogeneous block sizes are not supported with host buffer" ) kv_cache_layout = ( self.kv_cache_layout @@ -1777,7 +1806,22 @@ def _validate_remote_agent_handshake( # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. - if not self._has_mamba: + if self._has_mamba and self.use_mla: + # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are + # kernel-granularity views of the mamba-unified page. The MLA + # per-token page is TP-independent, so the block lengths must + # match up to the kernel block size ratio even under + # heterogeneous TP (remote kernel blocks may be smaller). + # SSM geometry is validated via ssm_sizes/conv offsets instead. + assert self.block_len_per_layer == [ + block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens + ], ( + "Hybrid MLA kernel-granularity block lengths must match " + f"between P and D (block_size_ratio={block_size_ratio}): " + f"local={self.block_len_per_layer}, " + f"remote={nixl_agent_meta.block_lens}." + ) + elif not self._has_mamba: assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) @@ -1871,27 +1915,57 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): "d2h", ) + @cached_property + def _attention_kv_caches(self) -> list[torch.Tensor]: + """Device KV caches of attention layers (mamba states excluded), + as consumed by the receive post-process.""" + assert self.device_kv_caches, ( + "_attention_kv_caches accessed before register_kv_caches" + ) + mamba_layers = { + name + for g, group in enumerate(self.kv_cache_config.kv_cache_groups) + if _is_ssm_spec(self._group_spec_types[g]) + for name in group.layer_names + } + kv_caches = self.device_kv_caches + return [cache for name, cache in kv_caches.items() if name not in mamba_layers] + def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[list[int]], + block_ids_list: list[tuple[list[int], int]], + convert: bool = True, ): """ Post process device kv cache after receiving from remote. - 3 types of post processing supported: + 3 types of conversion supported (``convert``): * kv_cache_postprocess_layout => convert from HND to NHD * kv_cache_postprocess_blksize => convert from small block size to large block size * kv_cache_postprocess_blksize_and_layout => convert from small block size to large block size and convert from HND to NHD + The transfer only covers ``covered_sub_blocks`` remote-sized + sub-blocks of each request's local attention blocks; the rest was + clipped, either by remote-block pairing (block-size ratio) or by the + hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were + excluded from the scheduler's alloc-time KV zeroing (which would race + the RDMA write), so everything past the covered range is zeroed here. + Stale bytes would otherwise surface as garbage or NaNs once decode + grows into the untransferred tail. """ if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." assert self.transfer_topo is not None - if self.enable_permute_local_kv and block_size_ratio > 1: + if not convert: + logger.debug( + "Post-processing device kv cache on receive by zeroing " + "untransferred blocks." + ) + elif self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " "block_size with %sx bigger and permuting layout from HND" @@ -1910,18 +1984,45 @@ def post_process_device_kv_on_receive( block_size_ratio, ) - for block_ids in block_ids_list: - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + attn_caches = self._attention_kv_caches + device = attn_caches[0].device + for block_ids, covered_sub_blocks in block_ids_list: + # Blocks the transfer didn't write: the token tail of the last + # partially covered block, then everything beyond it. + covered_blocks, sub_blocks_in_last = divmod( + covered_sub_blocks, block_size_ratio + ) + first_stale = covered_blocks + (1 if sub_blocks_in_last else 0) + has_stale = first_stale < len(block_ids) + indices = None + if convert or has_stale: + indices = async_tensor_h2d(block_ids, device, torch.long) + + if convert: + for cache in attn_caches: + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive( + cache, indices, block_size_ratio + ) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive( + cache, indices, block_size_ratio + ) - for cache in self.device_kv_caches.values(): - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive( - cache, indices, block_size_ratio - ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) + if sub_blocks_in_last: + last_block_id = block_ids[covered_blocks] + for cache in attn_caches: + # Both post-processed layouts leave tokens on dim 1. + sub_block_tokens = cache.shape[1] // block_size_ratio + zero_from = sub_blocks_in_last * sub_block_tokens + cache[last_block_id, zero_from:].zero_() + if has_stale: + assert indices is not None + stale_ids = indices[first_stale:] + for cache in attn_caches: + cache.index_fill_(0, stale_ids, 0) def post_process_device_kv_on_receive_heterogeneous_attn( self, block_ids: list[int] @@ -1935,13 +2036,8 @@ def post_process_device_kv_on_receive_heterogeneous_attn( indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) for _, cache_or_caches in self.device_kv_caches.items(): - blocks_to_update = cache_or_caches.index_select(1, indices) current_platform.pack_kv_cache( - key=blocks_to_update[0], - value=blocks_to_update[1], - key_cache=cache_or_caches[0], - value_cache=cache_or_caches[1], - block_ids=block_ids, + kv_cache=cache_or_caches, indices=indices, ) @@ -1996,18 +2092,33 @@ def get_finished(self) -> tuple[set[str], set[str]]: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) - # post processing for heteroblocksize + # Post processing for heteroblocksize/layout, and for blocks the + # transfer clipped. The latter happens either at remote-block + # granularity (block_size_ratio > 1) or at kernel-block + # granularity, when equal kernel pages meet differing logical + # block sizes and _apply_prefix_caching front-trims to the + # minimum count (hybrid heterogeneous TP). remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) block_size_ratio = self.transfer_topo.block_size_ratio( remote_info.remote_block_size ) - if not self.use_mla and ( - block_size_ratio > 1 or self.enable_permute_local_kv - ): - assert not self._is_hma_required - block_ids_for_blocksize_post_process[block_size_ratio].append( - meta.local_physical_block_ids[0] - ) + hetero_ppl = ( + remote_info.remote_physical_blocks_per_logical + != self._physical_blocks_per_logical_kv_block + ) + if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl: + for g, local_group in enumerate(meta.local_physical_block_ids): + if not local_group or _is_ssm_spec(self._group_spec_types[g]): + continue + # Number of remote-sized sub-blocks the transfer covered; + # everything past this was clipped and must be zeroed. + covered_sub_blocks = min( + len(local_group) * block_size_ratio, + len(meta.remote.block_ids[g]), + ) + block_ids_for_blocksize_post_process[block_size_ratio].append( + (local_group, covered_sub_blocks) + ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: block_ids_for_heterogeneous_attn_post_process.append( @@ -2017,7 +2128,14 @@ def get_finished(self) -> tuple[set[str], set[str]]: block_size_ratio, block_ids_list, ) in block_ids_for_blocksize_post_process.items(): - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + # MLA never needs the block-size/layout conversion, but its + # clipped blocks still need zeroing. + convert = not self.use_mla and ( + block_size_ratio > 1 or self.enable_permute_local_kv + ) + self.post_process_device_kv_on_receive( + block_size_ratio, block_ids_list, convert + ) for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) @@ -2207,6 +2325,45 @@ def get_mapped_blocks( return mapped_2d.flatten().astype(np.int64) + def _map_block_ids_for_block_size_ratio( + self, + local_block_ids: BlockIds, + remote_block_ids: BlockIds, + block_size_ratio: int, + ) -> tuple[BlockIds, BlockIds]: + """Map attention-group block ids to remote-block granularity. + + Each local attention block is split into ``block_size_ratio`` + sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the + remote list — the untransferred tail of the last local block — are + clipped here and zeroed in the receive post-process. Mamba state + blocks are indivisible and transfer 1:1, unexpanded. + + ex: remote (prefill) block ids with block_size 4: + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + Local (decode) block ids with block_size 16: [1, 2, 3] expand to + [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote. + """ + mapped_local: list[list[int]] = [] + mapped_remote: list[list[int]] = [] + for i, remote_group in enumerate(remote_block_ids): + local_group = local_block_ids[i] if local_block_ids else [] + if _is_ssm_spec(self._group_spec_types[i]): + mapped_local.append(list(local_group)) + mapped_remote.append(list(remote_group)) + continue + mapped = self.get_mapped_blocks( + np.asarray(local_group), block_size_ratio + ).tolist() + if len(mapped) > len(remote_group): + mapped = mapped[: len(remote_group)] + mapped_local.append(mapped) + mapped_remote.append(list(remote_group)) + if not any(mapped_local): + # Full prefix cache hit is indicated with an empty list. + return [], mapped_remote + return mapped_local, mapped_remote + def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: """ Convert block ids to kernel physical block ids. @@ -2282,16 +2439,25 @@ def _apply_prefix_caching( for i, remote_group in enumerate(remote_block_ids): num_local_blocks = len(local_block_ids[i]) num_remote_blocks = len(remote_group) - if ( - _is_ssm_spec(self._group_spec_types[i]) - and num_local_blocks < num_remote_blocks - ): - # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks - # prior to the last one are placeholders (null blocks). Mind that - # this doesn't really impact transfer, as we only still care about - # the last "block", the full in-place state. - assert num_local_blocks == 1, "SSM can only have one local block" - remote_block_ids[i] = remote_group[-num_local_blocks:] + if _is_ssm_spec(self._group_spec_types[i]): + if num_local_blocks == num_remote_blocks: + continue + # Only state-bearing slots reach here, single-state modes + # just one (see get_exchange_clipped_blocks), so differing + # counts mean position-indexed "all"-mode lists. A longer + # remote list carries earlier positions the local side + # already has (prefix hit) -> read its tail; a longer local + # list holds the position D recomputes itself, which gets + # no remote state. + assert num_local_blocks - num_remote_blocks <= 1, ( + f"Group {i}: unpairable SSM state slots, " + f"local={num_local_blocks} remote={num_remote_blocks}" + ) + num_blocks = min(num_local_blocks, num_remote_blocks) + if num_local_blocks < num_remote_blocks: + remote_block_ids[i] = remote_group[-num_blocks:] + else: + local_block_ids[i] = local_block_ids[i][:num_blocks] elif ( self._physical_blocks_per_logical_kv_block == remote_physical_per_logical @@ -2301,59 +2467,24 @@ def _apply_prefix_caching( remote_block_ids[i] = remote_group[-num_local_blocks:] else: # TODO Handle prefix caching with different block_sizes - max_padding = max( - self._physical_blocks_per_logical_kv_block, - remote_physical_per_logical, + # Allocation rounding legitimately leaves up to + # ppl - 1 trailing dead kernel blocks per side (plus one + # extra local block for the recomputed final token), so + # the counts may differ by up to the sum of the two + # ratios; anything larger indicates mismatched lists. + max_padding = ( + self._physical_blocks_per_logical_kv_block + + remote_physical_per_logical ) - assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( + assert abs(num_local_blocks - num_remote_blocks) <= max_padding, ( f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| >= {max_padding}" + f"{num_remote_blocks}| > {max_padding}" ) num_blocks = min(num_local_blocks, num_remote_blocks) local_block_ids[i] = local_block_ids[i][:num_blocks] remote_block_ids[i] = remote_group[:num_blocks] return local_block_ids, remote_block_ids - def get_backend_aware_kv_block_len( - self, layer_idx: int, first_split: bool = True, mamba_view: bool = False - ) -> int: - """ - Get the block length for one K/V element (K and V have the same size). - - For FA and other backends, this is equal to the length of the whole - block, as K and V are in separate regions. - For FlashInfer, this is half the length of the whole block, as K and V - share the same region. - Similarly, for SSM-based models, state and conv are interleaved, but crucially - the their size differs. - Reference diagram: - KVCacheTensor (Shared) - / \\ - / \\ - / \\ - Attention (FlashInfer) View Mamba View - | | - | | - +-------------------+ +-------------------+ - | KVCacheTensor | | KVCacheTensor | - | | | | - |<----- page ------>| |<----- page ------->| - | size | | size | - | Key 0 | Val 0 | |Conv 0 | SSM 0 | - | Key 1 | Val 1 | |Conv 1 | SSM 1 | - | ... | ... | | ... | ... | - | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | - | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | - +-------------------+ +--------------------+ - |1st_split-2nd_split| |1st_split-2nd_split | - """ - assert self.transfer_topo is not None - if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view: - block_len = self._mamba_ssm_size[not first_split] - else: - block_len = self.block_len_per_layer[layer_idx] - return block_len - def get_kv_connector_stats(self) -> KVConnectorStats | None: """ Get the KV transfer stats for the connector. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index f6f3de0a1f29..27eab825a103 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -180,6 +180,12 @@ def __init__(self): self.reqs_to_recv: dict[ReqId, ReqMeta] = {} self.reqs_to_save: dict[ReqId, ReqMeta] = {} self.reqs_to_send: dict[ReqId, float] = {} + # The scheduler process's time.perf_counter() when this metadata was + # built. reqs_to_send deadlines are stamped with the scheduler's + # clock, which is NOT comparable across processes (perf_counter is + # process/boot-local): workers must rebase the remaining TTL onto + # their own clock via this reference. 0.0 = unset (legacy metadata). + self.scheduler_clock: float = 0.0 self.reqs_in_batch: set[ReqId] = set() self.reqs_not_processed: set[ReqId] = set() # Heartbeat data grouped by remote engine, sent by D worker to P. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 51a54d06f078..9beceeabada6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -155,7 +155,7 @@ def update_state_after_alloc( if num_external_tokens > 0 else () ) - local_block_ids = self.get_sw_clipped_blocks( + local_block_ids = self.get_exchange_clipped_blocks( unhashed_local_block_ids ) @@ -262,7 +262,7 @@ def request_finished( # trimming down after allocating for the whole sequence length. Empty # blocks are always at the start of the list. # Here we "unpad" blocks to send the actual remote blocks to be read. - block_ids = self.get_sw_clipped_blocks(block_ids) + block_ids = self.get_exchange_clipped_blocks(block_ids) remote_num_tokens = request.num_computed_tokens diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 63969382dff8..cc32845a4a03 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -5,8 +5,6 @@ import time from typing import TYPE_CHECKING -import numpy as np - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -94,8 +92,20 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): assert req_id not in self._reqs_to_send # Add to requests that are waiting to be read and track expiration. + # Deadlines are stamped with the scheduler process's perf_counter, + # which is not comparable to ours when the worker runs in another + # process on another node (perf_counter epochs differ by boot time). + # Rebase the remaining TTL onto our clock; broadcast latency only + # lengthens the lease, which is the safe direction. A cross-node + # epoch gap larger than the TTL otherwise expires the lease on + # arrival and the blocks are freed before D reads them. + now_local = time.perf_counter() for req_id, expiration_time in metadata.reqs_to_send.items(): if req_id in self._reqs_to_process: + if metadata.scheduler_clock: + expiration_time = now_local + ( + expiration_time - metadata.scheduler_clock + ) self._reqs_to_send[req_id] = expiration_time # Send heartbeats to P-side engines to keep KV blocks alive while @@ -161,9 +171,9 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): ] # D may have to perform multiple reads from different remote ranks. - # MLA opt: when P TP > D TP, only a single read is executed for - # the first remote rank (cache is duplicated).. - if self.use_mla and tp_ratio < 0: + # Pure MLA reads once because its cache is replicated. Hybrid + # MLA+SSM still needs one read per SSM source rank. + if self.use_mla and tp_ratio < 0 and not self._has_mamba: assert len(read_specs) == 1 for i, spec in enumerate(read_specs): @@ -177,11 +187,11 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): req_id, ) # Get side handles. - if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size + if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1): # Remote tp_size > local tp_size: we must perform multiple # reads. Get the memory chunk onto which we will write to. - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: # Single read from remote, we write to the whole memory region. # Also handle remote block size different from local block size. @@ -203,7 +213,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_xfer_side_handle=remote_xfer_side_handle, ) - if self.use_mla and tp_ratio < 0 and read_specs: + if self.use_mla and tp_ratio < 0 and len(read_specs) == 1: # ..but we still need to notify the other remote ranks that we # have the blocks we need so they can update the request state. notif_id = f"{meta.remote.request_id}:{self.world_size}".encode() @@ -235,30 +245,11 @@ def _read_blocks( remote_info.remote_block_size ) if block_size_ratio > 1: - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - # NOTE: - # get_mapped_blocks will always expand block_ids for n times. - # ex: - # prefill block_ids with block_size as 4: - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - # Local decode block_ids with block_size as 16: [1, 2, 3] - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # Then we clip local to align with prefill - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) # NOTE(rob): having the staging blocks be on the READER side is # not going to work well (since we will have to call rearrange tensors). # after we detect the txn is complete (which means we cannot make the diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index 413ebdcdb09d..a02491e4872e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -170,7 +170,7 @@ def update_state_after_alloc( request.request_id, ) local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups() - local_block_ids = self.get_sw_clipped_blocks(local_block_ids) + local_block_ids = self.get_exchange_clipped_blocks(local_block_ids) # ``remote_*`` fields are P's coordinates (from D's perspective). # ``decode_*`` fields are D's own info that P needs for the @@ -271,7 +271,7 @@ def request_finished( time.perf_counter() + self._kv_lease_duration ) - block_ids = self.get_sw_clipped_blocks(block_ids) + block_ids = self.get_exchange_clipped_blocks(block_ids) remote_num_tokens = request.num_computed_tokens # Store finished blocks for worker-level matching with D diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 859dbae00b6f..13af0314f0e2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -8,10 +8,13 @@ ``send_notif``, matches D registrations with P finished blocks, and issues WRITE transfers via ``make_prepped_xfer`` / ``transfer``. -The engine main thread feeds the writer through three queues: +The engine main thread feeds the writer through queues: ``_reg_send_inbox`` (D-side regs to send), ``_finished_blocks_inbox`` (P-side blocks from metadata) and ``_pending_completion_notifs`` (non-PUSH_REG notifs forwarded back for HB / completion accounting). +The handshake-completion callback feeds ``_deferred_push_inbox`` with +matched pushes whose P→D handshake has finished so the writer can +(re-)issue the WRITE without ever blocking on the network. Wake model: the writer self-polls every ``_PUSH_WRITER_POLL_INTERVAL_MS`` only while it has unmatched @@ -21,8 +24,8 @@ ``start_load_kv`` (when handing it new work) and from ``get_finished`` (so each engine step gives the writer a chance to drain NIXL notifs); the handshake-completion callback sets the same event after a deferred -PUSH_REG send has been queued. When a request's lease expires (the base -worker reports it via ``done_sending``) or the WRITE completes, +PUSH_REG send or a deferred push WRITE has been queued. When a request's +lease expires (the base worker reports it via ``done_sending``) or the WRITE completes, ``get_finished`` enqueues an eviction onto ``_evict_finished_inbox`` so the writer drops any leftover ``_push_finished_blocks`` / ``_pending_d_registrations`` and stops self-polling. @@ -36,7 +39,6 @@ from typing import TYPE_CHECKING, Any import msgspec -import numpy as np from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( @@ -50,7 +52,10 @@ ReqMeta, TransferHandle, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + ReadSpec, + _is_attention_spec, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id from vllm.logger import init_logger @@ -109,6 +114,8 @@ def __init__( # ``_push_finished_blocks`` so an unmatched entry doesn't keep the # writer busy-polling forever. self._evict_finished_inbox: queue.Queue[str] = queue.Queue() + # Handshakes that have just completed and are ready for the WRITE on wthread + self._deferred_push_inbox = queue.Queue[tuple[str, BlockIds, dict[str, Any]]]() # Wake signal from engine main thread (start_load_kv / get_finished). # Writer self-polls at _PUSH_WRITER_POLL_INTERVAL_MS while it has @@ -184,8 +191,15 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): for req_id in metadata.reqs_not_processed: self._reqs_to_process.discard(req_id) assert req_id not in self._reqs_to_send + # Rebase scheduler-clock deadlines onto this worker's clock — see the + # equivalent block in pull_worker.start_load_kv for the rationale. + now_local = time.perf_counter() for req_id, expiration_time in metadata.reqs_to_send.items(): if req_id in self._reqs_to_process: + if metadata.scheduler_clock: + expiration_time = now_local + ( + expiration_time - metadata.scheduler_clock + ) self._reqs_to_send[req_id] = expiration_time # Heartbeats still leave from the main thread (base worker behaviour). @@ -206,7 +220,15 @@ def _push_writer_loop(self) -> None: break self._send_registration_to_p(rid, rd) - # 2. P-side finished blocks; match against pending regs. + # 2. Deferred P→D pushes whose handshake just completed; do xfer now + while True: + try: + rid, blocks, rd = self._deferred_push_inbox.get_nowait() + except queue.Empty: + break + self._do_start_push_kv(rid, blocks, rd) + + # 3. P-side finished blocks; match against pending regs. while True: try: rid, blocks = self._finished_blocks_inbox.get_nowait() @@ -218,7 +240,7 @@ def _push_writer_loop(self) -> None: else: self._push_finished_blocks[rid] = blocks - # 2b. Evict finished blocks for requests that have either + # 3b. Evict finished blocks for requests that have either # completed (WRITE acknowledged) or whose lease expired # without a D registration. Drop pending registrations # for the same reason so we don't leak state. @@ -230,7 +252,7 @@ def _push_writer_loop(self) -> None: self._push_finished_blocks.pop(rid, None) self._pending_d_registrations.pop(rid, None) - # 3. NIXL notifs: route PUSH_REG; forward the rest. + # 4. NIXL notifs: route PUSH_REG; forward the rest. for notifs in self.nixl_wrapper.get_new_notifs().values(): for notif in notifs: if notif.startswith(PUSH_REG_NOTIF_PREFIX): @@ -286,8 +308,10 @@ def _send_registration_to_p( reg_data["remote_port"], reg_data["remote_tp_size"], pp_size=remote_pp_size, - # D never addresses P memory in push mode; just load P's agents. - notif_agents_only=remote_pp_size > 1, + # D only ever sends PUSH_REG notifs to P and never reads or writes + # P's memory in push mode, so it never needs the transfer + # descriptors set up by the full add_remote_agent path. + notif_agents_only=True, ) if fut is None: self._do_send_reg_notif(req_id, reg_data) @@ -384,34 +408,53 @@ def _do_start_push_kv( ) -> None: """Start push-based KV transfer from P worker to D node. - ``local_block_ids`` are P's *logical* block IDs (from the P - scheduler's metadata). ``registration_data["local_block_ids"]`` - are D's *logical* block IDs (from D's scheduler, sent over the - PUSH_REG notif). All conversion to physical block IDs is - deferred to ``_xfer_blocks_for_req`` so each side uses its own - physical-blocks-per-logical ratio (P uses - ``self._physical_blocks_per_logical_kv_block``; D's ratio is - learned during the NIXL handshake).""" - decode_engine_id = registration_data["decode_engine_id"] - remote_block_ids = registration_data["local_block_ids"] - decode_host = registration_data["decode_host"] - decode_port = registration_data["decode_port"] - decode_request_id = registration_data["request_id"] + The P→D handshake runs on the base worker's background executor. + If it isn't ready yet we register a completion callback, defer the + WRITE, and re-drive this request via ``_deferred_push_inbox`` once + the handshake resolves -- so the writer thread never blocks on the + network (mirrors ``_send_registration_to_p``). + """ if not local_block_ids: logger.warning("No local blocks to push for request %s", request_id) return - if not self._ensure_d_handshake( + # ``local_block_ids`` are P's logical block IDs; ``remote_block_ids`` + # (D's, from the PUSH_REG notif) are also logical. + decode_engine_id = registration_data["decode_engine_id"] + remote_block_ids = registration_data["local_block_ids"] + decode_request_id = registration_data["request_id"] + + # Runs on the background executor; defer the WRITE until it's ready. + fut = self._ensure_handshake( decode_engine_id, - decode_host, - decode_port, + registration_data["decode_host"], + registration_data["decode_port"], registration_data["decode_tp_size"], - request_id, - ): + ) + if fut is not None: + + def _on_handshake( + f: Future[tuple[dict[tuple[int, int], str], float]], + rid: str = request_id, + blocks: BlockIds = local_block_ids, + rd: dict[str, Any] = registration_data, + ) -> None: + if (e := f.exception()) is not None: + # The engine reclaims the blocks via the TTL so we dont free here + self._log_failure( + failure_type="push_handshake_failed", req_id=rid, error=e + ) + return + self._deferred_push_inbox.put((rid, blocks, rd)) + self._push_writer_wake.set() + + fut.add_done_callback(_on_handshake) return - # Both sides are kept in logical form here; ``_xfer_blocks_for_req`` - # expands each side using the appropriate ratio. + # Both sides stay logical here; ``_xfer_blocks_for_req`` converts each + # to physical with its own physical-blocks-per-logical ratio -- P uses + # ``self._physical_blocks_per_logical_kv_block``, D's is learned during + # the NIXL handshake. logical_local = self._as_grouped_block_ids(local_block_ids) logical_remote = self._as_grouped_block_ids(remote_block_ids) physical_local = self._logical_to_kernel_block_ids( @@ -441,45 +484,6 @@ def _do_start_push_kv( elapsed_ms, ) - def _ensure_d_handshake( - self, - decode_engine_id: str, - decode_host: str, - decode_port: int, - decode_tp_size: int, - request_id: str, - ) -> bool: - """First-time P→D handshake. Blocking call on the writer thread. - - Returns True iff the handshake succeeded (or had already been - completed). Returns False if the handshake raised; the request is - skipped in that case (the engine layer will reschedule or fail it - via the standard lease/timeout path).""" - if decode_engine_id in self._remote_agents: - return True - try: - remote_agents, _ = self._nixl_handshake( - decode_host, - decode_port, - decode_tp_size, - decode_engine_id, - ) - except Exception: - logger.exception( - "Failed handshake to D %s for push %s", - decode_engine_id, - request_id, - ) - return False - with self._handshake_lock: - self._remote_agents[decode_engine_id] = remote_agents - logger.info( - "Push handshake to D %s done (%d agents)", - decode_engine_id, - len(remote_agents), - ) - return True - @staticmethod def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds: """Normalise a sequence of block IDs to a tuple-of-groups shape. @@ -511,41 +515,37 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): local_block_ids = meta.local_physical_block_ids num_groups = len(local_block_ids) - if self.use_mla and tp_ratio < 0: - # MLA latent is replicated across D's TP ranks: the tp-mapping - # collapses to one rank (fine for reads), but push must WRITE every - # D rank or the rest decode stale KV; only the dst differs per rank. + # MLA latent is replicated across D's TP ranks: the tp-mapping + # collapses it to one rank (fine for reads), but push must WRITE every + # D rank or the rest decode stale KV. For hybrid MLA+SSM the sharded + # SSM state already targets every covered D rank, so only the + # attention groups need widening; pure MLA writes to all handshaked + # ranks (only the dst differs per rank). + replicate_attn = self.use_mla and tp_ratio < 0 + if replicate_attn and not self._has_mamba: assert len(plan.all_source_ranks) == 1 - mla_local_ids = [list(ids) for ids in local_block_ids] - mla_remote_ids = [list(ids) for ids in remote_block_ids] - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=mla_local_ids, - remote_block_ids=mla_remote_ids, - ) - for rank in self.dst_xfer_side_handles[engine_id] - ] + write_ranks = sorted(self.dst_xfer_side_handles[engine_id]) else: - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=[ - list(local_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - remote_block_ids=[ - list(remote_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - ) - for rank in plan.all_source_ranks + write_ranks = list(plan.all_source_ranks) + + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds: + return [ + list(block_ids[g]) + if (replicate_attn and _is_attention_spec(self._group_spec_types[g])) + or rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) ] + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=group_ids(local_block_ids, rank), + remote_block_ids=group_ids(remote_block_ids, rank), + ) + for rank in write_ranks + ] + handles: list[int] = [] for i, spec in enumerate(read_specs): remote_block_size = remote_info.remote_block_size @@ -557,9 +557,12 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_block_size, req_id, ) - if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + if tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1): + # Multiple targets: write each rank its chunk of local memory. + # Hybrid MLA+SSM also lands here: its split handles replicate + # the attention descriptors and chunk only the SSM state. + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: local_xfer_side_handle = self.src_xfer_handles_by_block_size[ remote_block_size @@ -611,18 +614,11 @@ def _xfer_blocks( remote_info.remote_block_size ) if block_size_ratio > 1: - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) notif_id = f"{remote_request_id}:{self.world_size}".encode() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py new file mode 100644 index 000000000000..581f7fdbc8e3 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py @@ -0,0 +1,444 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Derivation of canonical page mappings for KV offloading. + +The only place in the offloading stack that reasons about parallelism +(TP/DCP/PCP); everything downstream consumes byte mappings. The canonical +page of a layer is the full offloaded block without parallelism: all KV +heads, all block_size * dcp * pcp tokens, in the worker's page encoding. +Uncertifiable layers get an opaque fallback mapping (fail closed). +""" + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +import torch + +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheConfig, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun + +if TYPE_CHECKING: + from vllm.config import VllmConfig + + +@dataclass(frozen=True) +class _RankContext: + """Sharding parameters of one worker rank within the offload group.""" + + tp_size: int + dcp_size: int + pcp_size: int + interleave: int + total_kv_heads: int + rank: int + + @property + def cp_size(self) -> int: + return self.dcp_size * self.pcp_size + + @property + def tp_rank(self) -> int: + return self.rank % self.tp_size + + @property + def total_cp_rank(self) -> int: + pcp_rank = self.rank // self.tp_size + return pcp_rank * self.dcp_size + self.tp_rank % self.dcp_size + + +@dataclass(frozen=True) +class ByteRegion: + """A byte region within a page that repeats once per token.""" + + local_offset: int + canonical_offset: int + bytes_per_token: int + canonical_token_stride: int + + +def _coalesce_runs(runs: list[CopyRun]) -> tuple[CopyRun, ...]: + """Collapse contiguous fragments within and across runs to minimize the + number of copy ops (e.g. a single-rank mapping becomes one whole-page run). + """ + out: list[CopyRun] = [] + for run in runs: + if ( + run.num_fragments > 1 + and run.local_stride == run.fragment_size + and run.canonical_stride == run.fragment_size + ): + size = run.fragment_size * run.num_fragments + run = CopyRun(run.local_offset, run.canonical_offset, size, 1, size, size) + prev = out[-1] if out else None + if ( + prev is not None + and prev.num_fragments == 1 + and run.num_fragments == 1 + and prev.local_offset + prev.fragment_size == run.local_offset + and prev.canonical_offset + prev.fragment_size == run.canonical_offset + ): + size = prev.fragment_size + run.fragment_size + out[-1] = CopyRun( + prev.local_offset, prev.canonical_offset, size, 1, size, size + ) + else: + out.append(run) + return tuple(out) + + +def _local_to_canonical_token(local_idx: int, ctx: _RankContext) -> int: + """Canonical position of one of this rank's local token indices.""" + chunk, pos_in_chunk = divmod(local_idx, ctx.interleave) + return (chunk * ctx.cp_size + ctx.total_cp_rank) * ctx.interleave + pos_in_chunk + + +def _interleave_cp_tokens( + regions: list[ByteRegion], + num_tokens: int, + ctx: _RankContext, +) -> tuple[CopyRun, ...]: + """Place each region's num_tokens rows at their canonical token positions, + one run per chunk of interleaved tokens.""" + runs: list[CopyRun] = [] + for region in regions: + if ctx.cp_size == 1: + runs.append( + CopyRun( + region.local_offset, + region.canonical_offset, + region.bytes_per_token, + num_tokens, + region.bytes_per_token, + region.canonical_token_stride, + ) + ) + continue + for chunk_start in range(0, num_tokens, ctx.interleave): + canonical_token = _local_to_canonical_token(chunk_start, ctx) + runs.append( + CopyRun( + region.local_offset + chunk_start * region.bytes_per_token, + region.canonical_offset + + canonical_token * region.canonical_token_stride, + region.bytes_per_token, + ctx.interleave, + region.bytes_per_token, + region.canonical_token_stride, + ) + ) + return _coalesce_runs(runs) + + +def _packed_kv_regions( + kv_cache: torch.Tensor, + spec: AttentionSpec, + head_shard: int, + num_head_shards: int, + cp_size: int, +) -> list[ByteRegion] | None: + """K and V adjacent per (token, head), in NHD or HND stride order.""" + bs, heads = spec.block_size, spec.num_kv_heads + elem = kv_cache.element_size() + head_elems = 2 * spec.head_size + if heads * bs * head_elems * elem != spec.real_page_size_bytes: + return None + _, head_stride, token_stride, inner_stride = kv_cache.stride() + if inner_stride != 1: + return None + head_bytes = head_elems * elem + token_row_bytes = heads * head_bytes + + if head_stride == head_elems and token_stride == heads * head_elems: # NHD + return [ + ByteRegion( + local_offset=0, + canonical_offset=head_shard * token_row_bytes, + bytes_per_token=token_row_bytes, + canonical_token_stride=num_head_shards * token_row_bytes, + ) + ] + + if head_stride == bs * head_elems and token_stride == head_elems: # HND + canonical_span = bs * cp_size # canonical tokens per offloaded block + return [ + ByteRegion( + local_offset=head * bs * head_bytes, + canonical_offset=(head_shard * heads + head) + * canonical_span + * head_bytes, + bytes_per_token=head_bytes, + canonical_token_stride=head_bytes, + ) + for head in range(heads) + ] + return None + + +def _split_kv_regions( + kv_cache: torch.Tensor, + spec: AttentionSpec, + head_shard: int, + num_head_shards: int, + cp_size: int, +) -> list[ByteRegion] | None: + """K and V in separate page halves, in NHD or HND stride order.""" + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size + elem = kv_cache.element_size() + if 2 * bs * heads * head_size * elem != spec.real_page_size_bytes: + return None + _, half_stride, token_stride, head_stride, inner_stride = kv_cache.stride() + if inner_stride != 1 or half_stride != bs * heads * head_size: + return None + head_bytes = head_size * elem + token_row_bytes = heads * head_bytes + canonical_span = bs * cp_size # canonical tokens per offloaded block + + if token_stride == heads * head_size and head_stride == head_size: # NHD + canonical_half_bytes = canonical_span * num_head_shards * token_row_bytes + return [ + ByteRegion( + local_offset=half * bs * token_row_bytes, + canonical_offset=half * canonical_half_bytes + + head_shard * token_row_bytes, + bytes_per_token=token_row_bytes, + canonical_token_stride=num_head_shards * token_row_bytes, + ) + for half in range(2) # K, then V + ] + + if head_stride == bs * head_size and token_stride == head_size: # HND + local_half_bytes = bs * heads * head_bytes + total_heads = num_head_shards * heads + return [ + ByteRegion( + local_offset=half * local_half_bytes + head * bs * head_bytes, + canonical_offset=(half * total_heads + head_shard * heads + head) + * canonical_span + * head_bytes, + bytes_per_token=head_bytes, + canonical_token_stride=head_bytes, + ) + for half in range(2) + for head in range(heads) + ] + return None + + +def _attention_byte_regions( + kv_cache: torch.Tensor, + spec: AttentionSpec, + num_blocks: int, + head_shard: int, + num_head_shards: int, + cp_size: int, +) -> list[ByteRegion] | None: + """Byte regions of an attention page, given this rank's head shard. + None when the physical layout is not recognized (fail closed).""" + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size + if tuple(kv_cache.shape) == (num_blocks, heads, bs, 2 * head_size): + return _packed_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size) + if tuple(kv_cache.shape) == (num_blocks, 2, bs, heads, head_size): + return _split_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size) + return None + + +def _layer_mapping( + spec: KVCacheSpec, + kv_cache: torch.Tensor | list[torch.Tensor] | None, + num_blocks: int, + ctx: _RankContext, +) -> CanonicalPageMapping | None: + """Certified mapping for one layer at one rank, or None (fail closed).""" + if not isinstance(spec, AttentionSpec): + return None + bs = spec.block_size + page = spec.real_page_size_bytes + if ctx.cp_size > 1 and (ctx.interleave > bs or bs % ctx.interleave): + return None + + if isinstance(spec, MLAAttentionSpec): + # TP-replicated latent; CP shards its tokens across the DCP groups + if ( + spec.compress_ratio != 1 + or page % bs + or ctx.tp_size % ctx.dcp_size + or spec.kv_quant_mode.is_per_token_head + ): + return None + row = page // bs + return CanonicalPageMapping( + canonical_page_size_bytes=ctx.cp_size * page, + local_page_size_bytes=page, + runs=_interleave_cp_tokens([ByteRegion(0, 0, row, row)], bs, ctx), + num_writers=ctx.tp_size // ctx.dcp_size, + writer_index=ctx.tp_rank // ctx.dcp_size, + parallelism_agnostic=ctx.cp_size == 1, + ) + + if spec.kv_quant_mode.is_per_token_head or not isinstance(kv_cache, torch.Tensor): + return None + total, tp = ctx.total_kv_heads, ctx.tp_size + if spec.num_kv_heads != max(1, total // tp): + return None + if total >= tp: + if total % tp: + return None + num_head_shards, replication = tp, 1 + else: + if tp % total: + return None + num_head_shards, replication = total, tp // total + # DCP shards tokens across ranks holding replicated KV + if replication % ctx.dcp_size: + return None + + head_shard = ctx.tp_rank // replication + regions = _attention_byte_regions( + kv_cache, spec, num_blocks, head_shard, num_head_shards, ctx.cp_size + ) + if regions is None: + return None + return CanonicalPageMapping( + canonical_page_size_bytes=ctx.cp_size * num_head_shards * page, + local_page_size_bytes=page, + runs=_interleave_cp_tokens(regions, bs, ctx), + num_writers=replication // ctx.dcp_size, + writer_index=(ctx.tp_rank % replication) // ctx.dcp_size, + parallelism_agnostic=ctx.cp_size == 1, + ) + + +def _opaque_fallback_mapping( + page_size_bytes: int, num_ranks: int, rank: int +) -> CanonicalPageMapping: + """Fallback: place the worker's page whole at a worker-exclusive offset.""" + run = CopyRun( + 0, rank * page_size_bytes, page_size_bytes, 1, page_size_bytes, page_size_bytes + ) + return CanonicalPageMapping( + canonical_page_size_bytes=num_ranks * page_size_bytes, + local_page_size_bytes=page_size_bytes, + runs=(run,), + num_writers=1, + writer_index=0, + parallelism_agnostic=False, + ) + + +def _run_intervals(runs: tuple[CopyRun, ...], canonical: bool) -> list[tuple[int, int]]: + intervals = [] + for run in runs: + offset = run.canonical_offset if canonical else run.local_offset + stride = run.canonical_stride if canonical else run.local_stride + for i in range(run.num_fragments): + start = offset + i * stride + intervals.append((start, start + run.fragment_size)) + return sorted(intervals) + + +def _is_exact_partition(intervals: list[tuple[int, int]], size: int) -> bool: + return ( + bool(intervals) + and intervals[0][0] == 0 + and intervals[-1][1] == size + and all(a[1] == b[0] for a, b in zip(intervals, intervals[1:])) + ) + + +def _verify_tiling(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None: + """Whichever ranks a block elects as writers must tile the canonical page + exactly once, and each rank's runs must cover exactly its local page.""" + size = per_rank[0].canonical_page_size_bytes + num_writers = per_rank[0].num_writers + for mapping in per_rank: + assert mapping.canonical_page_size_bytes == size + assert mapping.num_writers == num_writers + local = _run_intervals(mapping.runs, canonical=False) + assert _is_exact_partition(local, mapping.local_page_size_bytes), ( + f"runs do not cover the local page of layer {layer_name}" + ) + for block_id in range(num_writers): + stored: list[tuple[int, int]] = [] + for mapping in per_rank: + if mapping.is_writer(block_id): + stored += _run_intervals(mapping.runs, canonical=True) + stored.sort() + assert _is_exact_partition(stored, size), ( + f"writers of block {block_id} do not tile the canonical page " + f"of layer {layer_name}" + ) + + +def _unpadded_page_size(spec: KVCacheSpec) -> int | None: + if isinstance(spec, AttentionSpec): + return spec.unpadded_page_size_bytes + if isinstance(spec, MambaSpec): + return replace(spec, page_size_padded=None).page_size_bytes + return None + + +def derive_canonical_mappings( + vllm_config: "VllmConfig", + kv_cache_config: KVCacheConfig, + kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], +) -> dict[str, CanonicalPageMapping]: + """Per-layer canonical page mappings for this worker. + + Empty when the worker group is not exactly the TP x PCP grid; layers + absent from the result have no canonical representation. + """ + parallel_config = vllm_config.parallel_config + tp_size = parallel_config.tensor_parallel_size + pcp_size = parallel_config.prefill_context_parallel_size + group_size = tp_size * pcp_size + if parallel_config.world_size != group_size: + return {} + + def ctx(rank: int) -> _RankContext: + return _RankContext( + tp_size=tp_size, + dcp_size=parallel_config.decode_context_parallel_size, + pcp_size=pcp_size, + interleave=parallel_config.cp_kv_cache_interleave_size, + total_kv_heads=vllm_config.model_config.get_total_num_kv_heads(), + rank=rank, + ) + + my_rank = parallel_config.rank + num_blocks = kv_cache_config.num_blocks + + mappings: dict[str, CanonicalPageMapping] = {} + for kv_cache_group in kv_cache_config.kv_cache_groups: + group_kv_cache_spec = kv_cache_group.kv_cache_spec + if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs): + per_layer_specs = group_kv_cache_spec.kv_cache_specs + else: + per_layer_specs = {} + for layer_name in kv_cache_group.layer_names: + spec = per_layer_specs.get(layer_name, group_kv_cache_spec) + per_rank: list[CanonicalPageMapping] = [] + for rank in range(group_size): + mapping = _layer_mapping( + spec, kv_caches.get(layer_name), num_blocks, ctx(rank) + ) + if mapping is None: + break + per_rank.append(mapping) + if len(per_rank) != group_size: + page = _unpadded_page_size(spec) + if page is None: + continue + per_rank = [ + _opaque_fallback_mapping(page, group_size, rank) + for rank in range(group_size) + ] + _verify_tiling(layer_name, per_rank) + mappings[layer_name] = per_rank[my_rank] + return mappings diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b807b480f02d..b9837bcb4b07 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -5,7 +5,11 @@ from typing import TYPE_CHECKING from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, + MLAAttentionSpec, +) from vllm.v1.kv_offload.config import ( OffloadingCacheConfig, OffloadingConfig, @@ -40,7 +44,11 @@ def build_offloading_config( OffloadingGroupConfig( tokens_per_block=( group.kv_cache_spec.block_size - * parallel_config.decode_context_parallel_size + * ( + parallel_config.decode_context_parallel_size + if isinstance(group.kv_cache_spec, AttentionSpec) + else 1 + ) ), layer_names=tuple(group.layer_names), ) @@ -102,22 +110,48 @@ def build_offloading_config( ) worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks - # Only a single non-MLA full-attention group is parallelism-invariant: - # MLA latent KV is replicated per rank (never head-sharded), and the V2 - # model runner's KV layout is not known to be parallelism-invariant. - single_group = ( + single_group_spec = ( kv_cache_config.kv_cache_groups[0].kv_cache_spec if len(kv_cache_config.kv_cache_groups) == 1 else None ) + replicated_layout = ( + vllm_config.model_config.use_mla + # Exact type: fail closed on wrappers and sliding-window variants. + and type(single_group_spec) is MLAAttentionSpec + # Page accounting: one MLA page per layer, no packed/mixed rows. + and worker_kv_bytes_per_block > 0 + and worker_kv_bytes_per_block + == single_group_spec.page_size_bytes + * len(kv_cache_config.kv_cache_groups[0].layer_names) + # Safe MVP boundary: TP-only, no other parallel axes. + and parallel_config.tensor_parallel_size > 1 + and parallel_config.pipeline_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.world_size == parallel_config.tensor_parallel_size + # Shared /dev/shm mmap layout is single-node mp only. + and parallel_config.distributed_executor_backend == "mp" + and parallel_config.nnodes_within_dp == 1 + ) + + # Only a single non-MLA full-attention group is parallelism-invariant: + # MLA latent KV is replicated per rank (never head-sharded), and the V2 + # model runner's KV layout is not known to be parallelism-invariant. is_parallelism_agnostic = ( not vllm_config.use_v2_model_runner - and single_group is not None - and isinstance(single_group, FullAttentionSpec) - and not isinstance(single_group, MLAAttentionSpec) + and single_group_spec is not None + and isinstance(single_group_spec, FullAttentionSpec) + and not isinstance(single_group_spec, MLAAttentionSpec) ) kv_events_config = vllm_config.kv_events_config + cache_dtype = ( + vllm_config.model_config.dtype + if vllm_config.cache_config.cache_dtype == "auto" + else vllm_config.cache_config.cache_dtype + ) + return OffloadingConfig( groups=groups, worker_kv_bytes_per_block=worker_kv_bytes_per_block, @@ -128,7 +162,7 @@ def build_offloading_config( engine_id=engine_id, model=OffloadingModelConfig( name=vllm_config.model_config.model, - dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""), + dtype=str(cache_dtype).removeprefix("torch."), ), cache=OffloadingCacheConfig( tokens_per_hash=tokens_per_hash, @@ -144,4 +178,5 @@ def build_offloading_config( data_parallel_index=parallel_config.data_parallel_index, is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 4543cd6dd367..67ada00e03f6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -19,7 +19,13 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, NamedTuple -from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent +from vllm.distributed.kv_events import ( + MEDIUM_CPU, + MEDIUM_STORAGE, + BlockRemoved, + BlockStored, + KVCacheEvent, +) from vllm.logger import init_logger from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( @@ -28,6 +34,7 @@ get_kv_cache_spec_sliding_window, ) from vllm.v1.kv_offload.base import ( + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -43,6 +50,11 @@ logger = init_logger(__name__) +_MEDIUM_TO_EVENT_STR: dict[Medium, str] = { + Medium.CPU: MEDIUM_CPU, + Medium.STORAGE: MEDIUM_STORAGE, +} + class OffloadingEventGroupSpec(NamedTuple): kv_cache_spec_kind: str | None @@ -61,9 +73,9 @@ def get_offloading_event_group_spec( @dataclass(slots=True) class _OffloadEventMetadata: - """BlockStored payload snapshot for one OffloadKey, captured at store - time and kept until the matching eviction event. ``medium`` is forwarded - from the OffloadingEvent.""" + """BlockStored payload snapshot for one OffloadKey, captured while the + Request is available and kept until the matching eviction event. ``medium`` + is forwarded from the OffloadingEvent.""" # The chunk's constituent block hashes; the last one is the OffloadKey. block_hashes: tuple[BlockHash, ...] @@ -81,10 +93,11 @@ class _OffloadEventMetadata: class OffloadingEventsTracker: """Tracks offloaded chunks' KV event payloads from store to eviction. - The scheduler calls :meth:`record_store` from ``_build_store_jobs`` - while the ``Request`` is available, and routes the manager's raw - :class:`OffloadingEvent` stream through :meth:`take_events`. All state - is bounded by the CPU pool capacity and cleared by :meth:`reset`. + The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and + :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is + available. Deferred and missing lookups add no state. Under the connector's + supported success-only transfer model, entries follow primary allocations + until CPU removal translation or :meth:`reset`. """ def __init__(self, config: OffloadingKVEventsConfig): @@ -93,8 +106,7 @@ def __init__(self, config: OffloadingKVEventsConfig): config.enable_kv_cache_events and config.self_describing_kv_events ) - # OffloadKey -> payload snapshot, kept until the eviction event so - # BlockRemoved can fan out. Bounded: one entry per offloaded chunk. + # OffloadKey -> payload snapshot, kept until CPU removal or reset. self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {} def record_store( @@ -116,6 +128,23 @@ def record_store( meta = self._build_event_metadata(req, group_config, chunk_idx) self._pending_event_metadata[offload_key] = meta + def record_lookup( + self, + req: Request, + group_config: "GroupOffloadConfig", + chunk_idx: int, + offload_key: OffloadKey, + ) -> None: + """Snapshot metadata for a ready primary-tier lookup hit.""" + if not self.self_describing_enabled: + return + if group_config.sliding_window_size_in_chunks is not None: + return + if offload_key not in self._pending_event_metadata: + self._pending_event_metadata[offload_key] = self._build_event_metadata( + req, group_config, chunk_idx + ) + def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]: """Translate raw OffloadingEvents into self-describing KV events. @@ -165,7 +194,7 @@ def _build_event_metadata( assert len(chunk_hashes) == hbf if group_config.sliding_window_size_in_chunks is not None: - # record_store filters these out before calling this helper. + # The recording methods filter these out before calling this helper. raise AssertionError("self-describing events only support full attention") parent_block_hash: BlockHash | None @@ -201,7 +230,7 @@ def _build_event_metadata( def _placeholder_stored( self, key: OffloadKey, - medium: str, + medium: Medium, locality: str | None, ) -> BlockStored: return BlockStored( @@ -212,7 +241,7 @@ def _placeholder_stored( token_ids=[], lora_id=None, block_size=0, - medium=medium, + medium=_MEDIUM_TO_EVENT_STR[medium], lora_name=None, group_idx=get_offload_group_idx(key), locality=locality, @@ -232,7 +261,8 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: "OffloadingEventsTracker: no event metadata for " "offload key during BlockStored emission; emitting a " "placeholder payload. Expected for non-full-attention " - "groups; otherwise indicates a missing populate path." + "groups and promotions not observed as a primary-tier " + "hit before translation." ) yield self._placeholder_stored(key, event.medium, locality) continue @@ -249,7 +279,7 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: token_ids=list(meta.token_ids), block_size=meta.block_size, lora_id=meta.lora_id, - medium=event.medium, + medium=_MEDIUM_TO_EVENT_STR[event.medium], lora_name=meta.lora_name, extra_keys=( list(meta.extra_keys) if meta.extra_keys is not None else None @@ -290,7 +320,7 @@ def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: for group_idx, hashes in by_group.items(): yield BlockRemoved( block_hashes=hashes, - medium=event.medium, + medium=_MEDIUM_TO_EVENT_STR[event.medium], group_idx=group_idx, locality=locality, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 5e98c1266e20..fecce3d28fec 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -39,7 +39,9 @@ ) from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, + Locality, LookupResult, + Medium, OffloadingManager, OffloadingSpec, OffloadKey, @@ -47,13 +49,19 @@ ReqContext, RequestOffloadingContext, ScheduleEndContext, + TierFilter, + TierMatcher, make_offload_key, ) from vllm.v1.outputs import KVConnectorOutput -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) +KV_LOAD_TIERS_KEY = "kv_load_tiers" +MATCHER_MEDIUM_KEY = "medium" +MATCHER_LOCALITY_KEY = "locality" + @dataclass(slots=True) class TransferJobStatus: @@ -111,6 +119,26 @@ def get_sliding_window_size_in_chunks( return None +def is_store_reachable_swa_chunk( + absolute_chunk_index: int, + storable_chunk_count: int, + alignment_chunk_count: int | None, + sliding_window_chunks: int | None, + is_eagle_group: bool, +) -> bool: + """Return whether an SWA chunk can participate in an external-cache hit.""" + if alignment_chunk_count is None: + return True + assert sliding_window_chunks is not None + position_in_segment = absolute_chunk_index % alignment_chunk_count + segment_start = absolute_chunk_index - position_in_segment + actual_segment_length = min( + alignment_chunk_count, storable_chunk_count - segment_start + ) + reachable_tail = sliding_window_chunks + int(is_eagle_group) + return position_in_segment >= actual_segment_length - reachable_tail + + def resolve_mamba_align_size( spec: "OffloadingSpec", kv_cache_config: KVCacheConfig ) -> int | None: @@ -258,6 +286,8 @@ class RequestOffloadState: # time.monotonic() of this request's first deferred offload lookup; # None once consumed (observed) or while no lookup is pending. deferred_lookup_start_time: float | None = None + # True once on_request_finished has been signaled to the manager. + finished_signaled: bool = False def __post_init__(self) -> None: self.group_states = tuple( @@ -306,9 +336,12 @@ def update_block_id_groups( group_state.block_ids.extend(new_blocks) def storable_chunks( - self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int + self, + group_config: "GroupOffloadConfig", + group_state: RequestGroupState, + num_offloadable_tokens: int, ) -> int: - """Number of leading offloaded chunks eligible for store. + """Number of allocated leading offloaded chunks eligible for store. For eagle/MTP groups the volatile trailing chunk of the offloadable range is excluded while decoding: the draft-layer KV of the last @@ -325,7 +358,10 @@ def storable_chunks( is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens if group_config.is_eagle_group and is_decoding: num_chunks = max(0, num_chunks - 1) - return num_chunks + num_allocated_chunks = ( + len(group_state.block_ids) // self.config.blocks_per_chunk + ) + return min(num_chunks, num_allocated_chunks) def advance_stored_idx(self, num_offloadable_tokens: int) -> None: # max(): at the prefill->decode transition of a chunk-aligned prompt, @@ -336,7 +372,7 @@ def advance_stored_idx(self, num_offloadable_tokens: int) -> None: ): group_state.next_stored_chunk_idx = max( group_state.next_stored_chunk_idx, - self.storable_chunks(group_config, num_offloadable_tokens), + self.storable_chunks(group_config, group_state, num_offloadable_tokens), ) def update_num_hit_chunks(self, num_cached_tokens: int) -> None: @@ -348,10 +384,61 @@ def update_num_hit_chunks(self, num_cached_tokens: int) -> None: ) +def _parse_tier_filter(raw: Any) -> TierFilter: + """Parse raw kv_transfer_params tier matchers into a TierFilter.""" + if not isinstance(raw, list): + logger.warning( + "_parse_tier_filter: expected list, got %s; ignoring", + type(raw).__name__, + ) + return TierFilter.ALL + matchers: list[TierMatcher] = [] + for entry in raw: + if not isinstance(entry, dict): + logger.warning("_parse_tier_filter: entry is not a dict; skipping") + continue + medium: Medium | None = None + locality: Locality | None = None + raw_medium = entry.get(MATCHER_MEDIUM_KEY) + if raw_medium is not None: + try: + medium = Medium(raw_medium.upper()) + except (ValueError, AttributeError): + logger.warning( + "_parse_tier_filter: unknown medium %r; skipping entry", + raw_medium, + ) + continue + raw_locality = entry.get(MATCHER_LOCALITY_KEY) + if raw_locality is not None: + try: + locality = Locality(raw_locality.upper()) + except (ValueError, AttributeError): + logger.warning( + "_parse_tier_filter: unknown locality %r; skipping entry", + raw_locality, + ) + continue + matchers.append(TierMatcher(medium=medium, locality=locality)) + if not matchers: + if not raw: # input was [] — user explicitly wants nothing + return TierFilter(matchers=()) + # all entries were invalid — fall back to ALL + return TierFilter.ALL + return TierFilter(matchers=tuple(matchers)) + + def _create_req_context(req: Request) -> ReqContext: + params = req.kv_transfer_params + load_filter = TierFilter.ALL + if params: + raw = params.get(KV_LOAD_TIERS_KEY) + if raw is not None: + load_filter = _parse_tier_filter(raw) return ReqContext( req_id=req.request_id, - kv_transfer_params=req.kv_transfer_params, + kv_transfer_params=params, + load_tier_filter=load_filter, ) @@ -455,23 +542,28 @@ def _calc_num_offloadable_tokens( num = min(num, req_status.req.num_prompt_tokens) return num - def _maybe_cleanup_finished_req( - self, req_id: str, req_status: RequestOffloadState - ) -> None: - """Clean up req_status if finished and no in-flight jobs.""" - if req_status.req.is_finished() and not req_status.transfer_jobs: - del self._req_status[req_id] - def _maximal_prefix_lookup( - self, keys: Iterable[OffloadKey], req_context: ReqContext + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + req: Request, + group_config: GroupOffloadConfig, + start_chunk_idx: int, ) -> int | None: """Return the number of consecutive offloaded chunks from the start, or None if the backend deferred a lookup.""" hit_count = 0 defer_lookup = False - for key in keys: - match self.manager.lookup(key, req_context): + for local_idx, key in enumerate(keys): + result = self.manager.lookup(key, req_context) + match result: case LookupResult.HIT: + self._events_tracker.record_lookup( + req, + group_config, + start_chunk_idx + local_idx, + key, + ) hit_count += 1 case LookupResult.HIT_PENDING: defer_lookup = True @@ -616,7 +708,11 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: num_hit_chunks: int | None if sliding_window_size_in_chunks is None: num_hit_chunks = self._maximal_prefix_lookup( - offload_keys, req_status.req_context + offload_keys, + req_status.req_context, + req_status.req, + group_config, + start_chunk_idx, ) else: required_window = sliding_window_size_in_chunks @@ -823,6 +919,7 @@ def update_state_after_alloc( num_pending_gpu_blocks <= group_config.sliding_window_size_in_chunks * self.config.blocks_per_chunk + + 1 ) num_chunks = cdiv(num_cached_tokens, tokens_per_chunk) @@ -938,7 +1035,9 @@ def _build_store_jobs( continue req = req_status.req - if req.is_finished(): + if req.status is RequestStatus.FINISHED_ABORTED: + num_tokens_after_batch = req.num_computed_tokens + elif req.is_finished(): num_tokens_after_batch = req.num_tokens else: num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] @@ -955,7 +1054,7 @@ def _build_store_jobs( self.config.kv_group_configs, req_status.group_states ): num_chunks = req_status.storable_chunks( - group_config, num_offloadable_tokens + group_config, group_state, num_offloadable_tokens ) start_chunk_idx = group_state.next_stored_chunk_idx @@ -974,9 +1073,6 @@ def _build_store_jobs( ] assert len(offload_keys) == len(offload_block_ids) - alignment_chunk_count = group_config.alignment_chunk_count - tail = group_config.sliding_window_size_in_chunks - for key_idx, (offload_key, block_id) in enumerate( zip(offload_keys, offload_block_ids) ): @@ -984,20 +1080,22 @@ def _build_store_jobs( continue # Skip SWA chunks that can never serve a load hit: # within each full-attention alignment segment, only the - # trailing `tail` chunks are reachable by - # _sliding_window_lookup. For DeepSeek V4 with 100K - # tokens this reduces SWA stores by ~78%. - if alignment_chunk_count is not None: - assert tail is not None - abs_chunk_idx = start_chunk_idx + key_idx - pos_in_segment = abs_chunk_idx % alignment_chunk_count - if pos_in_segment < alignment_chunk_count - tail: - continue + # trailing chunks queried by _sliding_window_lookup are + # reachable. EAGLE/MTP requires one additional chunk that + # lookup later drops as its volatile draft tail. + abs_chunk_idx = start_chunk_idx + key_idx + if not is_store_reachable_swa_chunk( + abs_chunk_idx, + num_chunks, + group_config.alignment_chunk_count, + group_config.sliding_window_size_in_chunks, + group_config.is_eagle_group, + ): + continue new_offload_keys.append(offload_key) if not new_offload_keys: req_status.advance_stored_idx(num_offloadable_tokens) - self._maybe_cleanup_finished_req(req_id, req_status) continue store_output = self.manager.prepare_store( @@ -1008,12 +1106,10 @@ def _build_store_jobs( _ConnectorMetricName.ALLOCATION_FAILURE ) logger.warning("Request %s: cannot store chunks", req_id) - self._maybe_cleanup_finished_req(req_id, req_status) continue if not store_output.keys_to_store: req_status.advance_stored_idx(num_offloadable_tokens) - self._maybe_cleanup_finished_req(req_id, req_status) continue self._touch(req_status) @@ -1032,7 +1128,7 @@ def _build_store_jobs( group_config.sliding_window_size_in_chunks is not None ) num_chunks = req_status.storable_chunks( - group_config, num_offloadable_tokens + group_config, group_state, num_offloadable_tokens ) start_chunk_idx = group_state.next_stored_chunk_idx block_ids = group_state.block_ids @@ -1157,6 +1253,17 @@ def build_connector_meta( store_jobs=self._build_store_jobs(scheduler_output), jobs_to_flush=self._current_batch_jobs_to_flush, ) + + # All prepare_store calls for finished requests have been issued. + # Signal on_request_finished and clean up state where possible. + for req_id in scheduler_output.finished_req_ids or (): + req_status = self._req_status.get(req_id) + if req_status is None: + continue + req_status.finished_signaled = True + self.manager.on_request_finished(req_status.req_context) + if not req_status.transfer_jobs: + del self._req_status[req_id] self._current_batch_load_jobs = {} self._current_batch_jobs_to_flush = set() self._current_batch_allocated_block_ids = set() @@ -1247,7 +1354,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput): del self._jobs[job_id] req_status.transfer_jobs.remove(job_id) - if not req_status.transfer_jobs and req_status.req.is_finished(): + if req_status.finished_signaled and not req_status.transfer_jobs: del self._req_status[job_status.req_id] def get_stats(self) -> OffloadingConnectorStats | None: @@ -1289,7 +1396,6 @@ def request_finished( self.manager.on_request_finished(req_context) return False, None - self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) # Update offload keys with final block hash so _build_store_jobs can @@ -1333,6 +1439,8 @@ def reset_cache(self) -> None: for req_id, status in list(self._req_status.items()): if status.req.is_finished(): + if not status.finished_signaled: + self.manager.on_request_finished(status.req_context) del self._req_status[req_id] # Reset offloading manager cache diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 045e513c3f49..211903242edb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -5,6 +5,10 @@ import torch +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import ( + derive_canonical_mappings, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingConnectorMetadata, OffloadingWorkerMetadata, @@ -40,11 +44,17 @@ class OffloadingConnectorWorker: def __init__( self, spec: OffloadingSpec, + vllm_config: "VllmConfig", kv_cache_config: KVCacheConfig, ): self.spec = spec + self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config self.worker: OffloadingWorker | None = None + # Non-writers still ack: pending_count waits for world_size per job. + self._is_store_writer = ( + not self.spec.replicated_layout or self.spec.config.parallel.rank == 0 + ) # job_id -> req_id for in-flight loads. self._load_jobs: dict[int, ReqId] = {} @@ -56,11 +66,12 @@ def __init__( def _init_worker(self, kv_caches: CanonicalKVCaches) -> None: self.worker = self.spec.get_worker(kv_caches) - def register_kv_caches( - self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]] - ): + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): kv_cache_config = self.kv_cache_config num_blocks = kv_cache_config.num_blocks + mappings = derive_canonical_mappings( + self.vllm_config, kv_cache_config, kv_caches + ) # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use # stride(0) as the manager-block stride (equals total_num_bytes_per_block). @@ -120,24 +131,13 @@ def register_kv_caches( ) elif isinstance(layer_kv_cache_spec, MambaSpec): - state_tensors = kv_caches[layer_name] - assert isinstance(state_tensors, list) - - # re-construct the raw (num_blocks, page_size) tensor - # from the first state tensor - assert len(state_tensors) > 0 - first_state_tensor = state_tensors[0] - assert first_state_tensor.storage_offset() == 0 - tensor = ( - torch.tensor( - [], - dtype=torch.int8, - device=first_state_tensor.device, - ) - .set_(first_state_tensor.untyped_storage()) - .view((num_blocks, layer_kv_cache_spec.page_size_bytes)) + layer_kv_cache = kv_caches[layer_name] + assert layer_kv_cache.dtype == torch.int8 + tensors_per_block[layer_name] = ( + layer_kv_cache.view( + num_blocks, layer_kv_cache_spec.page_size_bytes + ), ) - tensors_per_block[layer_name] = (tensor,) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes unpadded_page_size_bytes[layer_name] = replace( @@ -210,10 +210,21 @@ def register_kv_caches( curr_tensor_idx = len(block_tensors) - 1 for layer_name in tensor_layer_names: + mapping = ( + mappings.get(layer_name) + if len(tensors_per_block[first_layer_name]) == 1 + else None + ) + assert ( + mapping is None + or mapping.local_page_size_bytes + == unpadded_page_size_bytes[layer_name] + ) block_data_refs[layer_name].append( CanonicalKVCacheRef( tensor_idx=curr_tensor_idx, page_size_bytes=(unpadded_page_size_bytes[layer_name]), + mapping=mapping, ) ) @@ -287,6 +298,9 @@ def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata) for job_id in kv_connector_metadata.jobs_to_flush: entry = kv_connector_metadata.store_jobs.pop(job_id, None) if entry is not None: + if not self._is_store_writer: + self._connector_worker_meta.mark_completed(job_id) + continue assert isinstance(entry.src_spec, GPULoadStoreSpec) self._unsubmitted_store_jobs.append( (job_id, entry.src_spec, entry.dst_spec) @@ -317,6 +331,10 @@ def start_kv_transfers(self, metadata: OffloadingConnectorMetadata): def prepare_store_kv(self, metadata: OffloadingConnectorMetadata): for job_id, entry in metadata.store_jobs.items(): + if not self._is_store_writer: + # Gate before queueing: no _unsubmitted_store_jobs entry. + self._connector_worker_meta.mark_completed(job_id) + continue # NOTE(orozery): defer the store to the beginning of the next # engine step, so that offloading starts AFTER transfers related # to token sampling, thereby avoiding delays to token generation. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 2fe4bf6a5a7c..a32d5d9cc8ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -51,6 +51,12 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): def prefer_cross_layer_blocks(self) -> bool: return True + @property + def requires_kv_delivery(self) -> bool: + # Runs as kv_both, but is a best-effort cache: a dropped save is just a + # future cache miss, so opt out of the producer-role default. + return False + def __init__( self, vllm_config: VllmConfig, @@ -69,7 +75,9 @@ def __init__( spec, vllm_config, kv_cache_config ) elif role == KVConnectorRole.WORKER: - self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config) + self.connector_worker = OffloadingConnectorWorker( + spec, vllm_config, kv_cache_config + ) def shutdown(self) -> None: if self.connector_worker is not None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index f1dac13ca517..3ed34f169d76 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -45,6 +45,12 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): """CPU KV cache offloading with custom kernel transfers and BlockPool LRU.""" + @property + def requires_kv_delivery(self) -> bool: + # Runs as kv_both, but is a best-effort cache: a dropped save is just a + # future cache miss, so opt out of the producer-role default. + return False + def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py index 634d59976f70..3879bee0ce29 100644 --- a/vllm/distributed/nixl_utils.py +++ b/vllm/distributed/nixl_utils.py @@ -21,7 +21,7 @@ def _maybe_set_ucx_rcache_limit() -> None: if "UCX_RCACHE_MAX_UNRELEASED" in os.environ: return - if "nixl" in sys.modules or "rixl" in sys.modules: + if "nixl" in sys.modules or "nixl_rocm" in sys.modules: logger.warning_once( "NIXL was already imported, we can't reset " "UCX_RCACHE_MAX_UNRELEASED. " @@ -36,8 +36,12 @@ def _maybe_set_ucx_rcache_limit() -> None: os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" +def _get_nixl_package_name() -> str: + return "nixl_rocm" if current_platform.is_rocm() else "nixl" + + def _get_nixl_module_name(name: str) -> str: - package_name = "rixl" if current_platform.is_rocm() else "nixl" + package_name = _get_nixl_package_name() if name == "nixlXferTelemetry": return f"{package_name}._bindings" return f"{package_name}._api" @@ -80,11 +84,11 @@ def __getattr__(name: str) -> Any: def is_nixl_available() -> bool: - """Lightweight check for nixl/rixl package without importing it.""" + """Lightweight check for the platform's NIXL package without importing it.""" import importlib.util - pkg = "rixl" if current_platform.is_rocm() else "nixl" - return importlib.util.find_spec(pkg) is not None + pkg = _get_nixl_package_name() + return pkg in sys.modules or importlib.util.find_spec(pkg) is not None __all__ = [ diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index f0647323e61d..a90e8acbcad8 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -127,6 +127,28 @@ def _register_group(group: "GroupCoordinator") -> None: _groups[group.unique_name] = weakref.ref(group) +def _apply_to_device_comms( + action: Callable[[DeviceCommunicatorBase], None], +) -> None: + """Apply ``action`` to every group's device communicator. + + Walks the registered parallel groups and skips those without a device + communicator (absent at ``world_size == 1``). + """ + comms = [] + for group_ref in _groups.values(): + group = group_ref() + if group is None: + continue + dc = group.device_communicator + if dc is None: + continue + comms.append(dc) + + for dc in comms: + action(dc) + + def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: assert group_name in _groups, f"Group {group_name} is not found." group = _groups[group_name]() @@ -392,6 +414,7 @@ def __init__( use_device_communicator: bool, # whether to use device communicator use_message_queue_broadcaster: bool = False, group_name: str | None = None, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -400,13 +423,10 @@ def __init__( self.rank = torch.distributed.get_rank() self.local_rank = local_rank self.device_index: int - if _WORLD is not None: - self.device_index = _WORLD.device_index - else: - assert local_rank >= 0, ( - "local_rank must be provided when creating the world group" - ) - self.device_index = local_rank + assert local_rank >= 0, ( + "local_rank must be provided when creating the world group" + ) + self.device_index = local_rank self_device_group = None self_cpu_group = None @@ -489,6 +509,7 @@ def __init__( device=self.device, device_group=self.device_group, unique_name=self.unique_name, + use_all2all=use_all2all, ) from vllm.distributed.device_communicators.shm_broadcast import MessageQueue @@ -1302,6 +1323,7 @@ def init_model_parallel_group( use_message_queue_broadcaster: bool = False, group_name: str | None = None, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> GroupCoordinator: return GroupCoordinator( group_ranks=group_ranks, @@ -1310,6 +1332,7 @@ def init_model_parallel_group( use_device_communicator=use_device_communicator, use_message_queue_broadcaster=use_message_queue_broadcaster, group_name=group_name, + use_all2all=use_all2all, ) @@ -1320,6 +1343,7 @@ def _init_stateless_group( backend: str, coord_store: Store, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> "StatelessGroupCoordinator": """Create a StatelessGroupCoordinator with the given parameters.""" from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator @@ -1335,6 +1359,7 @@ def _init_stateless_group( coord_store=coord_store, global_rank=world.rank, global_world_size=world.world_size, + use_all2all=use_all2all, ) @@ -1905,6 +1930,7 @@ def initialize_model_parallel( .unbind(0) ) group_ranks = [x.tolist() for x in group_ranks] + use_all2all = parallel_config.use_all2all if enable_elastic_ep: _EP = _init_stateless_group( group_ranks, @@ -1912,10 +1938,15 @@ def initialize_model_parallel( parallel_config.data_parallel_master_ip, backend, coord_store=coord_store, + use_all2all=use_all2all, ) else: _EP = init_model_parallel_group( - group_ranks, get_world_group().local_rank, backend, group_name="ep" + group_ranks, + get_world_group().local_rank, + backend, + group_name="ep", + use_all2all=use_all2all, ) # Create EPLB group with the same ranks as EP if EPLB is enabled. @@ -2031,6 +2062,20 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module): _EPLB.prepare_communication_buffer_for_model(model) +def checkpoint_prepare_distributed_state() -> None: + """Prepare every device communicator for a process checkpoint.""" + torch.accelerator.synchronize() + _apply_to_device_comms(lambda comm: comm.checkpoint_prepare()) + torch.accelerator.synchronize() + + +def checkpoint_restore_distributed_state() -> None: + """Restore every device communicator after a process checkpoint.""" + torch.accelerator.synchronize() + _apply_to_device_comms(lambda comm: comm.checkpoint_restore()) + torch.accelerator.synchronize() + + def model_parallel_is_initialized(): """Check if tensor and pipeline parallel groups are initialized.""" return _TP is not None and _PP is not None diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 5f4597d07cbb..38c74a97c55c 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -79,6 +79,7 @@ def __init__( host: str = "127.0.0.1", global_rank: int = 0, global_world_size: int = 1, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -191,6 +192,7 @@ def __init__( global_ranks=self.ranks, global_world_size=global_world_size, tcp_store_group=self.tcp_store_group, + use_all2all=use_all2all, ) self.mq_broadcaster = None diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2e377e292536..3fd21101e744 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -5,7 +5,15 @@ from abc import ABC, abstractmethod from collections.abc import Iterator from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + Protocol, + TypeVar, + runtime_checkable, +) import torch from typing_extensions import Self @@ -18,7 +26,7 @@ TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") -TConfig = TypeVar("TConfig", bound="WeightTransferConfig") +TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo") # A trainer supplies its parameters as a `WeightSource` (defined below): a # re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` @@ -108,7 +116,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024 @dataclass -class TrainerInitInfo(WeightTransferInitInfo): +class TrainerInitInfo: """Base trainer-side init info: which trainer rank drives the transfer. `rank` is this trainer process's rank, provided **explicitly** by the @@ -118,10 +126,25 @@ class TrainerInitInfo(WeightTransferInitInfo): while every rank still runs the trainer-side collectives. Backend subclasses add their own (positional) fields; `rank` is keyword-only so that ordering never conflicts. + + Every concrete subclass sets a class-level `backend` string (the same key it + registers under in `WeightTransferTrainerFactory`). The factory reads it to + dispatch, so callers pass only the init info/ It is a `ClassVar` + (a fixed per-backend constant), so it is not an ``__init__`` field. """ + backend: ClassVar[str] + rank: int = field(kw_only=True) + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if not getattr(cls, "backend", None): + raise TypeError( + f"{cls.__name__} must set a class-level `backend` string " + "(the WeightTransferTrainerFactory registry key)." + ) + @property def is_sender(self) -> bool: return self.rank == 0 @@ -370,10 +393,10 @@ def start_weight_update(self) -> None: ... def update_weights(self, update_info: dict[str, Any]) -> None: ... - def finish_weight_update(self) -> None: ... + def finish_weight_update(self, weight_version: str | None = None) -> None: ... -class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): +class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): """Trainer-side weight transfer engine. Symmetric to `WeightTransferEngine` but lives in the training process. @@ -382,6 +405,13 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): plans) on `self`. The `WeightSource` is required at `trainer_init`, then replayed each round by the no-argument `send_weights()`. + Unlike the worker engine, the trainer side does not take a + `WeightTransferConfig`: the backend is selected from the init info's + `backend` `ClassVar` (so callers pass only the init info), and the static + wire params (packed, buffer sizes) ride the backend-specific + `TrainerInitInfo`, which the sender also propagates to the worker at the init + handshake. + Multi-rank trainers: `trainer_init` and `send_weights` are called on *every* trainer rank. Rank 0 is the sender, resolved once at `trainer_init` into `is_sender`. Non-sender ranks still run every @@ -392,22 +422,18 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): Subclasses should define: init_info_cls: Type of backend-specific trainer init info - config_cls: Type of backend-specific config """ - # Subclasses should override these class attributes - init_info_cls: type[TInitInfo] - config_cls: type[TConfig] + # Subclasses should override this class attribute + init_info_cls: type[TTrainerInitInfo] def __init__( self, - config: TConfig, *, client: "VLLMWeightSyncClient", source: "WeightSource", is_sender: bool = True, ) -> None: - self.config = config self.is_sender = is_sender # The real client is held on every rank; each engine only *calls* it when # `is_sender`, so non-sender ranks never touch the wire. @@ -418,8 +444,7 @@ def __init__( @abstractmethod def trainer_init( cls, - config: TConfig, - init_info: TInitInfo, + init_info: TTrainerInitInfo, *, client: "VLLMWeightSyncClient", source: "WeightSource", diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 4f54a6e291e3..12dd0c9eacc8 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -77,8 +77,11 @@ def update_weights(self, update_info: dict[str, Any]) -> None: "update_weights", {"update_info": _json_safe_update_info(update_info)} ) - def finish_weight_update(self) -> None: - self._post("finish_weight_update") + def finish_weight_update(self, weight_version: str | None = None) -> None: + json = ( + {"weight_version": weight_version} if weight_version is not None else None + ) + self._post("finish_weight_update", json) class RayVLLMWeightSyncClient: @@ -108,7 +111,11 @@ def update_weights(self, update_info: dict[str, Any]) -> None: request = WeightTransferUpdateRequest(update_info=update_info) ray.get([h.update_weights.remote(request) for h in self.handles]) - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: import ray ray.get([h.finish_weight_update.remote() for h in self.handles]) + if weight_version is not None: + ray.get( + [h.update_weight_version.remote(weight_version) for h in self.handles] + ) diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 4ea27c5ef584..1172b4648025 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -18,9 +18,9 @@ from vllm.config import VllmConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, VLLMWeightSyncClient, WeightSource, - WeightTransferInitInfo, ) logger = init_logger(__name__) @@ -164,9 +164,7 @@ def loader() -> type[TrainerWeightTransferEngine]: @classmethod def trainer_init( cls, - backend: str, - config: "WeightTransferConfig", - init_info: "WeightTransferInitInfo", + init_info: "TrainerInitInfo", *, client: "VLLMWeightSyncClient", source: "WeightSource", @@ -176,16 +174,22 @@ def trainer_init( Called on every trainer rank (multi-rank trainers construct on all ranks; the sender is resolved inside the engine's ``trainer_init``). + The trainer side takes no `WeightTransferConfig` and no separate + `backend` argument: the backend is read from ``init_info.backend`` (a + `ClassVar` on each `TrainerInitInfo` subclass), and the static wire + params ride `init_info`. + Args: - backend: Backend name (must be registered). - config: Backend-specific weight transfer config. - init_info: Backend-specific trainer init info. + init_info: Backend-specific trainer init info. Its `backend` + selects the engine; it also carries the wire params (e.g. + `packed`). client: Inference-side control-plane client. source: `WeightSource` of `(name, tensor)` pairs to send each round. Raises: - ValueError: If the backend is not registered. + ValueError: If `init_info.backend` is not registered. """ + backend = init_info.backend if backend not in cls._registry: available = list(cls._registry.keys()) raise ValueError( @@ -200,7 +204,6 @@ def trainer_init( ) return engine_cls.trainer_init( - config=config, init_info=init_info, client=client, source=source, @@ -228,3 +231,12 @@ def trainer_init( "vllm.distributed.weight_transfer.sparse_nccl_engine", "SparseNCCLWeightTransferEngine", ) + + +# Trainer-side engines. Backends register here as they migrate to the stateful +# trainer engine; NCCL / sparse NCCL keep their static trainer path until then. +WeightTransferTrainerFactory.register_engine( + "ipc", + "vllm.distributed.weight_transfer.ipc_engine", + "IPCTrainerWeightTransferEngine", +) diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index f1b6070893bf..d512c9898cfc 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -3,19 +3,21 @@ """IPC-based weight transfer engine using CUDA IPC for communication.""" import pickle -from collections.abc import Callable, Iterator from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar import pybase64 as base64 -import ray -import requests import torch from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor +from typing_extensions import Self from vllm import envs from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, WeightTransferEngine, WeightTransferInitInfo, WeightTransferUpdateInfo, @@ -31,47 +33,35 @@ @dataclass -class IPCTrainerSendWeightsArgs: - """Arguments for IPC trainer_send_weights method.""" - - send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None] - """How to send updates to vLLM. Either a string ('ray' or 'http') for - built-in transports, or a callable that receives an - IPCWeightTransferUpdateInfo and performs the send.""" - llm_handle: Any = None - """Ray actor handle or list of handles (required for 'ray' send_mode).""" - url: str | None = None - """Base URL for HTTP endpoint (required for 'http' send_mode).""" - packed: bool = False - """Whether to use packed tensor transfer for bounded-memory chunking.""" - packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES - """Size in bytes for each packed tensor buffer when packed=True.""" +class IPCWeightTransferInitInfo(WeightTransferInitInfo): + """Worker-side init info for IPC weight transfer. No rendezvous needed. - def __post_init__(self): - """Validate that required arguments are provided for the selected mode.""" - if callable(self.send_mode): - return - if self.send_mode == "ray" and self.llm_handle is None: - raise ValueError("llm_handle is required for 'ray' send_mode") - if self.send_mode == "http" and self.url is None: - raise ValueError("url is required for 'http' send_mode") - if self.send_mode not in ("ray", "http"): - raise ValueError( - f"send_mode must be 'ray', 'http', or a callable, " - f"got {self.send_mode!r}" - ) + `packed` is a must-agree wire param: the trainer ships it here at the init + handshake so the worker decodes with the same setting the trainer encoded + with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does + not need the buffer size (producer-only).""" + + packed: bool = False @dataclass -class IPCWeightTransferInitInfo(WeightTransferInitInfo): - """Initialization info for IPC weight transfer backend. No init needed for IPC.""" +class IPCTrainerInitInfo(TrainerInitInfo): + """Trainer-side init info for IPC weight transfer. No rendezvous needed; + `rank` (from `TrainerInitInfo`) identifies this trainer process — rank 0 + ships the merged IPC handles. All ranks still join the handle all-gather. - pass + `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The + trainer propagates `packed` to the worker at `trainer_init` so the two sides + cannot disagree. `backend` is the factory dispatch key.""" + + backend: ClassVar[str] = "ipc" + packed: bool = False + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES @dataclass class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): - """Update info for IPC weight transfer backend.""" + """Per-round update info for the IPC weight transfer backend.""" names: list[str] dtype_names: list[str] @@ -85,8 +75,6 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): tensor_sizes: list[int] | None = None """Per-parameter sizes in bytes within the packed buffer. Required when packed=True, unused otherwise.""" - packed: bool = False - """Whether this update uses packed tensor format.""" def __post_init__(self): if self.ipc_handles_pickled is not None: @@ -119,17 +107,14 @@ def __post_init__(self): f"`shapes` should be of the same size as `names`: " f"got {len(self.shapes)} and {len(self.names)}" ) - if ( - not self.packed - and isinstance(self.ipc_handles, list) - and len(self.ipc_handles) != num_params - ): + # Unpacked transfers carry a per-parameter handle dict (a list); packed + # transfers carry a single dict for the whole buffer, so the list check + # only applies to the list case. + if isinstance(self.ipc_handles, list) and len(self.ipc_handles) != num_params: raise ValueError( f"`ipc_handles` should be of the same size as `names`: " f"got {len(self.ipc_handles)} and {len(self.names)}" ) - if self.packed and self.tensor_sizes is None: - raise ValueError("`tensor_sizes` is required when packed=True") class IPCWeightTransferEngine( @@ -154,57 +139,21 @@ def __init__( device: torch.device, model: torch.nn.Module, ) -> None: - """ - Initialize the IPC weight transfer engine. - - Args: - config: The configuration for the weight transfer engine - vllm_config: The full vLLM config - device: The device this worker's model lives on - model: The local model instance which will receive the weights - """ super().__init__(config, vllm_config, device, model) - - def parse_update_info( - self, update_dict: dict[str, Any] - ) -> IPCWeightTransferUpdateInfo: - """Parse update dict, deserializing pickled IPC handles if present. - - HTTP transport sends IPC handles as a base64-encoded pickle under the - key ``ipc_handles_pickled``. This method deserializes them back into - ``ipc_handles`` before constructing the typed dataclass, keeping - serialization concerns out of the dataclass itself. - - Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the - payload is deserialized via ``pickle.loads``. - """ - pickled = update_dict.pop("ipc_handles_pickled", None) - if pickled is not None: - if update_dict.get("ipc_handles") is not None: - raise ValueError( - "Cannot specify both `ipc_handles` and `ipc_handles_pickled`" - ) - - if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: - raise ValueError( - "Refusing to deserialize `ipc_handles_pickled` without " - "VLLM_ALLOW_INSECURE_SERIALIZATION=1" - ) - - update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled)) - - return super().parse_update_info(update_dict) + # Set from the trainer-supplied init info at the handshake; defaults are + # only for the (unreachable) receive-before-init case. + self.packed = False def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None: """ - Initialize the weight transfer mechanism. - This is called once at the beginning of training. - No initialization needed for IPC backend. + Initialize the weight transfer mechanism. No data-plane rendezvous is + needed for IPC; this just records the trainer-supplied wire params so the + worker decodes exactly as the trainer encoded. Args: - init_info: IPC initialization info (empty) + init_info: IPC initialization info (carries `packed`). """ - pass + self.packed = init_info.packed def start_weight_update(self) -> None: """Initialize layerwise reloading for the incoming checkpoint weights.""" @@ -226,6 +175,10 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: """ Receive weights from the trainer via CUDA IPC handles and load them. + Whether the transfer is packed is read from `self.packed`, set at the + init handshake from the trainer's init info, so it is guaranteed to + match how the trainer encoded. + Args: update_info: IPC update info containing parameter names, dtypes, shapes, and IPC handles. Each IPC handle is a mapping between physical @@ -238,8 +191,9 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: # rebuilt on the device the model lives on. device_index = self.device.index - if update_info.packed: - assert update_info.tensor_sizes is not None + if self.packed: + if update_info.tensor_sizes is None: + raise ValueError("`tensor_sizes` is required when packed=True") assert isinstance(update_info.ipc_handles, dict) weights = packed_ipc_consumer( ipc_handle=update_info.ipc_handles, @@ -274,71 +228,119 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: weight = rebuild_cuda_tensor(*list_args) weights.append((name, weight)) - self.model.load_weights(weights) + from vllm.model_executor.model_loader.mtp_validation import ( + disable_mtp_completeness_check, + ) + + with disable_mtp_completeness_check(): + self.model.load_weights(weights) def shutdown(self) -> None: pass @staticmethod - def trainer_send_weights( - iterator: Iterator[tuple[str, torch.Tensor]], - trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs, - ) -> None: - """Send weights from trainer to inference workers via CUDA IPC. + def trainer_send_weights(*args: Any, **kwargs: Any) -> None: + """Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead. - Supports two transport modes ('ray' and 'http') and two transfer - strategies: - - Non-packed (default): all weights in a single API call. - - Packed (packed=True): chunked transfer with bounded GPU memory. + Transitional stub kept only to satisfy the (still abstract) + `WeightTransferEngine.trainer_send_weights`; that member is dropped from + the worker ABC once every backend has migrated to the trainer engine. + """ + raise NotImplementedError( + "The static IPC trainer path has been replaced by " + "IPCTrainerWeightTransferEngine. Build it via " + "WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), " + "client=..., source=...) and drive it with send_weights()." + ) - For multi-GPU training, all ranks must call this method in - parallel. IPC handles are all-gathered across ranks and merged - so that each vLLM worker can find its own GPU UUID. Only rank 0 - sends the payload to vLLM. - .. note:: - This method calls ``update_weights`` internally. The caller must - call ``start_weight_update`` before and ``finish_weight_update`` - after this method. +class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]): + """Trainer-side CUDA IPC weight transfer engine. - Args: - iterator: Iterator of (name, tensor) pairs. For multi-GPU, - each rank should yield the full tensor on its own GPU - (e.g. via FSDP full_tensor()). - trainer_args: IPCTrainerSendWeightsArgs or equivalent dict. - """ - args = ( - IPCTrainerSendWeightsArgs(**trainer_args) - if isinstance(trainer_args, dict) - else trainer_args - ) - device_index = torch.accelerator.current_device_index() - gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid) - if args.packed: - IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid) - else: - IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid) + Called on every trainer rank. For multi-rank (e.g. FSDP) trainers all ranks + iterate the source (materializing each tensor) and contribute to the + IPC-handle all-gather; only the sender (rank 0) ships the merged handles to + the inference side. IPC transfer is straight-line (no concurrent broadcast + like NCCL): `update_weights` *is* the transfer, and it rides the client, so + it no-ops on non-senders. - @staticmethod - def _is_rank_zero() -> bool: - """Return True if this is rank 0 or no distributed group exists.""" - if not torch.distributed.is_initialized(): - return True - return torch.distributed.get_rank() == 0 + `packed` / `packed_buffer_size_bytes` come from `IPCTrainerInitInfo`; the + sender propagates `packed` to the worker at `trainer_init`. + """ + + init_info_cls = IPCTrainerInitInfo + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + is_sender: bool = True, + packed: bool = False, + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + ) -> None: + super().__init__(client=client, source=source, is_sender=is_sender) + self.packed = packed + self.packed_buffer_size_bytes = packed_buffer_size_bytes + self.device_index = torch.accelerator.current_device_index() + self.gpu_uuid = str(torch.cuda.get_device_properties(self.device_index).uuid) + + @classmethod + def trainer_init( + cls, + init_info: IPCTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + ) -> Self: + engine = cls( + client=client, + source=source, + is_sender=init_info.is_sender, + packed=init_info.packed, + packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, + ) + # IPC needs no data-plane rendezvous. The sender ships the must-agree + # `packed` flag so the worker decodes exactly as this trainer encodes. + if engine.is_sender: + engine.client.init_weight_transfer_engine({"packed": init_info.packed}) + return engine + + def send_weights(self) -> None: + source = self.source + if self.is_sender: + self.client.start_weight_update() + # Unpacked returns strong refs to its IPC-shared copies; they must stay + # alive across `finish_weight_update` too. + weight_refs = self._send(source) + if self.is_sender: + self.client.finish_weight_update() + self._post_send_sync() + del weight_refs + + # ---- data plane (runs on all ranks; only the sender ships) ---- + + def _send(self, source: WeightSource) -> list[torch.Tensor] | None: + if self.packed: + self._send_packed(source) + return None + return self._send_unpacked(source) - @staticmethod def _all_gather_and_merge_handles( + self, handles: list[dict[str, tuple]], ) -> list[dict[str, tuple]]: """All-gather and merge IPC handle dicts across ranks in one call. Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one per parameter or one per chunk). A single all_gather_object - collects every rank's full list, then rank 0 merges per-index so + collects every rank's full list, then the sender merges per-index so each dict maps every GPU UUID to its args. - Non-rank-0 returns a list of empty dicts. - No-op (returns handles unchanged) when no distributed group exists. + The all-gather runs over the *default* process group; this assumes the + default group is exactly the set of colocated trainer ranks and that the + sender is a member. No-op (returns handles unchanged) when no distributed + group exists. """ if ( not torch.distributed.is_initialized() @@ -352,7 +354,7 @@ def _all_gather_and_merge_handles( torch.distributed.barrier() torch.cuda.synchronize() - if torch.distributed.get_rank() == 0: + if self.is_sender: merged: list[dict[str, tuple]] = [] for param_idx in range(len(handles)): m: dict[str, tuple] = {} @@ -373,25 +375,23 @@ def _post_send_sync() -> None: torch.distributed.barrier() torch.cuda.ipc_collect() - @staticmethod - def _send_unpacked( - iterator: Iterator[tuple[str, torch.Tensor]], - args: IPCTrainerSendWeightsArgs, - gpu_uuid: str, - ) -> None: - """Send all weights in a single API call (non-packed mode).""" + def _send_unpacked(self, source: WeightSource) -> list[torch.Tensor]: + """Iterate the source, build one IPC handle per param, all-gather the + handles across ranks, and (sender) ship them in one update call. + + Returns the strong refs to every contiguous copy. reduce_tensor's args + do NOT keep storage alive, and non-contiguous inputs allocate fresh + storage in .contiguous(); the caller must keep these alive until the + post-send barrier (past `finish`) so the consumer's IPC views stay valid. + """ names: list[str] = [] dtype_names: list[str] = [] shapes: list[list[int]] = [] ipc_handles: list[dict[str, tuple]] = [] - # Hold strong refs to every contiguous copy until the send + post-send - # sync completes. reduce_tensor's returned args do NOT keep storage - # alive, and non-contiguous inputs allocate fresh storage in - # .contiguous() that would otherwise be GC'd before the consumer opens - # the IPC handle. weight_refs: list[torch.Tensor] = [] - for name, tensor in iterator: + for name, tensor in source: + # FSDP shards were gathered by the source; ensure contiguity here. names.append(name) dtype_names.append(str(tensor.dtype).split(".")[-1]) shapes.append(list(tensor.shape)) @@ -399,98 +399,64 @@ def _send_unpacked( weight = tensor.detach().contiguous() weight_refs.append(weight) _, ipc_args = reduce_tensor(weight) - ipc_handles.append({gpu_uuid: ipc_args}) - - ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles) - - if IPCWeightTransferEngine._is_rank_zero(): - IPCWeightTransferEngine._do_send( - args=args, - names=names, - dtype_names=dtype_names, - shapes=shapes, - ipc_handles=ipc_handles, - ) - - IPCWeightTransferEngine._post_send_sync() + ipc_handles.append({self.gpu_uuid: ipc_args}) + + ipc_handles = self._all_gather_and_merge_handles(ipc_handles) + self._do_send( + names=names, + dtype_names=dtype_names, + shapes=shapes, + ipc_handles=ipc_handles, + ) + return weight_refs - @staticmethod - def _send_packed( - iterator: Iterator[tuple[str, torch.Tensor]], - args: IPCTrainerSendWeightsArgs, - gpu_uuid: str, - ) -> None: + def _send_packed(self, source: WeightSource) -> None: """Send weights in bounded-memory chunks (packed mode).""" - post_iter_func: Callable = lambda item: item[1] - for chunk in packed_ipc_producer( - iterator=iterator, - gpu_uuid=gpu_uuid, - post_iter_func=post_iter_func, - buffer_size_bytes=args.packed_buffer_size_bytes, + iterator=iter(source), + gpu_uuid=self.gpu_uuid, + post_iter_func=lambda item: item[1], + buffer_size_bytes=self.packed_buffer_size_bytes, ): - ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles( - [chunk.ipc_handle] - )[0] - - if IPCWeightTransferEngine._is_rank_zero(): - IPCWeightTransferEngine._do_send( - args=args, - names=chunk.names, - dtype_names=chunk.dtype_names, - shapes=chunk.shapes, - ipc_handles=ipc_handle, - tensor_sizes=chunk.tensor_sizes, - packed=True, - ) - - IPCWeightTransferEngine._post_send_sync() + ipc_handle = self._all_gather_and_merge_handles([chunk.ipc_handle])[0] + self._do_send( + names=chunk.names, + dtype_names=chunk.dtype_names, + shapes=chunk.shapes, + ipc_handles=ipc_handle, + tensor_sizes=chunk.tensor_sizes, + ) + # Per-chunk barrier: the producer reuses a single IPC buffer across + # chunks, but only the sender waits for the consumers (via _do_send). + # Without syncing every rank here, non-sender ranks race ahead and + # overwrite their buffer while their colocated worker is still + # reading the current chunk, silently corrupting the transfer. + self._post_send_sync() - @staticmethod def _do_send( - args: IPCTrainerSendWeightsArgs, + self, names: list[str], dtype_names: list[str], shapes: list[list[int]], ipc_handles: list[dict[str, tuple]] | dict[str, tuple], tensor_sizes: list[int] | None = None, - packed: bool = False, ) -> None: - """Send a single update payload via the configured transport.""" + """Build one update payload and ship it via the client. Only the sender + ships (non-sender ranks already contributed to the handle all-gather). + + Emits raw `ipc_handles`; transports that cannot carry them natively + (HTTP/JSON) pickle them in their client (see `HTTPVLLMWeightSyncClient`). + """ + if not self.is_sender: + return update_fields: dict[str, Any] = { "names": names, "dtype_names": dtype_names, "shapes": shapes, - "packed": packed, + "ipc_handles": ipc_handles, } if tensor_sizes is not None: update_fields["tensor_sizes"] = tensor_sizes - update_fields["ipc_handles"] = ipc_handles update_info = IPCWeightTransferUpdateInfo(**update_fields) - - if callable(args.send_mode): - args.send_mode(update_info) - elif args.send_mode == "ray": - handles = ( - args.llm_handle - if isinstance(args.llm_handle, list) - else [args.llm_handle] - ) - ray.get( - [ - h.update_weights.remote(dict(update_info=asdict(update_info))) - for h in handles - ] - ) - elif args.send_mode == "http": - pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode( - "utf-8" - ) - http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"} - http_fields["ipc_handles_pickled"] = pickled_handles - - url = f"{args.url}/update_weights" - payload = {"update_info": http_fields} - response = requests.post(url, json=payload, timeout=300) - response.raise_for_status() + self.client.update_weights(asdict(update_info)) diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 5838a8ba8a1d..0ba44f3486e8 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -168,36 +168,41 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None: "Call init_transfer_engine() first." ) - if update_info.packed: - # Build iterator of (name, (shape, dtype)) from update_info - def state_dict_info_iterator(): + from vllm.model_executor.model_loader.mtp_validation import ( + disable_mtp_completeness_check, + ) + + with disable_mtp_completeness_check(): + if update_info.packed: + # Build iterator of (name, (shape, dtype)) from update_info + def state_dict_info_iterator(): + for name, dtype_name, shape in zip( + update_info.names, update_info.dtype_names, update_info.shapes + ): + dtype = getattr(torch, dtype_name) + yield (name, (shape, dtype)) + + packed_nccl_broadcast_consumer( + iterator=state_dict_info_iterator(), + group=self.model_update_group, + src=0, + post_unpack_func=self.model.load_weights, + buffer_size_bytes=update_info.packed_buffer_size_bytes, + num_buffers=update_info.packed_num_buffers, + device=self.device, + ) + else: + # Use simple one-by-one broadcasting for name, dtype_name, shape in zip( update_info.names, update_info.dtype_names, update_info.shapes ): dtype = getattr(torch, dtype_name) - yield (name, (shape, dtype)) - - packed_nccl_broadcast_consumer( - iterator=state_dict_info_iterator(), - group=self.model_update_group, - src=0, - post_unpack_func=self.model.load_weights, - buffer_size_bytes=update_info.packed_buffer_size_bytes, - num_buffers=update_info.packed_num_buffers, - device=self.device, - ) - else: - # Use simple one-by-one broadcasting - for name, dtype_name, shape in zip( - update_info.names, update_info.dtype_names, update_info.shapes - ): - dtype = getattr(torch, dtype_name) - weight = torch.empty(shape, dtype=dtype, device=self.device) - self.model_update_group.broadcast( - weight, src=0, stream=torch.cuda.current_stream() - ) - self.model.load_weights([(name, weight)]) - del weight + weight = torch.empty(shape, dtype=dtype, device=self.device) + self.model_update_group.broadcast( + weight, src=0, stream=torch.cuda.current_stream() + ) + self.model.load_weights([(name, weight)]) + del weight def shutdown(self) -> None: if self.model_update_group is not None: diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index ce42c691f49b..b34c58d53d76 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -10,8 +10,7 @@ import torch from torch.multiprocessing.reductions import reduce_tensor -# Default values for packed tensor configuration. -# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights. +# Default values for packed tensor transfer. DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB DEFAULT_PACKED_NUM_BUFFERS = 2 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 72dfc28688d6..219f8ffb082a 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -42,6 +42,7 @@ DiffusionConfig, ECTransferConfig, EPLBConfig, + FaultToleranceConfig, KernelConfig, KVEventsConfig, KVTransferConfig, @@ -76,7 +77,7 @@ from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend from vllm.config.load import SafetensorsLoadStrategy from vllm.config.lora import MaxLoRARanks -from vllm.config.mamba import MambaBackendEnum +from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm from vllm.config.model import ( ConvertOption, HfOverrides, @@ -85,7 +86,12 @@ RunnerOption, TokenizerMode, ) -from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MMTensorIPC +from vllm.config.multimodal import ( + MMCacheType, + MMEncoderTPMode, + MMHasherAlgorithm, + MMTensorIPC, +) from vllm.config.observability import DetailedTraceModules from vllm.config.parallel import ( All2AllBackend, @@ -101,10 +107,7 @@ from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins from vllm.ray.lazy_utils import is_in_ray_actor, is_ray_initialized -from vllm.transformers_utils.config import ( - is_interleaved, - maybe_override_with_speculators, -) +from vllm.transformers_utils.config import maybe_override_with_speculators from vllm.transformers_utils.repo_utils import get_model_path from vllm.transformers_utils.utils import is_cloud_storage from vllm.utils.argparse_utils import ( @@ -249,17 +252,19 @@ def get_type_hints(type_hint: TypeHint) -> set[TypeHint]: NEEDS_HELP = ( any("--help" in arg for arg in sys.argv) # vllm SUBCOMMAND --help - or (argv0 := sys.argv[0]).endswith("mkdocs") # mkdocs SUBCOMMAND - or argv0.endswith("mkdocs/__main__.py") # python -m mkdocs SUBCOMMAND + or "mkdocs" in sys.modules # mkdocs SUBCOMMAND ) def _maybe_add_docs_url(cls: Any) -> str: """Generate API docs URL for a vllm config class.""" - if not cls.__module__.startswith("vllm.config"): + import vllm.config + + name = cls.__name__ + if getattr(vllm.config, name, None) is not cls: return "" version = f"v{VLLM_VERSION}" if "dev" not in VLLM_VERSION else "latest" - return f"\n\nAPI docs: https://docs.vllm.ai/en/{version}/api/vllm/config/#vllm.config.{cls.__name__}" + return f"\n\nAPI docs: https://docs.vllm.ai/en/{version}/api/vllm/config/#vllm.config.{name}" def _expand_json_human_readable_numbers(val: str) -> str: @@ -527,8 +532,6 @@ class EngineArgs: kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None max_num_scheduled_tokens: int | None = None - max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills - max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold max_num_seqs: int | None = None max_logprobs: int = ModelConfig.max_logprobs @@ -568,6 +571,9 @@ class EngineArgs: mm_processor_cache_type: MMCacheType | None = ( MultiModalConfig.mm_processor_cache_type ) + mm_hasher_algorithm: MMHasherAlgorithm = get_field( + MultiModalConfig, "mm_hasher_algorithm" + ) mm_shm_cache_max_object_size_mb: int = ( MultiModalConfig.mm_shm_cache_max_object_size_mb ) @@ -588,6 +594,7 @@ class EngineArgs: renderer_num_workers: int = 1 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate + video_pruning_method: str = MultiModalConfig.video_pruning_method mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc mm_ipc_gpu_memory_gb: float = MultiModalConfig.mm_ipc_gpu_memory_gb # LoRA fields @@ -697,8 +704,11 @@ class EngineArgs: mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len + use_replayssm: bool = CacheConfig.use_replayssm mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON + mamba_ssu_algorithm: MambaSSUAlgorithm | None = None enable_mamba_cache_stochastic_rounding: bool = ( MambaConfig.enable_stochastic_rounding ) @@ -722,6 +732,11 @@ class EngineArgs: optimization_level: OptimizationLevel = VllmConfig.optimization_level performance_mode: PerformanceMode = VllmConfig.performance_mode + fault_tolerance_config: FaultToleranceConfig = get_field( + ParallelConfig, "fault_tolerance_config" + ) + enable_fault_tolerance: bool = ParallelConfig.enable_fault_tolerance + kv_offloading_size: float | None = CacheConfig.kv_offloading_size kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend tokens_only: bool = False @@ -735,6 +750,7 @@ class EngineArgs: fail_on_environ_validation: bool = False gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None + kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None def __post_init__(self): # support `EngineArgs(compilation_config={...})` @@ -754,6 +770,16 @@ def __post_init__(self): self.weight_transfer_config = WeightTransferConfig( **self.weight_transfer_config ) + if isinstance(self.fault_tolerance_config, dict): + if not self.enable_fault_tolerance: + logger.warning( + "--fault-tolerance-config was passed. Fault tolerance is being " + "automatically enabled." + ) + self.enable_fault_tolerance = True + self.fault_tolerance_config = FaultToleranceConfig( + **self.fault_tolerance_config + ) if isinstance(self.ir_op_priority, dict): self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority) @@ -937,6 +963,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: description=MambaConfig.__doc__, ) mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"]) + mamba_group.add_argument( + "--mamba-ssu-algorithm", **mamba_kwargs["ssu_algorithm"] + ) mamba_group.add_argument( "--enable-mamba-cache-stochastic-rounding", **mamba_kwargs["enable_stochastic_rounding"], @@ -1151,6 +1180,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: parallel_group.add_argument( "--worker-extension-cls", **parallel_kwargs["worker_extension_cls"] ) + parallel_group.add_argument( + "--enable-fault-tolerance", **parallel_kwargs["enable_fault_tolerance"] + ) + parallel_group.add_argument( + "--fault-tolerance-config", **parallel_kwargs["fault_tolerance_config"] + ) # KV cache arguments cache_kwargs = get_kwargs(CacheConfig) @@ -1203,6 +1238,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) + cache_group.add_argument( + "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"] + ) + cache_group.add_argument("--use-replayssm", **cache_kwargs["use_replayssm"]) cache_group.add_argument( "--kv-offloading-size", **cache_kwargs["kv_offloading_size"] ) @@ -1268,6 +1307,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: multimodal_group.add_argument( "--mm-processor-cache-type", **multimodal_kwargs["mm_processor_cache_type"] ) + multimodal_group.add_argument( + "--mm-hasher-algorithm", **multimodal_kwargs["mm_hasher_algorithm"] + ) multimodal_group.add_argument( "--mm-shm-cache-max-object-size-mb", **multimodal_kwargs["mm_shm_cache_max_object_size_mb"], @@ -1308,6 +1350,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: multimodal_group.add_argument( "--video-pruning-rate", **multimodal_kwargs["video_pruning_rate"] ) + multimodal_group.add_argument( + "--video-pruning-method", + **multimodal_kwargs["video_pruning_method"], + ) multimodal_group.add_argument( "--mm-tensor-ipc", **multimodal_kwargs["mm_tensor_ipc"] ) @@ -1440,13 +1486,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "default": None, }, ) - scheduler_group.add_argument( - "--max-num-partial-prefills", **scheduler_kwargs["max_num_partial_prefills"] - ) - scheduler_group.add_argument( - "--max-long-partial-prefills", - **scheduler_kwargs["max_long_partial_prefills"], - ) scheduler_group.add_argument( "--long-prefill-token-threshold", **scheduler_kwargs["long_prefill_token_threshold"], @@ -1614,12 +1653,20 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: default=None, help="Select GDN prefill backend.", ) + parser.add_argument( + "--kda-prefill-backend", + dest="kda_prefill_backend", + choices=["auto", "triton", "flashkda"], + default=None, + help="Select KDA prefill backend.", + ) return parser @classmethod def from_cli_args(cls, args: argparse.Namespace): # Get the list of attributes of this dataclass. attrs = [attr.name for attr in dataclasses.fields(cls)] + # Set the attributes from the parsed arguments. engine_args = cls( **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)} @@ -1680,6 +1727,7 @@ def create_model_config(self) -> ModelConfig: mm_processor_cache_gb=self.mm_processor_cache_gb, mm_processor_cache_type=self.mm_processor_cache_type, mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb, + mm_hasher_algorithm=self.mm_hasher_algorithm, mm_encoder_only=self.mm_encoder_only, mm_encoder_tp_mode=self.mm_encoder_tp_mode, mm_encoder_attn_backend=self.mm_encoder_attn_backend, @@ -1696,6 +1744,7 @@ def create_model_config(self) -> ModelConfig: override_attention_dtype=self.override_attention_dtype, logits_processors=self.logits_processors, video_pruning_rate=self.video_pruning_rate, + video_pruning_method=self.video_pruning_method, mm_tensor_ipc=self.mm_tensor_ipc, mm_ipc_gpu_memory_gb=self.mm_ipc_gpu_memory_gb, io_processor_plugin=self.io_processor_plugin, @@ -1886,7 +1935,8 @@ def create_engine_config( self._set_default_chunked_prefill_and_prefix_caching_args(model_config) self._set_default_reasoning_config_args() sliding_window: int | None = None - if not is_interleaved(model_config.hf_text_config): + layer_types = getattr(model_config.hf_text_config, "layer_types", None) + if layer_types is None or all(lt == "sliding_attention" for lt in layer_types): # Only set CacheConfig.sliding_window if the model is all sliding # window. Otherwise CacheConfig.sliding_window will override the # global layers in interleaved sliding window models. @@ -1919,6 +1969,8 @@ def create_engine_config( mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, + replayssm_buffer_len=self.replayssm_buffer_len, + use_replayssm=self.use_replayssm, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, ) @@ -1964,12 +2016,21 @@ def create_engine_config( assert not headless or not self.data_parallel_hybrid_lb, ( "data_parallel_hybrid_lb is not applicable in headless mode" ) - assert not (self.data_parallel_hybrid_lb and self.data_parallel_external_lb), ( - "data_parallel_hybrid_lb and data_parallel_external_lb cannot both be True." - ) - assert self.data_parallel_backend == "mp" or self.nnodes == 1, ( - "nnodes > 1 is only supported with data_parallel_backend=mp" - ) + if self.data_parallel_hybrid_lb and self.data_parallel_external_lb: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-hybrid-lb` and " + "`--data-parallel-external-lb` cannot be enabled together. " + "Enable only one load-balancing mode." + ) + if self.nnodes > 1 and self.data_parallel_backend != "mp": + raise ValueError( + "Invalid data-parallel launch options: " + f"`--nnodes {self.nnodes}` requires " + "`--data-parallel-backend mp`; got " + f"`--data-parallel-backend {self.data_parallel_backend}`. " + "Use the MP backend or set `--nnodes 1`." + ) inferred_data_parallel_rank = 0 if self.nnodes > 1: world_size = ( @@ -1980,13 +2041,22 @@ def create_engine_config( world_size_within_dp = ( self.pipeline_parallel_size * self.tensor_parallel_size ) + if world_size % self.nnodes != 0: + raise ValueError( + "Invalid data-parallel launch options: " + f"`--nnodes {self.nnodes}` must evenly divide the total " + f"world size ({world_size}). Adjust `--nnodes`, " + "`--data-parallel-size`, `--pipeline-parallel-size`, or " + "`--tensor-parallel-size`." + ) + if not 0 <= self.node_rank < self.nnodes: + raise ValueError( + "Invalid data-parallel launch options: `--node-rank` must " + f"be between 0 and {self.nnodes - 1}; got " + f"`--node-rank {self.node_rank}`. Set it to this node's " + "zero-based index." + ) local_world_size = world_size // self.nnodes - assert world_size % self.nnodes == 0, ( - f"world_size={world_size} must be divisible by nnodes={self.nnodes}." - ) - assert self.node_rank < self.nnodes, ( - f"node_rank={self.node_rank} must be less than nnodes={self.nnodes}." - ) inferred_data_parallel_rank = ( self.node_rank * local_world_size ) // world_size_within_dp @@ -2005,6 +2075,12 @@ def create_engine_config( data_parallel_external_lb = ( self.data_parallel_external_lb or self.data_parallel_rank is not None ) + if self.enable_fault_tolerance and not data_parallel_external_lb: + raise ValueError( + "Fault tolerance requires external load balancer mode " + "(--data-parallel-external-lb or --data-parallel-rank). " + "Internal LB mode is not supported." + ) if ( self.data_parallel_size > 1 and data_parallel_external_lb @@ -2017,14 +2093,21 @@ def create_engine_config( ) # Local DP rank = 1, use pure-external LB. if data_parallel_external_lb: - assert self.data_parallel_rank is not None, ( - "data_parallel_rank or node_rank must be specified if " - "data_parallel_external_lb is enable." - ) - assert self.data_parallel_size_local in (1, None), ( - "data_parallel_size_local must be 1 or None when data_parallel_rank " - "is set" - ) + if self.data_parallel_rank is None: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-external-lb` requires a data-parallel " + "rank. Set `--data-parallel-rank`, or set " + "`--data-parallel-size` greater than 1 and use `--nnodes` " + "with `--node-rank` so the rank can be inferred." + ) + if self.data_parallel_size_local not in (1, None): + raise ValueError( + "Invalid data-parallel launch options: an external " + "data-parallel rank requires `--data-parallel-size-local " + f"1`; got {self.data_parallel_size_local}. Set it to 1 or " + "omit it." + ) data_parallel_size_local = 1 # Use full external lb if we have local_size of 1. self.data_parallel_hybrid_lb = False @@ -2059,9 +2142,13 @@ def create_engine_config( self.node_rank, ) else: - assert not self.data_parallel_hybrid_lb, ( - "data_parallel_size_local must be set to use data_parallel_hybrid_lb." - ) + if self.data_parallel_hybrid_lb: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-hybrid-lb` requires " + "`--data-parallel-size-local`. Set it to the number of " + "data-parallel ranks on this node." + ) if self.data_parallel_backend == "ray" and ( envs.VLLM_RAY_DP_PACK_STRATEGY == "span" @@ -2151,6 +2238,8 @@ def create_engine_config( _api_process_count=self._api_process_count, _api_process_rank=self._api_process_rank, assigned_physical_gpu_ids=self._resolve_device_ids(), + enable_fault_tolerance=self.enable_fault_tolerance, + fault_tolerance_config=self.fault_tolerance_config, numa_bind=self.numa_bind, numa_bind_nodes=self.numa_bind_nodes, numa_bind_cpus=self.numa_bind_cpus, @@ -2190,8 +2279,6 @@ def create_engine_config( is_encoder_decoder=model_config.is_encoder_decoder, policy=self.scheduling_policy, scheduler_cls=self.scheduler_cls, - max_num_partial_prefills=self.max_num_partial_prefills, - max_long_partial_prefills=self.max_long_partial_prefills, long_prefill_token_threshold=self.long_prefill_token_threshold, scheduler_reserve_full_isl=self.scheduler_reserve_full_isl, watermark=self.watermark, @@ -2278,6 +2365,8 @@ def create_engine_config( mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()] else: mamba_config.backend = self.mamba_backend + if self.mamba_ssu_algorithm is not None: + mamba_config.ssu_algorithm = self.mamba_ssu_algorithm if self.enable_mamba_cache_stochastic_rounding: mamba_config.enable_stochastic_rounding = ( self.enable_mamba_cache_stochastic_rounding @@ -2286,6 +2375,7 @@ def create_engine_config( mamba_config.stochastic_rounding_philox_rounds = ( self.mamba_cache_philox_rounds ) + mamba_config.validate_ssu_algorithm() # Kernel config overrides kernel_config = copy.deepcopy(self.kernel_config) @@ -2368,6 +2458,8 @@ def create_engine_config( if self.gdn_prefill_backend is not None: self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend + if self.kda_prefill_backend is not None: + self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend config = VllmConfig( model_config=model_config, @@ -2402,14 +2494,6 @@ def create_engine_config( def _check_feature_supported(self): """Raise an error if the feature is not supported.""" - # No Concurrent Partial Prefills so far. - if ( - self.max_num_partial_prefills != SchedulerConfig.max_num_partial_prefills - or self.max_long_partial_prefills - != SchedulerConfig.max_long_partial_prefills - ): - _raise_unsupported_error(feature_name="Concurrent Partial Prefill") - if self.pipeline_parallel_size > 1: supports_pp = getattr( self.distributed_executor_backend, "supports_pp", False diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index c54123bea9e5..5a9b9f96d2c4 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -20,6 +20,7 @@ from vllm.tasks import SupportedTask from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult if TYPE_CHECKING: from vllm.v1.engine import PauseMode @@ -234,6 +235,16 @@ async def collective_rpc( """Perform a collective RPC call to the given path.""" raise NotImplementedError + async def handle_fault( + self, fault_tolerance_request: FaultToleranceRequest + ) -> FaultToleranceResult: + """send fault tolerance instruction to the engine""" + raise NotImplementedError + + async def get_status(self): + """Get fault tolerance status of all engines.""" + raise NotImplementedError + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: """Get supported tasks""" raise NotImplementedError @@ -256,6 +267,14 @@ async def update_weights(self, request: WeightTransferUpdateRequest) -> None: """Batched weight update for RL training.""" raise NotImplementedError - async def finish_weight_update(self) -> None: - """Finish the current weight update.""" + async def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + raise NotImplementedError + + async def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + raise NotImplementedError + + async def get_weight_version(self) -> str: + """Return the latest committed weight version.""" raise NotImplementedError diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index a470ab654094..14c3e4aca7a4 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -133,6 +133,18 @@ class AnthropicMessagesRequest(BaseModel): top_p: float | None = None # vLLM-specific fields that are not in Anthropic spec + cache_salt: str | None = Field( + default=None, + min_length=1, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters used for disaggregated serving.", diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index d516a7530b97..7c3b0597b5c4 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -486,6 +486,7 @@ def _build_base_request( temperature=anthropic_request.temperature, top_p=anthropic_request.top_p, top_k=anthropic_request.top_k, + cache_salt=anthropic_request.cache_salt, kv_transfer_params=anthropic_request.kv_transfer_params, ec_transfer_params=anthropic_request.ec_transfer_params, chat_template_kwargs=anthropic_request.chat_template_kwargs, diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index c89e9fa79d0e..2f8e11714bc9 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -467,7 +467,7 @@ def _merge_embeds( first_keys = set(data_items[0].keys()) if any(set(item.keys()) != first_keys for item in data_items[1:]): - raise ValueError( + raise VLLMValidationError( "All dictionaries in the list of embeddings must have the same keys." ) @@ -746,9 +746,15 @@ def _resolve_items( modality requires a processor, enforced by the guard below. """ if "image" in items_by_modality and "image_embeds" in items_by_modality: - raise ValueError("Mixing raw image and embedding inputs is not allowed") + raise VLLMValidationError( + "Mixing raw image and embedding inputs is not allowed", + parameter="image_embeds", + ) if "audio" in items_by_modality and "audio_embeds" in items_by_modality: - raise ValueError("Mixing raw audio and embedding inputs is not allowed") + raise VLLMValidationError( + "Mixing raw audio and embedding inputs is not allowed", + parameter="audio_embeds", + ) # `prompt_embeds` bypasses HF MM processors. Every other modality requires one. processor_modalities = items_by_modality.keys() - {"prompt_embeds"} if processor_modalities and mm_processor is None: @@ -945,16 +951,19 @@ def __init__( super().__init__() self._tracker = tracker + self._mm_processor_kwargs = mm_processor_kwargs - self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load( + @cached_property + def _connector(self) -> MediaConnector: + # Connector setup may probe VLLM_MEDIA_CACHE. Defer it until a request + # actually contains media so text-only parsing never blocks on that I/O. + return MEDIA_CONNECTOR_REGISTRY.load( envs.VLLM_MEDIA_CONNECTOR, - media_io_kwargs=tracker.media_io_kwargs, - allowed_local_media_path=tracker.allowed_local_media_path, - allowed_media_domains=tracker.allowed_media_domains, + media_io_kwargs=self._tracker.media_io_kwargs, + allowed_local_media_path=self._tracker.allowed_local_media_path, + allowed_media_domains=self._tracker.allowed_media_domains, ) - self._mm_processor_kwargs = mm_processor_kwargs - @property def model_config(self) -> ModelConfig: return self._tracker.model_config @@ -968,7 +977,9 @@ def parse_prompt_embeds(self, data: str) -> None: `tensor.shape[0]` placeholder tokens after tokenization. """ if not self.model_config.enable_prompt_embeds: - raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR) + raise VLLMValidationError( + _ENABLE_PROMPT_EMBEDS_ERROR, parameter="prompt_embeds" + ) tensor = safe_load_prompt_embeds(self.model_config, data.encode()) self._tracker.add("prompt_embeds", (tensor, None)) @@ -987,8 +998,9 @@ def parse_image_embeds( ) -> None: mm_config = self.model_config.get_multimodal_config() if not mm_config.enable_mm_embeds: - raise ValueError( - "You must set `--enable-mm-embeds` to input `image_embeds`" + raise VLLMValidationError( + "You must set `--enable-mm-embeds` to input `image_embeds`", + parameter="image_embeds", ) if isinstance(image_embeds, dict): @@ -1014,8 +1026,9 @@ def parse_audio_embeds( ) -> None: mm_config = self.model_config.get_multimodal_config() if not mm_config.enable_mm_embeds: - raise ValueError( - "You must set `--enable-mm-embeds` to input `audio_embeds`" + raise VLLMValidationError( + "You must set `--enable-mm-embeds` to input `audio_embeds`", + parameter="audio_embeds", ) if isinstance(audio_embeds, dict): @@ -1093,13 +1106,18 @@ def __init__( super().__init__() self._tracker = tracker - self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load( + self._mm_processor_kwargs: dict[str, Any] | None = mm_processor_kwargs + + @cached_property + def _connector(self) -> MediaConnector: + # Connector setup may probe VLLM_MEDIA_CACHE. Defer it until a request + # actually contains media so text-only parsing never blocks on that I/O. + return MEDIA_CONNECTOR_REGISTRY.load( envs.VLLM_MEDIA_CONNECTOR, - media_io_kwargs=tracker.media_io_kwargs, - allowed_local_media_path=tracker.allowed_local_media_path, - allowed_media_domains=tracker.allowed_media_domains, + media_io_kwargs=self._tracker.media_io_kwargs, + allowed_local_media_path=self._tracker.allowed_local_media_path, + allowed_media_domains=self._tracker.allowed_media_domains, ) - self._mm_processor_kwargs: dict[str, Any] | None = mm_processor_kwargs @property def model_config(self) -> ModelConfig: @@ -1117,7 +1135,9 @@ def parse_prompt_embeds(self, data: str) -> None: thread-pool executor via `safe_load_prompt_embeds_async`. """ if not self.model_config.enable_prompt_embeds: - raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR) + raise VLLMValidationError( + _ENABLE_PROMPT_EMBEDS_ERROR, parameter="prompt_embeds" + ) self._tracker.add( "prompt_embeds", partial(self._load_prompt_embeds_async, data.encode()) @@ -1151,25 +1171,35 @@ def parse_image_embeds( ) -> None: mm_config = self.model_config.get_multimodal_config() if not mm_config.enable_mm_embeds: - raise ValueError( - "You must set `--enable-mm-embeds` to input `image_embeds`" + raise VLLMValidationError( + "You must set `--enable-mm-embeds` to input `image_embeds`", + parameter="image_embeds", ) + placeholder = self._tracker.add( + "image_embeds", + partial(self._image_embeds_with_uuid_async, image_embeds, uuid), + ) + self._add_placeholder("image", placeholder) + + async def _image_embeds_with_uuid_async( + self, + image_embeds: str | dict[str, str] | None, + uuid: str | None, + ): if isinstance(image_embeds, dict): - embeds = { - k: self._connector.fetch_image_embedding(v) - for k, v in image_embeds.items() - } + tensors = await asyncio.gather( + *( + self._connector.fetch_image_embedding_async(v) + for v in image_embeds.values() + ) + ) + embeds = dict(zip(image_embeds, tensors)) elif isinstance(image_embeds, str): - embedding = self._connector.fetch_image_embedding(image_embeds) - embeds = embedding + embeds = await self._connector.fetch_image_embedding_async(image_embeds) else: embeds = None - - placeholder = self._tracker.add( - "image_embeds", partial(self._item_with_uuid_async, embeds, uuid) - ) - self._add_placeholder("image", placeholder) + return embeds, uuid def parse_audio_embeds( self, @@ -1178,25 +1208,35 @@ def parse_audio_embeds( ) -> None: mm_config = self.model_config.get_multimodal_config() if not mm_config.enable_mm_embeds: - raise ValueError( - "You must set `--enable-mm-embeds` to input `audio_embeds`" + raise VLLMValidationError( + "You must set `--enable-mm-embeds` to input `audio_embeds`", + parameter="audio_embeds", ) + placeholder = self._tracker.add( + "audio_embeds", + partial(self._audio_embeds_with_uuid_async, audio_embeds, uuid), + ) + self._add_placeholder("audio", placeholder) + + async def _audio_embeds_with_uuid_async( + self, + audio_embeds: str | dict[str, str] | None, + uuid: str | None, + ): if isinstance(audio_embeds, dict): - embeds = { - k: self._connector.fetch_audio_embedding(v) - for k, v in audio_embeds.items() - } + tensors = await asyncio.gather( + *( + self._connector.fetch_audio_embedding_async(v) + for v in audio_embeds.values() + ) + ) + embeds = dict(zip(audio_embeds, tensors)) elif isinstance(audio_embeds, str): - embedding = self._connector.fetch_audio_embedding(audio_embeds) - embeds = embedding + embeds = await self._connector.fetch_audio_embedding_async(audio_embeds) else: embeds = None - - placeholder = self._tracker.add( - "audio_embeds", partial(self._item_with_uuid_async, embeds, uuid) - ) - self._add_placeholder("audio", placeholder) + return embeds, uuid def parse_image_pil( self, @@ -1414,7 +1454,7 @@ def _get_full_multimodal_text_prompt( interleave_strings, ) logger.debug("Input prompt: %s", text_prompt) - raise ValueError( + raise VLLMValidationError( f"Found more '{placeholder}' placeholders in input prompt than " "actual multimodal data items." ) @@ -1592,10 +1632,14 @@ def _parse_chat_message_content_mm_part( tool_reference = tool_reference_params.get("name", None) return "tool_reference", tool_reference # Raise an error if no 'type' or direct URL is found. - raise ValueError("Missing 'type' field in multimodal part.") + raise VLLMValidationError( + "Missing 'type' field in multimodal part.", parameter="type" + ) if not isinstance(part_type, str): - raise ValueError("Invalid 'type' field in multimodal part.") + raise VLLMValidationError( + "Invalid 'type' field in multimodal part.", parameter="type" + ) return part_type, "unknown part_type content" @@ -1658,7 +1702,7 @@ def _reject_reserved_placeholder_in_text(text: str, model_config: ModelConfig) - caller move or inject splice positions via plain text content. """ if model_config.enable_prompt_embeds and PROMPT_EMBEDS_PLACEHOLDER_TOKEN in text: - raise ValueError( + raise VLLMValidationError( _RESERVED_PLACEHOLDER_IN_TEXT_ERROR.format( token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN ) @@ -1732,7 +1776,9 @@ def _parse_chat_message_content_part( modality = "audio" elif part_type == "prompt_embeds": if not content: - raise ValueError(_PROMPT_EMBEDS_MISSING_DATA_ERROR) + raise VLLMValidationError( + _PROMPT_EMBEDS_MISSING_DATA_ERROR, parameter="prompt_embeds" + ) mm_parser.parse_prompt_embeds(cast(str, content)) modality = "prompt_embeds" elif part_type == "audio_url": @@ -1992,13 +2038,19 @@ def get_history_tool_calls_cnt(conversation: list[ConversationMessage]): return idx -_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25") +_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25", "kimi_k3") def get_tool_call_id_type(model_config: ModelConfig) -> str: """Return the tool-call ID type for a given model configuration.""" hf_overrides = getattr(model_config, "hf_overrides", None) - if model_config.hf_text_config.model_type in _KIMI_MODEL_TYPES or ( + hf_config = getattr(model_config, "hf_config", None) + hf_text_config = getattr(model_config, "hf_text_config", None) + model_types = ( + getattr(hf_config, "model_type", None), + getattr(hf_text_config, "model_type", None), + ) + if any(model_type in _KIMI_MODEL_TYPES for model_type in model_types) or ( isinstance(hf_overrides, dict) and hf_overrides.get("model_type") in _KIMI_MODEL_TYPES ): diff --git a/vllm/entrypoints/cli/benchmark/main.py b/vllm/entrypoints/cli/benchmark/main.py index 1afac64b148d..9ea499870919 100644 --- a/vllm/entrypoints/cli/benchmark/main.py +++ b/vllm/entrypoints/cli/benchmark/main.py @@ -2,18 +2,38 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import os import sys import typing +from vllm import envs from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.logger import init_logger if typing.TYPE_CHECKING: from vllm.utils.argparse_utils import FlexibleArgumentParser else: FlexibleArgumentParser = argparse.ArgumentParser +logger = init_logger(__name__) + + +def maybe_exec_rust_bench() -> None: + if sys.argv[1:3] != ["bench", "serve"] or not envs.VLLM_USE_RUST_BENCH: + return + + rust_cli = envs.VLLM_RUST_FRONTEND_PATH + if rust_cli is None: + raise RuntimeError( + "VLLM_USE_RUST_BENCH=1 requires VLLM_RUST_FRONTEND_PATH " + "to resolve to the vllm-rs binary." + ) + + logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli) + os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]]) + def _import_bench_subcommand_modules() -> None: # Imported lazily so `BenchmarkSubcommandBase` subclasses register only diff --git a/vllm/entrypoints/cli/benchmark/mm_processor.py b/vllm/entrypoints/cli/benchmark/mm_processor.py index 26b93aacdc50..de0d62e743a8 100644 --- a/vllm/entrypoints/cli/benchmark/mm_processor.py +++ b/vllm/entrypoints/cli/benchmark/mm_processor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import inspect from vllm.benchmarks.mm_processor import add_cli_args, main from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase @@ -8,13 +9,59 @@ class BenchmarkMMProcessorSubcommand(BenchmarkSubcommandBase): - """The `mm-processor` subcommand for `vllm bench`.""" + r"""`vllm bench mm-processor` profiles the multimodal input processor pipeline of + vision-language models. It measures per-stage latency from the HuggingFace + processor through to the encoder forward pass, helping you identify + preprocessing bottlenecks and understand how different image resolutions or + item counts affect end-to-end request time. + + The benchmark supports two data sources: synthetic random multimodal inputs + (`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before + measurement to ensure stable results. + + ## Quick Start + + ```bash + vllm bench mm-processor \ + --model Qwen/Qwen2-VL-7B-Instruct \ + --dataset-name random-mm \ + --num-prompts 50 \ + --random-input-len 300 \ + --random-output-len 40 \ + --random-mm-base-items-per-request 2 \ + --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \ + --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' + ``` + + ## Measured Stages + + | Stage | Description | + | ----- | ----------- | + | `get_mm_hashes_secs` | Time spent hashing multimodal inputs | + | `get_cache_missing_items_secs` | Time spent looking up the processor cache | + | `apply_hf_processor_secs` | Time spent in the HuggingFace processor | + | `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs | + | `apply_prompt_updates_secs` | Time spent updating prompt tokens | + | `preprocessor_total_secs` | Total preprocessing time | + | `encoder_forward_secs` | Time spent in the encoder model forward pass | + | `num_encoder_calls` | Number of encoder invocations per request | + + The benchmark also reports end-to-end latency (TTFT + decode time) per + request. Use `--metric-percentiles` to select which percentiles to report + (default: p99) and `--output-json` to save results. + + For more examples (HF datasets, warmup, JSON output), see the + [multimodal processor benchmark guide](https://docs.vllm.ai/en/latest/benchmarking/cli/#multimodal-processor-benchmark). + """ name = "mm-processor" help = "Benchmark multimodal processor latency across different configurations." @classmethod def add_cli_args(cls, parser: FlexibleArgumentParser) -> None: + # The class docstring is the page overview / `--help` description. + if cls.__doc__: + parser.description = inspect.cleandoc(cls.__doc__) add_cli_args(parser) @staticmethod diff --git a/vllm/entrypoints/cli/launch.py b/vllm/entrypoints/cli/launch.py index d13a0c67c83c..91d4bf094b7d 100644 --- a/vllm/entrypoints/cli/launch.py +++ b/vllm/entrypoints/cli/launch.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import inspect import signal import uvloop @@ -36,9 +37,12 @@ class LaunchSubcommandBase(CLISubcommand): def add_cli_args(cls, parser: FlexibleArgumentParser) -> None: """Add the CLI arguments to the parser. - By default, adds the standard vLLM serving arguments. + By default, uses the subcommand's docstring as the description and adds + the standard vLLM serving arguments. Subclasses can override to add component-specific arguments. """ + if cls.__doc__: + parser.description = inspect.cleandoc(cls.__doc__) make_arg_parser(parser) @staticmethod @@ -47,7 +51,17 @@ def cmd(args: argparse.Namespace) -> None: class RenderSubcommand(LaunchSubcommandBase): - """The `render` subcommand for `vllm launch`.""" + """`vllm launch render` starts a GPU-less rendering server for preprocessing + and postprocessing only. + + ```bash + vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 + ``` + + This command reuses the standard serving parser, so model, frontend, + networking, and related CLI options follow the same conventions as + [`vllm serve`](https://docs.vllm.ai/en/latest/cli/serve/). + """ name = "render" help = "Launch a GPU-less rendering server (preprocessing and postprocessing only)." @@ -93,7 +107,6 @@ def subparser_init( cmd_subparser = launch_subparsers.add_parser( cmd_cls.name, help=cmd_cls.help, - description=cmd_cls.help, usage=f"vllm {self.name} {cmd_cls.name} [options]", ) cmd_subparser.set_defaults(launch_command=cmd_cls.cmd) diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index fe0b339b3ed7..3dc69dd3ad2b 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -54,6 +54,8 @@ def main(): logger.info("Delegating entrypoint handling to vllm-omni") omni_main() else: + vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench() + # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default if len(sys.argv) > 1 and sys.argv[1] == "bench": logger.debug( diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index d5e9b2bc874a..e8f72ab230a2 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -58,6 +58,10 @@ def cmd(args: argparse.Namespace) -> None: uvloop.run(serve_grpc(args)) return + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) + if args.headless: if args.api_server_count is not None and args.api_server_count > 0: raise ValueError( @@ -103,7 +107,7 @@ def cmd(args: argparse.Namespace) -> None: # - Hybrid LB: Use local DP size (internal LB for local ranks only) # - Internal LB: Use full DP size if args.api_server_count is None: - if is_multi_port or is_external_lb or envs.VLLM_RUST_FRONTEND_PATH: + if is_multi_port or is_external_lb or rust_frontend_path: args.api_server_count = 1 elif is_hybrid_lb: args.api_server_count = args.data_parallel_size_local or 1 @@ -120,7 +124,7 @@ def cmd(args: argparse.Namespace) -> None: "Defaulting api_server_count to data_parallel_size (%d).", args.api_server_count, ) - elif envs.VLLM_RUST_FRONTEND_PATH and args.api_server_count > 1: + elif rust_frontend_path and args.api_server_count > 1: logger.warning( "Ignoring --api-server-count=%d when using rust front-end process", args.api_server_count, @@ -140,7 +144,7 @@ def cmd(args: argparse.Namespace) -> None: run_dp_supervisor(args) elif args.api_server_count < 1: run_headless(args) - elif args.api_server_count > 1 or envs.VLLM_RUST_FRONTEND_PATH: + elif args.api_server_count > 1 or rust_frontend_path: run_multi_api_server(args) else: # Single API server (this process). @@ -256,7 +260,9 @@ def signal_handler(signum, frame): def run_multi_api_server(args: argparse.Namespace): assert not args.headless - rust_frontend_path = envs.VLLM_RUST_FRONTEND_PATH + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) num_api_servers: int = args.api_server_count assert num_api_servers > 0 @@ -320,9 +326,12 @@ def signal_handler(signum, frame): defer_api_server_ports=not (rust_frontend_path or is_ray_dp), ) - with launch_core_engines( - vllm_config, executor_class, log_stats, addresses, num_api_servers - ) as (local_engine_manager, coordinator, addresses, tensor_queue): + with launch_core_engines(vllm_config, executor_class, log_stats, addresses) as ( + local_engine_manager, + coordinator, + addresses, + tensor_queue, + ): stats_update_address = ( coordinator.get_stats_publish_address() if coordinator else None ) diff --git a/vllm/entrypoints/cohere/__init__.py b/vllm/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/entrypoints/cohere/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py new file mode 100644 index 000000000000..b532e60920a8 --- /dev/null +++ b/vllm/entrypoints/cohere/api_router.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FastAPI router for the Cohere Chat v2 API (``POST /cohere/v2/chat``). + +The Cohere v2 protocol models are sourced from the official ``cohere`` +Python SDK (``pip install cohere``). To keep that an *optional* +dependency for vLLM, the SDK-dependent imports - and the route handler +itself - are gated on a one-shot probe at module load. If the SDK isn't +installed, :func:`attach_router` becomes a no-op (with an info log) and +vLLM continues to boot normally. + +Even when the SDK is installed, :func:`attach_router` also requires +``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose +the route. This keeps non-Cohere deployments that pull in the SDK for +unrelated reasons (e.g. test dependencies) from accidentally exposing +the api. + +Note: the handler must live at module scope (not inside +``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves +the ``CohereChatV2Request`` body annotation against the module's +globals. Defining it locally inside ``attach_router`` would hide the +type from ``get_type_hints``, causing FastAPI to silently degrade the +body parameter into a query parameter and reject every request with +422. +""" + +import json +from http import HTTPStatus + +from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse +from starlette.middleware.base import BaseHTTPMiddleware + +import vllm.envs as envs +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + sanitize_message, + validate_json_request, + with_cancellation, +) +from vllm.logger import init_logger + +_COHERE_PATH_PREFIX = "/cohere/" + +logger = init_logger(__name__) + + +try: + import cohere # noqa: F401 -- dependency probe +except ImportError: + _SDK_AVAILABLE = False +else: + _SDK_AVAILABLE = True + + +if _SDK_AVAILABLE: + from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, + CohereError, + ) + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + + router = APIRouter() + + def _serving(request: Request) -> CohereServingChatV2 | None: + return getattr(request.app.state, "cohere_serving_chat_v2", None) + + def _request_id(raw_request: Request | None) -> str | None: + """Best-effort lookup of the active request id. + + Prefers the id the underlying chat handler stamped onto + ``raw_request.state.request_metadata`` (if it got that far before + failing), falling back to the ``X-Request-Id`` HTTP header. May + return ``None`` if neither is available, in which case the field + is omitted from the response. + """ + if raw_request is None: + return None + meta = getattr(raw_request.state, "request_metadata", None) + if meta is not None and getattr(meta, "request_id", None): + return meta.request_id + return raw_request.headers.get("X-Request-Id") + + def _error_response( + error: ErrorResponse, + raw_request: Request | None, + *, + fallback_status: int = HTTPStatus.BAD_REQUEST, + ) -> JSONResponse: + """Translate vLLM's internal error envelope into Cohere's shape.""" + info = error.error + status = info.code or fallback_status + return JSONResponse( + status_code=status, + content=CohereError( + message=sanitize_message(info.message), + id=_request_id(raw_request), + ).model_dump(exclude_none=True), + ) + + @router.post( + "/cohere/v2/chat", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, + HTTPStatus.NOT_FOUND.value: {"model": CohereError}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, + }, + ) + @with_cancellation + @load_aware_call + async def chat_v2(request: CohereChatV2Request, raw_request: Request): + handler = _serving(raw_request) + if handler is None: + return JSONResponse( + status_code=HTTPStatus.NOT_IMPLEMENTED.value, + content=CohereError( + message="The model does not support the Cohere v2 chat API.", + id=_request_id(raw_request), + ).model_dump(exclude_none=True), + ) + + try: + result = await handler.create_chat_v2(request, raw_request) + except Exception as e: # noqa: BLE001 - report as 500 for parity + logger.exception("Error in /cohere/v2/chat: %s", e) + return JSONResponse( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + content=CohereError( + message=sanitize_message(str(e)), + id=_request_id(raw_request), + ).model_dump(exclude_none=True), + ) + + match result: + case ErrorResponse(): + return _error_response(result, raw_request) + case CohereChatV2Response(): + return JSONResponse(content=result.model_dump(exclude_none=True)) + case _: + return StreamingResponse(content=result, media_type="text/event-stream") + + class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware): + """Rewrite vLLM error bodies into the Cohere ``{message, id}`` shape. + + The endpoint handler above already returns :class:`CohereError` for + errors it owns, but globally-registered exception handlers (e.g. + :func:`validation_exception_handler` for pydantic body errors, + :func:`http_exception_handler`, engine error handlers) fire *before* + the handler runs and produce vLLM's internal + ``ErrorResponse`` shape (``{"error": {"message": ...}}``). That + shape doesn't match the ``CohereError`` schema advertised on the + route's OpenAPI ``responses``, so clients (and schema-conformance + tests like ``test_openai_schema.py``) would see a mismatch on + those paths. This middleware normalises any error body on + ``/cohere/*`` responses to :class:`CohereError`. + """ + + async def dispatch(self, request: Request, call_next): + response = await call_next(request) + if not request.url.path.startswith(_COHERE_PATH_PREFIX): + return response + if response.status_code < 400: + return response + content_type = response.headers.get("content-type", "") + if not content_type.startswith("application/json"): + return response + + body = b"".join([chunk async for chunk in response.body_iterator]) + translated = _translate_vllm_error_body(body, request) + if translated is not None: + return translated + passthrough_headers = { + k: v + for k, v in response.headers.items() + if k.lower() != "content-length" + } + return Response( + content=body, + status_code=response.status_code, + headers=passthrough_headers, + media_type=content_type, + ) + + def _translate_vllm_error_body(raw: bytes, request: Request) -> JSONResponse | None: + """Translate a vLLM ``ErrorResponse`` body to a ``CohereError`` body. + + Returns ``None`` if ``raw`` does not match the vLLM error envelope + (which signals the middleware to pass the body through unchanged). + """ + try: + data = json.loads(raw) + except (json.JSONDecodeError, TypeError, ValueError): + return None + if not ( + isinstance(data, dict) + and isinstance(data.get("error"), dict) + and "message" in data["error"] + ): + return None + try: + err = ErrorResponse.model_validate(data) + except Exception: # noqa: BLE001 - malformed envelope; pass through + return None + return _error_response(err, request) + + +def attach_router(app: FastAPI) -> None: + """Register ``POST /cohere/v2/chat`` on ``app``. + + No-op when either: + + * the ``VLLM_ENABLE_COHERE_API`` env var isn't set to ``1``. The + Cohere v2 endpoint is opt-in because it carries Cohere-specific + request/response semantics (grounding citations, tool_plan, + PLAN/THINKING_CONTENT blocks) that are only meaningful when + serving a Cohere Command-family model. + * the optional ``cohere`` SDK isn't installed (the v2 protocol + models live there) + + The two skip paths log at different levels: an operator who set + ``VLLM_ENABLE_COHERE_API=1`` but forgot to install ``cohere`` sees + a WARNING (they explicitly asked for the endpoint and it's silently + absent), whereas the default-off skip logs at debug. + """ + enabled = envs.VLLM_ENABLE_COHERE_API + if not enabled: + logger.debug( + "VLLM_ENABLE_COHERE_API is not set; /cohere/v2/chat endpoint " + "disabled. Set VLLM_ENABLE_COHERE_API=1 to enable it." + ) + return + if not _SDK_AVAILABLE: + logger.warning( + "VLLM_ENABLE_COHERE_API=1 but the `cohere` SDK is not " + "installed; /cohere/v2/chat will not be exposed. Install " + "with `pip install cohere` to enable the endpoint." + ) + return + app.include_router(router) + app.add_middleware(CohereErrorEnvelopeMiddleware) diff --git a/vllm/entrypoints/cohere/cohere_chat_message.py b/vllm/entrypoints/cohere/cohere_chat_message.py new file mode 100644 index 000000000000..6d221adce3e2 --- /dev/null +++ b/vllm/entrypoints/cohere/cohere_chat_message.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM chat-protocol extensions for grounded (citation-carrying) models. + +This module keeps the OpenAI chat completion protocol classes +(:class:`ChatMessage` / :class:`DeltaMessage`) free of Cohere-specific fields. +It provides two things: + +* :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation + representation produced by the Cohere reasoning parser + (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by + :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere + SDK shape -- that conversion happens in the serving layer. +* :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage` + / :class:`DeltaMessage` subclasses that add a ``citations`` field. They + are only instantiated by the citation-aware serving handler + (:class:`vllm.entrypoints.cohere.serving.CohereServingChatV2`) and by the + Cohere reasoning parser's streaming path. Non-Cohere code paths continue + to construct plain :class:`ChatMessage` / :class:`DeltaMessage`. + +The module intentionally does *not* import the ``cohere`` Python SDK so the +reasoning parser (which is Cohere-model-specific but ships without the SDK +dependency) can import from here freely. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import Field, model_serializer + +from vllm.entrypoints.openai.chat_completion.protocol import ChatMessage +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + OpenAIBaseModel, +) + + +class CitationSource(OpenAIBaseModel): + """Source attribution for a :class:`Citation`. + + Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source + discriminator (``document`` or ``tool``); ``id`` is the citing + document/tool-output identifier; ``document`` and ``tool_output`` carry + the original payload that produced the citation. + + Sources are fully resolved by the reasoning parser at emit time + (see :func:`_melody_sources_to_vllm` in + :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the + ``POSITION_TO_SOURCE_KEY`` map forwarded through + ``chat_template_kwargs`` by + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`. + That means every instance that reaches the serving / wire layer + already has ``type`` / ``id`` populated; the serving layer only has + to coerce to :class:`cohere.types.Citation` and rewrite + ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning models. + """ + + type: Literal["document", "tool"] | None = None + id: str | None = None + document: dict[str, Any] | None = None + tool_output: dict[str, Any] | None = None + + +class Citation(OpenAIBaseModel): + """A citation grounding a span of generated text in source material. + + vLLM-internal representation used by the Cohere reasoning parser and + the Cohere v2 serving layer. This is not the on-wire Cohere SDK shape: + conversion to the SDK's ``cohere.types.Citation`` happens in + :mod:`vllm.entrypoints.cohere.serving`. + """ + + start: int | None = None + """Start character offset in the surrounding text content.""" + end: int | None = None + """End character offset (exclusive) in the surrounding text content.""" + text: str | None = None + """The cited text snippet.""" + sources: list[CitationSource] = Field(default_factory=list) + """Source documents / tool outputs that ground this citation.""" + content_index: int | None = None + """Index of the content block this citation refers to (when the message + has multiple content blocks).""" + type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None + """Which kind of content block this citation grounds: the user-visible + text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a + tool-plan block (``PLAN``). ``None`` means unspecified.""" + + +class CohereChatMessage(ChatMessage): + """:class:`ChatMessage` extension carrying grounding citations. + + Only instantiated by :class:`CohereServingChatV2` (and any other + :class:`OpenAIServingChat` subclass that opts in by overriding + ``_create_chat_message``). Regular OpenAI-compatible handlers keep + emitting plain :class:`ChatMessage` so their response schema is + unchanged. + + The response envelope declares ``message: SerializeAsAny[ChatMessage]``, + so pydantic serializes this subclass with its own schema (including + ``citations``) when it flows through + :class:`ChatCompletionResponseChoice`. + """ + + citations: list[Citation] | None = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + # ``mode="wrap"`` fully overrides (rather than chains) the + # parent's ``@model_serializer``, so we explicitly delegate via + # ``super()._serialize(handler)`` to preserve the parent's + # cleanup (e.g. stripping empty ``tool_calls``). Then we drop an + # unset ``citations`` field so the wire matches the OpenAI-style + # contract that optional vLLM extensions are omitted rather than + # serialized as ``null``. + data = super()._serialize(handler) + if not data.get("citations"): + data.pop("citations", None) + return data + + +class CohereDeltaMessage(DeltaMessage): + """:class:`DeltaMessage` extension carrying grounding citations for streaming. + + Emitted by + :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser` + on delta events whose payload includes citations. Non-Cohere parsers + return plain :class:`DeltaMessage`, so their streamed shape is + unchanged. + + The response envelope declares + ``delta: SerializeAsAny[DeltaMessage]``, so pydantic serializes this + subclass with its own schema when it flows through + :class:`ChatCompletionResponseStreamChoice`. + """ + + citations: list[Citation] | None = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + # See ``CohereChatMessage._serialize`` for why we delegate via + # ``super()`` instead of calling ``handler(self)`` directly. + data = super()._serialize(handler) + if not data.get("citations"): + data.pop("citations", None) + return data diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py new file mode 100644 index 000000000000..2850e8eddae8 --- /dev/null +++ b/vllm/entrypoints/cohere/protocol.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere Chat v2 API protocol. + +The bulk of the wire types come straight from the official ``cohere`` +Python SDK so we stay in sync with the upstream specification and +avoid re-declaring the schema. We only own three things locally: + +1. The top-level request body model (the SDK doesn't ship one — its + ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific + extensions (``kv_transfer_params`` / ``chat_template_kwargs``). +2. The non-streaming response envelope (the SDK exposes the message + shape via :class:`AssistantMessageResponse` but no full response + wrapper). +3. The streaming discriminated union (the SDK exports each event type + individually but not as a combined ``Annotated[Union[...], + discriminator]``). + +Importing this module pulls in the ``cohere`` package. The router that +mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still +boots without the SDK installed. + +See https://docs.cohere.com/reference/chat for the upstream spec. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from cohere import types as _sdk +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + ChatMessageV2, + ChatRequestSafetyMode, + Citation, + CitationOptions, + Document, + ResponseFormatV2, + SystemChatMessageV2, + Thinking, + ToolCallV2, + ToolChatMessageV2, + ToolV2, + UserChatMessageV2, +) +from pydantic import BaseModel, Field, field_validator + +# Re-export the SDK wire-format types alongside our local extensions so +# ``vllm.entrypoints.cohere.serving`` and friends can import everything +# they need from this module. +__all__ = [ + "AssistantChatMessageV2", + "AssistantMessageResponse", + "ChatMessageV2", + "Citation", + "CitationEndEvent", + "CitationOptions", + "CitationStartEvent", + "CohereChatV2Request", + "CohereChatV2Response", + "CohereError", + "CohereFinishReason", + "CohereLogprobItem", + "CohereUsage", + "CohereUsageBilledUnits", + "CohereUsageTokens", + "ContentDeltaEvent", + "ContentEndEvent", + "ContentStartEvent", + "Document", + "MessageEndEvent", + "MessageStartEvent", + "ResponseFormatV2", + "SystemChatMessageV2", + "Thinking", + "ToolCallDeltaEvent", + "ToolCallEndEvent", + "ToolCallStartEvent", + "ToolCallV2", + "ToolChatMessageV2", + "ToolPlanDeltaEvent", + "ToolV2", + "UserChatMessageV2", +] + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class CohereError(BaseModel): + """Top-level error body returned by ``/cohere/v2/chat`` error responses. + + Cohere's documented error schemas are uniform: ``{message, id}``. + """ + + message: str + id: str | None = None + + +# --------------------------------------------------------------------------- +# Tool choice / finish reasons +# --------------------------------------------------------------------------- +# +# These literals aren't first-class enums in the SDK but are documented at +# https://docs.cohere.com/reference/chat. We declare them here so the +# request/response models can validate them. + +CohereToolChoice = Literal["REQUIRED", "NONE"] + +CohereFinishReason = Literal[ + "COMPLETE", + "STOP_SEQUENCE", + "MAX_TOKENS", + "TOOL_CALL", + "ERROR", + "TIMEOUT", +] + + +# --------------------------------------------------------------------------- +# Request +# --------------------------------------------------------------------------- + + +class CohereChatV2Request(BaseModel): + """Cohere Chat v2 request body. + + Mirrors the schema documented at https://docs.cohere.com/reference/chat. + All structured fields delegate to the official SDK types so the body + schema stays in sync with the upstream spec. + """ + + model: str + messages: list[ChatMessageV2] + stream: bool | None = False + + # Tooling + tools: list[ToolV2] | None = None + strict_tools: bool | None = None + tool_choice: CohereToolChoice | None = None + + # Grounding + documents: list[str | Document] | None = None + citation_options: CitationOptions | None = None + + # Output + response_format: ResponseFormatV2 | None = None + safety_mode: ChatRequestSafetyMode | None = None + max_tokens: int | None = None + stop_sequences: list[str] | None = None + + # Sampling + temperature: float | None = None + seed: int | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + k: int | None = None + p: float | None = None + logprobs: bool | None = None + + # Reasoning + thinking: Thinking | None = None + + # Scheduling + priority: int | None = None + + # vLLM-specific extensions (not in Cohere spec). These mirror what the + # Anthropic and OpenAI surfaces already expose so V2 callers can reach + # the same engine knobs when needed. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, + description="KVTransfer parameters used for disaggregated serving.", + ) + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the template." + ), + ) + + @field_validator("model") + @classmethod + def _validate_model(cls, v: str) -> str: + if not v: + raise ValueError("model is required") + return v + + @field_validator("max_tokens") + @classmethod + def _validate_max_tokens(cls, v: int | None) -> int | None: + if v is not None and v < 0: + raise ValueError("max_tokens must be non-negative") + return v + + @field_validator("messages", mode="before") + @classmethod + def _normalize_message_roles(cls, v: Any) -> Any: + """Rewrite OpenAI-style ``developer`` roles to ``system``. + + Cohere's v2 ``ChatMessageV2`` discriminated union only admits + the four literal roles ``user`` / ``assistant`` / ``system`` / + ``tool``. OpenAI's ``developer`` role is documented as + high-priority system instructions, so we alias it onto + ``system`` *before* the SDK discriminator runs (otherwise it + rejects the message with a ``literal_error`` against each union + member). Mirrors ``_role_to_melody`` in the renderer so the + same alias is honoured no matter which surface the message + arrives through. + + ``mode="before"`` is required so the rewrite happens before the + ``list[ChatMessageV2]`` coercion runs the SDK's discriminated + union; a default-mode validator would never see ``developer`` + because validation would have already failed. On any + structural malformation (non-iterable input, items without a + dict-shaped ``role`` field, etc.) we hand ``v`` back unchanged + and let Pydantic's normal coercion surface a precise error. + """ + try: + return [ + {**msg, "role": "system"} + if msg.get("role", "").lower() == "developer" + else msg + for msg in v + ] + except (AttributeError, TypeError): + return v + + @field_validator("messages") + @classmethod + def _validate_messages(cls, v: list[ChatMessageV2]) -> list[ChatMessageV2]: + if not v: + raise ValueError("messages must contain at least one message") + return v + + +# --------------------------------------------------------------------------- +# Usage / Logprobs +# --------------------------------------------------------------------------- +# +# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage +# envelope is documented separately in the OpenAPI spec and we declare it +# here. + + +class CohereUsageBilledUnits(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + search_units: float | None = None + classifications: float | None = None + + +class CohereUsageTokens(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + + +class CohereUsage(BaseModel): + billed_units: CohereUsageBilledUnits | None = None + tokens: CohereUsageTokens | None = None + cached_tokens: float | None = None + + +class CohereLogprobItem(BaseModel): + text: str | None = None + token_ids: list[int] + logprobs: list[float] | None = None + + +# --------------------------------------------------------------------------- +# Non-streaming response +# --------------------------------------------------------------------------- + + +class CohereChatV2Response(BaseModel): + """Cohere Chat v2 non-streaming response body. + + Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in + the documented v2 response envelope (``id``, ``finish_reason``, + ``usage``, ``logprobs``). The single constructor in + :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for + supplying a non-empty ``id`` (falling back to a synthesized one if + the upstream response is missing it) to this model. + """ + + id: str + finish_reason: CohereFinishReason + message: AssistantMessageResponse + usage: CohereUsage | None = None + logprobs: list[CohereLogprobItem] | None = None + + # vLLM-specific extension. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, description="KVTransfer parameters." + ) + + +# --------------------------------------------------------------------------- +# Streaming events +# --------------------------------------------------------------------------- +# +# Cohere V2 streams a sequence of typed JSON events delivered as Server- +# Sent Events; each event's ``type`` field carries the discriminator. The +# SDK exposes a Pydantic model per event but none of them declare ``type`` +# as a field (the SDK relies on its own deserializer for discrimination), +# so a naive ``model_dump_json()`` would silently drop the discriminator +# and break clients that demux on ``type``. +# +# We therefore subclass each SDK event and bake the wire-format ``type`` +# string in as a ``Literal`` field with a default. ``model_dump()`` now +# emits ``type`` for free, and surfaces can simply construct the event +# class and serialize it -- no manual ``type`` parameter needed. +# +# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI +# ``StreamedChatResponseV2`` schema for the wire-format reference. + + +class MessageStartEvent(_sdk.ChatMessageStartEvent): + type: Literal["message-start"] = "message-start" + + +class ContentStartEvent(_sdk.ChatContentStartEvent): + type: Literal["content-start"] = "content-start" + + +class ContentDeltaEvent(_sdk.ChatContentDeltaEvent): + type: Literal["content-delta"] = "content-delta" + + +class ContentEndEvent(_sdk.ChatContentEndEvent): + type: Literal["content-end"] = "content-end" + + +class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent): + type: Literal["tool-plan-delta"] = "tool-plan-delta" + + +class ToolCallStartEvent(_sdk.ChatToolCallStartEvent): + type: Literal["tool-call-start"] = "tool-call-start" + + +class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent): + type: Literal["tool-call-delta"] = "tool-call-delta" + + +class ToolCallEndEvent(_sdk.ChatToolCallEndEvent): + type: Literal["tool-call-end"] = "tool-call-end" + + +class CitationStartEvent(_sdk.CitationStartEvent): + type: Literal["citation-start"] = "citation-start" + + +class CitationEndEvent(_sdk.CitationEndEvent): + type: Literal["citation-end"] = "citation-end" + + +class MessageEndEvent(_sdk.ChatMessageEndEvent): + type: Literal["message-end"] = "message-end" diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py new file mode 100644 index 000000000000..04180b5cfadc --- /dev/null +++ b/vllm/entrypoints/cohere/serving.py @@ -0,0 +1,1678 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py +"""Cohere Chat v2 API serving handler. + +Implements ``POST /cohere/v2/chat`` by translating the incoming Cohere v2 request +into a standard :class:`ChatCompletionRequest` and delegating to +:class:`OpenAIServingChat`. The actual prompt rendering is handled by +vLLM's renderer pipeline (``vllm.renderers``): + +- For Cohere Command-family models, set ``--tokenizer-mode cohere`` and the + :class:`vllm.renderers.cohere.CohereRenderer` will template the request + via the ``cohere_melody`` library (``render_cmd3`` / ``render_cmd4``) + and surface citations through the Cohere-scoped + :class:`CohereChatMessage` / :class:`CohereDeltaMessage` subclasses + (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`). +- For any other model the default Jinja-based :class:`HfRenderer` is used, + and the endpoint behaves as a plain v2-shaped wrapper around chat + completions. +""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import AsyncGenerator +from enum import Enum +from http import HTTPStatus +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + Citation, + Document, + SystemChatMessageV2, + ToolCallV2, + ToolCallV2Function, + ToolChatMessageV2, + UserChatMessageV2, +) +from fastapi import Request +from pydantic import BaseModel + +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.cohere.cohere_chat_message import ( + CitationSource as InternalCitationSource, +) +from vllm.entrypoints.cohere.cohere_chat_message import ( + CohereChatMessage, +) +from vllm.entrypoints.cohere.protocol import ( + CitationEndEvent, + CitationStartEvent, + CohereChatV2Request, + CohereChatV2Response, + CohereFinishReason, + CohereUsage, + CohereUsageBilledUnits, + CohereUsageTokens, + ContentDeltaEvent, + ContentEndEvent, + ContentStartEvent, + MessageEndEvent, + MessageStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolPlanDeltaEvent, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionStreamResponse, + ChatCompletionToolsParam, + ChatMessage, +) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.engine.protocol import ( + ErrorResponse, + JsonSchemaResponseFormat, + ResponseFormat, + StreamOptions, +) +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.utils.api_utils import sanitize_message +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.parser.abstract_parser import Parser +from vllm.renderers.cohere import MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY + +if TYPE_CHECKING: + from vllm.renderers.online_renderer import OnlineRenderer + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# SSE helpers +# --------------------------------------------------------------------------- + + +def _sse(data: str) -> str: + """Wrap a JSON payload in a Server-Sent Event frame. + + Cohere's stream uses bare ``data:`` lines (no ``event:`` prefix); each + JSON object's ``type`` field carries the event discriminator. + """ + return f"data: {data}\n\n" + + +# Cohere v2 SSE stream terminator. Declared by the upstream OpenAPI spec +# (``x-fern-streaming.terminator: "[DONE]"`` on the ``/v2/chat`` stream +# operation) and observed by the cohere-python / Fern-generated clients +# (the Python SDK breaks its read loop on ``_sse.data == "[DONE]"``). +# Must be emitted after the closing ``message-end`` event. +_DONE_FRAME = _sse("[DONE]") + + +def _emit(event: BaseModel) -> str: + """Serialize a typed stream event into an SSE frame. + + The typed event classes in ``vllm.entrypoints.cohere.protocol`` + (``MessageStartEvent``, ``ContentStartEvent``, ``CitationStartEvent``, + ...) bake the wire-format ``type`` field into the model definition, + so a plain ``model_dump_json()`` already carries the discriminator. + """ + return _sse(event.model_dump_json(exclude_none=True)) + + +class ContentBlockType(str, Enum): + """Wire-format / internal discriminator for chat content blocks. + + ``THINKING`` and ``TEXT`` are the documented Cohere v2 content-block + type discriminators on the wire. ``TOOL_CALL`` is reserved for the + internal stream state machine (see :class:`_StreamState`) when a + tool call is the currently open block; it is never serialized. + """ + + THINKING = "thinking" + TEXT = "text" + TOOL_CALL = "tool_call" + + +# Mapping of vLLM/OpenAI finish reasons to Cohere's enum. +_FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = { + "stop": "COMPLETE", + "length": "MAX_TOKENS", + "tool_calls": "TOOL_CALL", + "stop_sequence": "STOP_SEQUENCE", + "error": "ERROR", + None: "COMPLETE", +} + + +def _map_finish_reason(reason: str | None) -> CohereFinishReason: + return _FINISH_REASON_MAP.get(reason, "COMPLETE") + + +class _CitablePosition(NamedTuple): + """One entry in melody's citation numbering scheme, resolved to wire shape. + + Populated by :meth:`CohereServingChatV2._walk_citable_positions` + (the single source of truth for melody's numbering rule) and + consumed by both the inbound + (:meth:`CohereServingChatV2._build_doc_id_to_prompt_position`, + which drives citation resolution on the request side) and outbound + (:meth:`CohereServingChatV2._build_position_to_source`, whose + result is forwarded to the reasoning parser via + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``) helpers -- so + the numbering rule only lives in one place. + + ``ids`` lists every id under which this position is addressable: + top-level documents contribute both their explicit id (if set) and + the ``doc_{idx}`` synthetic fallback, in that order so + ``setdefault``-style resolution prefers the explicit id; + tool-message documents contribute their single ``document.id``. + ``source`` is the fully-resolved wire-shape source (``type`` / + ``id`` / ``document`` or ``tool_output`` payload) for the + outbound direction. + """ + + bucket: int + result_idx: int + ids: tuple[str, ...] + source: InternalCitationSource + + +# --------------------------------------------------------------------------- +# Serving class +# --------------------------------------------------------------------------- + + +class CohereServingChatV2(OpenAIServingChat): + """Handler for the Cohere Chat v2 API (``POST /cohere/v2/chat``). + + The handler is intentionally thin: it converts the v2 request into a + :class:`ChatCompletionRequest` (preserving Cohere-specific fields such + as ``documents``, ``safety_mode`` and ``citation_options`` via + ``chat_template_kwargs``) and delegates to the underlying chat + completion machinery. All Cohere-specific templating happens in the + renderer (:class:`vllm.renderers.cohere.CohereRenderer`) when the + engine is started with ``--tokenizer-mode cohere``. + + Citations flow through ``CohereChatMessage`` / ``CohereDeltaMessage`` + subclasses (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`) so + the OpenAI-shared :class:`ChatMessage` / :class:`DeltaMessage` keep + their declared schemas unchanged. Non-streaming responses pick up + the subclass by overriding :meth:`_create_chat_message`; streaming + deltas are emitted directly by the Cohere reasoning parser. + """ + + def __init__( + self, + engine_client: EngineClient, + models: OpenAIServingModels, + response_role: str, + *, + online_renderer: OnlineRenderer, + request_logger: RequestLogger | None, + chat_template: str | None, + chat_template_content_format: ChatTemplateContentFormatOption, + return_tokens_as_token_ids: bool = False, + reasoning_parser: str = "", + enable_auto_tools: bool = False, + tool_parser: str | None = None, + enable_prompt_tokens_details: bool = False, + enable_force_include_usage: bool = False, + default_chat_template_kwargs: dict[str, Any] | None = None, + is_reasoning_model: bool = True, + ) -> None: + super().__init__( + engine_client=engine_client, + models=models, + response_role=response_role, + online_renderer=online_renderer, + request_logger=request_logger, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + return_tokens_as_token_ids=return_tokens_as_token_ids, + reasoning_parser=reasoning_parser, + enable_auto_tools=enable_auto_tools, + tool_parser=tool_parser, + enable_prompt_tokens_details=enable_prompt_tokens_details, + enable_force_include_usage=enable_force_include_usage, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + # Controls how the assistant's chain-of-thought is surfaced on + # turns that also contain tool calls. + # + # - ``True`` (default): the model is a reasoning Command-family + # model; reasoning is always surfaced as a ``thinking`` content + # block (non-streaming) or as ``content-start`` / ``content- + # delta`` events for a thinking block (streaming), regardless of + # whether tool calls also appear. + # - ``False``: the model is an older non-reasoning Command model + # that uses Cohere's ``tool_plan`` field for its chain-of- + # thought before tool calls; reasoning is surfaced as + # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events + # (streaming) on tool-call turns, and the thinking content block + # is dropped. + # + # TODO: replace this manual flag with automatic detection from the + # model's capabilities config once that exists. + self._is_reasoning_model = is_reasoning_model + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + + async def create_chat_v2( + self, + request: CohereChatV2Request, + raw_request: Request | None = None, + ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse: + """Implements ``POST /cohere/v2/chat``.""" + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Received Cohere v2 chat request %s", request.model_dump_json() + ) + + chat_req = self._convert_v2_to_chat_completion(request) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Converted Cohere v2 -> ChatCompletion: %s", + chat_req.model_dump_json(), + ) + + generator = await self.create_chat_completion(chat_req, raw_request) + + match generator: + case ErrorResponse(): + return generator + case ChatCompletionResponse(): + return self._chat_completion_to_v2(generator, request) + case _: + return self._chat_completion_stream_to_v2(generator, request) + + # ================================================================== + # Request conversion: Cohere V2 -> ChatCompletionRequest + # ================================================================== + + @classmethod + def _convert_v2_to_chat_completion( + cls, request: CohereChatV2Request + ) -> ChatCompletionRequest: + openai_messages: list[dict[str, Any]] = [] + cls._convert_messages(request.messages, openai_messages) + + chat_req = cls._build_base_chat_completion(request, openai_messages) + cls._apply_streaming_options(chat_req, request) + cls._apply_response_format(chat_req, request) + cls._apply_tools(chat_req, request) + cls._apply_tool_choice(chat_req, request) + cls._apply_cohere_template_kwargs(chat_req, request) + return chat_req + + @classmethod + def _convert_messages( + cls, + messages: list, + openai_messages: list[dict[str, Any]], + ) -> None: + for msg in messages: + match msg: + case SystemChatMessageV2(): + openai_messages.append( + { + "role": "system", + "content": cls._coerce_text_content(msg.content), + } + ) + case UserChatMessageV2(): + openai_messages.append(cls._convert_user_message(msg)) + case AssistantChatMessageV2(): + openai_messages.append(cls._convert_assistant_message(msg)) + case ToolChatMessageV2(): + openai_messages.append(cls._convert_tool_message(msg)) + case _: # pragma: no cover - guarded by Pydantic discriminator + raise ValueError(f"Unsupported Cohere v2 message: {msg!r}") + + @staticmethod + def _coerce_text_content(content: str | list[Any]) -> str: + if isinstance(content, str): + return content + parts: list[str] = [] + for block in content: + text = getattr(block, "text", None) + if text: + parts.append(text) + return "".join(parts) + + @classmethod + def _convert_user_message(cls, msg: UserChatMessageV2) -> dict[str, Any]: + if isinstance(msg.content, str): + return {"role": "user", "content": msg.content} + + # Discriminate by ``type`` rather than isinstance so we don't have + # to import every individual ``*Content`` variant from the SDK - + # the union (``UserMessageV2Content``) covers both ``TextContent`` + # and ``ImageUrlContent`` and both expose ``type`` as a Literal. + content_parts: list[dict[str, Any]] = [] + for block in msg.content: + if block.type == "text": + content_parts.append({"type": "text", "text": block.text}) + elif block.type == "image_url": + image_url: dict[str, Any] = {"url": block.image_url.url} + if getattr(block.image_url, "detail", None) is not None: + image_url["detail"] = block.image_url.detail + content_parts.append( + { + "type": "image_url", + "image_url": image_url, + } + ) + if len(content_parts) == 1 and content_parts[0]["type"] == "text": + return {"role": "user", "content": content_parts[0]["text"]} + return {"role": "user", "content": content_parts} + + @classmethod + def _convert_assistant_message(cls, msg: AssistantChatMessageV2) -> dict[str, Any]: + out: dict[str, Any] = {"role": "assistant"} + + # Cohere splits reasoning out into ``thinking`` content blocks. We + # collapse them back into the OpenAI ``reasoning`` field for the + # downstream chat template, while text blocks become ``content``. + text_parts: list[str] = [] + thinking_parts: list[str] = [] + if isinstance(msg.content, str): + text_parts.append(msg.content) + elif msg.content is not None: + for block in msg.content: + if block.type == "text": + text_parts.append(block.text) + elif block.type == "thinking": + thinking_parts.append(block.thinking) + + # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool + # calls; preserve it as reasoning so templates that expect a + # planning block still see it. + if msg.tool_plan: + thinking_parts.append(msg.tool_plan) + + if text_parts: + out["content"] = "".join(text_parts) + if thinking_parts: + out["reasoning"] = "".join(thinking_parts) + + if msg.tool_calls: + out["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": (tc.function.name if tc.function else "") or "", + "arguments": (tc.function.arguments if tc.function else None) + or "{}", + }, + } + for tc in msg.tool_calls + ] + return out + + @classmethod + def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: + if isinstance(msg.content, str): + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": msg.content, + } + + # When the tool result is text-only, flatten to a string for + # maximum compatibility with standard chat templates. When it + # includes documents, preserve them as structured content parts + # so the cohere renderer can surface them as grounding sources + # (the :class:`CohereRenderer` understands + # ``{type: document, document: {...}}`` blocks). Non-cohere + # renderers may not honor document blocks, but that matches + # the broader "documents are no-op for OSS models" contract + # documented on this endpoint. + # + # Tool message content uses ``ToolMessageV2Content``, which is a + # union of ``TextToolContent`` and ``DocumentToolContent`` - + # distinct from the user-message text/image union. We discriminate + # on the ``type`` literal so we don't have to import each variant. + has_documents = any(block.type == "document" for block in msg.content) + if not has_documents: + text = "\n".join( + block.text for block in msg.content if block.type == "text" + ) + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": text, + } + + parts: list[dict[str, Any]] = [] + for block in msg.content: + if block.type == "text": + parts.append({"type": "text", "text": block.text}) + elif block.type == "document": + parts.append( + { + "type": "document", + "document": block.document.model_dump(exclude_none=True), + } + ) + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": parts, + } + + @classmethod + def _build_base_chat_completion( + cls, + request: CohereChatV2Request, + openai_messages: list[dict[str, Any]], + ) -> ChatCompletionRequest: + return ChatCompletionRequest( + model=request.model, + messages=openai_messages, + max_tokens=request.max_tokens, + max_completion_tokens=request.max_tokens, + stop=request.stop_sequences, + temperature=request.temperature, + top_p=request.p, + top_k=request.k, + seed=request.seed, + frequency_penalty=request.frequency_penalty, + presence_penalty=request.presence_penalty, + logprobs=request.logprobs, + priority=request.priority or 0, + kv_transfer_params=request.kv_transfer_params, + chat_template_kwargs=request.chat_template_kwargs, + ) + + @classmethod + def _apply_streaming_options( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + if request.stream: + chat_req.stream = True + chat_req.stream_options = StreamOptions.model_validate( + {"include_usage": True} + ) + + @classmethod + def _apply_response_format( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + rf = request.response_format + if rf is None or rf.type == "text": + return + chat_req.response_format = ResponseFormat( + type="json_schema" if rf.json_schema else "json_object", + json_schema=( + JsonSchemaResponseFormat( + name="cohere_v2_json_schema", + json_schema=rf.json_schema, + ) + if rf.json_schema + else None + ), + ) + + @classmethod + def _apply_tools( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + if not request.tools: + return + # Cohere's ``strict_tools`` is the spec-equivalent of OpenAI's + # per-function ``strict`` flag: when true the API guarantees tool + # call arguments match the declared JSON schema. We surface it in + # both places so OpenAI-shaped consumers see it on the function + # definition and the cohere renderer can read it back from + # ``chat_template_kwargs`` for cmd3/cmd4 preamble selection. + strict = bool(request.strict_tools) + chat_req.tools = [ + ChatCompletionToolsParam.model_validate( + { + "type": "function", + "function": { + "name": tool.function.name, + "description": tool.function.description, + "parameters": tool.function.parameters, + **({"strict": True} if strict else {}), + }, + } + ) + for tool in request.tools + ] + + @classmethod + def _apply_tool_choice( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + # Cohere v2's ``tool_choice`` only admits ``REQUIRED`` / ``NONE`` + # (see :data:`CohereToolChoice`); named-tool selection is not + # part of the v2 spec. Map both onto OpenAI's equivalents. + if request.tool_choice == "REQUIRED": + chat_req.tool_choice = "required" + elif request.tool_choice == "NONE": + chat_req.tool_choice = "none" + elif chat_req.tools: + # Mirrors Cohere's "free choice" default when tools are present. + chat_req.tool_choice = "auto" + + @classmethod + def _apply_cohere_template_kwargs( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + """Forward Cohere-specific request fields into ``chat_template_kwargs``. + + The :class:`vllm.renderers.cohere.CohereRenderer` consumes these + kwargs to drive ``cohere_melody.render_cmd3`` / ``render_cmd4``. + Other renderers ignore unknown kwargs, so this is also a no-op for + non-cohere ``--tokenizer-mode`` settings. + """ + kwargs = dict(chat_req.chat_template_kwargs or {}) + + if request.documents: + documents: list[dict[str, Any]] = [] + for idx, doc in enumerate(request.documents): + if isinstance(doc, str): + documents.append({"id": f"doc_{idx}", "data": {"text": doc}}) + else: + documents.append( + { + "id": doc.id or f"doc_{idx}", + "data": ( + doc.data + if isinstance(doc.data, dict) + else {"text": doc.data} + ), + } + ) + kwargs.setdefault("documents", documents) + + if request.safety_mode is not None: + kwargs.setdefault("safety_mode", str(request.safety_mode).lower()) + if request.citation_options is not None: + kwargs.setdefault( + "citation_options", + request.citation_options.model_dump(exclude_none=True), + ) + if request.thinking is not None: + kwargs.setdefault( + "thinking", + request.thinking.model_dump(exclude_none=True), + ) + + # ``strict_tools`` is intentionally NOT forwarded here. It's a + # decoder-guidance flag that ``_apply_tools`` already maps onto + # per-function OpenAI ``strict: true`` on ``chat_req.tools``; + # ``cohere_melody.render_cmd3``/``render_cmd4`` have no + # ``strict_tools`` knob, and the renderer's residual-forward path + # would otherwise surface it as a Jinja variable ``{{ strict_tools }}`` + # that templates have no defined use for. + + # Citations on assistant messages in the request don't fit + # OpenAI's ``ChatMessage`` / ``ConversationMessage`` schemas, so + # we forward them under the reserved ``MESSAGES_CITATIONS_KEY`` + # chat_template_kwargs entry keyed by request-message index. + # See ``_sdk_citation_to_melody`` for the id-resolution rules. + citations_by_index = cls._collect_message_citations(request) + if citations_by_index: + kwargs.setdefault(MESSAGES_CITATIONS_KEY, citations_by_index) + + # Forward the ``(bucket, result_idx) -> resolved wire source`` + # map to the reasoning parser via chat_template_kwargs. The + # parser (constructed with the same kwargs) uses it to attach + # real document ids / types / payloads to every source it + # emits, so citation deltas already carry SDK-shape sources by + # the time they reach the streaming loop. See + # ``_build_position_to_source`` for the numbering rule and + # ``_melody_sources_to_vllm`` in the reasoning parser for the + # consumer. + position_to_source = cls._build_position_to_source(request) + if position_to_source: + kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source) + + if kwargs: + chat_req.chat_template_kwargs = kwargs + + @classmethod + def _collect_message_citations( + cls, request: CohereChatV2Request + ) -> dict[int, list[dict[str, Any]]]: + """Extract ``AssistantChatMessageV2.citations`` keyed by message index. + + Returns an empty dict when no assistant message in the request + carries any resolvable citations, so callers can use truthiness + as the "should I emit the ``MESSAGES_CITATIONS_KEY`` entry?" + test. Citations whose sources can't be resolved are dropped + individually; messages that end up with zero surviving + citations are omitted from the result. + """ + doc_positions = cls._build_doc_id_to_prompt_position( + request.messages, request.documents + ) + out: dict[int, list[dict[str, Any]]] = {} + for idx, msg in enumerate(request.messages): + if not isinstance(msg, AssistantChatMessageV2) or not msg.citations: + continue + resolved: list[dict[str, Any]] = [] + for c in msg.citations: + converted = cls._sdk_citation_to_melody(c, doc_positions) + if converted is not None: + resolved.append(converted) + if resolved: + out[idx] = resolved + return out + + @classmethod + def _walk_citable_positions( + cls, messages: list, documents: list | None + ) -> list[_CitablePosition]: + """Single-pass implementation of melody's citation numbering rule. + + Mirrors ``PromptRenderIds::from_messages`` in + ``melody/src/templating/util.rs``. Emits one + :class:`_CitablePosition` per document-bearing slot melody + assigns while rendering the prompt: + + * **Bucket numbering.** Bucket ``0`` is reserved for the + top-level ``documents`` array when present, and each unique + ``tool_call_id`` claims the next integer on first-seen basis + -- registered from whichever of the assistant's + ``tool_calls[].id`` or the tool message's ``tool_call_id`` + appears first in the history. + * **Per-bucket result-index slots.** Consumed by *every* + content item, not just documents -- a text block occupies a + slot too, and the source at that slot carries an id-less + ``tool_output`` payload synthesized via + :meth:`_text_to_tool_output_payload`, matching the cohere + api's behavior for text-only tool content. Strings and + empty contents each consume exactly one slot because the + cohere renderer wraps them into a single text block before + handing the message off to melody. Multiple tool messages + sharing a ``tool_call_id`` accumulate into the same bucket, + so later docs' slots are offset by the sum of previous + messages' content lengths. + * **Doc ids attached to each position.** Top-level docs + contribute both the explicit ``id`` (if set) and the + ``doc_{idx}`` fallback used for inbound addressing; tool + document blocks contribute their single ``document.id`` + when present. The ``ids`` tuple is empty for slots that + are only reachable positionally (text-only tool content, + docs without a client-provided id) -- those slots still + have a fully-populated :class:`InternalCitationSource`, + just no client-visible id to cite them by. + + Ids alone key the inbound helper + (:meth:`_build_doc_id_to_prompt_position`, which drives + request-side citation resolution); the resolved + :class:`InternalCitationSource` (with ``type`` / ``id`` / + payload) keys the outbound helper + (:meth:`_build_position_to_source`, whose result is forwarded + to the reasoning parser via + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``). + + Cohere's on-wire semantics -- ``ToolSource.id`` and + ``DocumentSource.id`` both carry a document id, with ``type`` + acting only as a payload-shape hint -- are the invariant this + walk is mirroring. + """ + out: list[_CitablePosition] = [] + tool_call_id_to_bucket: dict[str, int] = {} + # Next unassigned ``tool_result_index`` per bucket. Advances + # by ``len(content)`` (or 1 for string / missing content) for + # every tool message, whether or not the content had any doc + # ids in it -- matching melody's per-slot accounting. + bucket_result_len: dict[int, int] = {} + next_bucket = 0 + + if documents: + for idx, doc in enumerate(documents): + # Every top-level doc is addressable inbound by a + # synthetic ``doc_{idx}`` fallback (so clients can cite + # positional docs). On the *outbound* wire we prefer + # the client's explicit id when present, and omit the + # ``id`` field entirely when it's not -- matching + # cohere's ``DocumentSource`` where ``id`` is optional + # and preserving "no id in, no id out" symmetry. + fallback_id = f"doc_{idx}" + if isinstance(doc, str): + payload: dict[str, Any] = {"text": doc} + explicit_id: str | None = None + else: + payload = ( + dict(doc.data) + if isinstance(doc.data, dict) + else {"text": doc.data} + ) + explicit_id = doc.id or None + if explicit_id: + payload.setdefault("id", explicit_id) + # Register explicit id first so it wins over the + # fallback for ``setdefault``-style inbound lookup. + ids: tuple[str, ...] = ( + (explicit_id, fallback_id) if explicit_id else (fallback_id,) + ) + out.append( + _CitablePosition( + bucket=0, + result_idx=idx, + ids=ids, + source=InternalCitationSource( + type="document", id=explicit_id, document=payload + ), + ) + ) + next_bucket = 1 + + def register_bucket(tool_call_id: str) -> int: + nonlocal next_bucket + bucket = tool_call_id_to_bucket.get(tool_call_id) + if bucket is None: + bucket = next_bucket + tool_call_id_to_bucket[tool_call_id] = bucket + next_bucket += 1 + return bucket + + for msg in messages: + if isinstance(msg, ToolChatMessageV2): + if not msg.tool_call_id: + continue + bucket = register_bucket(msg.tool_call_id) + base = bucket_result_len.get(bucket, 0) + if isinstance(msg.content, list): + for offset, block in enumerate(msg.content): + info = cls._tool_content_block_to_source(block) + if info is None: + continue + # Only real document ids get registered in the + # inbound-facing ``ids`` tuple; text-only tool + # content produces a source with ``id=None`` + # to match the cohere api -- the text is + # exposed as ``tool_output.content`` but + # there's no id for clients to cite by. + pos_ids: tuple[str, ...] = (info.id,) if info.id else () + out.append( + _CitablePosition( + bucket=bucket, + result_idx=base + offset, + ids=pos_ids, + source=info, + ) + ) + bucket_result_len[bucket] = base + len(msg.content) + else: + # A string or missing content is wrapped by the + # cohere renderer into a single text block before + # melody sees it, so it consumes one slot. For + # string content, emit a synthetic id-less tool + # source carrying the text as its payload so a + # model citation against this slot still has + # something to attach on the wire (matching the + # id-less tool_output behavior for text-only + # ``ToolChatMessageV2`` content blocks). Missing + # content still advances the slot cursor but + # produces no source -- a citation against it + # would be dropped by ``_to_wire_citation``. + if isinstance(msg.content, str): + payload = cls._text_to_tool_output_payload(msg.content) + out.append( + _CitablePosition( + bucket=bucket, + result_idx=base, + ids=(), + source=InternalCitationSource( + type="tool", id=None, tool_output=payload + ), + ) + ) + bucket_result_len[bucket] = base + 1 + elif isinstance(msg, AssistantChatMessageV2) and msg.tool_calls: + for tc in msg.tool_calls: + if tc.id: + register_bucket(tc.id) + + return out + + @classmethod + def _build_doc_id_to_prompt_position( + cls, messages: list, documents: list | None + ) -> dict[str, tuple[int, int]]: + """Map every citable document id to its ``(bucket, result_index)``. + + Inbound helper: converts request-side citation ``source.id`` + refs into melody's numeric ``FilterCitation`` addressing (see + :meth:`_sdk_citation_to_melody`). First-seen-id wins so an + explicit id on a top-level document takes precedence over the + ``doc_{idx}`` fallback that :meth:`_walk_citable_positions` + also emits for the same slot. + + Thin wrapper over :meth:`_walk_citable_positions`; the + numbering rule lives there. + """ + doc_positions: dict[str, tuple[int, int]] = {} + for pos in cls._walk_citable_positions(messages, documents): + for pos_id in pos.ids: + doc_positions.setdefault(pos_id, (pos.bucket, pos.result_idx)) + return doc_positions + + @staticmethod + def _sdk_citation_to_melody( + citation: Any, + doc_positions: dict[str, tuple[int, int]], + ) -> dict[str, Any] | None: + """Convert a Cohere SDK ``Citation`` to melody's ``FilterCitation``. + + Melody's ``Source`` uses a two-level numeric address: + ``tool_call_index`` picks a bucket (``0`` = top-level + ``documents`` when present, ``1..N`` = tool results in history + order) and ``tool_result_indices`` picks positions inside that + bucket. The Cohere ``Citation.sources`` schema is per-document: + each entry's ``id`` is the id of the specific document that + grounds the citation (``type`` distinguishes payload shape, not + id semantics). We resolve each source id against + ``doc_positions`` and aggregate hits sharing a bucket into one + melody ``Source`` so the renderer emits compact + ``text`` markers. + + Returns ``None`` when any source id fails to resolve or when + the input has no sources at all: emitting a citation with an + empty ``sources=[]`` would render a malformed + ``text`` marker, and silently attributing a + citation to the wrong bucket is worse than dropping the + ```` markup. This matches the cohere api's behavior. + + ``Source.document_ids`` is intentionally NOT populated: it is a + parser-*output* lookup table (see ``PromptRenderIds`` in + melody), and melody's deserializer ignores it on input. + + The Cohere ``type`` literal (``TEXT_CONTENT`` / ``THINKING_CONTENT`` + / ``PLAN``) collapses to melody's boolean ``is_thinking``. + ``PLAN`` is treated as a thinking-style citation because cmd3 / + cmd4 templates render both inline with the same ``...`` + markup. + """ + raw_sources = getattr(citation, "sources", None) or [] + if not raw_sources: + return None + + grouped: dict[int, list[int]] = {} + for src in raw_sources: + src_id = getattr(src, "id", None) + if not src_id: + return None + position = doc_positions.get(src_id) + if position is None: + return None + bucket, result_idx = position + grouped.setdefault(bucket, []).append(result_idx) + + sources = [ + {"tool_call_index": bucket, "tool_result_indices": indices} + for bucket, indices in grouped.items() + ] + + cite_type = str(getattr(citation, "type", "") or "").upper() + is_thinking = cite_type in ("THINKING_CONTENT", "PLAN") + return { + "start_index": getattr(citation, "start", None) or 0, + "end_index": getattr(citation, "end", None) or 0, + "text": getattr(citation, "text", None) or "", + "sources": sources, + "is_thinking": is_thinking, + } + + # ================================================================== + # Response conversion: ChatCompletion -> Cohere V2 + # ================================================================== + + def _chat_completion_to_v2( + self, + response: ChatCompletionResponse, + request: CohereChatV2Request, + ) -> CohereChatV2Response: + choice = response.choices[0] + msg = choice.message + + # Build content blocks as dicts; ``AssistantMessageResponse`` + # validates them into the proper discriminated union variants + # (``Text/ThinkingAssistantMessageResponseContentItem``). + content_blocks: list[dict[str, Any]] = [] + if msg.reasoning: + content_blocks.append( + {"type": ContentBlockType.THINKING, "thinking": msg.reasoning} + ) + if msg.content: + content_blocks.append({"type": ContentBlockType.TEXT, "text": msg.content}) + + tool_calls: list[ToolCallV2] | None = None + if msg.tool_calls: + tool_calls = [ + ToolCallV2( + id=tc.id, + function=ToolCallV2Function( + name=tc.function.name, + arguments=tc.function.arguments, + ), + ) + for tc in msg.tool_calls + ] + + # Cohere's ``tool_plan`` is the planning text emitted before tool + # calls on older, non-reasoning Command models. For those models + # we surface ``reasoning`` as ``tool_plan`` and drop the thinking + # block. Reasoning Command models emit a regular thinking block + # alongside tool calls, so for them we leave the thinking block + # in place and never set ``tool_plan``. + tool_plan: str | None = None + if not self._is_reasoning_model and tool_calls and msg.reasoning: + tool_plan = msg.reasoning + content_blocks = [ + blk + for blk in content_blocks + if blk.get("type") != ContentBlockType.THINKING + ] + + assistant_msg = AssistantMessageResponse( + content=content_blocks or None, + tool_calls=tool_calls, + tool_plan=tool_plan, + citations=self._extract_citations_if_any(msg), + ) + + usage = self._build_usage(response) + + return CohereChatV2Response( + id=response.id or f"chat_{int(time.time() * 1000)}", + finish_reason=_map_finish_reason(choice.finish_reason), + message=assistant_msg, + usage=usage, + kv_transfer_params=response.kv_transfer_params, + ) + + def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage: + """Route response construction through the citation-carrying subclass. + + Overrides :meth:`OpenAIServingChat._create_chat_message` so every + non-streaming construction site in the base full-generator + produces a :class:`CohereChatMessage`. That lets + :meth:`_finalize_response_message` below set citations on + ``message.citations`` directly, without relying on + ``extra="allow"`` attribute injection. + """ + return CohereChatMessage(*args, **kwargs) + + def _finalize_response_message( + self, + message: ChatMessage, + *, + parser: Parser | None, + ) -> ChatMessage: + """Copy grounding citations off the reasoning parser onto the message. + + The Cohere reasoning parser + (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the + citations produced by its most recent unary ``extract_reasoning`` + call on ``last_unary_citations``. We surface them here on + :class:`CohereChatMessage` so downstream response conversion can + pick them up without the base :class:`OpenAIServingChat` having to + know about citations. + """ + # ``_create_chat_message`` above guarantees the concrete type at + # runtime. ``cast`` narrows the declared ``ChatMessage`` for + # mypy without adding a runtime check we don't need: even in the + # invariant-violated case (a subclass reset the factory back to + # plain ``ChatMessage``), ``OpenAIBaseModel``'s ``extra="allow"`` + # config lets the ``.citations`` write survive into + # ``model_dump`` via the extras bucket, so the wire is correct + # either way. + message = cast(CohereChatMessage, message) + citations = getattr( + getattr(parser, "reasoning_parser", None), + "last_unary_citations", + None, + ) + if citations: + message.citations = citations + return message + + def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None: + """Coerce ``CohereChatMessage.citations`` into the Cohere v2 wire shape. + + :class:`CohereChatMessage` carries a + ``citations: list[vllm...Citation] | None`` field populated by + :meth:`_finalize_response_message` (which reads it off the + reasoning parser's ``last_unary_citations`` cache). Sources are + already fully resolved by the parser (see + :func:`_melody_sources_to_vllm` in + :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using + the position map forwarded via ``chat_template_kwargs``. This + method's remaining job is: + + * Drop citations whose sources all failed to resolve (empty + ``sources`` list) -- fail-closed policy matching the inbound + ``_sdk_citation_to_melody``. + * Rewrite ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning + models, since ``_chat_completion_to_v2`` surfaces + ``msg.reasoning`` as ``tool_plan`` on the same + ``is_reasoning_model`` flag. + * Coerce to :class:`cohere.types.Citation` for the wire. + """ + raw = getattr(msg, "citations", None) + if not raw: + return None + out: list[Citation] = [] + for c in raw: + resolved = self._to_wire_citation(c) + if resolved is not None: + out.append(resolved) + return out or None + + def _to_wire_citation(self, citation: Any) -> Citation | None: + """Coerce one parser-produced citation into the SDK wire shape. + + Accepts both :class:`InternalCitation` instances (what the + unary path receives directly from the parser via + ``last_unary_citations``) and plain ``dict`` payloads (what + the streaming path sees, because ``DeltaMessage.citations`` is + an untyped extras field on the OpenAI wire protocol and dict + shapes survive ``model_validate_json`` unchanged). Both flows + end up carrying the same key names, so a small ``_get`` helper + handles both without a shape-specific code path. + """ + + def _get(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + raw_sources = _get(citation, "sources") or [] + wire_sources: list[dict[str, Any]] = [] + for src in raw_sources: + # Sources arrive pre-resolved from the parser (or as trusted + # dicts on the streaming path). Match the cohere api's + # behavior of preserving id-less sources (e.g. text-only + # tool results): cohere's ``Source`` types treat ``id`` as + # optional, so we simply drop the field via + # ``exclude_none=True`` instead of dropping the source. + if isinstance(src, dict): + wire = {k: v for k, v in src.items() if v is not None} + if wire: + wire_sources.append(wire) + elif isinstance(src, InternalCitationSource): + wire_sources.append(src.model_dump(exclude_none=True)) + if not wire_sources: + return None + + cite_type = _get(citation, "type") + if cite_type == "THINKING_CONTENT" and not self._is_reasoning_model: + # Non-reasoning Command models surface the reasoning block + # as ``tool_plan`` (see ``_chat_completion_to_v2``), so any + # ``is_thinking`` citation from melody is actually a PLAN + # citation on the wire. + cite_type = "PLAN" + + payload = { + "start": _get(citation, "start"), + "end": _get(citation, "end"), + "text": _get(citation, "text"), + "sources": wire_sources, + "type": cite_type, + } + try: + return Citation.model_validate(payload) + except Exception: # pragma: no cover - defensive + logger.warning( + "Skipping malformed citation payload: %r", payload, exc_info=True + ) + return None + + @classmethod + def _build_position_to_source( + cls, request: CohereChatV2Request + ) -> dict[tuple[int, int], InternalCitationSource]: + """Invert the request-side numbering into ``(bucket, idx) -> source``. + + Outbound helper: forwarded to the reasoning parser via + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]`` so + :func:`_melody_sources_to_vllm` can attach real document ids / + types / payloads to each melody source it emits. Each entry + carries the fully-resolved wire-shape info: ``type`` + (``"document"`` for the reserved top-level bucket 0, ``"tool"`` + for any tool-call bucket), ``id`` (explicit doc id, the + ``doc_{idx}`` fallback, or ``None`` for id-less sources), and + ``document`` / ``tool_output`` payload -- matching the cohere + api's ``Source`` shape. + + Thin wrapper over :meth:`_walk_citable_positions`; the + numbering rule lives there. + """ + return { + (pos.bucket, pos.result_idx): pos.source + for pos in cls._walk_citable_positions(request.messages, request.documents) + } + + @staticmethod + def _text_to_tool_output_payload(text: str) -> dict[str, Any]: + """Wrap a bare tool-content text into a ``tool_output`` payload. + + Matches the cohere api's handling of text-only tool content: + attempt to interpret ``text`` as a JSON object first; fall + back to ``{"content": }`` so a text-only tool result + still has a structured payload we can attach to a citation + source. + """ + if not text: + return {} + try: + parsed = json.loads(text) + except (json.JSONDecodeError, TypeError): + return {"content": text} + if isinstance(parsed, dict): + return parsed + return {"content": text} + + @classmethod + def _tool_content_block_to_source( + cls, + block: Any, + ) -> InternalCitationSource | None: + """Extract the wire-shape source for one tool-message content block. + + Handles both cohere ``DocumentToolContent`` and + ``TextToolContent`` blocks: + + * ``document`` blocks with an explicit id produce a + ``tool`` source keyed by that id (payload is the document's + ``data`` dict, or ``{"text": ...}`` for scalar data). + * ``document`` blocks *without* an id and ``text`` blocks + produce an id-less ``tool`` source: the cohere api emits + the payload with an empty id in this case, and cohere's + SDK treats ``ToolSource.id`` as optional -- so we simply + omit ``id`` on the wire instead of inventing a synthetic + one. + + Returns ``None`` only for shapes we don't recognise at all; + that slot still gets counted (callers advance the bucket + cursor over the block regardless) so numbering stays aligned + with melody. + """ + block_type = getattr(block, "type", None) + if block_type == "document": + doc: Document | None = getattr(block, "document", None) + if doc is None: + return None + doc_id = doc.id or None + data = doc.data + if isinstance(data, dict): + payload: dict[str, Any] = dict(data) + elif data is None: + payload = {} + else: + # Match the top-level ``documents`` handling (see + # ``_apply_cohere_template_kwargs``): non-dict payloads + # get wrapped as ``{"text": ...}`` so clients get a + # renderable body instead of a bare id. + payload = {"text": data} + if doc_id: + payload.setdefault("id", doc_id) + return InternalCitationSource( + type="tool", + id=doc_id, + tool_output=payload, + ) + if block_type == "text": + text = getattr(block, "text", None) or "" + payload = cls._text_to_tool_output_payload(text) + return InternalCitationSource( + type="tool", + id=None, + tool_output=payload, + ) + return None + + @staticmethod + def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None: + if response.usage is None: + return None + prompt = response.usage.prompt_tokens + completion = response.usage.completion_tokens or 0 + cached: int | None = None + if response.usage.prompt_tokens_details is not None: + cached = response.usage.prompt_tokens_details.cached_tokens + return CohereUsage( + billed_units=CohereUsageBilledUnits( + input_tokens=prompt, + output_tokens=completion, + ), + tokens=CohereUsageTokens( + input_tokens=prompt, + output_tokens=completion, + ), + cached_tokens=cached, + ) + + # ================================================================== + # Stream conversion: chat completion stream -> Cohere V2 SSE events + # ================================================================== + + async def _chat_completion_stream_to_v2( + self, + generator: AsyncGenerator[str, None], + request: CohereChatV2Request, + ) -> AsyncGenerator[str, None]: + """Translate an OpenAI-style chat completion SSE stream into Cohere's + v2 stream-event format. + + Cohere's v2 stream lifecycle is: + + message-start + [content-start, content-delta..., content-end]* + [tool-plan-delta]* + [tool-call-start, tool-call-delta..., tool-call-end]* + message-end + """ + state = _StreamState() + + try: + async for item in generator: + if not item.startswith("data:"): + continue + data_str = item[len("data:") :].strip().rstrip("\n") + if not data_str: + continue + if data_str == "[DONE]": + # OpenAI's stream terminator. Fall through to the + # post-loop cleanup so we always emit ``message-end`` + # even if the usage-only chunk was skipped. + break + + chunk = ChatCompletionStreamResponse.model_validate_json(data_str) + state.last_chunk_id = chunk.id + + if not state.started: + yield _emit( + MessageStartEvent( + id=chunk.id, + delta={"message": {"role": "assistant"}}, + ) + ) + state.started = True + + # The final OpenAI chunk has no choices and only carries usage. + if not chunk.choices: + for ev in self._close_open_blocks(state): + yield ev + yield self._build_message_end_event( + chunk_id=chunk.id, + finish_reason=state.finish_reason, + usage_chunk=chunk, + ) + state.ended = True + continue + + choice = chunk.choices[0] + if choice.finish_reason is not None: + state.finish_reason = choice.finish_reason + + delta = choice.delta + + # Reasoning -> thinking content block + reasoning = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if reasoning: + for ev in self._handle_thinking_delta(state, reasoning): + yield ev + + if delta.content: + for ev in self._handle_text_delta(state, delta.content): + yield ev + + if delta.tool_calls: + for ev in self._handle_tool_call_deltas(state, delta.tool_calls): + yield ev + + # Citations: a Cohere-specific extension on DeltaMessage that + # the cohere renderer/parsers may populate. + delta_citations = getattr(delta, "citations", None) + if delta_citations: + for ev in self._handle_citation_deltas(state, delta_citations): + yield ev + + except Exception as exc: + logger.exception("Error converting chat completion stream to v2") + if state.started and not state.ended: + yield _sse( + json.dumps( + { + "type": "message-end", + "delta": { + "error": sanitize_message(str(exc)), + "finish_reason": "ERROR", + }, + } + ) + ) + state.ended = True + yield _DONE_FRAME + return + + # Normal completion or ``[DONE]``: ensure ``message-end`` is always + # emitted. Upstream may close the stream without sending the final + # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only + # terminator); without this fallback Cohere clients would hang + # waiting for the closing event. + if state.started and not state.ended: + for ev in self._close_open_blocks(state): + yield ev + yield self._build_message_end_event( + chunk_id=state.last_chunk_id, + finish_reason=state.finish_reason, + usage_chunk=None, + ) + state.ended = True + + # Stream terminator. Cohere's v2 SSE protocol ends every stream + # with ``data: [DONE]\n\n`` after ``message-end``; Fern-generated + # clients (Go/Java) and cohere-python all key their read loop off + # this sentinel. + yield _DONE_FRAME + + # -- per-delta helpers -------------------------------------------- + + def _handle_thinking_delta(self, state: _StreamState, delta_text: str) -> list[str]: + # Non-reasoning Command models: emit ``tool-plan-delta`` events + # directly instead of opening a thinking content block. The + # ``tool-plan-delta`` event has no start/end pair around it. + if not self._is_reasoning_model: + events: list[str] = list(self._close_open_blocks(state)) + events.append( + _emit( + ToolPlanDeltaEvent( + delta={"message": {"tool_plan": delta_text}}, + ) + ) + ) + return events + + # Reasoning model (default): open / continue a thinking block. + events = [] + if state.active_block != ContentBlockType.THINKING: + events.extend(self._close_open_blocks(state)) + idx = state.next_content_index() + state.active_block = ContentBlockType.THINKING + state.active_block_index = idx + events.append( + _emit( + ContentStartEvent( + index=idx, + delta={ + "message": { + "content": { + "type": ContentBlockType.THINKING, + "thinking": "", + } + } + }, + ) + ) + ) + events.append( + _emit( + ContentDeltaEvent( + index=state.active_block_index, + delta={"message": {"content": {"thinking": delta_text}}}, + ) + ) + ) + return events + + def _handle_text_delta(self, state: _StreamState, delta_text: str) -> list[str]: + events: list[str] = [] + if state.active_block != ContentBlockType.TEXT: + events.extend(self._close_open_blocks(state)) + idx = state.next_content_index() + state.active_block = ContentBlockType.TEXT + state.active_block_index = idx + events.append( + _emit( + ContentStartEvent( + index=idx, + delta={ + "message": { + "content": { + "type": ContentBlockType.TEXT, + "text": "", + } + } + }, + ) + ) + ) + events.append( + _emit( + ContentDeltaEvent( + index=state.active_block_index, + delta={"message": {"content": {"text": delta_text}}}, + ) + ) + ) + return events + + def _handle_tool_call_deltas(self, state: _StreamState, deltas: list) -> list[str]: + events: list[str] = [] + for tc in deltas: + tc_index = tc.index + fn = tc.function + + if tc_index not in state.tool_calls_seen: + # New tool call. Close any open content/tool block first. + events.extend(self._close_open_blocks(state)) + state.tool_calls_seen.add(tc_index) + state.active_tool_index = tc_index + state.active_block = ContentBlockType.TOOL_CALL + events.append( + _emit( + ToolCallStartEvent( + index=tc_index, + delta={ + "message": { + "tool_calls": { + "id": tc.id or "", + "type": "function", + "function": { + "name": (fn.name if fn else "") or "", + "arguments": (fn.arguments if fn else None) + or "", + }, + } + } + }, + ) + ) + ) + continue + + if fn and fn.arguments: + events.append( + _emit( + ToolCallDeltaEvent( + index=tc_index, + delta={ + "message": { + "tool_calls": { + "function": { + "arguments": fn.arguments, + } + } + } + }, + ) + ) + ) + return events + + def _handle_citation_deltas( + self, + state: _StreamState, + citations: list, + ) -> list[str]: + """Emit ``citation-start`` / ``citation-end`` events for a delta. + + The reasoning parser populates + :class:`CohereDeltaMessage.citations` (see + :mod:`vllm.entrypoints.cohere.cohere_chat_message`) with + sources whose ``id`` / ``type`` / ``document`` / ``tool_output`` + are already resolved against the request's document context + (see :func:`_melody_sources_to_vllm` in the parser). This + method only coerces to the SDK wire shape, applies the + ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning + models, and drops citations whose sources didn't resolve -- + the shared :meth:`_to_wire_citation` handles all three. + """ + events: list[str] = [] + for c in citations: + citation = self._to_wire_citation(c) + if citation is None: + continue + idx = state.next_citation_index() + events.append( + _emit( + CitationStartEvent( + index=idx, + delta={ + "message": { + "citations": citation.model_dump(exclude_none=True) + } + }, + ) + ) + ) + events.append(_emit(CitationEndEvent(index=idx))) + return events + + # -- block lifecycle helpers -------------------------------------- + + def _close_open_blocks(self, state: _StreamState) -> list[str]: + """Emit ``content-end`` / ``tool-call-end`` for the currently open + block (if any) and reset the corresponding ``_StreamState`` slots. + """ + events: list[str] = [] + if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING): + events.append(_emit(ContentEndEvent(index=state.active_block_index))) + elif state.active_block == ContentBlockType.TOOL_CALL: + events.append(_emit(ToolCallEndEvent(index=state.active_tool_index))) + state.active_block = None + state.active_block_index = None + state.active_tool_index = None + return events + + def _build_message_end_event( + self, + chunk_id: str, + finish_reason: str | None, + usage_chunk: ChatCompletionStreamResponse | None = None, + ) -> str: + delta: dict[str, Any] = { + "finish_reason": _map_finish_reason(finish_reason), + } + if usage_chunk is not None and usage_chunk.usage is not None: + prompt = usage_chunk.usage.prompt_tokens + completion = usage_chunk.usage.completion_tokens or 0 + usage_block: dict[str, Any] = { + "billed_units": { + "input_tokens": prompt, + "output_tokens": completion, + }, + "tokens": { + "input_tokens": prompt, + "output_tokens": completion, + }, + } + if usage_chunk.usage.prompt_tokens_details is not None: + cached = usage_chunk.usage.prompt_tokens_details.cached_tokens + if cached is not None: + usage_block["cached_tokens"] = cached + delta["usage"] = usage_block + return _emit(MessageEndEvent(id=chunk_id, delta=delta)) + + # ================================================================== + # Helpers for the router + # ================================================================== + + @staticmethod + def create_error_response( + message: str | Exception, + err_type: str = "bad_request", + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, + param: str | None = None, + ) -> ErrorResponse: + # Override of :meth:`BaseServing.create_error_response` that uses + # Cohere-flavored defaults (``type="bad_request"``, ``code=400``) + # so the router can translate the envelope uniformly. ``param`` + # is accepted for signature parity with the base class but is + # not surfaced in the Cohere wire format. + from vllm.entrypoints.openai.engine.protocol import ErrorInfo + + del param # unused; kept for signature compatibility + return ErrorResponse( + error=ErrorInfo( + message=str(message), + type=err_type, + code=int(status_code), + ) + ) + + +# --------------------------------------------------------------------------- +# Stream state +# --------------------------------------------------------------------------- + + +class _StreamState: + """Tracks which Cohere v2 stream block (if any) is currently open.""" + + def __init__(self) -> None: + self.started: bool = False + self.ended: bool = False + self.finish_reason: str | None = None + self.last_chunk_id: str = "" + self.active_block: ContentBlockType | None = None + self.active_block_index: int | None = None + self.active_tool_index: int | None = None + self._next_index: int = 0 + self._next_citation_index: int = 0 + self.tool_calls_seen: set[int] = set() + + def next_content_index(self) -> int: + idx = self._next_index + self._next_index += 1 + return idx + + def next_citation_index(self) -> int: + idx = self._next_citation_index + self._next_citation_index += 1 + return idx diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 062a7947e790..92c564ac31da 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -4,6 +4,8 @@ from fastapi import FastAPI +import vllm.envs as envs + if TYPE_CHECKING: from argparse import Namespace @@ -41,6 +43,12 @@ def register_generate_api_routers(app: FastAPI): register_anthropic_api_router(app) + from vllm.entrypoints.cohere.api_router import ( + attach_router as register_cohere_api_router, + ) + + register_cohere_api_router(app) + from .generative_scoring.api_router import register_generative_scoring_api_router register_generative_scoring_api_router(app) @@ -55,6 +63,23 @@ async def init_generate_state( ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import load_chat_template + + # The Cohere serving handler depends on the optional `cohere` SDK for + # its wire-format protocol models, and is additionally gated on the + # `VLLM_ENABLE_COHERE_API` env flag (see + # `vllm.entrypoints.cohere.api_router.attach_router`). Skip the import + # entirely when the endpoint isn't going to be exposed, both because + # the SDK may not be installed and because the serving object holds + # nontrivial state (chat handler, warmup) that would otherwise be + # unused. + if envs.VLLM_ENABLE_COHERE_API: + try: + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + except ImportError: + CohereServingChatV2 = None # type: ignore[assignment,misc] + else: + CohereServingChatV2 = None # type: ignore[assignment,misc] + from vllm.entrypoints.mcp.tool_server import ( DemoToolServer, MCPToolServer, @@ -86,6 +111,19 @@ async def init_generate_state( tool_server = None resolved_chat_template = load_chat_template(args.chat_template) + # Fold the dedicated ``--cohere-format`` CLI flag into the renderer's + # default chat-template kwargs. The cohere renderer reads + # ``chat_template_kwargs["cohere_format"]`` to pick cmd3 vs cmd4 + # rendering; making this a first-class flag keeps the right format + # discoverable for ``vllm serve --tokenizer-mode cohere`` users + # without forcing them to hand-construct a JSON dict for + # ``--default-chat-template-kwargs``. Per-request overrides still + # take precedence (see ``merge_kwargs`` in + # ``ChatCompletionRequest.build_chat_params``). + default_chat_template_kwargs = dict(args.default_chat_template_kwargs or {}) + if getattr(args, "cohere_format", None): + default_chat_template_kwargs.setdefault("cohere_format", args.cohere_format) + # Render endpoints are always backed by OnlineRenderer so that # /v1/chat/completions/render and /v1/completions/render work on both # generate-mode and render-only servers. Created in init_app_state. @@ -106,7 +144,7 @@ async def init_generate_state( enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, enable_log_outputs=args.enable_log_outputs, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None @@ -119,7 +157,7 @@ async def init_generate_state( request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, @@ -140,8 +178,6 @@ async def init_generate_state( if "generate" in supported_tasks else None ) - if state.openai_serving_chat is not None: - state.openai_serving_chat.warmup() state.openai_serving_completion = ( OpenAIServingCompletion( engine_client, @@ -171,11 +207,32 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None ) + state.cohere_serving_chat_v2 = ( + CohereServingChatV2( + engine_client, + state.openai_serving_models, + args.response_role, + online_renderer=state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + default_chat_template_kwargs=default_chat_template_kwargs, + is_reasoning_model=args.cohere_is_reasoning_model, + ) + if CohereServingChatV2 is not None and "generate" in supported_tasks + else None + ) from .generative_scoring.serving import ServingGenerativeScoring diff --git a/vllm/entrypoints/launcher.py b/vllm/entrypoints/launcher.py index 08a3ab58c780..3a903af2b1fd 100644 --- a/vllm/entrypoints/launcher.py +++ b/vllm/entrypoints/launcher.py @@ -92,6 +92,18 @@ async def serve_http( ) ) + # make sure uvicorn server signal handler has been registered. + # Otherwise the app signal handler below will be overwritten in server.serve. + # server.started is set to True after server complete server.startup. + # We need to constraint the time sequence. + logger.info("API server: waiting for HTTP server to start") + while not server.started and not server_task.done(): + await asyncio.sleep(0.1) + if server_task.done(): + # Propagate startup failure (e.g. port bind error) instead of hanging. + await server_task + logger.info("API server: HTTP server started") + shutdown_event = asyncio.Event() def signal_handler() -> None: diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index b3205728e496..478aa80357db 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -55,6 +55,7 @@ from vllm.v1.engine.llm_engine import LLMEngine from vllm.v1.sample.logits_processor import LogitsProcessor +from ..renderers import ChatParams from .offline_utils import _O, _R, OfflineInferenceMixin if TYPE_CHECKING: @@ -352,6 +353,8 @@ def _make_config(value: Any, cls: type[_R]) -> _R: self.chat_template = load_chat_template(chat_template) self.input_processor = self.llm_engine.input_processor + self.renderer.warmup(ChatParams(chat_template=self.chat_template)) + # The renderer thread pool is only consumed by the async renderer # path; the synchronous `LLM` entrypoint runs multimodal # preprocessing serially. Warn so the setting is not a silent @@ -885,9 +888,19 @@ def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: "update_weights", kwargs={"update_info": update_info_dict} ) - def finish_weight_update(self) -> None: - """Finish the current weight update.""" + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 59c7ee84caef..5a6dd8c83dd0 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -28,7 +28,6 @@ from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware @@ -42,20 +41,15 @@ ) from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.server_utils import ( - engine_error_handler, exception_handler, - generation_error_handler, get_uvicorn_log_config, http_exception_handler, lifespan, log_response, validation_exception_handler, + vllm_error_handler, ) -from vllm.exceptions import ( - VLLMNotFoundError, - VLLMUnprocessableEntityError, - VLLMValidationError, -) +from vllm.exceptions import VLLMError from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_derenderer import OnlineDerenderer @@ -67,7 +61,6 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import is_valid_ipv6_address from vllm.utils.system_utils import decorate_logs, set_ulimit -from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.version import __version__ as VLLM_VERSION prometheus_multiproc_dir: tempfile.TemporaryDirectory @@ -269,6 +262,13 @@ def build_app( register_pooling_api_routers(app, supported_tasks, model_config) + if args.enable_fault_tolerance: + from vllm.entrypoints.serve.fault_tolerance.api_router import ( + register_fault_tolerance_api_router, + ) + + register_fault_tolerance_api_router(app) + # Endpoint plugins are attached last so their routes are registered after all core # routers. This runs even for the CPU only render server. A plugin eligible for # the `render` task still gets its routes registered. It receives @@ -284,23 +284,26 @@ def build_app( allow_headers=args.allowed_headers, ) + # Exception handlers are registered in four layers: + # 1. framework errors raised by FastAPI/Starlette + # 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler + # 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` + # 4. the raw ``Exception`` handler as a safety net + # Registering specific exception types (rather than only ``Exception``) + # ensures they are handled by ``ExceptionMiddleware`` (inside the Prometheus + # middleware) rather than ``ServerErrorMiddleware`` (outside it), so their + # status codes are recorded correctly. app.exception_handler(HTTPException)(http_exception_handler) app.exception_handler(RequestValidationError)(validation_exception_handler) - app.exception_handler(EngineGenerateError)(engine_error_handler) - app.exception_handler(EngineDeadError)(engine_error_handler) - app.exception_handler(GenerationError)(generation_error_handler) - # Register specific exception types so they are handled by - # ExceptionMiddleware (inside the Prometheus middleware) rather than - # ServerErrorMiddleware (outside it). Without this, these exceptions - # propagate through Prometheus as unhandled and get recorded as 5xx - # even though they result in 4xx responses to the client. - app.exception_handler(VLLMValidationError)(exception_handler) - app.exception_handler(VLLMUnprocessableEntityError)(exception_handler) - app.exception_handler(VLLMNotFoundError)(exception_handler) + + app.exception_handler(VLLMError)(vllm_error_handler) + + # TODO(zqzten): remove these fallback handlers after migration to VLLMError app.exception_handler(ValueError)(exception_handler) app.exception_handler(TypeError)(exception_handler) app.exception_handler(OverflowError)(exception_handler) app.exception_handler(NotImplementedError)(exception_handler) + app.exception_handler(Exception)(exception_handler) # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY @@ -423,6 +426,7 @@ async def init_app_state( default_chat_template_kwargs=args.default_chat_template_kwargs, log_error_stack=args.log_error_stack, ) + state.online_renderer.warmup() state.online_derenderer = OnlineDerenderer( model_config=engine_client.model_config, @@ -526,6 +530,7 @@ async def init_render_app_state( default_chat_template_kwargs=args.default_chat_template_kwargs, log_error_stack=args.log_error_stack, ) + state.online_renderer.warmup() state.online_derenderer = OnlineDerenderer( model_config=vllm_config.model_config, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 8c1694bbf73e..1cdfd2f698f9 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -10,7 +10,13 @@ ChatCompletionAudio as OpenAIChatCompletionAudio, ) from openai.types.chat.chat_completion_message import Annotation as OpenAIAnnotation -from pydantic import Field, PrivateAttr, model_serializer, model_validator +from pydantic import ( + Field, + PrivateAttr, + SerializeAsAny, + model_serializer, + model_validator, +) from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ( @@ -91,7 +97,12 @@ class ChatCompletionLogProbs(OpenAIBaseModel): class ChatCompletionResponseChoice(OpenAIBaseModel): index: int - message: ChatMessage + # ``SerializeAsAny`` lets pydantic honor subclasses of ``ChatMessage`` + # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereChatMessage``) + # so that added fields like ``citations`` survive JSON serialization + # instead of being stripped down to the base schema. Plain + # ``ChatMessage`` instances serialize identically to before. + message: SerializeAsAny[ChatMessage] logprobs: ChatCompletionLogProbs | None = None # per OpenAI spec this is the default finish_reason: str | None = "stop" @@ -139,7 +150,11 @@ class ChatCompletionResponse(OpenAIBaseModel): class ChatCompletionResponseStreamChoice(OpenAIBaseModel): index: int - delta: DeltaMessage + # ``SerializeAsAny`` lets pydantic honor subclasses of ``DeltaMessage`` + # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereDeltaMessage``) + # so streaming ``citations`` survive JSON serialization. Plain + # ``DeltaMessage`` instances serialize identically to before. + delta: SerializeAsAny[DeltaMessage] logprobs: ChatCompletionLogProbs | None = None finish_reason: str | None = None stop_reason: int | str | None = None @@ -179,6 +194,7 @@ def _propagate_defer_loading(self) -> "ChatCompletionToolsParam": @model_serializer(mode="wrap") def _serialize(self, handler): data = handler(self) + data = {k: v for k, v in data.items() if k in type(self).model_fields} if self.defer_loading is None: data.pop("defer_loading", None) return data @@ -476,6 +492,16 @@ class ChatCompletionRequest(OpenAIBaseModel): "can detect such behavior and terminate early, saving time and tokens.", ) + stream_interval: Annotated[int, Field(ge=1)] | None = Field( + default=None, + description=( + "Number of tokens to batch into each streamed chunk. Raises the " + "server's `--stream-interval` for this request. Values below the " + "server setting are clamped up to it. The first and last chunks " + "are always sent immediately. Ignored for non-streaming requests." + ), + ) + # --8<-- [end:chat-completion-extra-params] @model_validator(mode="before") @@ -523,8 +549,8 @@ def _materialize_tool_calls_after(self) -> "ChatCompletionRequest": msg["tool_calls"] = list(tool_calls) return self - _grammar_from_tool_parser: bool = PrivateAttr(default=False) - """CAUTION: Should only be set by ``ToolParser.adjust_request``.""" + _grammar_from_parser: bool = PrivateAttr(default=False) + """CAUTION: Should only be set by the parser-engine adapter's adjust_request.""" def build_chat_params( self, @@ -555,6 +581,12 @@ def build_chat_params( ), media_io_kwargs=self.media_io_kwargs, return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask), + # No-tools requests default to tool_choice="none" at the API + # layer. Collapse that default before rendering, so K3 emits a + # model-visible tool-choice instruction only for requests with a + # tools block. + tool_choice=self.tool_choice if self.tools else None, + response_format=self.response_format, ) def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: @@ -687,9 +719,10 @@ def to_sampling_params( skip_special_tokens=self.skip_special_tokens, spaces_between_special_tokens=self.spaces_between_special_tokens, include_stop_str_in_output=self.include_stop_str_in_output, - output_kind=RequestOutputKind.DELTA - if self.stream - else RequestOutputKind.FINAL_ONLY, + output_kind=( + RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY + ), + stream_interval=self.stream_interval, structured_outputs=self.extract_structured_outputs(), logit_bias=self.logit_bias, bad_words=self.bad_words, @@ -757,6 +790,18 @@ def check_logprobs(cls, data): parameter="logprob_token_ids", ) + # These fields are integers, but `mode="before"` runs on the raw + # request data, so a non-numeric value (e.g. a JSON string) would + # reach the comparisons below and raise TypeError -> HTTP 500. Reject + # it here so the client gets a clean 400 instead. + for field_name in ("prompt_logprobs", "top_logprobs"): + field_value = data.get(field_name) + if field_value is not None and not isinstance(field_value, (int, float)): + raise VLLMValidationError( + f"`{field_name}` must be an integer.", + parameter=field_name, + value=field_value, + ) if (prompt_logprobs := data.get("prompt_logprobs")) is not None: if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1): raise VLLMValidationError( @@ -837,9 +882,10 @@ def check_tool_usage(cls, data): # Reject empty tools array, matching OpenAI API behavior if data.get("tools") == []: - raise ValueError( + raise VLLMValidationError( "`tools` must not be an empty array. " - "Either provide at least one tool or omit the field entirely." + "Either provide at least one tool or omit the field entirely.", + parameter="tools", ) # if "tool_choice" is not specified but tools are provided, @@ -1063,13 +1109,15 @@ def check_batch_mode(cls, data: Any) -> Any: if isinstance(data, BatchChatCompletionRequest): data = data.model_dump(exclude_unset=True) if data.get("use_beam_search"): - raise ValueError( + raise VLLMValidationError( "Batch chat completions do not support beam search. " - "Please set `use_beam_search` to False." + "Please set `use_beam_search` to False.", + parameter="use_beam_search", ) if data.get("logprob_token_ids") and not data.get("logprobs"): - raise ValueError( - "when using `logprob_token_ids`, `logprobs` must be set to true." + raise VLLMValidationError( + "when using `logprob_token_ids`, `logprobs` must be set to true.", + parameter="logprob_token_ids", ) response_format = data.get("response_format") rf_type = ( @@ -1083,8 +1131,10 @@ def check_batch_mode(cls, data: Any) -> Any: validate_structured_outputs_structural_tag(structured_outputs) n = data.get("n", 1) if n is not None and n != 1: - raise ValueError( - "Batch chat completions do not support `n > 1`. Please set `n` to 1." + raise VLLMValidationError( + "Batch chat completions do not support `n > 1`. Please set `n` to 1.", + parameter="n", + value=n, ) return data diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 4c05db0b8a48..7b3ce551b209 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -60,12 +60,10 @@ from vllm.outputs import RequestOutput from vllm.parser import ParserManager from vllm.parser.abstract_parser import Parser -from vllm.renderers import ChatParams from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.utils.collection_utils import as_list -from vllm.utils.mistral import is_mistral_tool_parser logger = init_logger(__name__) @@ -157,15 +155,6 @@ def __init__( model_name=self.model_config.model, is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) - if ( - self.parser_cls is not None - and is_mistral_tool_parser(self.parser_cls.tool_parser_cls) - and self.parser_cls.reasoning_parser_cls is not None - ): - from vllm.tool_parsers.mistral_tool_parser import MistralToolParser - - MistralToolParser.model_can_reason = True - self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none self.enable_prompt_tokens_details = enable_prompt_tokens_details @@ -188,15 +177,6 @@ def __init__( self.supports_code_interpreter = False self.python_tool = None - def warmup(self) -> None: - self.renderer.warmup( - ChatParams( - chat_template=self.chat_template, - chat_template_content_format=self.chat_template_content_format, - chat_template_kwargs=self.default_chat_template_kwargs, - ) - ) - def _effective_chat_template_kwargs( self, request: ChatCompletionRequest ) -> dict[str, Any]: @@ -350,7 +330,7 @@ async def _create_chat_completion( else: if not request.include_reasoning: reasoning_ended = True - elif request._grammar_from_tool_parser: + elif request._grammar_from_parser: # The Mistral grammar already includes an optional # `think?` rule that handles both reasoning and # non-reasoning outputs. @@ -411,6 +391,34 @@ def get_chat_request_role(self, request: ChatCompletionRequest) -> str: return self.response_role return request.messages[-1]["role"] + def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage: + """Construct the response :class:`ChatMessage` for the non-streaming path. + + The full-generator calls this at every construction site so + subclasses can swap in a specialized :class:`ChatMessage` + subclass (e.g. :class:`CohereServingChatV2` returning + :class:`CohereChatMessage`) without duplicating the branchy + tool-choice / auto-tools logic that decides which fields are + populated. The default returns a plain :class:`ChatMessage`. + """ + return ChatMessage(*args, **kwargs) + + def _finalize_response_message( + self, + message: ChatMessage, + *, + parser: Parser | None, + ) -> ChatMessage: + """Subclass hook to enrich a fully-constructed :class:`ChatMessage`. + + Default is a no-op. Subclasses that need to surface parser-side + extras (e.g. :class:`CohereServingChatV2` reading grounding + citations off the reasoning parser and populating + :class:`CohereChatMessage.citations`) override this to inspect + ``parser`` and mutate/replace ``message``. + """ + return message + async def chat_completion_stream_generator( self, request: ChatCompletionRequest, @@ -914,13 +922,19 @@ async def chat_completion_full_generator( ) is_required_tool_choice = request.tool_choice == "required" + # All six construction sites route through ``self._create_chat_message`` + # so subclasses can swap in a specialized :class:`ChatMessage` + # (e.g. the Cohere v2 handler's ``CohereChatMessage``) without + # having to duplicate this branch logic. if (not self.enable_auto_tools or not tool_parser_cls) and ( not is_named_tool_choice and not is_required_tool_choice ): - message = ChatMessage(role=role, reasoning=reasoning, content=content) + message = self._create_chat_message( + role=role, reasoning=reasoning, content=content + ) elif is_named_tool_choice or is_required_tool_choice: - message = ChatMessage( + message = self._create_chat_message( role=role, reasoning=reasoning, content=content or "", @@ -933,7 +947,9 @@ async def chat_completion_full_generator( # if the request doesn't use tool choice # OR specifies to not use a tool elif not request.tool_choice or request.tool_choice == "none": - message = ChatMessage(role=role, reasoning=reasoning, content=content) + message = self._create_chat_message( + role=role, reasoning=reasoning, content=content + ) # handle when there are tools and tool choice is auto elif ( @@ -944,7 +960,7 @@ async def chat_completion_full_generator( ): auto_tools_called = tool_calls is not None and len(tool_calls) > 0 if tool_calls: - message = ChatMessage( + message = self._create_chat_message( role=role, reasoning=reasoning, content=content, @@ -955,7 +971,7 @@ async def chat_completion_full_generator( ) else: - message = ChatMessage( + message = self._create_chat_message( role=role, reasoning=reasoning, content=content, @@ -968,7 +984,17 @@ async def chat_completion_full_generator( " if tools should be extracted. Returning a standard chat " "completion." ) - message = ChatMessage(role=role, reasoning=reasoning, content=content) + message = self._create_chat_message( + role=role, reasoning=reasoning, content=content + ) + + # Subclass hook: enrich the constructed message with any + # parser-side extras that don't fit through the plain + # ``(reasoning, content, tool_calls)`` tuple. Base is a no-op; + # citation-aware handlers use this to surface grounding + # metadata cached on the reasoning parser. + message = self._finalize_response_message(message, parser=parser) + # In OpenAI's API, when a tool is called, the finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 8dbb69943903..fbb5a6df5f50 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -148,6 +148,25 @@ class BaseFrontendArgs: """If set to False, output deltas will not be logged. Relevant only if --enable-log-outputs is set. """ + cohere_is_reasoning_model: bool = True + """Cohere ``/cohere/v2/chat`` only. Whether the served model is a + reasoning Command-family model. When True (default), the assistant's + chain-of-thought is surfaced as a ``thinking`` content block (or + ``content-*`` events on the stream). When False, reasoning is + surfaced as Cohere's ``tool_plan`` field (or ``tool-plan-delta`` + events) whenever the model emits tool calls, matching older non- + reasoning Command models. Has no effect on the non-Cohere + endpoints.""" + cohere_format: str = "cmd4" + """Cohere ``--tokenizer-mode cohere`` only. Which Cohere prompt + format to render: ``cmd4`` (current Command A+ models; default) or + ``cmd3`` (earlier Cmd-A and Cmd-A reasoning models). Selecting the + wrong format silently produces a prompt the model wasn't trained + on, which most commonly manifests as the model emitting text but no + citations / tool calls / thinking blocks. Equivalent to passing + ``--default-chat-template-kwargs '{"cohere_format": "..."}'`` -- any + explicit request-level ``chat_template_kwargs.cohere_format`` takes + priority.""" log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE """If set to True, log the stack trace of error responses""" tokens_only: bool = False diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 73677b16af9b..79ad9799e18d 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -233,6 +233,16 @@ class CompletionRequest(OpenAIBaseModel): ), ) + stream_interval: Annotated[int, Field(ge=1)] | None = Field( + default=None, + description=( + "Number of tokens to batch into each streamed chunk. Raises the " + "server's `--stream-interval` for this request. Values below the " + "server setting are clamped up to it. The first and last chunks " + "are always sent immediately. Ignored for non-streaming requests." + ), + ) + # --8<-- [end:completion-extra-params] def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: @@ -365,6 +375,7 @@ def to_sampling_params( output_kind=RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY, + stream_interval=self.stream_interval, structured_outputs=self.extract_structured_outputs(), logit_bias=self.logit_bias, allowed_token_ids=self.allowed_token_ids, @@ -468,6 +479,18 @@ def check_logprobs(cls, data): parameter="logprob_token_ids", ) + # These fields are integers, but `mode="before"` runs on the raw + # request data, so a non-numeric value (e.g. a JSON string) would + # reach the comparisons below and raise TypeError -> HTTP 500. Reject + # it here so the client gets a clean 400 instead. + for field_name in ("prompt_logprobs", "logprobs"): + field_value = data.get(field_name) + if field_value is not None and not isinstance(field_value, (int, float)): + raise VLLMValidationError( + f"`{field_name}` must be an integer.", + parameter=field_name, + value=field_value, + ) if (prompt_logprobs := data.get("prompt_logprobs")) is not None: if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1): raise VLLMValidationError( diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index d669ec4d1d58..8ce6233c1ba1 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -257,7 +257,7 @@ def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: name = f"APIServer_DP{child_args.data_parallel_rank}" set_process_title(name) decorate_logs(name) - if envs.VLLM_RUST_FRONTEND_PATH: + if envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH: _run_rust_vllm_dp_server(child_args) else: _run_python_vllm_dp_server(child_args) diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 05536f14217c..805c639d7d16 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -19,7 +19,7 @@ from vllm.config.utils import replace from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMServerError, VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import StructuredOutputsParams from vllm.utils import random_uuid @@ -296,6 +296,7 @@ class FunctionDefinition(OpenAIBaseModel): @model_serializer(mode="wrap") def _serialize(self, handler): data = handler(self) + data = {k: v for k, v in data.items() if k in type(self).model_fields} if self.strict is None: data.pop("strict", None) if self.defer_loading is None: @@ -404,7 +405,7 @@ def _serialize(self, handler): return data -class GenerationError(Exception): +class GenerationError(VLLMServerError): """raised when finish_reason indicates internal server error (500)""" def __init__(self, message: str = "Internal server error"): diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 3f6857dcc323..ad92766cb7b1 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -337,6 +337,7 @@ def build_chat_params( extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, + tool_choice=self.tool_choice if self.tools else None, ) def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 8590afe07fbd..344ce0a209c1 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -610,7 +610,13 @@ async def _make_request( request: ResponsesRequest, prev_response: ResponsesResponse | None, ): - tool_dicts = construct_tool_dicts(request.tools, request.tool_choice) + tool_dicts = construct_tool_dicts( + request.tools, + request.tool_choice, + exclude_tools_when_tool_choice_none=( + self.online_renderer.exclude_tools_when_tool_choice_none + ), + ) # Construct the input messages. messages = construct_input_messages( request_instructions=request.instructions, @@ -743,11 +749,14 @@ def _make_request_with_harmony( request: ResponsesRequest, prev_response: ResponsesResponse | None, ): - if request.tool_choice not in ("auto", "none"): - raise NotImplementedError( - "Only 'auto' or 'none' tool_choice is supported " - "in response API with Harmony" - ) + if self.parser is not None: + # HarmonyParser doesn't need chat_template_kwargs + # TODO: Unify adjust_request() call with non-harmony branch + self.parser( + self.renderer.get_tokenizer(), + request.tools, + model_config=self.model_config, + ).adjust_request(request=request) arrival_time = time.time() messages = self._construct_input_messages_with_harmony(request, prev_response) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 07a9704f9b54..d3de64844a0f 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -31,8 +31,11 @@ from vllm import envs from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionMessageParam -from vllm.entrypoints.openai.engine.protocol import FunctionCall +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionMessageParam, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.engine.protocol import FunctionCall, FunctionDefinition from vllm.entrypoints.openai.responses.protocol import ResponseInputOutputItem from vllm.logger import init_logger from vllm.tool_parsers.utils import ( @@ -363,27 +366,30 @@ def extract_tool_types(tools: list[Tool]) -> set[str]: return tool_types -def convert_tool_responses_to_completions_format(tool: dict) -> dict: +def convert_tool_responses_to_completions_format( + tool: dict, +) -> ChatCompletionToolsParam: """ - Convert a flat tool schema: + Convert a flat Responses tool schema: {"type": "function", "name": "...", "description": "...", "parameters": {...}} - into: - {"type": "function", "function": {...}} + into a Chat Completions tool param for chat-template rendering. """ - return { - "type": "function", - "function": tool, - } + return ChatCompletionToolsParam( + type="function", + function=FunctionDefinition.model_validate( + {k: v for k, v in tool.items() if k != "type"} + ), + ) def construct_tool_dicts( - tools: list[Tool], tool_choice: ToolChoice + tools: list[Tool], + tool_choice: ToolChoice, + exclude_tools_when_tool_choice_none: bool = False, ) -> list[dict[str, Any]] | None: - if not tools or (tool_choice == "none"): - tool_dicts = None - else: - tool_dicts = [ - convert_tool_responses_to_completions_format(tool) - for tool in iter_response_function_tool_dicts(tools) - ] - return tool_dicts + if not tools or (tool_choice == "none" and exclude_tools_when_tool_choice_none): + return None + return [ + convert_tool_responses_to_completions_format(tool).model_dump() + for tool in iter_response_function_tool_dicts(tools) + ] diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index cdc4c16cacb1..f74b740276e8 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from concurrent.futures import Executor from typing import Any, Final, cast from vllm import ( @@ -10,21 +11,17 @@ ) from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, ChatTemplateConfig, - ChatTemplateContentFormatOption, - ConversationMessage, ) -from vllm.entrypoints.serve.engine.typing import RendererChatRequest, RendererRequest -from vllm.inputs import EngineInput, SingletonPrompt from vllm.lora.request import LoRARequest from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq -from vllm.tool_parsers import ToolParser +from vllm.utils.async_utils import make_async from vllm.utils.mistral import is_mistral_tokenizer from ..typing import ( - ALLOfflineInputsContext, + AnyOfflineInputsContext, + AnyRenderParam, EncodeChatRenderParams, EncodeCMPLRenderParams, OfflineEncodeInputsContext, @@ -66,14 +63,25 @@ def __init__( chat_template_config.trust_request_chat_template ) + self.template_kwargs = None + self.tool_dicts = None + + # Shared thread pool executor for preprocessing + self._executor: Executor = self.renderer._executor + self.render_async = make_async(self.render, executor=self._executor) + ####################################### # online APIs def create_pooling_params(self, request): return request.to_pooling_params() - def pre_process_online(self, ctx: PoolingServeContext): + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: request = ctx.request + renderer = self.renderer + requests: Sequence[AnyRenderParam] if isinstance(request, PoolingChatLikeRequest): self._validate_chat_template( @@ -81,24 +89,87 @@ def pre_process_online(self, ctx: PoolingServeContext): chat_template_kwargs=request.chat_template_kwargs, trust_request_chat_template=self.trust_request_chat_template, ) - _, engine_inputs = self._preprocess_chat_online( - request, - request.messages, - default_template=self.chat_template, - default_template_content_format=self.chat_template_content_format, - default_template_kwargs=None, + + num_requests = 1 + default_template_kwargs = merge_kwargs( + self.template_kwargs, + dict( + tools=self.tool_dicts, + tokenize=is_mistral_tokenizer(renderer.tokenizer), + ), + ) + + mm_config = self.model_config.multimodal_config + tok_params = request.build_tok_params(self.model_config) + chat_params = request.build_chat_params( + self.chat_template, self.chat_template_content_format + ).with_defaults( + default_template_kwargs, + default_media_io_kwargs=( + mm_config.media_io_kwargs if mm_config else None + ), + ) + + params_seq = self._params_to_seq(ctx.pooling_params, num_requests) + seq_lora_requests = self._lora_request_to_seq( + ctx.lora_request, num_requests ) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + requests = [ + EncodeChatRenderParams( + conversations=request.messages, + chat_params=chat_params, + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] + + return requests + elif isinstance(request, PoolingCompletionLikeRequest): - engine_inputs = self._preprocess_cmpl_online( - request, - prompt_input=request.input, - prompt_embeds=None, + model_config = self.model_config + prompts_seq = prompt_to_seq(request.input) + num_requests = len(prompts_seq) + + parsed_prompts = [ + ( + prompt + if isinstance(prompt, bytes) + else parse_model_prompt(model_config, prompt) + ) + for prompt in prompts_seq + ] + tok_params = request.build_tok_params(model_config) + + params_seq = self._params_to_seq(ctx.pooling_params, num_requests) + seq_lora_requests = self._lora_request_to_seq( + ctx.lora_request, num_requests ) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + requests = [ + EncodeCMPLRenderParams( + prompts=parsed_prompts[i], + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] + + return requests else: raise ValueError(f"Invalid {self.name} request type") - ctx.engine_inputs = engine_inputs - def post_process_online( self, ctx: PoolingServeContext, @@ -109,7 +180,7 @@ def post_process_online( # offline APIs def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext + self, ctx: AnyOfflineInputsContext ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineEncodeInputsContext) @@ -209,84 +280,6 @@ def render( priorities=render_params["priorities"], ) - def _preprocess_cmpl_online( - self, - request: RendererRequest, - prompt_input: str | list[str] | list[int] | list[list[int]] | None, - prompt_embeds: bytes | list[bytes] | None, - ) -> list[EngineInput]: - renderer = self.renderer - model_config = self.model_config - - prompts = list[SingletonPrompt | bytes]() - if prompt_embeds is not None: # embeds take higher priority - prompts.extend(prompt_to_seq(prompt_embeds)) - if prompt_input is not None: - prompts.extend(prompt_to_seq(prompt_input)) - - parsed_prompts = [ - ( - prompt - if isinstance(prompt, bytes) - else parse_model_prompt(model_config, prompt) - ) - for prompt in prompts - ] - tok_params = request.build_tok_params(model_config) - - return renderer.render_cmpl( - parsed_prompts, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt") - if (v := getattr(request, k, None)) is not None - }, - ) - - def _preprocess_chat_online( - self, - request: RendererChatRequest, - messages: list[ChatCompletionMessageParam], - default_template: str | None, - default_template_content_format: ChatTemplateContentFormatOption, - default_template_kwargs: dict[str, Any] | None, - tool_dicts: list[dict[str, Any]] | None = None, - tool_parser: type[ToolParser] | None = None, - ) -> tuple[list[ConversationMessage], list[EngineInput]]: - renderer = self.renderer - - default_template_kwargs = merge_kwargs( - default_template_kwargs, - dict( - tools=tool_dicts, - tokenize=is_mistral_tokenizer(renderer.tokenizer), - ), - ) - - mm_config = self.model_config.multimodal_config - - tok_params = request.build_tok_params(self.model_config) - chat_params = request.build_chat_params( - default_template, default_template_content_format - ).with_defaults( - default_template_kwargs, - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), - ) - - (conversation,), (engine_input,) = renderer.render_chat( - [messages], - chat_params, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt") - if (v := getattr(request, k, None)) is not None - }, - ) - - return conversation, [engine_input] - def _validate_chat_template( self, request_chat_template: str | None, @@ -341,10 +334,13 @@ def _lora_request_to_seq( def _priority_to_seq( self, - priority: Sequence[int] | None, + priority: int | Sequence[int] | None, num_requests: int, ) -> Sequence[int]: if priority is not None: + if isinstance(priority, int): + return [priority] * num_requests + if len(priority) != num_requests: raise ValueError( f"The lengths of prompts ({num_requests}) " diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 79f93a148049..36b1c22e61a5 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - +import asyncio from abc import ABC, abstractmethod from collections.abc import AsyncGenerator, Mapping from concurrent.futures import Executor @@ -63,9 +63,6 @@ def __init__( # Shared thread pool executor for preprocessing and postprocessing. self._executor: Executor = self.renderer._executor - self._preprocessing_async = make_async( - self._preprocessing, executor=self._executor - ) self._postprocessing_async = make_async( self._postprocessing, executor=self._executor ) @@ -77,7 +74,7 @@ async def __call__( ) -> Response: io_processor = self.get_io_processor(request) ctx = await self._init_ctx(io_processor, request, raw_request) - await self._preprocessing_async(io_processor, ctx) + await self._preprocessing(io_processor, ctx) await self._prepare_generators(ctx) await self._collect_batch(ctx) return await self._postprocessing_async(io_processor, ctx) @@ -86,11 +83,17 @@ async def __call__( def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: raise NotImplementedError - @torch.inference_mode() - def _preprocessing( + async def _preprocessing( self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext ): - return io_processor.pre_process_online(ctx) + requests = io_processor.get_request_factory_online(ctx) + + if len(requests) == 0: + raise ValueError("You must pass at least one prompt") + + ctx.engine_inputs = await asyncio.gather( + *[io_processor.render_async(request) for request in requests] + ) @torch.inference_mode() def _postprocessing( @@ -110,16 +113,26 @@ async def _init_ctx( await self._check_model(request) pooling_params = io_processor.create_pooling_params(request) + lora_request = self._maybe_get_adapters(request) + priorities = getattr(request, "priority", 0) + prompt_extras = { + k: v + for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") + if (v := getattr(request, k, None)) is not None + } + ctx = PoolingServeContext( request=request, raw_request=raw_request, model_name=model_name, pooling_params=pooling_params, request_id=request_id, + lora_request=lora_request, + priorities=priorities, + prompt_extras=prompt_extras, ) self._validate_request(ctx) - ctx.lora_request = self._maybe_get_adapters(ctx.request) return ctx async def _prepare_generators( @@ -127,7 +140,7 @@ async def _prepare_generators( ctx: PoolingServeContext, ): if ctx.engine_inputs is None: - raise ValueError("Engine prompts not available") + raise ValueError("Engine inputs not available") generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] @@ -139,12 +152,7 @@ async def _prepare_generators( assert ctx.pooling_params is not None pooling_params = ctx.pooling_params - - if isinstance(pooling_params, list): - for params in pooling_params: - params.verify(self.model_config) - else: - pooling_params.verify(self.model_config) + pooling_params.verify(self.model_config) for i, engine_input in enumerate(ctx.engine_inputs): prompt_request_id = ( @@ -153,26 +161,20 @@ async def _prepare_generators( else ctx.prompt_request_ids[i] ) - params = ( - pooling_params[i] - if isinstance(pooling_params, list) - else pooling_params - ) - self._log_inputs( prompt_request_id, - engine_input, - params=params, + engine_input["prompts"], + params=engine_input["params"], lora_request=ctx.lora_request, ) generator = self.engine_client.encode( - engine_input, - params, - prompt_request_id, - lora_request=ctx.lora_request, + request_id=prompt_request_id, + prompt=engine_input["prompts"], + pooling_params=engine_input["params"], + lora_request=engine_input["lora_requests"], + priority=engine_input["priorities"], trace_headers=trace_headers, - priority=getattr(ctx.request, "priority", 0), ) generators.append(generator) diff --git a/vllm/entrypoints/pooling/classify/api_router.py b/vllm/entrypoints/pooling/classify/api_router.py index 9e016a72e843..fc86f49d6fd1 100644 --- a/vllm/entrypoints/pooling/classify/api_router.py +++ b/vllm/entrypoints/pooling/classify/api_router.py @@ -16,8 +16,11 @@ router = APIRouter() -def classify(request: Request) -> ServingClassification | None: - return request.app.state.serving_classification +def classify(request: Request) -> ServingClassification: + handler = getattr(request.app.state, "serving_classification", None) + if handler is None: + raise NotImplementedError("The model does not support Classification API") + return handler @router.post("/classify", dependencies=[Depends(validate_json_request)]) @@ -27,7 +30,4 @@ async def create_classify( request: ClassificationRequest, raw_request: Request ) -> Response: handler = classify(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Classification API") - return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/embed/api_router.py b/vllm/entrypoints/pooling/embed/api_router.py index 7ffb5840d5bd..adee8a94310b 100644 --- a/vllm/entrypoints/pooling/embed/api_router.py +++ b/vllm/entrypoints/pooling/embed/api_router.py @@ -18,8 +18,11 @@ router = APIRouter() -def embedding(request: Request) -> ServingEmbedding | None: - return request.app.state.serving_embedding +def embedding(request: Request) -> ServingEmbedding: + handler = getattr(request.app.state, "serving_embedding", None) + if handler is None: + raise NotImplementedError("The model does not support Embeddings API") + return handler @router.post( @@ -37,9 +40,6 @@ async def create_embedding( raw_request: Request, ): handler = embedding(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Embeddings API") - return await handler(request, raw_request) @@ -58,7 +58,4 @@ async def create_cohere_embedding( raw_request: Request, ): handler = embedding(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Embeddings API") - return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 5c2f825a9675..57728c88377c 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -17,7 +17,7 @@ ChatCompletionMessageParam, CustomChatCompletionMessageParam, ) -from vllm.inputs import EngineInput, tokens_input +from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput from vllm.renderers import merge_kwargs @@ -28,11 +28,14 @@ from ..base.io_processor import PoolingIOProcessor from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( - ALLOfflineInputsContext, + AnyOfflineInputsContext, + AnyRenderParam, ChunkedEmbeddingMetadata, + EncodeChatRenderParams, OfflineEncodeInputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, + PoolingEngineInput, PoolingServeContext, RequestFactory, ) @@ -76,9 +79,11 @@ def __init__(self, *args, **kwargs): list(self.task_instructions.keys()), ) - def pre_process_online(self, ctx: PoolingServeContext): + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: if isinstance(ctx.request, CohereEmbedRequest): - self._pre_process_cohere_online(ctx) + requests = self._get_request_factory_cohere_online(ctx) elif isinstance( ctx.request, ( @@ -88,12 +93,11 @@ def pre_process_online(self, ctx: PoolingServeContext): EmbeddingBatchChatInputRequest, ), ): - self._pre_process_openai_chat_online(ctx) + requests = self._get_request_factory_chat_input_online(ctx) else: - super().pre_process_online(ctx) + requests = super().get_request_factory_online(ctx) - if self.enable_chunked_processing: - self._pre_process_chunked(ctx) + return requests def post_process_online( self, @@ -114,18 +118,21 @@ def post_process_online( # PTAL: examples/pooling/embed/openai_embedding_long_text ################################################################# - def _pre_process_chunked(self, ctx: PoolingServeContext) -> None: + def maybe_pre_process_chunked(self, ctx: PoolingServeContext) -> None: + if not self.enable_chunked_processing: + return None + if ctx.engine_inputs is None: raise ValueError("Engine prompts not available") ctx.original_engine_inputs = ctx.engine_inputs request_id = ctx.request_id max_model_len = self.model_config.max_model_len - chunked_engine_inputs: list[EngineInput] = [] + chunked_engine_inputs: list[PoolingEngineInput] = [] prompt_request_ids: list[str] = [] chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] = [] for prompt_idx, engine_input in enumerate(ctx.engine_inputs): - token_ids = engine_input.get("prompt_token_ids", None) + token_ids = engine_input["prompts"].get("prompt_token_ids", None) if token_ids is None: raise NotImplementedError( "Long Text Embedding with Chunked Processing does " @@ -138,7 +145,12 @@ def _pre_process_chunked(self, ctx: PoolingServeContext) -> None: chunk_list(prompt_token_ids, max_model_len) ): chunked_engine_inputs.append( - tokens_input(prompt_token_ids=chunk_tokens) + PoolingEngineInput( + prompts=tokens_input(prompt_token_ids=chunk_tokens), + params=engine_input["params"], + lora_requests=engine_input["lora_requests"], + priorities=engine_input["priorities"], + ) ) prompt_request_ids.append( f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}" @@ -234,7 +246,7 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None: # Get original prompt token IDs for this prompt original_prompt = original_engine_inputs[prompt_idx] - token_ids = original_prompt.get("prompt_token_ids", None) + token_ids = original_prompt["prompts"].get("prompt_token_ids", None) if token_ids is None: raise NotImplementedError( "Long Text Embedding with Chunked Processing does " @@ -262,6 +274,68 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None: return None + ################################################################# + # Chat input Request Preprocessing & Postprocessing + ################################################################# + + def _get_request_factory_chat_input_online( + self, + ctx: PoolingServeContext, + ) -> Sequence[AnyRenderParam]: + request = ctx.request + renderer = self.renderer + + self._validate_chat_template( + request_chat_template=request.chat_template, + chat_template_kwargs=request.chat_template_kwargs, + trust_request_chat_template=self.trust_request_chat_template, + ) + + if isinstance( + request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest) + ): + all_messages = request.messages + else: + all_messages = [request.messages] + num_requests = len(all_messages) + + default_template_kwargs = merge_kwargs( + self.template_kwargs, + dict( + tools=self.tool_dicts, + tokenize=is_mistral_tokenizer(renderer.tokenizer), + ), + ) + + mm_config = self.model_config.multimodal_config + tok_params = request.build_tok_params(self.model_config) + chat_params = request.build_chat_params( + self.chat_template, self.chat_template_content_format + ).with_defaults( + default_template_kwargs, + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), + ) + + params_seq = self._params_to_seq(ctx.pooling_params, num_requests) + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + requests = [ + EncodeChatRenderParams( + conversations=all_messages[i], + chat_params=chat_params, + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] + + return requests + ################################################################# # Cohere Request Preprocessing & Postprocessing ################################################################# @@ -378,71 +452,9 @@ def create_pooling_params(self, request): ) return super().create_pooling_params(request) - def _pre_process_openai_chat_online( - self, - ctx: PoolingServeContext[ - EmbeddingChatRequest - | EmbeddingBatchChatRequest - | EmbeddingChatInputRequest - | EmbeddingBatchChatInputRequest - ], - ) -> None: - request = ctx.request - self._validate_chat_template( - request_chat_template=request.chat_template, - chat_template_kwargs=request.chat_template_kwargs, - trust_request_chat_template=self.trust_request_chat_template, - ) - - if isinstance( - request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest) - ): - all_messages = request.messages - else: - all_messages = [request.messages] - ctx.engine_inputs = self._batch_render_openai_chat(request, all_messages) - - def _batch_render_openai_chat( - self, - request: ( - EmbeddingChatRequest - | EmbeddingBatchChatRequest - | EmbeddingChatInputRequest - | EmbeddingBatchChatInputRequest - ), - all_messages: Sequence[list[ChatCompletionMessageParam]], - ) -> list[EngineInput]: - renderer = self.renderer - mm_config = self.model_config.multimodal_config - - tok_params = request.build_tok_params(self.model_config) - chat_params = request.build_chat_params( - self.chat_template, - self.chat_template_content_format, - ).with_defaults( - merge_kwargs( - None, - dict( - tools=None, - tokenize=is_mistral_tokenizer(renderer.tokenizer), - ), - ), - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), - ) - - _, engine_inputs = renderer.render_chat( - all_messages, - chat_params, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt") - if (v := getattr(request, k, None)) is not None - }, - ) - return engine_inputs - - def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None: + def _get_request_factory_cohere_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: """Convert a ``CohereEmbedRequest`` into engine prompts. If a model has a chat template the task instruction are rendered @@ -479,13 +491,12 @@ def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None: task_prefix = self._get_task_instruction_prefix(input_type) if task_prefix is None: - ctx.engine_inputs = self._preprocess_cohere_text_completion( - request, + return self._get_request_factory_cohere_text_completion( + ctx, texts, truncate_prompt_tokens, truncation_side, ) - return all_messages = [ self._mixed_input_to_messages( @@ -497,27 +508,30 @@ def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None: for text in texts ] if self._has_chat_template(): - ctx.engine_inputs = self._batch_render_chat( - request, + return self._batch_render_chat( + ctx, all_messages, truncate_prompt_tokens, truncation_side, ) else: - ctx.engine_inputs = self._preprocess_cohere_text_completion( - request, + return self._get_request_factory_cohere_text_completion( + ctx, self._apply_task_instruction(texts, input_type), truncate_prompt_tokens, truncation_side, ) - return task_prefix = self._get_task_instruction_prefix(input_type) all_messages = [ self._mixed_input_to_messages(inp, task_prefix=task_prefix) for inp in input ] - ctx.engine_inputs = self._batch_render_chat( - request, all_messages, truncate_prompt_tokens, truncation_side + + return self._batch_render_chat( + ctx, + all_messages, + truncate_prompt_tokens, + truncation_side, ) def _has_chat_template(self) -> bool: @@ -531,13 +545,14 @@ def _has_chat_template(self) -> bool: is not None ) - def _preprocess_cohere_text_completion( + def _get_request_factory_cohere_text_completion( self, - request: CohereEmbedRequest, + ctx: PoolingServeContext, texts: list[str], truncate_prompt_tokens: int | None, truncation_side: Literal["left", "right"] | None, - ) -> list[EngineInput]: + ) -> Sequence[AnyRenderParam]: + request = ctx.request proxy = EmbeddingCompletionRequest( model=request.model, input=texts, @@ -546,50 +561,32 @@ def _preprocess_cohere_text_completion( truncate_prompt_tokens=truncate_prompt_tokens, truncation_side=truncation_side, ) - return self._preprocess_cmpl_online( - proxy, prompt_input=proxy.input, prompt_embeds=None - ) + ctx.request = proxy + requests = super().get_request_factory_online(ctx) + ctx.request = request + return requests def _batch_render_chat( self, - request: CohereEmbedRequest, + ctx: PoolingServeContext, all_messages: Sequence[list[ChatCompletionMessageParam]], truncate_prompt_tokens: int | None, truncation_side: Literal["left", "right"] | None, - ) -> list[EngineInput]: + ) -> Sequence[AnyRenderParam]: """Batch-render multiple conversations through the chat template.""" - if not all_messages: - return [] - - proxy = EmbeddingChatRequest( + request = ctx.request + proxy = EmbeddingBatchChatRequest( model=request.model, - messages=list(all_messages[0]), + messages=all_messages, dimensions=request.output_dimension, encoding_format="float", truncate_prompt_tokens=truncate_prompt_tokens, truncation_side=truncation_side, ) - - renderer = self.renderer - mm_config = self.model_config.multimodal_config - - tok_params = proxy.build_tok_params(self.model_config) - chat_params = proxy.build_chat_params( - self.chat_template, - self.chat_template_content_format, - ).with_defaults( - merge_kwargs( - None, - dict( - tools=None, - tokenize=is_mistral_tokenizer(renderer.tokenizer), - ), - ), - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), - ) - - _, engine_inputs = renderer.render_chat(all_messages, chat_params, tok_params) - return engine_inputs + ctx.request = proxy + requests = self._get_request_factory_chat_input_online(ctx) + ctx.request = request + return requests def _validate_input_type(self, input_type: str | None) -> None: """Raise if *input_type* is not supported by this model.""" @@ -640,7 +637,9 @@ class TokenEmbedIOProcessor(PoolingIOProcessor): class JinaRankingTokenEmbedIOProcessor( TokenEmbedIOProcessor, JinaRankingIOProcessorMixin ): - def pre_process_online(self, ctx: PoolingServeContext): + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: request = ctx.request if isinstance(request, PoolingCompletionLikeRequest): prompts = request.input @@ -655,20 +654,15 @@ def pre_process_online(self, ctx: PoolingServeContext): query=text_prompts[-1], docs=text_prompts[:-1] ) - engine_inputs = self._preprocess_cmpl_online( - request, - prompt_input=prompt_input, - prompt_embeds=None, - ) + request.input = prompt_input + return super().get_request_factory_online(ctx) elif isinstance(request, PoolingChatLikeRequest): raise ValueError("The JinaForRanking does not support chat Request.") else: raise ValueError(f"Invalid {self.name} request type") - ctx.engine_inputs = engine_inputs - def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext + self, ctx: AnyOfflineInputsContext ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineEncodeInputsContext) if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2: diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index b2912e544f28..96a87fea2018 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -93,9 +93,9 @@ class EmbeddingBatchChatRequest( ``messages`` instead of introducing a separate batch-specific field. """ - messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = ( - Field(..., min_length=1) - ) + messages: Sequence[ + Annotated[list[ChatCompletionMessageParam], Field(min_length=1)] + ] = Field(..., min_length=1) def to_pooling_params(self): return PoolingParams( @@ -133,9 +133,9 @@ def normalize_input_messages(cls, data): class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest): """OpenAI embeddings request with batched chat conversations in ``input``.""" - input: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = ( - Field(..., min_length=1) - ) + input: Sequence[ + Annotated[list[ChatCompletionMessageParam], Field(min_length=1)] + ] = Field(..., min_length=1) @model_validator(mode="before") @classmethod diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index fd8140982ecf..a9f6dd310ec7 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -9,6 +9,7 @@ from vllm.outputs import PoolingRequestOutput from vllm.utils.serial_utils import EmbedDType, Endianness +from ..base.io_processor import PoolingIOProcessor from ..base.serving import PoolingServing from ..typing import PoolingServeContext from ..utils import ( @@ -53,6 +54,12 @@ def __init__(self, *args, **kwargs): def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: return EmbedIOProcessor(*args, **kwargs) + async def _preprocessing( + self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext + ): + await super()._preprocessing(io_processor, ctx) + self.io_processor.maybe_pre_process_chunked(ctx) + def _build_response( self, ctx: PoolingServeContext, diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index bb3be812c6b0..ade3c039849f 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -25,7 +25,7 @@ from .scoring.io_processor import ScoringIOProcessor from .scoring.typing import ScoreInput from .typing import ( - ALLOfflineInputsContext, + AnyOfflineInputsContext, OfflineEncodeInputsContext, OfflineOutputsContext, OfflinePluginInputsContext, @@ -104,7 +104,7 @@ def encode( io_processor = self.pooling_io_processors[pooling_task] - ctx: ALLOfflineInputsContext + ctx: AnyOfflineInputsContext if isinstance(prompts, dict) and "data" in prompts: ctx = OfflinePluginInputsContext( pooling_task=pooling_task, @@ -399,6 +399,9 @@ def _run_tiling_engine( num_requests: int, use_tqdm: bool | Callable[..., tqdm] = True, ): + if num_requests == 0: + raise ValueError("You must pass at least one prompt") + # Keeping max_num_seqs * 2 requests in the core can already saturate the core. # Therefore, keep most requests waiting outside the core. max_requests_in_core = ( diff --git a/vllm/entrypoints/pooling/pooling/api_router.py b/vllm/entrypoints/pooling/pooling/api_router.py index 653a36f699ac..40f7b9eb28a8 100644 --- a/vllm/entrypoints/pooling/pooling/api_router.py +++ b/vllm/entrypoints/pooling/pooling/api_router.py @@ -17,8 +17,11 @@ router = APIRouter() -def pooling(request: Request) -> ServingPooling | None: - return request.app.state.serving_pooling +def pooling(request: Request) -> ServingPooling: + handler = getattr(request.app.state, "serving_pooling", None) + if handler is None: + raise NotImplementedError("The model does not support Pooling API") + return handler @router.post( @@ -33,7 +36,4 @@ def pooling(request: Request) -> ServingPooling | None: @load_aware_call async def create_pooling(request: PoolingRequest, raw_request: Request): handler = pooling(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Pooling API") - return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index b0c38b63a32c..ecc356c477a8 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -10,7 +10,9 @@ from ..base.io_processor import PoolingIOProcessor from ..typing import ( - ALLOfflineInputsContext, + AnyOfflineInputsContext, + AnyRenderParam, + EncodeCMPLRenderParams, OfflineEncodeInputsContext, OfflineOutputsContext, OfflinePluginInputsContext, @@ -49,7 +51,9 @@ def __init__(self, *args, **kwargs): ####################################### # online APIs - def pre_process_online(self, ctx: PoolingServeContext): + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: assert isinstance(ctx.request, IOProcessorRequest) validated_prompt = self.io_processor.parse_data(ctx.request.data) @@ -66,24 +70,33 @@ def pre_process_online(self, ctx: PoolingServeContext): ) for prompt in prompt_to_seq(raw_prompts) ] + num_requests = len(parsed_prompts) tok_params = ctx.request.build_tok_params(self.model_config) - ctx.engine_inputs = self.renderer.render_cmpl( - parsed_prompts, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt") - if (v := getattr(ctx.request, k, None)) is not None - }, - ) - pooling_params = self.io_processor.merge_pooling_params() if pooling_params.task is None: pooling_params.task = "plugin" ctx.pooling_params = pooling_params + params_seq = self._params_to_seq(ctx.pooling_params, num_requests) + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + requests = [ + EncodeCMPLRenderParams( + prompts=parsed_prompts[i], + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] + return requests + def post_process_online( self, ctx: PoolingServeContext, @@ -114,7 +127,7 @@ def post_process_online( # offline APIs def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext + self, ctx: AnyOfflineInputsContext ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflinePluginInputsContext) assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index f67b5e912f30..a26b3ca6d0ac 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -20,12 +20,18 @@ logger = init_logger(__name__) -def score(request: Request) -> ServingScores | None: - return request.app.state.serving_scores +def score(request: Request) -> ServingScores: + handler = getattr(request.app.state, "serving_scores", None) + if handler is None: + raise NotImplementedError("The model does not support Score API") + return handler -def rerank(request: Request) -> ServingScores | None: - return request.app.state.serving_scores +def rerank(request: Request) -> ServingScores: + handler = getattr(request.app.state, "serving_scores", None) + if handler is None: + raise NotImplementedError("The model does not support Rerank (Score) API") + return handler @router.post( @@ -40,9 +46,6 @@ def rerank(request: Request) -> ServingScores | None: @load_aware_call async def create_score(request: ScoreRequest, raw_request: Request): handler = score(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Score API") - return await handler(request, raw_request) @@ -77,9 +80,6 @@ async def create_score_v1(request: ScoreRequest, raw_request: Request): @load_aware_call async def do_rerank(request: RerankRequest, raw_request: Request): handler = rerank(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Rerank (Score) API") - return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index e8b3efc32eae..5dfe8a3a136f 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -6,8 +6,7 @@ import torch.nn.functional as F -from vllm import PoolingParams, PoolingRequestOutput, PromptType, TokensPrompt -from vllm.inputs import EngineInput +from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.renderers import TokenizeParams from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( @@ -20,8 +19,11 @@ from ...chat_utils import ChatTemplateResolutionError from ..base.io_processor import PoolingIOProcessor +from ..pooling.protocol import PoolingCompletionRequest from ..typing import ( - ALLOfflineInputsContext, + AnyOfflineInputsContext, + AnyPoolingRequest, + AnyRenderParam, EncodeChatRenderParams, EncodeCMPLRenderParams, OfflineEncodeInputsContext, @@ -167,17 +169,7 @@ def valid_inputs( ) return scoring_data - -class BiEncoderIOProcessor(ScoringIOProcessor): - name = "bi-encoder" - pooling_task: PoolingTask = "embed" - - ####################################### - # online APIs - - def pre_process_online(self, ctx: ScoringServeContext): - request = ctx.request - + def valid_inputs_online(self, request: AnyPoolingRequest): if isinstance(request, ScoreRequest): data_1 = request.data_1 data_2 = request.data_2 @@ -185,9 +177,24 @@ def pre_process_online(self, ctx: ScoringServeContext): data_1 = request.query data_2 = request.documents else: - raise ValueError(f"Invalid {self.name} request type") + raise ValueError(f"Invalid {request.__class__.__name__} request type") scoring_data = self.valid_inputs(data_1, data_2) + return scoring_data + + +class BiEncoderIOProcessor(ScoringIOProcessor): + name = "bi-encoder" + pooling_task: PoolingTask = "embed" + + ####################################### + # online APIs + + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: + request = ctx.request + scoring_data = self.valid_inputs_online(request) max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( request=request @@ -197,19 +204,37 @@ def pre_process_online(self, ctx: ScoringServeContext): scoring_data, max_tokens_per_query, max_tokens_per_doc ) - tok_params = request.build_tok_params(self.model_config) - engine_inputs = self._pre_process( - scoring_data, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") - if (v := getattr(request, k, None)) is not None - }, + data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config) + data_2 = score_data_to_prompts( + scoring_data.data_2, "document", self.model_config ) + prompts = data_1 + data_2 + ctx.n_queries = len(data_1) + + prompts_seq = prompt_to_seq(prompts) + parsed_prompts = [ + parse_model_prompt(self.model_config, prompt) for prompt in prompts_seq + ] + num_requests = len(parsed_prompts) + + tok_params = request.build_tok_params(self.model_config) + params_seq = self._params_to_seq(ctx.pooling_params, num_requests) + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) - ctx.engine_inputs = engine_inputs - ctx.n_queries = len(scoring_data.data_1) + requests = [ + EncodeCMPLRenderParams( + prompts=parsed_prompts[i], + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] + return requests def post_process_online( self, @@ -226,7 +251,7 @@ def post_process_online( # offline APIs def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext + self, ctx: AnyOfflineInputsContext ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineScoringInputsContext) @@ -267,41 +292,6 @@ def post_process_offline( ####################################### # helpers - def _pre_process( - self, - scoring_data: ScoringData, - tok_params: TokenizeParams, - prompt_extras: dict[str, Any] | None = None, - ) -> Sequence[EngineInput]: - data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config) - data_2 = score_data_to_prompts( - scoring_data.data_2, "document", self.model_config - ) - - return self._preprocess_cmpl_offline( - prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras - ) - - def _preprocess_cmpl_offline( - self, - prompts: PromptType | Sequence[PromptType], - tok_params: TokenizeParams, - prompt_extras: dict[str, Any] | None = None, - ) -> Sequence[EngineInput]: - prompts = prompt_to_seq(prompts) - parsed_prompts = [ - ( - prompt - if isinstance(prompt, bytes) - else parse_model_prompt(self.model_config, prompt) - ) - for prompt in prompts - ] - - return self.renderer.render_cmpl( - parsed_prompts, tok_params, prompt_extras=prompt_extras - ) - def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): emb_data_1 = outputs[:n_queries] emb_data_2 = outputs[n_queries:] @@ -432,49 +422,50 @@ def __init__(self, *args, **kwargs): ####################################### # online APIs - def pre_process_online(self, ctx: ScoringServeContext): + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: request = ctx.request + scoring_data = self.valid_inputs_online(request) + data_1 = scoring_data.data_1 + data_2 = scoring_data.data_2 + num_requests = len(data_2) - if isinstance(request, ScoreRequest): - data_1 = request.data_1 - data_2 = request.data_2 - elif isinstance(request, RerankRequest): - data_1 = request.query - data_2 = request.documents - else: - raise ValueError(f"Invalid {self.name} request type") - - scoring_data = self.valid_inputs(data_1, data_2) + if len(data_1) == 1: + data_1 = data_1 * num_requests max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( request=request ) tok_params = request.build_tok_params(self.model_config) - pooling_params = self.create_pooling_params(request) - - engine_inputs, pooling_params_list = self._pre_process( - scoring_data, - tok_params, - pooling_params, - chat_template=self.chat_template, - max_tokens_per_query=max_tokens_per_query, - max_tokens_per_doc=max_tokens_per_doc, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") - if (v := getattr(request, k, None)) is not None - }, - ) + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + requests = [ + ScoringRenderParams( + data_1=data_1[i], + data_2=data_2[i], + chat_template=self.chat_template, + max_tokens_per_query=max_tokens_per_query, + max_tokens_per_doc=max_tokens_per_doc, + tok_params=tok_params, + prompt_extras=ctx.prompt_extras, + skip_mm_cache=False, + params=ctx.pooling_params, + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + for i in range(num_requests) + ] - ctx.engine_inputs = engine_inputs - ctx.pooling_params = pooling_params_list + return requests ####################################### # offline APIs def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext + self, ctx: AnyOfflineInputsContext ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineScoringInputsContext) @@ -485,10 +476,12 @@ def get_request_factory_offline( if len(data_1) == 1: data_1 = data_1 * num_requests + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + pooling_params=ctx.pooling_params + ) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) - prompt_extras = ctx.pooling_params.extra_kwargs seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) @@ -500,6 +493,8 @@ def request_factory() -> RequestGenerator: data_1=data_1[i], data_2=data_2[i], chat_template=ctx.chat_template, + max_tokens_per_query=max_tokens_per_query, + max_tokens_per_doc=max_tokens_per_doc, tok_params=tok_params, prompt_extras=prompt_extras, skip_mm_cache=False, @@ -531,17 +526,13 @@ def render( params = render_params["params"] prompt_extras = render_params["prompt_extras"] - max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( - pooling_params=params - ) - _, engine_prompt = self.get_score_prompt( data_1=render_params["data_1"], data_2=render_params["data_2"], encode_kwargs=tok_params.get_encode_kwargs(), chat_template=render_params["chat_template"], - max_tokens_per_query=max_tokens_per_query, - max_tokens_per_doc=max_tokens_per_doc, + max_tokens_per_query=render_params["max_tokens_per_query"], + max_tokens_per_doc=render_params["max_tokens_per_doc"], chat_template_kwargs=prompt_extras.get("chat_template_kwargs") if prompt_extras else None, @@ -551,29 +542,16 @@ def render( if token_type_ids := engine_prompt.pop("token_type_ids", None): params = params.clone() - compressed = compress_token_type_ids(token_type_ids) - params.extra_kwargs = {"compressed_token_type_ids": compressed} - - engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time) - - return PoolingEngineInput( - prompts=engine_input, - params=params, - lora_requests=render_params["lora_requests"], - priorities=render_params["priorities"], - ) + compressed = compress_token_type_ids( + _apply_post_tokenization_to_token_type_ids( + self.tokenizer, tok_params, token_type_ids + ) + ) + params.extra_kwargs = { + **(params.extra_kwargs or {}), + "compressed_token_type_ids": compressed, + } - def _pre_process( - self, - scoring_data: ScoringData, - tok_params: TokenizeParams, - pooling_params: PoolingParams | None, - chat_template: str | None = None, - max_tokens_per_query: int = 0, - max_tokens_per_doc: int = 0, - prompt_extras: dict[str, Any] | None = None, - ) -> tuple[Sequence[EngineInput], list[PoolingParams]]: - arrival_time = time.time() engine_prompt_extras = ( { k: v @@ -584,55 +562,18 @@ def _pre_process( else None ) - data_1 = scoring_data.data_1 - data_2 = scoring_data.data_2 + if engine_prompt_extras: + target_prompt = extract_target_prompt(self.model_config, engine_prompt) + target_prompt.update(engine_prompt_extras) - if len(data_1) == 1: - data_1 = data_1 * len(data_2) - - if pooling_params is None: - pooling_params = PoolingParams(task="classify") - - pooling_params_list = list[PoolingParams]() - engine_inputs = list[EngineInput]() - for q, d in zip(data_1, data_2): - _, engine_prompt = self.get_score_prompt( - data_1=q, - data_2=d, - encode_kwargs=tok_params.get_encode_kwargs(), - chat_template=chat_template, - max_tokens_per_query=max_tokens_per_query, - max_tokens_per_doc=max_tokens_per_doc, - chat_template_kwargs=prompt_extras.get("chat_template_kwargs") - if prompt_extras - else None, - ) - - token_type_ids = engine_prompt.pop("token_type_ids", None) - tok_params.apply_post_tokenization(self.tokenizer, engine_prompt) - - if token_type_ids is not None: - params = pooling_params.clone() - compressed = compress_token_type_ids( - _apply_post_tokenization_to_token_type_ids( - self.tokenizer, tok_params, token_type_ids - ) - ) - params.extra_kwargs = { - **(params.extra_kwargs or {}), - "compressed_token_type_ids": compressed, - } - pooling_params_list.append(params) - else: - pooling_params_list.append(pooling_params) + engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time) - if engine_prompt_extras: - target_prompt = extract_target_prompt(self.model_config, engine_prompt) - target_prompt.update(engine_prompt_extras) - engine_inputs.append( - self.renderer.process_for_engine(engine_prompt, arrival_time) - ) - return engine_inputs, pooling_params_list + return PoolingEngineInput( + prompts=engine_input, + params=params, + lora_requests=render_params["lora_requests"], + priorities=render_params["priorities"], + ) def get_score_prompt( self, @@ -837,16 +778,27 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM name = "jina-reranking-scoring" pooling_task: PoolingTask = "token_embed" - def get_request_factory_offline( - self, ctx: ALLOfflineInputsContext - ) -> tuple[RequestFactory, int]: - assert isinstance(ctx, OfflineScoringInputsContext) + def get_request_factory_online( + self, ctx: PoolingServeContext + ) -> Sequence[AnyRenderParam]: + request = ctx.request + ctx.n_queries = 1 - scoring_data = ctx.scoring_data - prompt_extras = ctx.pooling_params.extra_kwargs + prompt_extras = ctx.prompt_extras + scoring_data = self.valid_inputs_online(request) + + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + request=request + ) + + if max_tokens_per_query > 0 or max_tokens_per_doc > 0: + scoring_data = self._truncate_scoring_data( + scoring_data, max_tokens_per_query, max_tokens_per_doc + ) queries = self.ensure_str(scoring_data.data_1) docs = self.ensure_str(scoring_data.data_2) + chat_template_kwargs = ( prompt_extras.get("chat_template_kwargs") if prompt_extras else None ) @@ -868,24 +820,28 @@ def get_request_factory_offline( for q, d in zip(queries, docs) ] - return PoolingIOProcessor.get_request_factory_offline( - self, - OfflineEncodeInputsContext( - pooling_task=self.pooling_task, - prompts=prompts, - tokenization_kwargs=ctx.tokenization_kwargs, - pooling_params=ctx.pooling_params, - lora_request=ctx.lora_request, - priorities=ctx.priorities, - ), + # Forward truncation from the real request: the base factory reads + # these off ctx.request, so omitting them here silently drops + # truncate_prompt_tokens for Jina rerank/score (unlike the embed and + # bi/cross-encoder paths, which read them from the real request). + ctx.request = PoolingCompletionRequest( + task="token_embed", + input=prompts, + truncate_prompt_tokens=request.truncate_prompt_tokens, + truncation_side=request.truncation_side, ) + requests = PoolingIOProcessor.get_request_factory_online(self, ctx) + ctx.request = request + return requests + + def get_request_factory_offline( + self, ctx: AnyOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineScoringInputsContext) + + scoring_data = ctx.scoring_data + prompt_extras = ctx.pooling_params.extra_kwargs - def _pre_process( - self, - scoring_data: ScoringData, - tok_params: TokenizeParams, - prompt_extras: dict[str, Any] | None = None, - ) -> Sequence[EngineInput]: queries = self.ensure_str(scoring_data.data_1) docs = self.ensure_str(scoring_data.data_2) chat_template_kwargs = ( @@ -909,8 +865,16 @@ def _pre_process( for q, d in zip(queries, docs) ] - return self._preprocess_cmpl_offline( - prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras + return PoolingIOProcessor.get_request_factory_offline( + self, + OfflineEncodeInputsContext( + pooling_task=self.pooling_task, + prompts=prompts, + tokenization_kwargs=ctx.tokenization_kwargs, + pooling_params=ctx.pooling_params, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + ), ) def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index 5937664d5687..34a2887d5442 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -190,7 +190,7 @@ def _request_output_to_rerank_response( async def flash_late_interaction(self, *args, **kwargs) -> Response: ctx = await self._init_ctx(self.io_processor, *args, **kwargs) - await self._preprocessing_async(self.io_processor, ctx) + await self._preprocessing(self.io_processor, ctx) # stage 1: encode queries and cache token embeddings on workers. await self._flash_late_interaction_encode_queries(ctx) @@ -211,7 +211,6 @@ async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext) query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)] query_uses = [n_docs if n_queries == 1 else 1] * n_queries - query_pooling_params_list = [] for i in range(n_queries): pooling_params = ctx.pooling_params.clone() pooling_params.late_interaction_params = ( @@ -220,23 +219,21 @@ async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext) query_uses=query_uses[i], ) ) - query_pooling_params_list.append(pooling_params) + query_engine_inputs[i]["params"] = pooling_params - assert ( - n_queries - == len(query_pooling_params_list) - == len(query_engine_inputs) - == len(query_keys) - ) + assert n_queries == len(query_engine_inputs) == len(query_keys) query_ctx = ScoringServeContext( request=ctx.request, raw_request=ctx.raw_request, model_name=ctx.model_name, request_id=ctx.request_id, - pooling_params=query_pooling_params_list, + pooling_params=ctx.pooling_params, prompt_request_ids=query_keys, engine_inputs=query_engine_inputs, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + prompt_extras=ctx.prompt_extras, ) await self._prepare_generators(query_ctx) @@ -255,30 +252,27 @@ async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext): query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)] doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)] - doc_pooling_params_list = [] for i in range(n_docs): query_idx = 0 if n_queries == 1 else i pooling_params = ctx.pooling_params.clone() pooling_params.late_interaction_params = build_late_interaction_doc_params( query_key=query_keys[query_idx] ) - doc_pooling_params_list.append(pooling_params) + doc_engine_inputs[i]["params"] = pooling_params - assert ( - n_docs - == len(doc_pooling_params_list) - == len(doc_engine_inputs) - == len(doc_keys) - ) + assert n_docs == len(doc_engine_inputs) == len(doc_keys) doc_ctx = ScoringServeContext( request=ctx.request, raw_request=ctx.raw_request, model_name=ctx.model_name, request_id=ctx.request_id, - pooling_params=doc_pooling_params_list, + pooling_params=ctx.pooling_params, prompt_request_ids=doc_keys, engine_inputs=doc_engine_inputs, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + prompt_extras=ctx.prompt_extras, ) await self._prepare_generators(doc_ctx) diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index b44c9476e789..288f329d09a3 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -88,10 +88,13 @@ class PoolingServeContext(Generic[PoolingRequestT]): raw_request: Request | None = None model_name: str request_id: str - pooling_params: PoolingParams | list[PoolingParams] + pooling_params: PoolingParams + lora_request: LoRARequest | None + priorities: int | Sequence[int] | None + prompt_extras: dict[str, Any] | None + created_time: int = field(default_factory=lambda: int(time.time())) - lora_request: LoRARequest | None = None - engine_inputs: Sequence[EngineInput] | None = None + engine_inputs: Sequence["PoolingEngineInput"] | None = None prompt_request_ids: list[str] | None = None result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = ( @@ -100,7 +103,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): final_res_batch: list[PoolingRequestOutput] = field(default_factory=list) ## for Long Text Embedding with Chunked Processing - original_engine_inputs: Sequence[EngineInput] | None = None + original_engine_inputs: Sequence["PoolingEngineInput"] | None = None chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None ## for bi-encoder & late-interaction @@ -118,7 +121,7 @@ class OfflineInputsContext: pooling_task: PoolingTask tokenization_kwargs: dict[str, Any] | None lora_request: Sequence[LoRARequest | None] | None - priorities: Sequence[int] | None + priorities: int | Sequence[int] | None @dataclass @@ -140,7 +143,7 @@ class OfflinePluginInputsContext(OfflineInputsContext): pooling_params: PoolingParams | Sequence[PoolingParams] | None -ALLOfflineInputsContext: TypeAlias = ( +AnyOfflineInputsContext: TypeAlias = ( OfflineEncodeInputsContext | OfflineScoringInputsContext | OfflinePluginInputsContext @@ -178,6 +181,15 @@ class ScoringRenderParams(RenderParams): data_1: ScoreData data_2: ScoreData chat_template: str | None + max_tokens_per_query: int + max_tokens_per_doc: int + + +AnyRenderParam: TypeAlias = ( + EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams +) +RequestGenerator: TypeAlias = Generator[AnyRenderParam] +RequestFactory: TypeAlias = Callable[[], RequestGenerator] class PoolingEngineInput(TypedDict): @@ -185,9 +197,3 @@ class PoolingEngineInput(TypedDict): params: PoolingParams lora_requests: LoRARequest | None priorities: int - - -RequestGenerator: TypeAlias = Generator[ - EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams -] -RequestFactory: TypeAlias = Callable[[], RequestGenerator] diff --git a/vllm/entrypoints/scale_out/derender/api_router.py b/vllm/entrypoints/scale_out/derender/api_router.py index 3f88d51f0a92..4a139144db95 100644 --- a/vllm/entrypoints/scale_out/derender/api_router.py +++ b/vllm/entrypoints/scale_out/derender/api_router.py @@ -5,15 +5,17 @@ from fastapi import APIRouter, Depends, Request from fastapi.responses import JSONResponse -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse -from vllm.entrypoints.openai.completion.protocol import CompletionResponse from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import validate_json_request from vllm.logger import init_logger from ..token_in_token_out.protocol import ( - DerenderChatRequest, - DerenderCompletionRequest, + DerenderChatRequestUnion, + DerenderChatStreamRequest, + DerenderChatStreamResponse, + DerenderCompletionRequestUnion, + DerenderCompletionStreamRequest, + DerenderCompletionStreamResponse, ) from .serving import ServingDerender @@ -29,46 +31,96 @@ def derender(request: Request) -> ServingDerender | None: @router.post( "/v1/chat/completions/derender", dependencies=[Depends(validate_json_request)], - response_model=ChatCompletionResponse, responses={ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -async def derender_chat_completion(request: DerenderChatRequest, raw_request: Request): +async def derender_chat_completion( + request: DerenderChatRequestUnion, + raw_request: Request, +): + """Derender a generate response into a ChatCompletionResponse. + + Accepts both non-streaming (``stream=false``, default) and streaming + (``stream=true``) request bodies on the same path; FastAPI validates and + routes on the ``stream`` discriminator. + + Non-streaming: body is ``DerenderChatRequest`` (``generate_response`` with + the complete token list). Returns a ``ChatCompletionResponse``. + + Streaming: body is ``DerenderChatStreamRequest`` (one ``generate_chunk`` + delta + optional ``stream_state``). Returns a ``DerenderChatStreamResponse`` + (``chunk`` + ``stream_state``). The client carries ``stream_state`` between + successive calls, one per SSE chunk from ``/inference/v1/generate``. + """ handler = derender(raw_request) if handler is None: raise NotImplementedError( "The model does not support Chat Completions Derender API" ) - result = await handler.derender_chat_response(request) + if isinstance(request, DerenderChatStreamRequest): + stream_result = await handler.derender_chat_stream_response(request) + if isinstance(stream_result, ErrorResponse): + return JSONResponse( + content=stream_result.model_dump(), + status_code=stream_result.error.code, + ) + chunk, stream_state = stream_result + response = DerenderChatStreamResponse(chunk=chunk, stream_state=stream_state) + return JSONResponse(content=response.model_dump()) + result = await handler.derender_chat_response(request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) - return JSONResponse(content=result.model_dump()) @router.post( "/v1/completions/derender", dependencies=[Depends(validate_json_request)], - response_model=CompletionResponse, responses={ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -async def derender_completion(request: DerenderCompletionRequest, raw_request: Request): +async def derender_completion( + request: DerenderCompletionRequestUnion, + raw_request: Request, +): + """Derender a generate response into a CompletionResponse. + + Accepts both non-streaming (``stream=false``, default) and streaming + (``stream=true``) request bodies on the same path. + + Non-streaming: body is ``DerenderCompletionRequest``. Returns a + ``CompletionResponse``. + + Streaming: body is ``DerenderCompletionStreamRequest`` (one + ``generate_chunk`` + optional ``stream_state``). Returns a + ``DerenderCompletionStreamResponse`` (``chunk`` + ``stream_state``). + """ handler = derender(raw_request) if handler is None: raise NotImplementedError("The model does not support Completions Derender API") - result = await handler.derender_completion_response(request) + if isinstance(request, DerenderCompletionStreamRequest): + stream_result = await handler.derender_completion_stream_response(request) + if isinstance(stream_result, ErrorResponse): + return JSONResponse( + content=stream_result.model_dump(), + status_code=stream_result.error.code, + ) + chunk, stream_state = stream_result + response = DerenderCompletionStreamResponse( + chunk=chunk, stream_state=stream_state + ) + return JSONResponse(content=response.model_dump()) + result = await handler.derender_completion_response(request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) - return JSONResponse(content=result.model_dump()) diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index 613ff65ad0ea..4b72ce2e9fc4 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -4,8 +4,14 @@ from typing import cast import vllm.envs as envs -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse -from vllm.entrypoints.openai.completion.protocol import CompletionResponse +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponse, + ChatCompletionStreamResponse, +) +from vllm.entrypoints.openai.completion.protocol import ( + CompletionResponse, + CompletionStreamResponse, +) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, UsageInfo, @@ -28,7 +34,10 @@ from ..token_in_token_out.mm_serde import encode_mm_kwargs_item from ..token_in_token_out.protocol import ( DerenderChatRequest, + DerenderChatStreamRequest, DerenderCompletionRequest, + DerenderCompletionStreamRequest, + DerenderStreamState, GenerateResponse, MultiModalFeatures, PlaceholderRangeInfo, @@ -200,7 +209,9 @@ async def derender_completion_response( total_prompt_tokens, total_completion_tokens, ) = await self.online_derenderer.derender_completion( - request.generate_responses, request.prompt_tokens + request.generate_responses, + request.prompt_tokens, + completion_request=request.completion_request, ) first = request.generate_responses[0] @@ -237,6 +248,90 @@ async def derender_completion_response( kv_transfer_params=kv_params, ) + async def derender_chat_stream_response( + self, + request: DerenderChatStreamRequest, + ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState] | ErrorResponse: + """Streaming counterpart to ``derender_chat_response``. + + Processes one ``GenerateStreamResponse`` chunk and returns the + derendered chunk together with the updated client carried state. + + ``parser is None`` or no ``chat_request`` until reasoning/tool call + functionality added in future PR. + """ + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + return error_check_ret + + try: + chunk, updated_state = await self.online_derenderer.derender_chat_stream( + model=request.model, + generate_chunk=request.generate_chunk, + state=request.stream_state, + chat_request=request.chat_request, + prompt_tokens=request.prompt_tokens, + ) + except NotImplementedError as exc: + return self.create_error_response(exc) + except ValueError as exc: + return self.create_error_response(str(exc)) + except (KeyError, IndexError) as exc: + return self.create_error_response( + f"invalid stream_state: detokenization failed ({exc!r})" + ) + + logger.debug( + "derender_chat_stream request_id=%s model=%s delta_tokens=%d", + request.generate_chunk.request_id, + request.model, + sum( + len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids + ), + ) + return chunk, updated_state + + async def derender_completion_stream_response( + self, + request: DerenderCompletionStreamRequest, + ) -> tuple[CompletionStreamResponse, DerenderStreamState] | ErrorResponse: + """Streaming counterpart to ``derender_completion_response``. + + Processes one ``GenerateStreamResponse`` chunk (one output sequence's + delta) and returns the derendered chunk and updated state. + """ + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + return error_check_ret + + try: + ( + chunk, + updated_state, + ) = await self.online_derenderer.derender_completion_stream( + model=request.model, + generate_chunk=request.generate_chunk, + state=request.stream_state, + prompt_tokens=request.prompt_tokens, + completion_request=request.completion_request, + ) + except ValueError as exc: + return self.create_error_response(str(exc)) + except (KeyError, IndexError) as exc: + return self.create_error_response( + f"invalid stream_state: detokenization failed ({exc!r})" + ) + + logger.debug( + "derender_completion_stream request_id=%s model=%s delta_tokens=%d", + request.generate_chunk.request_id, + request.model, + sum( + len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids + ), + ) + return chunk, updated_state + @staticmethod def _extract_mm_features( engine_input: EngineInput, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index c22e70b014c4..8dad837613d7 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any +from typing import Any, Literal, TypeAlias from pydantic import ( BaseModel, @@ -14,8 +14,12 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProbs, ChatCompletionRequest, + ChatCompletionStreamResponse, +) +from vllm.entrypoints.openai.completion.protocol import ( + CompletionRequest, + CompletionStreamResponse, ) -from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.engine.protocol import StreamOptions, UsageInfo from vllm.logprobs import Logprob from vllm.renderers import TokenizeParams @@ -252,21 +256,16 @@ class GenerateResponse(BaseModel): ) -####### Derender (postprocessing) ####### - - class DerenderChatRequest(BaseModel): """Request for the /v1/chat/completions/derender endpoint (non-streaming). - Wraps a complete GenerateResponse and caller-supplied metadata needed to - produce a fully-formed ChatCompletionResponse without a GPU. - - Streaming derender would require a separate endpoint design with - incremental token delivery, ``OutputProcessor``-based detokenization, - and ``parser.parse_delta()`` instead of ``parser.parse()``. + Wraps a complete GenerateResponse and caller supplied metadata needed to + produce a fully formed ChatCompletionResponse without a GPU. """ # --8<-- [start:derender-chat-request] + stream: Literal[False] = False + model: str """Served model name.""" @@ -285,7 +284,7 @@ class DerenderChatRequest(BaseModel): Required by the parsing so that tool/reasoning parsers can receive the full request context they expect (request.tools, request.tool_choice, - request._grammar_from_tool_parser, etc.). + request._grammar_from_parser, etc.). """ # --8<-- [end:derender-chat-request] @@ -299,6 +298,8 @@ class DerenderCompletionRequest(BaseModel): """ # --8<-- [start:derender-completion-request] + stream: Literal[False] = False + model: str """Served model name.""" @@ -330,3 +331,156 @@ def _validate_prompt_tokens_length(self) -> "DerenderCompletionRequest": f"generate_responses length ({len(self.generate_responses)})" ) return self + + +class DerenderStreamState(BaseModel): + """Per sequence state for stateless streaming derender. + + The client carries this between successive per chunk HTTP calls to the + streaming derender endpoint. All fields are plain JSON serializable data. + No opaque tokenizer or parser internals are stored here. + + The detokenization strategy carries the incremental decode offsets + directly rather than re-sending the whole token history each chunk. + ``detokenize_incrementally`` only ever reads the trailing token window + ``prev_tokens[prefix_offset:]``, so we carry just that tail plus the two + offsets. Each chunk resumes exactly where the last one stopped, including + any partially processed multi-byte character (tracked by ``read_offset``), + then trims and rebases the window so it never grows with generation length. + + Performance: + - Compute per chunk is O(delta). One ``detokenize_incrementally`` call per + new token, independent of how many tokens preceded it. + - Transport per chunk is O(window). The carried tail is bounded by the + incremental detokenization offset, so cumulative bytes over the wire are + O(n) rather than the O(n^2) a full history round trip would incur. + """ + + prev_tokens: list[str] = Field(default_factory=list) + """Trailing decode window. Token strings from ``prefix_offset`` onward. + + Bounded, trimmed and rebased each chunk to the tail + ``detokenize_incrementally`` still reads, so it does not grow with the + number of chunks. + """ + + prefix_offset: int = Field(default=0, ge=0) + """Prefix offset into ``prev_tokens`` for incremental detokenization.""" + + read_offset: int = Field(default=0, ge=0) + """Read offset into ``prev_tokens`` for incremental detokenization.""" + + @field_validator("prev_tokens") + @classmethod + def _bound_prev_tokens(cls, v: list[str]) -> list[str]: + # INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET is small (5) and the trimmed + # window is O(offset). A generous limit rejects unusually large or malformed + # payloads without restricting legitimate multi-byte sequences. + limit = 1024 + if len(v) > limit: + raise ValueError(f"prev_tokens length ({len(v)}) exceeds maximum ({limit})") + return v + + role_sent: bool = False + """True once the initial ``role: "assistant"`` delta has been emitted. + + Prevents re-emitting the role on subsequent chunks even when the detok + window is transiently empty (e.g. usage only final chunk). + """ + + # TODO: Properties used in follow on PR for tool call parsing + last_content: str | None = None + """Last emitted cumulative assistant content text.""" + + last_reasoning: str | None = None + """Last emitted cumulative reasoning text.""" + + last_tool_call_ids: list[str] = Field(default_factory=list) + """Stable tool-call IDs, assigned once when each call first appears. + + Prevents ID regeneration across re-parsing. + """ + + +class DerenderChatStreamRequest(BaseModel): + """One chunk streaming derender request for /v1/chat/completions/derender. + + The client sends one request per SSE chunk received from + ``/inference/v1/generate``. Each request carries the generate chunk + plus the ``stream_state`` returned by the previous call (``None`` on the + first call). The response contains the derendered chunk and the updated + state to be passed to the next call. + + This implements stateless no server side session. All mutable state lives in + the client carried ``stream_state``. + """ + + stream: Literal[True] + + model: str + generate_chunk: GenerateStreamResponse + """One SSE chunk from ``/inference/v1/generate`` (``stream=True``).""" + + stream_state: DerenderStreamState | None = None + """Client carried detok state from the previous call. ``None`` on first.""" + + prompt_tokens: int | None = None + """Prompt token count for usage. Forwarded from the render step.""" + + chat_request: ChatCompletionRequest | None = None + """The original (post adjust_request) ChatCompletionRequest from /render.""" + + +class DerenderCompletionStreamRequest(BaseModel): + """One chunk streaming derender request for /v1/completions/derender. + + Parallel to ``DerenderChatStreamRequest`` for the completions endpoint. + Each call processes one SSE chunk (one output sequence's delta) and + returns the derendered chunk plus updated state. + """ + + stream: Literal[True] + + model: str + generate_chunk: GenerateStreamResponse + """One SSE chunk from ``/inference/v1/generate``.""" + + stream_state: DerenderStreamState | None = None + """Client-carried detok state. ``None`` on the first call.""" + + prompt_tokens: int | None = None + """Prompt token count for usage.""" + + completion_request: CompletionRequest | None = None + """The original (post adjust_request) CompletionRequest from /render.""" + + +class DerenderChatStreamResponse(BaseModel): + """Response for one streaming chat derender chunk. + + Pairs the derendered SSE chunk with the updated client carried state to + pass to the next call. + """ + + chunk: ChatCompletionStreamResponse + stream_state: DerenderStreamState + + +class DerenderCompletionStreamResponse(BaseModel): + """Response for one streaming completions derender chunk. + + Parallel to ``DerenderChatStreamResponse`` for the completions endpoint. + """ + + chunk: CompletionStreamResponse + stream_state: DerenderStreamState + + +# Determines the type by checking the ``stream`` field's literal value. A body without +# ``stream`` validates as the non-streaming member +# (``stream`` defaults to ``False`` there), so FastAPI can validate and dispatch both +# shapes on a single path. +DerenderChatRequestUnion: TypeAlias = DerenderChatRequest | DerenderChatStreamRequest +DerenderCompletionRequestUnion: TypeAlias = ( + DerenderCompletionRequest | DerenderCompletionStreamRequest +) diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py index 8a2494a59df2..392fcf567476 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -5,7 +5,7 @@ from http import HTTPStatus from typing import Annotated -from fastapi import APIRouter, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -203,11 +203,29 @@ async def update_weights(raw_request: Request): @router.post("/finish_weight_update") -async def finish_weight_update(raw_request: Request): - await engine_client(raw_request).finish_weight_update() +async def finish_weight_update( + raw_request: Request, + weight_version: Annotated[str | None, Body(embed=True)] = None, +): + await engine_client(raw_request).finish_weight_update(weight_version) return JSONResponse(content={"message": "Weight update finished"}) +@router.post("/update_weight_version") +async def update_weight_version( + raw_request: Request, + new_version: Annotated[str, Body(embed=True)], +): + await engine_client(raw_request).update_weight_version(new_version) + return JSONResponse(content={"success": True, "new_version": new_version}) + + +@router.get("/weight_info") +async def weight_info(raw_request: Request): + weight_version = await engine_client(raw_request).get_weight_version() + return JSONResponse(content={"weight_version": weight_version}) + + @router.get("/get_world_size") async def get_world_size( raw_request: Request, diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index e711a257ddd4..02a242509050 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -12,10 +12,7 @@ from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, ) -from vllm.entrypoints.serve.elastic_ep.middleware import ( - get_scaling_elastic_ep, - set_scaling_elastic_ep, -) +from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep from vllm.entrypoints.serve.utils.api_utils import validate_json_request from vllm.logger import init_logger @@ -64,8 +61,6 @@ async def scale_elastic_ep(raw_request: Request): status_code=400, detail="drain_timeout must be a positive integer" ) - # Set scaling flag to prevent new requests - set_scaling_elastic_ep(True) client = engine_client(raw_request) try: await client.scale_elastic_ep(new_data_parallel_size, drain_timeout) @@ -83,8 +78,6 @@ async def scale_elastic_ep(raw_request: Request): except Exception as e: logger.error("Scale failed: %s", e) raise HTTPException(status_code=500, detail="Scale failed") from e - finally: - set_scaling_elastic_ep(False) @router.post("/is_scaling_elastic_ep") diff --git a/vllm/entrypoints/serve/engine/typing.py b/vllm/entrypoints/serve/engine/typing.py index 2e01c092c7b2..a6878d500150 100644 --- a/vllm/entrypoints/serve/engine/typing.py +++ b/vllm/entrypoints/serve/engine/typing.py @@ -17,7 +17,9 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderChatRequest, + DerenderChatStreamRequest, DerenderCompletionRequest, + DerenderCompletionStreamRequest, GenerateRequest, GenerateResponse, ) @@ -54,6 +56,7 @@ def build_chat_params( | TokenizeCompletionRequest | DetokenizeRequest | DerenderCompletionRequest + | DerenderCompletionStreamRequest ) ChatLikeRequest: TypeAlias = ( @@ -61,6 +64,7 @@ def build_chat_params( | BatchChatCompletionRequest | TokenizeChatRequest | DerenderChatRequest + | DerenderChatStreamRequest ) SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest diff --git a/vllm/entrypoints/serve/fault_tolerance/__init__.py b/vllm/entrypoints/serve/fault_tolerance/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/fault_tolerance/api_router.py b/vllm/entrypoints/serve/fault_tolerance/api_router.py new file mode 100644 index 000000000000..960833af2155 --- /dev/null +++ b/vllm/entrypoints/serve/fault_tolerance/api_router.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import uuid +from http import HTTPStatus + +from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse + +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import validate_json_request +from vllm.logger import init_logger +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest + +logger = init_logger(__name__) + +router = APIRouter() + +_ALLOWED_INSTRUCTIONS = {"retry"} + + +def _validate_payload(body: dict) -> tuple[str, dict]: + if not isinstance(body, dict): + raise HTTPException(400, "Request body must be a JSON object.") + instruction = body.get("instruction") + if not instruction: + raise HTTPException(400, "'instruction' is required.") + if instruction not in _ALLOWED_INSTRUCTIONS: + raise HTTPException(400, f"Invalid instruction: '{instruction}'.") + params = body.get("params", {}) + if not isinstance(params, dict): + raise HTTPException(400, "'params' must be an object.") + return instruction, params + + +@router.post( + "/fault_tolerance/apply", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.ACCEPTED.value: {"model": dict}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + }, +) +async def process_fault_tolerance_instruction( + raw_request: Request, background_tasks: BackgroundTasks +): + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException(400, "Invalid JSON format") from e + + instruction, params = _validate_payload(body) + ft_request = FaultToleranceRequest( + instruction=instruction, + params=params, + request_id=str(uuid.uuid4()), + ) + + client: EngineClient = raw_request.app.state.engine_client + # Recovery runs cross-rank collective ops that only complete once every rank + # has been dispatched. Run it in the background and return immediately so the + # orchestrator can dispatch to all ranks without blocking; completion is + # observed by polling GET /fault_tolerance/status. + background_tasks.add_task(_run_fault_recovery, client, ft_request) + return JSONResponse( + status_code=HTTPStatus.ACCEPTED.value, + content={ + "message": "Request accepted; poll /fault_tolerance/status for updates.", + "request_id": ft_request.request_id, + }, + background=background_tasks, + ) + + +async def _run_fault_recovery( + client: EngineClient, ft_request: FaultToleranceRequest +) -> None: + """Drive recovery to completion after the 202 response is sent.""" + try: + result = await client.handle_fault(ft_request) + except Exception: + logger.exception("[FT] Recovery dispatch failed.") + return + if not result.success: + logger.error( + "[FT] Recovery failed for request %s: %s", + ft_request.request_id, + result.reason, + ) + + +@router.get("/fault_tolerance/status") +async def get_status(raw_request: Request): + client: EngineClient = raw_request.app.state.engine_client + return JSONResponse(content=await client.get_status()) + + +def register_fault_tolerance_api_router(app: FastAPI): + app.include_router(router) diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index 2de21f616485..5fc855d9b5a9 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -331,7 +331,7 @@ def log_version_and_model(lgr: Logger, version: str, model_name: str) -> None: " ${b}▀▀${r} ${w}▀▀▀▀▀ ▀▀▀▀▀ ▀ ▀${r}\n" ) colors = { - "w": "\033[97;1m", # white + "w": "\033[1m", # bold, default foreground "o": "\033[93m", # orange "b": "\033[94m", # blue "r": "\033[0m", # reset diff --git a/vllm/entrypoints/serve/utils/error_response.py b/vllm/entrypoints/serve/utils/error_response.py index fc17a75c75aa..2aa785c53bb1 100644 --- a/vllm/entrypoints/serve/utils/error_response.py +++ b/vllm/entrypoints/serve/utils/error_response.py @@ -28,7 +28,9 @@ def create_error_response( ) from vllm.exceptions import ( + VLLMClientError, VLLMNotFoundError, + VLLMServerError, VLLMUnprocessableEntityError, VLLMValidationError, ) @@ -45,8 +47,23 @@ def create_error_response( err_type = "NotFoundError" status_code = HTTPStatus.NOT_FOUND param = None + elif isinstance(exc, VLLMClientError): + # Any other client-caused error defaults to 400. + err_type = "BadRequestError" + status_code = HTTPStatus.BAD_REQUEST + param = None + elif isinstance(exc, GenerationError): + err_type = "InternalServerError" + status_code = exc.status_code + param = None + elif isinstance(exc, VLLMServerError): + # Any other server-caused error defaults to 500. + err_type = "InternalServerError" + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + param = None + # Fallback for raw exceptions not yet migrated to VLLMError. + # TODO(zqzten): remove these fallback handlers after migration to VLLMError elif isinstance(exc, (ValueError, TypeError, OverflowError)): - # Common validation errors from user input err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = None @@ -54,10 +71,6 @@ def create_error_response( err_type = "NotImplementedError" status_code = HTTPStatus.NOT_IMPLEMENTED param = None - elif isinstance(exc, GenerationError): - err_type = "InternalServerError" - status_code = exc.status_code - param = None elif any(cls.__name__ == "TemplateError" for cls in type(exc).__mro__): # jinja2.TemplateError and its subclasses (avoid importing jinja2) err_type = "BadRequestError" diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index c65206580905..03e02e34f743 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -31,7 +31,7 @@ create_error_response, sanitize_message, ) -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMError, VLLMValidationError from vllm.logger import init_logger from vllm.utils.gc_utils import freeze_gc_heap from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError @@ -39,7 +39,7 @@ logger = init_logger("vllm.entrypoints.openai.server_utils") -GUARDED_PREFIX = ("/v1", "/v2", "/inference") +GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") class AuthenticationMiddleware: @@ -325,6 +325,16 @@ async def log_response(request: Request, call_next): return response +async def vllm_error_handler(req: Request, exc: VLLMError): + """Dispatch a vLLM-specific error to the appropriate handler.""" + if isinstance(exc, (EngineGenerateError, EngineDeadError)): + return await engine_error_handler(req, exc) + elif isinstance(exc, GenerationError): + return await generation_error_handler(req, exc) + else: + return await exception_handler(req, exc) + + async def engine_error_handler( req: Request, exc: EngineDeadError | EngineGenerateError ): diff --git a/vllm/entrypoints/speech_to_text/base/protocol.py b/vllm/entrypoints/speech_to_text/base/protocol.py index e8cb61a41472..70d28726e4f1 100644 --- a/vllm/entrypoints/speech_to_text/base/protocol.py +++ b/vllm/entrypoints/speech_to_text/base/protocol.py @@ -8,4 +8,5 @@ ## Protocols for Audio AudioResponseFormat: TypeAlias = Literal["json", "text", "srt", "verbose_json", "vtt"] +TranscriptionResponseFormat: TypeAlias = AudioResponseFormat | Literal["diarized_json"] _LONG_INFO = torch.iinfo(torch.long) diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 7c703f2535c4..8065c13d43f5 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -43,6 +43,7 @@ from ..transcription.protocol import ( TranscriptionResponse, + TranscriptionResponseDiarized, TranscriptionResponseStreamChoice, TranscriptionResponseVerbose, TranscriptionSegment, @@ -56,7 +57,9 @@ TranslationStreamResponse, ) -SpeechToTextResponse: TypeAlias = TranscriptionResponse | TranslationResponse +SpeechToTextResponse: TypeAlias = ( + TranscriptionResponse | TranslationResponse | TranscriptionResponseDiarized +) SpeechToTextResponseVerbose: TypeAlias = ( TranscriptionResponseVerbose | TranslationResponseVerbose ) @@ -70,6 +73,7 @@ TranscriptionResponse | TranslationResponse | TranscriptionResponseVerbose + | TranscriptionResponseDiarized | TranslationResponseVerbose ) @@ -174,6 +178,7 @@ def _decode_and_chunk_speech( y, sr = load_audio( buf, sr=self.asr_config.sample_rate, + mono=True, max_duration_s=self.max_audio_decode_duration_s, ) except ValueError: @@ -258,7 +263,7 @@ async def _preprocess_speech_to_text( request: SpeechToTextRequest, audio_data: bytes, request_id: str, - ) -> tuple[list[EngineInput], float]: + ) -> tuple[list[EngineInput], float, list[float]]: # Validate request request.language = self.model_cls.validate_language(request.language) request.to_language = ( @@ -277,6 +282,13 @@ async def _preprocess_speech_to_text( # Run cpu intensive preprocess step in a separate thread pool executor. chunks, duration = await self._decode_and_chunk_speech_async(audio_data) + chunk_start_offsets: list[float] = [0.0] + + for chunk in chunks[:-1]: + chunk_start_offsets.append( + chunk_start_offsets[-1] + chunk.shape[-1] / self.asr_config.sample_rate + ) + if request.language is None and getattr( self.model_cls, "supports_explicit_language_detection", False ): @@ -306,7 +318,7 @@ async def _preprocess_speech_to_text( engine_inputs = await self.renderer.render_cmpl_async(parsed_prompts) - return engine_inputs, duration + return engine_inputs, duration, chunk_start_offsets def _preprocess_verbose_prompt(self, prompt: EncoderDecoderDictPrompt): dec_prompt = prompt["decoder_prompt"] @@ -391,7 +403,7 @@ def _get_verbose_segments( SpeechToTextSegment, segment_class( id=len(segments), - seek=start_time, + seek=int(start_time), start=start_time + BASE_OFFSET * start_timestamp, end=start_time + BASE_OFFSET * end_timestamp, temperature=request.temperature, @@ -441,10 +453,15 @@ async def _create_speech_to_text( if self.engine_client.errored: raise self.engine_client.dead_error - if request.response_format not in ["text", "json", "verbose_json"]: + if request.response_format not in [ + "text", + "json", + "verbose_json", + "diarized_json", + ]: return self.create_error_response( "Currently only support response_format: " - "`text`, `json` or `verbose_json`" + "`text`, `json`, `verbose_json` or `diarized_json`" ) if ( @@ -455,9 +472,20 @@ async def _create_speech_to_text( f"Currently do not support verbose_json for {request.model}" ) - if request.response_format == "verbose_json" and request.stream: + if ( + request.response_format == "diarized_json" + and not self.model_cls.supports_diarized_transcription + ): return self.create_error_response( - "verbose_json format doesn't support streaming case" + f"Currently do not support diarized_json for {request.model}" + ) + + if ( + request.response_format in {"verbose_json", "diarized_json"} + and request.stream + ): + return self.create_error_response( + f"{request.response_format} format doesn't support streaming case" ) request_id = f"{self.task_type}-{self._base_request_id(raw_request)}" @@ -467,7 +495,11 @@ async def _create_speech_to_text( lora_request = self._maybe_get_adapters(request) - engine_inputs, duration_s = await self._preprocess_speech_to_text( + ( + engine_inputs, + duration_s, + chunk_start_offsets, + ) = await self._preprocess_speech_to_text( request=request, audio_data=audio_data, request_id=request_id, @@ -587,11 +619,10 @@ async def _create_speech_to_text( assert len(list_result_generator) == 1, ( "`max_audio_clip_s` is set to None, audio cannot be chunked" ) + assert len(chunk_start_offsets) == len(list_result_generator) result_generator = merge_async_iterators(*list_result_generator) async for idx, op in result_generator: - start_time = ( - float(idx * chunk_size_in_s) if chunk_size_in_s is not None else 0.0 - ) + start_time = chunk_start_offsets[idx] if request.response_format == "verbose_json": assert op.outputs[0].logprobs segments: list[SpeechToTextSegment] = self._get_verbose_segments( @@ -624,7 +655,33 @@ async def _create_speech_to_text( # rounded up as per openAI specs "seconds": int(math.ceil(duration_s)), } - if request.response_format != "verbose_json": + if request.response_format == "diarized_json": + diarized_segments = self.model_cls.parse_diarized_transcript(text) + if not diarized_segments: + return self.create_error_response( + "Model output did not contain a valid diarized transcript" + ) + final_response = cast( + T, + TranscriptionResponseDiarized( + duration=duration_s, + text=separator.join( + segment.text for segment in diarized_segments + ), + segments=[ + { + "id": f"seg_{index}", + "start": segment.start, + "end": segment.end, + "text": segment.text, + "speaker": segment.speaker, + } + for index, segment in enumerate(diarized_segments) + ], + usage=usage, + ), + ) + elif request.response_format != "verbose_json": final_response = cast( T, TranscriptionResponse(text=text, usage=usage) ) @@ -634,7 +691,7 @@ async def _create_speech_to_text( TranscriptionResponseVerbose( text=text, language=request.language, - duration=str(duration_s), + duration=duration_s, segments=total_segments, ), ) @@ -648,7 +705,7 @@ async def _create_speech_to_text( TranslationResponseVerbose( text=text, language=request.language, - duration=str(duration_s), + duration=duration_s, segments=total_segments, ), ) diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index 3d6600fe3656..3220e1505099 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -27,7 +27,7 @@ ) from vllm.utils import random_uuid -from ..base.protocol import _LONG_INFO, AudioResponseFormat +from ..base.protocol import _LONG_INFO, TranscriptionResponseFormat if TYPE_CHECKING: import numpy as np @@ -88,7 +88,7 @@ class TranscriptionRequest(OpenAIBaseModel): should match the audio language. """ - response_format: AudioResponseFormat = Field(default="json") + response_format: TranscriptionResponseFormat = Field(default="json") """ The format of the output, in one of these options: `json`, `text`, `srt`, `verbose_json`, or `vtt`. @@ -384,7 +384,7 @@ class TranscriptionSegment(OpenAIBaseModel): class TranscriptionResponseVerbose(OpenAIBaseModel): - duration: str + duration: float """The duration of the input audio.""" language: str @@ -400,6 +400,27 @@ class TranscriptionResponseVerbose(OpenAIBaseModel): """Extracted words and their corresponding timestamps.""" +class TranscriptionDiarizedSegment(OpenAIBaseModel): + """A speaker-attributed transcription segment.""" + + type: Literal["transcript.text.segment"] = "transcript.text.segment" + id: str + start: float + end: float + text: str + speaker: str + + +class TranscriptionResponseDiarized(OpenAIBaseModel): + """OpenAI-compatible diarized transcription response.""" + + task: Literal["transcribe"] = "transcribe" + duration: float + text: str + segments: list[TranscriptionDiarizedSegment] + usage: TranscriptionUsageAudio + + TranscriptionResponseVariant: TypeAlias = ( - TranscriptionResponse | TranscriptionResponseVerbose + TranscriptionResponse | TranscriptionResponseVerbose | TranscriptionResponseDiarized ) diff --git a/vllm/entrypoints/speech_to_text/transcription/serving.py b/vllm/entrypoints/speech_to_text/transcription/serving.py index a59ae8aa313b..51898841fa0f 100644 --- a/vllm/entrypoints/speech_to_text/transcription/serving.py +++ b/vllm/entrypoints/speech_to_text/transcription/serving.py @@ -18,6 +18,7 @@ from .protocol import ( TranscriptionRequest, TranscriptionResponse, + TranscriptionResponseDiarized, TranscriptionResponseStreamChoice, TranscriptionResponseVerbose, TranscriptionStreamResponse, @@ -55,6 +56,7 @@ async def create_transcription( ) -> ( TranscriptionResponse | TranscriptionResponseVerbose + | TranscriptionResponseDiarized | AsyncGenerator[str, None] | ErrorResponse ): @@ -70,6 +72,8 @@ async def create_transcription( response_class=( TranscriptionResponseVerbose if request.response_format == "verbose_json" + else TranscriptionResponseDiarized + if request.response_format == "diarized_json" else TranscriptionResponse ), stream_generator_method=self.transcription_stream_generator, diff --git a/vllm/entrypoints/speech_to_text/translation/protocol.py b/vllm/entrypoints/speech_to_text/translation/protocol.py index 6e457682c2f6..d8836554c565 100644 --- a/vllm/entrypoints/speech_to_text/translation/protocol.py +++ b/vllm/entrypoints/speech_to_text/translation/protocol.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import time from typing import TYPE_CHECKING, Literal, TypeAlias @@ -78,7 +79,6 @@ class TranslationRequest(OpenAIBaseModel): `verbose_json`, or `vtt`. """ - # TODO support additional sampling parameters # --8<-- [start:translation-sampling-params] use_beam_search: bool = False """Whether or not beam search should be used.""" @@ -103,6 +103,28 @@ class TranslationRequest(OpenAIBaseModel): will use [log probability](https://en.wikipedia.org/wiki/Log_probability) to automatically increase the temperature until certain thresholds are hit. """ + + top_p: float | None = None + """Enables nucleus (top-p) sampling, where tokens are selected from the + smallest possible set whose cumulative probability exceeds `p`. + """ + + top_k: int | None = None + """Limits sampling to the `k` most probable tokens at each step.""" + + min_p: float | None = None + """Filters out tokens with a probability lower than `min_p`, ensuring a + minimum likelihood threshold during sampling. + """ + + frequency_penalty: float | None = 0.0 + """The frequency penalty to use for sampling.""" + + repetition_penalty: float | None = None + """The repetition penalty to use for sampling.""" + + presence_penalty: float | None = 0.0 + """The presence penalty to use for sampling.""" # --8<-- [end:translation-sampling-params] # --8<-- [start:translation-extra-params] @@ -139,11 +161,23 @@ class TranslationRequest(OpenAIBaseModel): max_completion_tokens: int | None = None """The maximum number of tokens to generate.""" + + vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( + default=None, + description=( + "Additional request parameters with (list of) string or " + "numeric values, used by custom extensions." + ), + ) # --8<-- [end:translation-extra-params] # Default sampling parameters for translation requests. _DEFAULT_SAMPLING_PARAMS: dict = { + "repetition_penalty": 1.0, "temperature": 0, + "top_p": 1.0, + "top_k": 0, + "min_p": 0.0, } def build_stt_params( @@ -199,14 +233,38 @@ def to_sampling_params( temperature = default_sampling_params.get( "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"] ) + if (top_p := self.top_p) is None: + top_p = default_sampling_params.get( + "top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"] + ) + if (top_k := self.top_k) is None: + top_k = default_sampling_params.get( + "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"] + ) + if (min_p := self.min_p) is None: + min_p = default_sampling_params.get( + "min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"] + ) + if (repetition_penalty := self.repetition_penalty) is None: + repetition_penalty = default_sampling_params.get( + "repetition_penalty", + self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"], + ) return SamplingParams.from_optional( temperature=temperature, max_tokens=max_tokens, seed=self.seed, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=self.frequency_penalty, + repetition_penalty=repetition_penalty, + presence_penalty=self.presence_penalty, output_kind=RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY, + extra_args=self.vllm_xargs, skip_clone=True, # Created fresh per request, safe to skip clone ) @@ -226,6 +284,16 @@ def validate_stream_options(cls, data): parameter=invalid_param, ) + xargs = data.get("vllm_xargs") + if isinstance(xargs, str): + try: + data["vllm_xargs"] = json.loads(xargs) + except json.JSONDecodeError as e: + raise VLLMValidationError( + f"Failed to parse vllm_xargs. Must be valid JSON: {e}", + parameter="vllm_xargs", + ) from e + return data @@ -289,7 +357,7 @@ class TranslationSegment(OpenAIBaseModel): class TranslationResponseVerbose(OpenAIBaseModel): - duration: str + duration: float """The duration of the input audio.""" language: str diff --git a/vllm/envs.py b/vllm/envs.py index 67d8fa5a2c65..4587139f1b9b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -131,6 +131,7 @@ VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 + AITER_SITUV2_A8W4: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_USE_AITER_MHA: bool = True @@ -153,6 +154,7 @@ K_SCALE_CONSTANT: int = 200 V_SCALE_CONSTANT: int = 100 VLLM_USE_RUST_FRONTEND: bool = False + VLLM_USE_RUST_BENCH: bool = False VLLM_RUST_FRONTEND_PATH: str | None = "auto" VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 @@ -186,13 +188,15 @@ VLLM_MOE_USE_DEEP_GEMM: bool = True VLLM_USE_DEEP_GEMM_E8M0: bool = True VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True + VLLM_DCP_Q_REPLICATE: bool = False VLLM_DEEP_GEMM_WARMUP: Literal[ "skip", "full", "relax", ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True - VLLM_MOE_SKIP_PADDING: bool = False + VLLM_MOE_SKIP_PADDING: bool = True + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -237,6 +241,7 @@ VLLM_LOOPBACK_IP: str = "" VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True VLLM_ENABLE_RESPONSES_API_STORE: bool = False + VLLM_ENABLE_COHERE_API: bool = False VLLM_HAS_FLASHINFER_CUBIN: bool = False VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True @@ -547,22 +552,24 @@ def _deprecated_triton_attn_use_td() -> None: return None -def _resolve_rust_frontend_path() -> str | None: - """Resolve the Rust frontend binary path. +def _resolve_rust_cli_path() -> str | None: + """Resolve the vllm-rs binary path. - Returns None if VLLM_USE_RUST_FRONTEND is not enabled. + Returns None unless VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH is enabled. When enabled, resolves VLLM_RUST_FRONTEND_PATH ("auto" by default) to the actual binary path. """ - use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) + use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) or bool( + int(os.environ.get("VLLM_USE_RUST_BENCH", "0")) + ) raw = os.environ.get("VLLM_RUST_FRONTEND_PATH", "auto") if not use_rust: if os.environ.get("VLLM_RUST_FRONTEND_PATH") is not None: logger.warning( - "VLLM_RUST_FRONTEND_PATH is set but VLLM_USE_RUST_FRONTEND " - "is not enabled. The Rust frontend will not be used. " - "Set VLLM_USE_RUST_FRONTEND=1 to enable it." + "VLLM_RUST_FRONTEND_PATH is set without enabling " + "VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH. " + "Set one of them to 1 to use the vllm-rs binary." ) return None @@ -999,7 +1006,6 @@ def _resolve_rust_frontend_path() -> str | None: # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. - # - "identity": returns raw video bytes for model processor to handle. # # Custom backend implementations can be registered # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and @@ -1206,6 +1212,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ROCM_USE_AITER_MOE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") ), + # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up- + # interleaved flydsl kernels instead of the default a16w4 separated path. + # Shared with the AITER runtime, which reads the same env var directly. + "AITER_SITUV2_A8W4": lambda: ( + os.getenv("AITER_SITUV2_A8W4", "0").lower() in ("true", "1") + ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass # for large batches @@ -1339,10 +1351,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_RUST_FRONTEND": lambda: bool( int(os.getenv("VLLM_USE_RUST_FRONTEND", "0")) ), - # Path to the Rust frontend binary. Defaults to "auto" which discovers - # the binary installed with the vllm package. Only used when - # VLLM_USE_RUST_FRONTEND=1. - "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_frontend_path(), + # If set, use the packaged Rust client for `vllm bench serve`. + "VLLM_USE_RUST_BENCH": lambda: bool(int(os.getenv("VLLM_USE_RUST_BENCH", "0"))), + # Path to the vllm-rs binary. Defaults to "auto" which discovers the + # binary installed with the vllm package. Used when VLLM_USE_RUST_FRONTEND=1 + # or VLLM_USE_RUST_BENCH=1. + "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_cli_path(), # If set, vllm will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` @@ -1490,6 +1504,8 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool( int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1")) ), + # Opt-in MLA DCP query replication: skip the decode query all-gather. + "VLLM_DCP_Q_REPLICATE": lambda: bool(int(os.getenv("VLLM_DCP_Q_REPLICATE", "0"))), # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm # JIT all the required kernels before model execution so there is no # JIT'ing in the hot-path. However, this warmup increases the engine @@ -1516,9 +1532,20 @@ def _resolve_rust_frontend_path() -> str | None: ), # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that - # treats topk_id == -1 as a skip sentinel; off by default because not all - # kernels support it yet. - "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))), + # treats topk_id == -1 as a skip sentinel + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))), + # Kimi-K3 only. Under sequence-parallel MoE the dense and shared-expert MLPs + # are replicated on every rank, so each rank streams the whole weight to + # serve its own token shard. Shard them across TP instead: the MLP then + # all-gathers the full token set, computes this rank's partial, and + # reduce-scatters (which both sums across TP and restores the sequence + # sharding). Trades weight bandwidth and resident memory for two collectives + # per layer, so it only wins at low token counts: intended for decode + # instances in a P/D disaggregated deployment, not for prefill or unified + # serving. + "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) + ), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( @@ -1545,7 +1572,7 @@ def _resolve_rust_frontend_path() -> str | None: # tensors above will instead be sent via a separate message. # While the sending side still actually copies the tensor # in all cases, on the receiving side, tensors above this - # limit will actually be zero-copy decoded. + # limit will actually be zero-copy decoded. The unit is bytes. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int( os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256") ), @@ -1745,6 +1772,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ENABLE_RESPONSES_API_STORE": lambda: bool( int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0")) ), + # If set to 1, expose the Cohere Chat v2 API at ``POST /cohere/v2/chat``. + # Default off + "VLLM_ENABLE_COHERE_API": lambda: bool( + int(os.getenv("VLLM_ENABLE_COHERE_API", "0")) + ), # If set, use the fp8 mfma in rocm paged attention. "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool( int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0")) @@ -2112,6 +2144,13 @@ def compile_factors() -> dict[str, object]: "VLLM_CACHE_ROOT", # Runtime memory-plan persistence; does not affect compiled graphs. "VLLM_ENABLE_STARTUP_PLAN", + # Location-only derived paths: where a cache/config directory lives + # cannot affect compiled artifacts, and hashing them means relocating + # HOME or the XDG roots silently invalidates every compile cache + # (VLLM_CACHE_ROOT above and VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR below + # are already ignored for the same reason). + "VLLM_XLA_CACHE_PATH", + "VLLM_CONFIG_ROOT", "LD_LIBRARY_PATH", "VLLM_SERVER_DEV_MODE", "VLLM_DP_MASTER_IP", diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 4112c3de24ba..4383e9a64415 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -6,7 +6,25 @@ from typing import Any -class VLLMValidationError(ValueError): +class VLLMError(Exception): + """Base class for all vLLM-specific errors. + + Subclasses are split into `VLLMClientError` (caused by the request, mapped + to 4xx) and `VLLMServerError` (caused by the server, mapped to 5xx). + Dispatching on this hierarchy lets the entrypoints decide the HTTP status + without relying on raw Python exception types such as `ValueError`. + """ + + +class VLLMClientError(VLLMError): + """Base class for errors caused by the client request (4xx).""" + + +class VLLMServerError(VLLMError): + """Base class for errors caused by the server (5xx).""" + + +class VLLMValidationError(VLLMClientError): """vLLM-specific validation error for request validation failures. Args: @@ -36,7 +54,7 @@ def __str__(self): return f"{base} ({', '.join(extras)})" if extras else base -class VLLMNotFoundError(Exception): +class VLLMNotFoundError(VLLMClientError): """vLLM-specific NotFoundError""" pass @@ -66,7 +84,7 @@ def __str__(self): return self.message -class VLLMUnprocessableEntityError(ValueError): +class VLLMUnprocessableEntityError(VLLMClientError): """vLLM-specific error for unprocessable entity requests. This exception is raised when the request content is invalid or cannot be diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index f997004d2fbd..bda40bebbc86 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -7,6 +7,8 @@ from typing_extensions import NotRequired, TypedDict, assert_never +from vllm.exceptions import VLLMValidationError + if TYPE_CHECKING: import torch @@ -284,7 +286,7 @@ class EncoderDecoderInput(TypedDict): def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: if enc_input["type"] == "embeds": - raise ValueError( + raise VLLMValidationError( "Embedding inputs are not supported for encoder-decoder models" ) @@ -302,7 +304,7 @@ def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: def _validate_dec_input(dec_input: SingletonInput) -> DecoderEngineInput: if dec_input["type"] == "embeds": - raise ValueError( + raise VLLMValidationError( "Embedding inputs are not supported for encoder-decoder models" ) diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index 33a71b8f853f..a40028a4d0bc 100644 --- a/vllm/ir/ops/layernorm.py +++ b/vllm/ir/ops/layernorm.py @@ -23,10 +23,14 @@ def rms_norm( @rms_norm.register_input_generator def _rms_norm_input_generator( - num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5 + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + epsilon: float = 1e-5, + device: torch.device | str | None = None, ) -> tuple: - x = torch.randn(num_tokens, hidden_size, dtype=dtype) - weight = torch.randn(hidden_size, dtype=dtype) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) return x, weight, epsilon @@ -64,9 +68,13 @@ def fused_add_rms_norm( @fused_add_rms_norm.register_input_generator def _fused_add_rms_norm_input_generator( - num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5 + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + epsilon: float = 1e-5, + device: torch.device | str | None = None, ) -> tuple: - x = torch.randn(num_tokens, hidden_size, dtype=dtype) - x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype) - weight = torch.randn(hidden_size, dtype=dtype) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) return x, x_residual, weight, epsilon diff --git a/vllm/kernels/__init__.py b/vllm/kernels/__init__.py index 075bc01f3ba3..e83e2772f06d 100644 --- a/vllm/kernels/__init__.py +++ b/vllm/kernels/__init__.py @@ -2,6 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kernel implementations for vLLM.""" -from . import aiter_ops, oink_ops, vllm_c, xpu_ops +from . import aiter_ops, oink_ops, vllm_c -__all__ = ["vllm_c", "aiter_ops", "oink_ops", "xpu_ops"] +__all__ = ["vllm_c", "aiter_ops", "oink_ops"] diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 273bc58935b7..de8a53ae8b77 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -75,7 +75,6 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor direct_register_aiter_op( op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake ) - rms_add_no_var_16bit_only = ( lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) @@ -112,6 +111,12 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import rmsnorm2d_fwd_with_add + x_shape = x.shape + residual_shape = residual.shape + hidden_size = x.shape[-1] + x = x.reshape(-1, hidden_size) + residual = residual.reshape(-1, hidden_size) + # TODO can out = x and residual_out = residual to save memory? # Need to check if the kernel supports in-place residual output # (if yes set mutates_args and inplace) @@ -125,7 +130,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( weight, variance_epsilon, ) - return out, residual_out + return out.reshape(x_shape), residual_out.reshape(residual_shape) def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( diff --git a/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json index afb30aab42d4..f6f018bb3506 100644 --- a/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json +++ b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json @@ -234,7 +234,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], @@ -1354,7 +1354,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], @@ -1564,7 +1564,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], @@ -1634,7 +1634,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], @@ -2546,7 +2546,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], @@ -2686,7 +2686,7 @@ ], "range_warp_specializes": [ null, - true, + false, null, null ], diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index 6ae5d9939e37..9bb6943a2cbc 100644 --- a/vllm/kernels/vllm_c.py +++ b/vllm/kernels/vllm_c.py @@ -12,6 +12,7 @@ """Most kernels in this file are supported on all CUDA-alike platforms.""" IS_ROCM = current_platform.is_rocm() """ROCm needs shape normalization before calling some vLLM C kernels.""" +GPGPU_DEVICE = CUDA_ALIKE or current_platform.is_xpu() rms_no_var_size = lambda x, weight, epsilon, variance_size=None: ( variance_size is None and (weight is None or weight.dtype == x.dtype) @@ -20,7 +21,7 @@ @ir.ops.rms_norm.register_impl( - "vllm_c", supports_args=rms_no_var_size, supported=CUDA_ALIKE + "vllm_c", supports_args=rms_no_var_size, supported=GPGPU_DEVICE ) def rms_norm( x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None @@ -32,7 +33,9 @@ def rms_norm( if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()): original_shape = x.shape x = x.reshape(-1, original_shape[-1]) - output = torch.empty_like(x) + # empty_like preserves the strides of transposed inputs, but the + # libtorch-stable kernel requires a contiguous output tensor. + output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output.reshape(original_shape) @@ -51,7 +54,7 @@ def rms_norm( @ir.ops.fused_add_rms_norm.register_impl( "vllm_c", supports_args=rms_add_no_var_size, - supported=CUDA_ALIKE, + supported=GPGPU_DEVICE, inplace=True, ) def fused_add_rms_norm( diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py deleted file mode 100644 index df82962d802f..000000000000 --- a/vllm/kernels/xpu_ops.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch -from torch import Tensor - -from vllm import ir -from vllm.platforms import current_platform - -current_platform.import_kernels() - - -def is_xpu_kernels_found() -> bool: - from importlib.util import find_spec - - return find_spec("vllm_xpu_kernels") is not None - - -XPU_KERNELS_SUPPORTED = is_xpu_kernels_found() -"""Kernels in this file are supported if vLLM XPU kernels are installed.""" - -rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and ( - weight is None or weight.dtype == x.dtype -) - - -@ir.ops.rms_norm.register_impl( - "xpu_kernels", supports_args=rms_no_var, supported=XPU_KERNELS_SUPPORTED -) -def rms_norm( - x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None -) -> Tensor: - assert variance_size is None - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) - output = torch.empty(x.shape, device=x.device, dtype=x.dtype) - torch.ops._C.rms_norm(output, x, weight, epsilon) - return output - - -rms_add_no_var_size = ( - lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None - and (weight is None or weight.dtype == x.dtype) -) - - -@ir.ops.fused_add_rms_norm.register_impl( - "xpu_kernels", - supports_args=rms_add_no_var_size, - supported=XPU_KERNELS_SUPPORTED, - inplace=True, -) -def fused_add_rms_norm( - x: Tensor, - x_residual: Tensor, - weight: Tensor | None, - epsilon: float, - variance_size: int | None = None, -) -> tuple[Tensor, Tensor]: - assert variance_size is None - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) - torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) - return x, x_residual diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index b447af87802e..d80848b14bf1 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -44,7 +44,7 @@ def __init__(self, base_layer: MoERunner) -> None: "Monolithic kernels are not supported for Fused MoE LoRA." ) - # Use the MoE-aware TP rank/size: when EP is active, FusedMoE collapses + # Use the MoE-aware TP rank/size: when EP is active, MoERunner collapses # moe_parallel_config.tp_size to 1 (experts are sharded across the # TP group instead). moe_parallel_config = self.moe_config.moe_parallel_config diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 01d049632834..f094f11f09e8 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -24,7 +24,7 @@ @dataclass(frozen=True) class MoEEPLoadSpec: - """Per-expert-parallel slicing metadata for one FusedMoE LoRA module. + """Per-expert-parallel slicing metadata for one FusedMoEFactory LoRA module. Threaded into the LoRA loader so per-expert weights from EP ranks other than this one can be skipped before they ever hit CPU memory. @@ -193,7 +193,7 @@ def from_local_checkpoint( skip_prefixes: List of module name prefixes to skip during loading. Models can define this to skip modules not used in inference (e.g., MTP layers). Format: ["mtp."] - moe_ep_spec: When 2D FusedMoE LoRA modules are present with + moe_ep_spec: When 2D FusedMoEFactory LoRA modules are present with expert parallelism enabled, the (ep_rank, local, global) slicing metadata shared across all MoE layers. Non-local expert weights are skipped at read time instead of being diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 2c4be5678630..ba545bc9c72d 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -739,7 +739,7 @@ def _register_packed_modules(self, module_full_name: str) -> None: def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: for module_name, new_module_names in self.packed_modules.items(): - # For 2D FusedMoE modules with EP, narrow the per-expert + # For 2D RoutedExperts modules with EP, narrow the per-expert # sub-module list to this rank's owned experts so pack_moe # produces a tensor sized to local_num_experts directly. packed_module_names = new_module_names @@ -1096,7 +1096,7 @@ def _restrict_to_local_experts( def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None: """ - Per-rank slicing metadata for 2D FusedMoE LoRA modules. + Per-rank slicing metadata for 2D RoutedEXperts LoRA modules. """ if not self._use_ep or not self._is_moe: return None diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index a1514c9206be..bf98e46e8546 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -284,7 +284,11 @@ def enabled(cls) -> bool: enabled = f"+{cls.name}" in custom_ops disabled = f"-{cls.name}" in custom_ops - assert not (enabled and disabled), f"Cannot enable and disable {cls.name}" + if enabled and disabled: + raise ValueError( + "custom_ops cannot both enable and disable the same operation: " + f"{cls.name}. Remove either the '+' or '-' directive" + ) return (CustomOp.default_on() or enabled) and not disabled @@ -299,7 +303,10 @@ def default_on() -> bool: compilation_config = get_cached_compilation_config() count_none = compilation_config.custom_ops.count("none") count_all = compilation_config.custom_ops.count("all") - assert count_none + count_all == 1 + if count_none + count_all != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) return not count_none > 0 or count_all > 0 diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index fcc50ffcb0ed..aaf46875e80e 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -74,6 +74,12 @@ MxFp4LinearKernel, MxFp4LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mxfp4.aiter import ( + AiterMxfp4LinearKernel, +) +from vllm.model_executor.kernels.linear.mxfp4.emulation import ( + EmulationMxfp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( FlashInferMxFp4LinearKernel, ) @@ -86,6 +92,13 @@ from vllm.model_executor.kernels.linear.mxfp4.xpu import ( XPUMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp6 import ( + MxFp6LinearKernel, + MxFp6LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mxfp6.emulation import ( + EmulationMxfp6LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8 import ( Mxfp8LinearKernel, Mxfp8LinearLayerConfig, @@ -274,6 +287,7 @@ def _get_linear_backend() -> str: AiterFp8BlockScaledMMKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterMxfp4LinearKernel, }, "machete": { MacheteLinearKernel, @@ -290,6 +304,8 @@ def _get_linear_backend() -> str: "emulation": { EmulationMxfp8LinearKernel, EmulationNvFp4LinearKernel, + EmulationMxfp6LinearKernel, + EmulationMxfp4LinearKernel, }, "xpu": { XPUW8A8FP8LinearKernel, @@ -463,11 +479,25 @@ def _filter_kernels_by_backend( ], } +_POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = { + PlatformEnum.CUDA: [ + EmulationMxfp6LinearKernel, + ], + PlatformEnum.ROCM: [ + EmulationMxfp6LinearKernel, + ], +} + _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { PlatformEnum.CUDA: [ FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel, + EmulationMxfp4LinearKernel, + ], + PlatformEnum.ROCM: [ + AiterMxfp4LinearKernel, + EmulationMxfp4LinearKernel, ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, @@ -805,9 +835,15 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel: ) -def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: +def init_mxfp4_linear_kernel( + activation_quant_key: QuantKey | None = None, +) -> MxFp4LinearKernel: """Select and instantiate the best MXFP4 linear kernel for the current platform.""" + config = MxFp4LinearLayerConfig( + activation_quant_key=activation_quant_key, + ) + linear_backend = _get_linear_backend() platform = current_platform._enum @@ -836,8 +872,13 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: failure_reasons.append(f"{kernel_cls.__name__}: {reason}") continue + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__) - return kernel_cls(MxFp4LinearLayerConfig()) + return kernel_cls(config) raise ValueError( "Failed to find a kernel that can implement the " @@ -845,6 +886,59 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: ) +def init_mxfp6_linear_kernel( + weight_quant_key: QuantKey, + activation_quant_key: QuantKey | None = None, +) -> MxFp6LinearKernel: + """Select and instantiate the best MXFP6 linear kernel for the + current platform.""" + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, + activation_quant_key=activation_quant_key, + ) + + linear_backend = _get_linear_backend() + + platform = current_platform._enum + possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, [])) + + # Apply --linear-backend filtering when set. + if linear_backend != "auto": + filtered = _filter_kernels_by_backend(linear_backend, possible) + if not filtered: + raise ValueError( + f"--linear-backend={linear_backend} was requested but no " + f"'{linear_backend}' kernel exists for MXFP6 layers." + ) + possible = filtered + + failure_reasons = [] + for kernel_cls in possible: + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: + failure_reasons.append( + f" {kernel_cls.__name__} disabled by environment variable" + ) + continue + + is_supported, reason = kernel_cls.is_supported() + if not is_supported: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + logger.info_once("Using %s for MXFP6 GEMM", kernel_cls.__name__) + return kernel_cls(config) + + raise ValueError( + "Failed to find a kernel that can implement the " + "MXFP6 linear layer. Reasons: \n" + "\n".join(failure_reasons) + ) + + def init_wfp8_a16_linear_kernel( weight_quant_key: QuantKey, activation_quant_key: QuantKey, @@ -884,6 +978,7 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: """Select and instantiate the best NVFP4 linear kernel for the current platform.""" config = NvFp4LinearLayerConfig() + a16_kernels = (MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel) # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the # batch-invariant CUTLASS implementation when available, otherwise fall @@ -924,6 +1019,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: force_kernel = MarlinNvFp4LinearKernel if force_kernel is not None: + if use_a16 and force_kernel not in a16_kernels: + raise ValueError(f"{force_kernel.__name__} does not support W4A16") is_supported, reason = force_kernel.is_supported() if not is_supported: raise ValueError( @@ -936,6 +1033,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: # Auto-select from registry (or --linear-backend filtered). platform = current_platform._enum possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, [])) + if use_a16: + possible = [kernel for kernel in possible if kernel in a16_kernels] # Apply --linear-backend filtering when set. if linear_backend != "auto": @@ -1025,6 +1124,10 @@ def register_linear_kernel( if platform not in _POSSIBLE_MXFP4_KERNELS: _POSSIBLE_MXFP4_KERNELS[platform] = [] _POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class) + elif kernel_type == "mxfp6": + if platform not in _POSSIBLE_MXFP6_KERNELS: + _POSSIBLE_MXFP6_KERNELS[platform] = [] + _POSSIBLE_MXFP6_KERNELS[platform].append(kernel_class) else: raise ValueError(f"Unrecognized kernel type: {kernel_type}") @@ -1079,6 +1182,12 @@ def register_linear_kernel( "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", + "MxFp6LinearKernel", + "MxFp6LinearLayerConfig", + "init_mxfp6_linear_kernel", + "EmulationMxfp6LinearKernel", + "AiterMxfp4LinearKernel", + "EmulationMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", "MarlinMxFp4LinearKernel", "FlashInferCutedslMxfp8LinearKernel", diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index f6931ceb1c83..704fa7da4b41 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -307,7 +307,3 @@ def kernel( ) if const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() - - -def make_host_bf16(k_val: int, bs: int = 128): - return LLBf16Dotprod(k=k_val, bs=bs) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py b/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py new file mode 100644 index 000000000000..44964f3e47bf --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr + + +class CuteSkinnyGemm: + """Shape-dynamic low-latency GEMM for small token counts. + + Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16 + inputs, FP32 accumulators, and an output matching the input dtype. The + residual term is optional and is added before the output conversion. N and + K are runtime values. The tiny M dimension is fully unrolled alongside a + small set of tuning parameters. + """ + + def __init__( + self, + *, + element_type, + num_rows: int, + block_size: int, + outputs_per_block: int, + vector_width: int = 8, + k_unroll: int = 1, + has_residual: bool = False, + use_pdl: bool = False, + ) -> None: + if block_size % cute.arch.WARP_SIZE != 0: + raise ValueError("block_size must be a multiple of the warp size") + self.element_type = element_type + self.num_rows = num_rows + self.block_size = block_size + self.outputs_per_block = outputs_per_block + self.vector_width = vector_width + self.k_unroll = k_unroll + self.has_residual = has_residual + self.use_pdl = use_pdl + self.num_warps = block_size // cute.arch.WARP_SIZE + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + stream: CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=self.num_rows * self.outputs_per_block * self.num_warps * 4, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + num_rows: cutlass.Constexpr = self.num_rows + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + + acc_layout = cute.make_layout( + (num_rows, outputs_per_block), stride=(outputs_per_block, 1) + ) + acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + # Layout after both divides is (M/N, K_TILE, K_LANE, K_VEC). + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout((num_rows, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, vector_width), stride=(vector_width, 1) + ), + self.element_type, + ) + + for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None]) + + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + tB = tB_all[n_idx, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to( + cutlass.Float32 + ) * b_regs[ni, vi].to(cutlass.Float32) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ) + if const_expr(self.has_residual): + total += gResidual[mi, n_idx].to(cutlass.Float32) + gC[mi, n_idx] = cutlass.Float32(total).to(self.element_type) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 16bf9632965b..b0471b565737 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -31,20 +31,71 @@ def is_available() -> bool: return _cutedsl_available +# Default configs _DEFAULT_DOTPROD_BS = 128 _DEFAULT_DOTPROD_MAX_M = 4 _DEFAULT_SPLITK_CONFIG = (6, 4) -_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = { - (7168, 256): 6, + +# ll_bf16 router shapes covered by warmup/tuning: +# (4096, 256) # DSV4-Flash +# (6144, 256) # GLM5.2 +# (7168, 256) # DSV3.2 +# (7168, 384) # DSV4-Pro + +# SM100f-specific tuned configs +_SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { + (6144, 256): {M: 256 for M in (1, 3, 4)}, +} +_SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + (4096, 256): { + **{M: (8, 5) for M in (5, 8)}, + 9: (8, 2), + }, + (7168, 256): {14: (8, 2)}, + (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, + (7168, 384): {M: (7, 5) for M in (13, 16)}, } -_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + +# SM90-specific tuned configs +_SM90_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { + (4096, 256): {M: 256 for M in (1, 3)}, + (7168, 384): {M: 256 for M in (1, 2)}, +} +_SM90_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + (4096, 256): { + **{M: (8, 2) for M in range(5, 8)}, + **{M: (8, 5) for M in range(10, 12)}, + **{M: (8, 2) for M in (13, 16)}, + **{M: (8, 5) for M in (14, 15)}, + }, + (7168, 256): {8: (6, 5)}, + (6144, 256): {M: (8, 2) for M in (9, 11)}, (7168, 384): { - 5: (4, 4), - **{M: (5, 4) for M in range(6, 17)}, + **{M: (8, 2) for M in (6, 8, 12)}, + **{M: (7, 5) for M in (7, 9, 10, 11, 13, 14, 15, 16)}, }, } +def _arch_tuned_configs() -> tuple[ + dict[tuple[int, int], dict[int, int]], + dict[tuple[int, int], dict[int, tuple[int, int]]], +]: + from vllm.platforms import current_platform + + if current_platform.is_device_capability_family(100): + return ( + _SM100F_TUNED_DOTPROD_BS, + _SM100F_TUNED_SPLITK_CONFIGS, + ) + if current_platform.is_device_capability(90): + return ( + _SM90_TUNED_DOTPROD_BS, + _SM90_TUNED_SPLITK_CONFIGS, + ) + return {}, {} + + _cute_ctx = None @@ -89,11 +140,12 @@ def __init__(self) -> None: self._splitk_cache: dict[tuple[int, int], Any] = {} def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: - dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M) - if dotprod_max_m >= M or K < 2048: - return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS) + tuned_bs, tuned_splitk = _arch_tuned_configs() + if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: + bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) + return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) - split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get( + split_k, num_stages = tuned_splitk.get((K, N), {}).get( M, _DEFAULT_SPLITK_CONFIG ) return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py new file mode 100644 index 000000000000..e1dd18920b75 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_cutedsl_available: bool | None = None + + +@dataclass(frozen=True, slots=True) +class SkinnyGemmConfig: + num_rows: int + block_size: int + outputs_per_block: int + k_unroll: int = 1 + vector_width: int = 8 + + +class ShapeDynamicSkinnyGemm: + def __init__(self) -> None: + self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {} + self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set() + self._warmup_registered = False + + @staticmethod + def is_available() -> bool: + global _cutedsl_available + if _cutedsl_available is not None: + return _cutedsl_available + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + + _cutedsl_available = True + except ImportError: + _cutedsl_available = False + logger.info("cuteDSL is not available; skinny GEMM is disabled") + return _cutedsl_available + + @staticmethod + def _config(m: int, n: int, k: int) -> SkinnyGemmConfig: + num_rows = m + wide_block = 224 + if m == 1 and k >= 7168 and k % (wide_block * 8) == 0: + if n % 3 == 0: + k_unroll = 2 if n <= 2304 else 4 + return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll) + if 2304 < n < 4096 and n % 2 == 0: + return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4) + + if k <= 2048 or k % (128 * 8) != 0: + outputs_per_block = 2 if k <= 2048 else 4 + if n % outputs_per_block: + outputs_per_block = 1 + if k % (64 * 8) == 0: + return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2) + if k % (32 * 8) == 0: + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2) + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4) + + block_size = 64 if 4096 <= n < 8192 else 128 + outputs_per_block = 1 if m == 1 and n <= 2304 else 2 + if n % outputs_per_block: + outputs_per_block = 1 + k_unroll = 2 if n <= 2304 or n >= 16384 else 1 + return SkinnyGemmConfig( + num_rows, + block_size, + outputs_per_block, + k_unroll=k_unroll, + ) + + @staticmethod + def _cutlass_dtype(dtype: torch.dtype): + from cutlass import BFloat16, Float16 + + return BFloat16 if dtype == torch.bfloat16 else Float16 + + @staticmethod + def _stream(): + from cuda.bindings.driver import CUstream + + from vllm.utils.torch_utils import current_stream + + return CUstream(current_stream().cuda_stream) + + @staticmethod + def _use_pdl() -> bool: + from vllm.platforms import current_platform + + return current_platform.is_arch_support_pdl() + + def _compile( + self, + dtype: torch.dtype, + config: SkinnyGemmConfig, + has_residual: bool, + ) -> None: + import cutlass.cute as cute + from quack.compile_utils import make_fake_tensor + + from ._skinny_gemm import CuteSkinnyGemm + + element_type = self._cutlass_dtype(dtype) + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor( + element_type, + (config.num_rows, k), + divisibility=config.vector_width, + ) + b = make_fake_tensor( + element_type, + (n, k), + divisibility=config.vector_width, + ) + c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + kernel = CuteSkinnyGemm( + element_type=element_type, + num_rows=config.num_rows, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=has_residual, + use_pdl=self._use_pdl(), + ) + self._compiled[(dtype, config, has_residual)] = cute.compile( + kernel, + a, + b, + residual, + c, + self._stream(), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", + ) + + def request_warmup( + self, + dtype: torch.dtype, + shapes: Iterable[tuple[int, int, int]], + ) -> None: + """Request compilation of ``(M, N, K)`` shapes before graph capture.""" + self.request_warmup_configs( + dtype, + (self._config(m, n, k) for m, n, k in shapes), + ) + + def request_warmup_configs( + self, + dtype: torch.dtype, + configs: Iterable[SkinnyGemmConfig], + *, + has_residual: bool = False, + ) -> None: + """Request compilation of explicit measured configs before capture.""" + self._warmup_configs.update((dtype, config, has_residual) for config in configs) + if self._warmup_registered: + return + from vllm.model_executor.warmup.cutedsl_warmup import ( + register_cutedsl_warmup_provider, + ) + + register_cutedsl_warmup_provider(self) + self._warmup_registered = True + + def get_cutedsl_warmup_compile_units(self): + from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit + + return tuple( + CuTeDSLCompileUnit( + name=( + "shape-dynamic skinny GEMM with residual" + if has_residual + else "shape-dynamic skinny GEMM" + ), + key=("shape-dynamic-skinny-gemm", dtype, config, has_residual), + compile=partial(self._compile, dtype, config, has_residual), + ) + for dtype, config, has_residual in sorted( + self._warmup_configs, + key=lambda item: ( + str(item[0]), + item[1].num_rows, + item[1].block_size, + item[1].outputs_per_block, + item[1].k_unroll, + item[1].vector_width, + item[2], + ), + ) + ) + + def __call__( + self, + a: torch.Tensor, + b: torch.Tensor, + config: SkinnyGemmConfig | None = None, + residual: torch.Tensor | None = None, + ) -> torch.Tensor: + if a.dim() != 2 or b.dim() != 2: + raise ValueError("a and b must be 2D tensors") + if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype: + raise ValueError("a and b must have the same BF16 or FP16 dtype") + if not a.is_cuda or not b.is_cuda or a.device != b.device: + raise ValueError("a and b must be CUDA tensors on the same device") + if not a.is_contiguous() or not b.is_contiguous(): + raise ValueError("a and b must be contiguous") + if a.shape[1] != b.shape[1]: + raise ValueError("a and b must have matching K dimensions") + if not 1 <= a.shape[0] <= 16: + raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16") + if residual is not None: + if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]): + raise ValueError("residual must have shape (M, N)") + if residual.dtype != a.dtype: + raise ValueError("residual must have the same dtype as a and b") + if residual.device != a.device or not residual.is_cuda: + raise ValueError("residual must be on the same CUDA device") + if not residual.is_contiguous(): + raise ValueError("residual must be contiguous") + + config = config or self._config(a.shape[0], b.shape[0], a.shape[1]) + if config.num_rows != a.shape[0]: + raise ValueError("config num_rows must match M") + if b.shape[0] % config.outputs_per_block != 0: + raise ValueError("N must be divisible by outputs_per_block") + if a.shape[1] % (config.block_size * config.vector_width) != 0: + raise ValueError( + "K must be divisible by block_size * vector_width for this config" + ) + has_residual = residual is not None + cache_key = (a.dtype, config, has_residual) + if cache_key not in self._compiled: + self._compile(a.dtype, config, has_residual) + output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device) + residual_arg = output if residual is None else residual + self._compiled[cache_key](a, b, residual_arg, output, self._stream()) + return output + + +shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm() diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 87ed8d1b582f..5a4f65297158 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -25,6 +25,7 @@ unpack_cols, ) from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -111,8 +112,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: else: padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) - # Allocate marlin workspace. - self.workspace = marlin_make_workspace_new(device) + # Allocate marlin workspace, reusing existing storage on reload. + self.workspace = marlin_make_workspace_new( + device, existing=getattr(self, "workspace", None) + ) # Default names since marlin requires empty parameters for these, # TODO: remove this requirement from marlin (allow optional tensors) @@ -174,7 +177,9 @@ def transform_w_s(x): getattr(layer, self.w_gidx_name) ) self._transform_param(layer, self.w_gidx_name, lambda _: g_idx) - layer.g_idx_sort_indices = g_idx_sort_indices + replace_parameter( + layer, "g_idx_sort_indices", g_idx_sort_indices, prefer_copy=True + ) else: setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device)) layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py new file mode 100644 index 000000000000..f1336dd008ae --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn.parameter import Parameter + +import vllm.envs as envs +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +logger = init_logger(__name__) + +# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP +# which can force the engine core to spawn instead of fork. +# is_aiter_found_and_supported() checks platform + arch + library availability via +# find_spec/amdsmi, so it stays HIP-free. +# Actual aiter imports are deferred to the functions/methods that need them, +# where HIP initialization is expected. +if is_aiter_found_and_supported(): + from vllm.utils.torch_utils import direct_register_custom_op + + def gemm_with_dynamic_quant( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + rocm_use_aiter_fp4_asm_gemm: bool = False, + out_dtype: torch.dtype | None = torch.bfloat16, + x_scales: torch.Tensor | None = None, + ) -> torch.Tensor: + from aiter.ops.triton.gemm_afp4wfp4 import ( + gemm_afp4wfp4, + gemm_afp4wfp4_preshuffled_weight_scales, + ) + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + if rocm_use_aiter_fp4_asm_gemm: + from aiter import gemm_a4w4, per_1x32_f4_quant_hip + + M = x.shape[0] + N = weight.shape[0] + K = weight.shape[1] + if rocm_use_aiter_fp4_asm_gemm: + if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): + if x_scales is None: + # use hip quant kernel for performance + if M >= 32: + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) + else: + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False) + else: + x_q = x + x_s = x_scales + + if M >= 32: + x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1) + else: + x_s = x_s[:M, ...].view(torch.uint8) + + y = torch.empty(M, N, device=x_q.device, dtype=out_dtype) + gemm_afp4wfp4_preshuffled_weight_scales( + x_q.view(torch.uint8), + weight.view(torch.uint8).view(weight.shape[0] // 16, -1), + x_s, + weight_scale.view(torch.uint8).view( + weight_scale.shape[0] // 32, -1 + ), + out_dtype, + y, + ) + else: + if x_scales is None: + # use hip quant kernel for performance + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) + else: + x_q = x + x_s = x_scales + + y = gemm_a4w4( + x_q, + weight.view(x_q.dtype), + x_s, + weight_scale.view(x_s.dtype), + dtype=out_dtype, + bpreshuffle=True, + ) + return y[:M] + else: + if x_scales is None: + x_q, x_s = dynamic_mxfp4_quant(x) + else: + x_q = x + x_s = x_scales + y = torch.empty( + x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype + ) + + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) + return y + + def gemm_with_dynamic_quant_fake( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_scales: torch.Tensor = None, + rocm_use_aiter_fp4_asm_gemm: bool = False, + out_dtype: torch.dtype | None = torch.bfloat16, + ) -> torch.Tensor: + return torch.empty( + (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device + ) + + direct_register_custom_op( + op_name="gemm_with_dynamic_quant", + op_func=gemm_with_dynamic_quant, + mutates_args=[], + fake_impl=gemm_with_dynamic_quant_fake, + dispatch_key=current_platform.dispatch_key, + ) + + +class AiterMxfp4LinearKernel(MxFp4LinearKernel): + """AITER-based native MXFP4 GEMM kernel for ROCm.""" + + def __init__(self, config: MxFp4LinearLayerConfig) -> None: + super().__init__(config) + self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() + self.out_dtype = torch.get_default_dtype() + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.supports_mx(): + return False, "current platform does not support native MXFP4 computation" + + from vllm._aiter_ops import is_aiter_found_and_supported + from vllm.model_executor.kernels.linear import _get_linear_backend + + linear_backend = _get_linear_backend() + + if ( + current_platform.is_rocm() + and current_platform.supports_mx() + and "AiterMxfp4LinearKernel" not in envs.VLLM_DISABLED_KERNELS + and linear_backend == "auto" + and not is_aiter_found_and_supported() + ): + logger.warning_once( + "This platform supports native MXFP4 W4A4 MOE " + "computation via AITER MOE backend, but AITER is not " + "found or not supported. Consider installing AITER: " + "https://github.com/ROCm/aiter." + ) + + if is_aiter_found_and_supported(): + return True, None + return False, "AITER not found or not supported on the current platform" + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.use_asm_gemm: + from aiter.ops.shuffle import shuffle_weight + + weight_scale = layer.weight_scale.data + sm, sn = weight_scale.shape + weight_scale = weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) + weight_scale = weight_scale.permute(0, 3, 5, 2, 4, 1, 6).contiguous() + weight_scale = weight_scale.view(sm, sn) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + layer.weight = Parameter( + shuffle_weight(layer.weight.data, layout=(16, 16)), + requires_grad=False, + ) + else: + layer.weight_scale = Parameter( + layer.weight_scale.data.T.contiguous(), requires_grad=False + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + y = torch.ops.vllm.gemm_with_dynamic_quant( + x, + layer.weight, + layer.weight_scale, + self.use_asm_gemm, + self.out_dtype, + ) + if bias is not None: + y = y + bias + return y diff --git a/vllm/model_executor/kernels/linear/mxfp4/base.py b/vllm/model_executor/kernels/linear/mxfp4/base.py index 868faa4731d5..9c3088df52a1 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/base.py +++ b/vllm/model_executor/kernels/linear/mxfp4/base.py @@ -6,6 +6,8 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + @dataclass class MxFp4LinearLayerConfig: @@ -13,9 +15,13 @@ class MxFp4LinearLayerConfig: All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per byte) and per-block weight scales (group size 32). + + Attributes: + activation_quant_key: Identifies the activation quantization format, + or `None` when activations must not be quantized. """ - pass + activation_quant_key: QuantKey | None = None class MxFp4LinearKernel(ABC): diff --git a/vllm/model_executor/kernels/linear/mxfp4/emulation.py b/vllm/model_executor/kernels/linear/mxfp4/emulation.py new file mode 100644 index 000000000000..66436521b8d9 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/emulation.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from functools import partial + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + dequant_mxfp4, + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( + quant_dequant_mxfp6, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E3M2Dynamic, +) +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +logger = init_logger(__name__) + +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[ + QuantKey, Callable[[torch.Tensor], torch.Tensor] +] = { + kMxfp4Dynamic: quant_dequant_mxfp4, + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + + +class EmulationMxfp4LinearKernel(MxFp4LinearKernel): + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear).""" + + def __init__(self, config: MxFp4LinearLayerConfig) -> None: + super().__init__(config) + if config.activation_quant_key is None: + # no input Q/DQ for weight-only + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( + lambda x: x + ) + else: + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[ + config.activation_quant_key + ] + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in ( + None, + kMxfp4Dynamic, + kMxfp6E3M2Dynamic, + kMxfp6E2M3Dynamic, + ): + return False, "only supports MXFP4 or MXFP6 or unquantized activations" + + if ( + current_platform.is_rocm() + and current_platform.supports_mx() + and config.activation_quant_key != kMxfp4Dynamic + ): + logger.warning_once( + "The current platform supports native MXFP4/MXFP6 computation, " + f"but kernels for activation_quant_key={config.activation_quant_key} " + f"are not yet integrated in vLLM. Using EmulationMxfp4LinearKernel, " + "with simulated weight dequantization and activation " + "QDQ (quantize and dequantize), with the linear " + "layers computed in high precision." + ) + + if not current_platform.supports_mx(): + logger.warning_once( + "The current platform does not support native MXFP4 " + "computation. Using EmulationMxfp4LinearKernel, with simulated weight " + "dequantization and activation QDQ (quantize and dequantize), with " + "the linear layers computed in high precision." + ) + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype) + qdq_x = self.quant_dequant_func(x) + return F.linear(qdq_x, dq_w, bias) diff --git a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py index 532a7d55e929..4491e855229d 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py @@ -7,6 +7,9 @@ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( swizzle_mxfp4_scales, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutedsl @@ -28,6 +31,8 @@ def is_supported( @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index d93f5d481580..f38b8ead63af 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -3,14 +3,18 @@ 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.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + class HummingMxFp4LinearKernel(MxFp4LinearKernel): """Humming GEMM Kernel for MXFP4.""" @@ -28,7 +32,15 @@ def is_supported( return True, None @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in (None, kMxfp4Dynamic): + return False, "only supports MXFP4 dynamic or unquantized activations" + if config.activation_quant_key is not None: + logger.warning_once( + "HummingMxFp4LinearKernel is a weight-only (A16) kernel; " + "the requested activation quantization (%s) is ignored.", + config.activation_quant_key, + ) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py index 38440752072e..30ff0d58d4c0 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py @@ -3,8 +3,13 @@ import torch +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic + from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + class MarlinMxFp4LinearKernel(MxFp4LinearKernel): @classmethod @@ -20,7 +25,15 @@ def is_supported( return False, "Marlin FP4 not available" @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in (None, kMxfp4Dynamic): + return False, "only supports MXFP4 dynamic or unquantized activations" + if config.activation_quant_key is not None: + logger.warning_once( + "MarlinMxFp4LinearKernel is a weight-only (A16) kernel; " + "the requested activation quantization (%s) is ignored.", + config.activation_quant_key, + ) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/xpu.py b/vllm/model_executor/kernels/linear/mxfp4/xpu.py index 8d33939d2ed9..c499d0a17a45 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/xpu.py +++ b/vllm/model_executor/kernels/linear/mxfp4/xpu.py @@ -6,6 +6,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( xpu_mxfp4_quantize as quant_mxfp4, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -24,7 +27,9 @@ def is_supported( return True, None @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp6/__init__.py b/vllm/model_executor/kernels/linear/mxfp6/__init__.py new file mode 100644 index 000000000000..bbdd6f0b52c8 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.kernels.linear.mxfp6.base import ( + MxFp6LinearKernel, + MxFp6LinearLayerConfig, +) + +__all__ = [ + "MxFp6LinearKernel", + "MxFp6LinearLayerConfig", +] diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py new file mode 100644 index 000000000000..ef00bc57647d --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + + +@dataclass +class MxFp6LinearLayerConfig: + """Configuration for an MXFP6 linear layer. + + All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per + byte) and per-block weight scales (group size 32). + + Attributes: + weight_quant_key: Identifies the weight quantization format. Can be + kMxfp6E2M3Static or kMxfp6E3M2Static. + activation_quant_key: Identifies the activation quantization format, + or `None` when activations must not be quantized. + """ + + weight_quant_key: QuantKey + activation_quant_key: QuantKey | None = None + + +class MxFp6LinearKernel(ABC): + """Base class for MXFP6 quantized linear kernels. + + Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc). + The kernel selection mechanism iterates over registered subclasses in + priority order,calling ``is_supported`` and ``can_implement`` to find the best + match for the current hardware. + """ + + def __init__(self, config: MxFp6LinearLayerConfig) -> None: + assert self.can_implement(config)[0] + assert self.is_supported()[0] + self.config = config + + @classmethod + @abstractmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + """Return whether this kernel can run on the current platform.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + """Return whether this kernel can handle *config*.""" + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Transform weights into the format required by this kernel. + + Called once after checkpoint weights have been loaded onto the + device. Implementations should repack / swizzle / pad weights + and scales in-place on *layer*. + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the quantized GEMM.""" + raise NotImplementedError diff --git a/vllm/model_executor/kernels/linear/mxfp6/emulation.py b/vllm/model_executor/kernels/linear/mxfp6/emulation.py new file mode 100644 index 000000000000..0ef68deb0d02 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/emulation.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from functools import partial + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( + dequant_mxfp6, + quant_dequant_mxfp6, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, +) + +from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig + +_WEIGHT_DEQUANT_FUNCS: dict[QuantKey, Callable[..., torch.Tensor]] = { + kMxfp6E3M2Static: partial(dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Static: partial(dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[ + QuantKey, Callable[[torch.Tensor], torch.Tensor] +] = { + kMxfp4Dynamic: quant_dequant_mxfp4, + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + + +class EmulationMxfp6LinearKernel(MxFp6LinearKernel): + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear).""" + + def __init__(self, config: MxFp6LinearLayerConfig) -> None: + super().__init__(config) + self.dequant_func = _WEIGHT_DEQUANT_FUNCS[config.weight_quant_key] + if config.activation_quant_key is None: + # no input Q/DQ for weight-only + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( + lambda x: x + ) + else: + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[ + config.activation_quant_key + ] + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + if config.weight_quant_key not in ( + kMxfp6E2M3Static, + kMxfp6E3M2Static, + ): + return False, "only supports MXFP6 weights" + if config.activation_quant_key not in ( + None, + kMxfp4Dynamic, + kMxfp6E3M2Dynamic, + kMxfp6E2M3Dynamic, + ): + return False, "only supports MXFP4 or MXFP6 or unquantized activations" + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) + qdq_x = self.quant_dequant_func(x) + return F.linear(qdq_x, dq_w, bias) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 1b39491ab346..ab396d026ad1 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -9,6 +9,9 @@ rocm_aiter_ops, ) from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, ) @@ -16,6 +19,7 @@ from vllm.platforms import current_platform from .BlockScaledMMLinearKernel import ( + FP8BlockParams, Fp8BlockScaledMMLinearKernel, ) from .cutlass import CutlassInt8ScaledMMLinearKernel @@ -370,11 +374,27 @@ def __init__(self, config: FP8ScaledMMLinearLayerConfig): super().__init__(config) n, k = config.weight_shape - self.use_triton = ( - not current_platform.is_fp8_fnuz() - and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) + _on_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + _on_gfx1250 = on_gfx1250() + + self.use_triton = not current_platform.is_fp8_fnuz() and ( + rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250 ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + + params = FP8BlockParams.from_layer(layer) + if params.weight_scale_inv is not None: + ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV + else: + ws, attr = params.weight_scale, params.WEIGHT_SCALE + if ws is not None and ws.dtype == torch.float8_e8m0fnu: + replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) + @classmethod def is_supported(cls, compute_capability=None): return ( @@ -406,19 +426,12 @@ def apply_block_scaled_mm( Bs: torch.Tensor, ) -> torch.Tensor: if As.dtype != Bs.dtype: - from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, - ) - if As.dtype == torch.float8_e8m0fnu: As = _upcast_e8m0_to_fp32(As).contiguous() else: As = As.to(torch.float32) - if Bs.dtype == torch.float8_e8m0fnu: - Bs = _upcast_e8m0_to_fp32(Bs).contiguous() - else: - Bs = Bs.to(torch.float32) + Bs = Bs.to(torch.float32) out_dtype = self.config.out_dtype if self.use_triton: diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 2b6d3ed73698..a04be2e097c0 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -104,10 +104,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if not on_mi3xx(): - return False, "requires MI3xx." + if get_cdna_version() <= 2: + return False, "requires CDNA3+" if compute_capability is not None and compute_capability < 94: return False, "requires compute capability 94 and above." diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index 64bc5b6c8bbe..4a35f49a22e3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,15 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_gfx12x, on_mi3xx + from vllm.platforms.rocm import get_cdna_version, on_gfx12x, on_gfx1250 - if not (on_mi3xx() or on_gfx12x()): - return False, "requires MI3xx or gfx12x" + # wvSplitKQ (skinny GEMM) is excluded from the gfx1250 build. + if on_gfx1250(): + return False, "wvSplitKQ (skinny GEMM) is not built on gfx1250" + + # Restore RDNA4 (gfx12x) dropped by the get_cdna_version()>2 refactor. + if get_cdna_version() <= 2 and not on_gfx12x(): + return False, "requires CDNA3+ (gfx942/gfx950) or RDNA4 (gfx12x)" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index f30d6ced1d35..e6ec7926c4bd 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -197,6 +197,45 @@ def is_supported( 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) + + # Checkpoint scale is [n_blocks, k_blocks] (one value per 128x128 tile). + # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout. + # We store the transposed contiguous buffer as a .t() view so that: + # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape + # - apply_block_scaled_mm recovers the contiguous buffer via .t() + scale_kn = scale.data.t().contiguous() # [k_blocks, n_blocks] + replace_parameter(layer, scale_attr, scale_kn.t()) # view: [n_blocks, k_blocks] + + if getattr(layer, "is_bmm", False): + self._prepare_bmm_params(layer, scale_kn) + + def _prepare_bmm_params( + self, layer: torch.nn.Module, scale_kn: torch.Tensor + ) -> None: + """Precompute batched weight and scale for grouped fp8_bmm (e.g. wo_a). + + Splits scale [k_blocks, n_blocks] into [G, k_blocks, n_blocks_per_group] + and weight [N_total, K] into [G, K, N_per_group] for batch GEMM. + """ + batch = layer.bmm_batch_size + k_blocks, n_blocks = scale_kn.shape + layer.bmm_scale = ( + scale_kn.reshape(k_blocks, batch, n_blocks // batch) + .permute(1, 0, 2) + .contiguous() + ) + w = layer.weight + N_total, K = w.shape + layer.bmm_weight = w.reshape(batch, N_total // batch, K).permute( + 0, 2, 1 + ) # [G, K, N_per_group] + def apply_block_scaled_mm( self, A: torch.Tensor, @@ -204,13 +243,14 @@ def apply_block_scaled_mm( As: torch.Tensor, 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. + # B is [N, K]; .t() gives [K, N] view (no copy). + # Bs is stored as [n_blocks, k_blocks] view; .t() recovers the + # contiguous [k_blocks, n_blocks] buffer that oneDNN expects. return torch.ops._xpu_C.fp8_gemm( A, B.t(), self.config.out_dtype, As, - Bs.t().contiguous(), + Bs.t(), torch.Tensor(), ) diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 50abac3e0101..3afd7069aa3e 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -856,21 +856,23 @@ def hc_prenorm_gemm_block_m_tilelang( T.sync_threads() if warp_id == 0: + reduced_acc = T.alloc_local((block_m,), T.float32) + reduced_sqr = T.alloc_local((block_m,), T.float32) + T.clear(reduced_acc) + T.clear(reduced_sqr) for i_m in T.unroll(block_m): token_idx = i_mt * block_m + i_m if token_idx < num_tokens: if lane < tile_n: - reduced_acc = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): - reduced_acc += warp_acc[i_w, i_m, lane] + reduced_acc[i_m] += warp_acc[i_w, i_m, lane] out_idx = i_t * tile_n + lane if out_idx < n_out: - out[0, token_idx, out_idx] = reduced_acc + out[0, token_idx, out_idx] = reduced_acc[i_m] if lane == 0 and i_t == 0: - reduced_sqr = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): - reduced_sqr += warp_sqr[i_w, i_m] - sqrsum[0, token_idx] = reduced_sqr + reduced_sqr[i_m] += warp_sqr[i_w, i_m] + sqrsum[0, token_idx] = reduced_sqr[i_m] if ENABLE_PDL: T.pdl_trigger() diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index a8c91af9adc4..953a47e0fae3 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -158,6 +158,51 @@ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) +@CustomOp.register("situ_and_mul") +class SituAndMul(CustomOp): + """SituGLU activation used by Kimi models. + + Computes beta * tanh(gate / beta) * sigmoid(gate) * up. When + ``linear_beta`` is set, the up projection is also softly clipped with + linear_beta * tanh(up / linear_beta). + """ + + def __init__( + self, + beta: float = 1.0, + linear_beta: float | None = None, + *, + compile_native: bool = True, + ): + super().__init__(compile_native=compile_native) + self.beta = float(beta) + self.linear_beta = None if linear_beta is None else float(linear_beta) + if current_platform.is_cuda_alike(): + self.op = torch.ops._C.situ_and_mul + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (gate * up).to(x.dtype) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries. + # linear_beta<=0 signals "unset" to the kernel (up passed through). + d = x.shape[-1] // 2 + out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device) + self.op( + out, x, self.beta, -1.0 if self.linear_beta is None else self.linear_beta + ) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + @CustomOp.register("silu_and_mul_with_clamp") class SiluAndMulWithClamp(CustomOp): """SwiGLU activation with input clamping (used by some MoE shared experts). @@ -613,8 +658,8 @@ class ReLUSquaredActivation(CustomOp): # --8<-- [end:relu2] - def __init__(self): - super().__init__() + def __init__(self, *, compile_native: bool = True): + super().__init__(compile_native=compile_native) if current_platform.is_cuda_alike(): self.op = torch.ops._C.relu_squared diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 375049d3b773..c8e5a9c6ac3d 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -680,6 +680,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: head_size=self.head_size, head_size_v=self.head_size, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=quant_mode, tq_slot_size=tq_config.slot_size_aligned, ) else: @@ -794,7 +795,7 @@ def unified_kv_cache_update( layer_slot_mapping, ) - return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype) + return key.new_empty(0) def unified_kv_cache_update_fake( diff --git a/vllm/model_executor/layers/attention/chunked_local_attention.py b/vllm/model_executor/layers/attention/chunked_local_attention.py index cb595438adef..72aad84781bb 100644 --- a/vllm/model_executor/layers/attention/chunked_local_attention.py +++ b/vllm/model_executor/layers/attention/chunked_local_attention.py @@ -20,7 +20,6 @@ ) from vllm.v1.attention.selector import get_attn_backend from vllm.v1.kv_cache_interface import ( - AttentionSpec, ChunkedLocalAttentionSpec, KVCacheSpec, get_kv_quant_mode, @@ -42,7 +41,7 @@ class ChunkedLocalAttentionBuilder(underlying_builder): # type: ignore def get_cudagraph_support( cls: type["AttentionMetadataBuilder"], vllm_config: VllmConfig, - kv_cache_spec: AttentionSpec, + kv_cache_spec: KVCacheSpec, ) -> AttentionCGSupport: # Explicit override in case the underlying builder specialized this getter. # @override omitted only because of mypy limitation due to type variable. diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 16420dbc65e7..3f737bed5821 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -272,6 +272,7 @@ ) from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, + get_num_attention_heads_from_layers, split_decodes_and_prefills, ) from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs @@ -367,6 +368,7 @@ def __init__( q_lora_rank: int | None, kv_lora_rank: int, kv_b_proj: ColumnParallelLinear, + dcp_q_replicate: bool = False, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -374,6 +376,7 @@ def __init__( use_sparse: bool = False, indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + non_causal_multi_token_decode: bool = False, **extra_impl_args, ): super().__init__() @@ -385,10 +388,12 @@ def __init__( self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.kv_b_proj = kv_b_proj + self.dcp_q_replicate = dcp_q_replicate + self.W_UK_T_dcp_qrep: torch.Tensor | None = None self.head_size = kv_lora_rank + qk_rope_head_dim self.layer_name = prefix self.indexer = indexer - + self.non_causal_multi_token_decode = non_causal_multi_token_decode self.num_kv_heads = 1 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim @@ -513,29 +518,36 @@ def __init__( compilation_config.static_forward_context[prefix] = self self.prefill_backend: MLAPrefillBackend | None - try: - prefill_backend_cls = get_mla_prefill_backend(vllm_config) - except ValueError: - if ( - not self.impl.is_sparse - or vllm_config.attention_config.mla_prefill_backend is not None - ): - raise + if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill: logger.warning_once( - "No MLA prefill backend supports this model; sparse MLA will use the " - "top-k MQA path only (no dense-MHA prefill)." + "Sparse MLA impl has no dense-MHA prefill path; using the top-k " + "MQA path only." ) self.prefill_backend = None else: - self.prefill_backend = prefill_backend_cls( - num_heads=self.num_heads, - scale=self.scale, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - vllm_config=vllm_config, - ) + try: + prefill_backend_cls = get_mla_prefill_backend(vllm_config) + except ValueError: + if ( + not self.impl.is_sparse + or vllm_config.attention_config.mla_prefill_backend is not None + ): + raise + logger.warning_once( + "No MLA prefill backend supports this model; sparse MLA will " + "use the top-k MQA path only (no dense-MHA prefill)." + ) + self.prefill_backend = None + else: + self.prefill_backend = prefill_backend_cls( + num_heads=self.num_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + vllm_config=vllm_config, + ) self.kv_cache = torch.tensor([]) @@ -592,6 +604,7 @@ def forward( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, output_shape: torch.Size | None = None, + q_dcp_replicated: torch.Tensor | None = None, ) -> torch.Tensor: if self.calculate_kv_scales: torch.ops.vllm.maybe_calc_kv_scales( @@ -647,6 +660,7 @@ def forward( self_kv_cache, attn_metadata, output=output, + q_dcp_replicated=q_dcp_replicated, ) return output else: @@ -666,6 +680,7 @@ def forward( output, encoded, kv_cache_dummy_dep=kv_cache_dummy_dep, + q_dcp_replicated=q_dcp_replicated, ) return output @@ -683,6 +698,7 @@ def forward_impl( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + q_dcp_replicated: torch.Tensor | None = None, ) -> torch.Tensor: assert output is not None, "Output tensor must be provided." @@ -739,6 +755,8 @@ def forward_impl( output_padded = output output = output[:num_actual_toks, ...] q = q[:num_actual_toks, ...] + if q_dcp_replicated is not None: + q_dcp_replicated = q_dcp_replicated[:num_actual_toks, ...] k_c_normed = k_c_normed[:num_actual_toks, ...] k_pe = k_pe[:num_actual_toks, ...] @@ -754,11 +772,23 @@ def forward_impl( num_mha_tokens = q.size(0) - num_mqa_tokens if self.impl.is_sparse and num_mha_tokens > 0: + prefill = getattr(attn_metadata, "prefill", None) + use_dense_mha = getattr(prefill, "use_dense_mha", False) prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - use_mha = ( + use_masked_mha = ( self.prefill_backend is not None - and prefill_max_seq_len <= attn_metadata.topk_tokens # type: ignore[attr-defined] - and not self._vllm_config.attention_config.sparse_mla_force_mqa + and self.impl.masked_mha_available # type: ignore[attr-defined] + and self.impl.dcp_world_size <= 1 + and prefill is not None + and _use_masked_mha( + backend_name=self.attn_backend.get_name(), + tensor_parallel_size=self._vllm_config.parallel_config.tensor_parallel_size, + query_len=prefill.max_query_len, + seq_len=prefill_max_seq_len, + ) + ) + use_mha = (use_dense_mha or use_masked_mha) and not ( + self._vllm_config.attention_config.sparse_mla_force_mqa ) if not use_mha: num_mqa_tokens = q.size(0) @@ -768,6 +798,11 @@ def forward_impl( quant_key is not None and self.prefill_backend is not None and self.prefill_backend.supports_quant_output(quant_key) + and ( + not self.impl.is_sparse + or attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] + <= attn_metadata.topk_tokens # type: ignore[attr-defined] + ) and attn_metadata is not None and attn_metadata.prefill is not None and attn_metadata.prefill.chunked_context is None @@ -794,7 +829,12 @@ def forward_impl( ) if num_mqa_tokens > 0: - mqa_q = q[:num_mqa_tokens] + if q_dcp_replicated is not None: + mqa_q = q_dcp_replicated[:num_mqa_tokens] + qrep_decode = True + else: + mqa_q = q[:num_mqa_tokens] + qrep_decode = False mqa_output_slice = output[:num_mqa_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( @@ -834,7 +874,9 @@ def forward_impl( else: # Pads the head_dim if necessary (for the underlying kernel) N, B, P = mqa_q_nope.shape - _, _, L = self.W_UK_T.shape + W_UK_T = self.W_UK_T_dcp_qrep if qrep_decode else self.W_UK_T + assert W_UK_T is not None + _, _, L = W_UK_T.shape if self.q_pad_num_heads is not None: mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L)) @@ -843,7 +885,7 @@ def forward_impl( mqa_ql_nope = mqa_q_nope.new_empty((N, B, L)) # Multiply (N, B, P) x (N, P, L) -> (N, B, L) - torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope) + torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope) # Convert from (N, B, L) to (B, N, L) mqa_ql_nope = mqa_ql_nope.transpose(0, 1) @@ -856,6 +898,7 @@ def forward_impl( ) else: mqa_q = (mqa_ql_nope, mqa_q_pe) + # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused) if self.impl.dcp_world_size > 1: if self.use_pcp: if self.impl.dcp_world_size > self.impl.pcp_world_size: @@ -864,8 +907,11 @@ def forward_impl( mqa_q = get_tp_group().all_gather(mqa_q, dim=1) else: if isinstance(mqa_q, tuple): + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) mqa_q = torch.cat(mqa_q, dim=-1) - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) + if not qrep_decode: + # mqa_q do allgather in head dim. + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) # call decode attn if not self.impl.is_sparse: @@ -953,6 +999,21 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): self.kv_b_proj, out_dtype=act_dtype ).T + if self.dcp_q_replicate: + # qrep wired here: validate unsupported decode backends once. + assert self.q_pad_num_heads in (None, self.num_heads), ( + "DCP query replication is unsupported on head-padding MLA " + "backends (q_pad_num_heads)." + ) + if ( + self.is_aiter_triton_fp4_bmm_enabled + or self.is_aiter_triton_fp8_bmm_enabled + ): + raise NotImplementedError( + "DCP query replication is not implemented for the aiter " + "FP4/FP8 MLA BMM paths." + ) + assert kv_b_proj_weight.shape == ( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), @@ -1033,6 +1094,10 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) # Convert from (L, N, P) to (N, P, L) replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) + if self.dcp_q_replicate: + self.W_UK_T_dcp_qrep = get_dcp_group().all_gather( + self.W_UK_T.contiguous(), dim=0 + ) # If we should not load quant weights, we initialize the scales to 1.0 # as the default value. See [Note: Register q/k/v/prob scales in state dict] @@ -1083,6 +1148,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): @@ -1177,6 +1243,7 @@ def unified_mla_attention_with_output( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + q_dcp_replicated: torch.Tensor | None = None, ) -> None: # kv_cache_dummy_dep is not used but accepting it creates a data dependency # that ensures torch.compile preserves ordering between KV cache update and @@ -1197,6 +1264,7 @@ def unified_mla_attention_with_output( quant_scale_ue8m0=quant_scale_ue8m0, quant_col_major=quant_col_major, quant_tma_aligned=quant_tma_aligned, + q_dcp_replicated=q_dcp_replicated, ) @@ -1213,6 +1281,7 @@ def unified_mla_attention_with_output_fake( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + q_dcp_replicated: torch.Tensor | None = None, ) -> None: return @@ -1340,6 +1409,7 @@ class ChunkedContextMetadata: seq_tot: list[int] max_seq_lens: list[int] seq_lens: torch.Tensor + context_lens: torch.Tensor workspace: torch.Tensor token_to_seq: torch.Tensor chunk_total_token: list[int] @@ -1361,6 +1431,9 @@ class ChunkedContextMetadata: q_data_type: torch.dtype | None = None output_dtype: torch.dtype | None = None prefill_backend: MLAPrefillBackend | None = None + query_lens_cpu: torch.Tensor | None = None + use_dense_mha: bool = False + topk_mask_workspace: torch.Tensor | None = None @dataclass @@ -1403,6 +1476,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]): num_decode_tokens: int num_prefills: int + causal: bool = True + # The dimension of the attention heads head_dim: int | None = None @@ -1455,6 +1530,39 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims: ) +_DSV32_MASKED_MHA_THRESHOLDS: dict[str, dict[int, tuple[int | None, ...]]] = { + "FLASHMLA_SPARSE": { + 1: (1536, 4096, None, None, None), + 2: (512, 1024, 4096, None, None), + 4: (512, 1024, 1536, 8192, None), + 8: (512, 512, 1024, 2048, 16384), + }, + "FLASHINFER_MLA_SPARSE": { + 8: (512, 1024, 1024, 2048, 32768), + }, +} +_DSV32_SEQ_LEN_BUCKETS = (2048, 4096, 8192, 16384, 32768) + + +def _use_masked_mha( + *, + backend_name: str, + tensor_parallel_size: int, + query_len: int, + seq_len: int, +) -> bool: + thresholds = _DSV32_MASKED_MHA_THRESHOLDS.get(backend_name, {}).get( + tensor_parallel_size + ) + if thresholds is None: + return False + for bucket_idx, bucket_seq_len in enumerate(_DSV32_SEQ_LEN_BUCKETS): + if seq_len <= bucket_seq_len: + min_query_len = thresholds[bucket_idx] + return min_query_len is not None and query_len >= min_query_len + return False + + @functools.cache def backend_supports_prefill_query_quantization() -> bool: """Check if the selected MLA prefill backend supports query quantization. @@ -1629,6 +1737,7 @@ def build_mla_chunked_context_metadata( seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(), max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(), seq_lens=chunk_seq_lens, + context_lens=context_lens_cpu.to(device, non_blocking=True), token_to_seq=token_to_seq_cpu.to(device, non_blocking=True), chunk_total_token=chunk_total_token.tolist(), workspace=chunked_prefill_workspace, @@ -1652,6 +1761,7 @@ def build_mla_chunked_context_metadata( seq_tot=chunk_seq_lens.sum(dim=1).tolist(), max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(), seq_lens=chunk_seq_lens, + context_lens=context_lens_cpu.to(device, non_blocking=True), token_to_seq=token_to_seq_cpu.to(device, non_blocking=True), chunk_total_token=chunk_total_token, workspace=chunked_prefill_workspace, @@ -1669,6 +1779,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): understand this class """ + kv_cache_spec: AttentionSpec + # Defines the level of query length support for this backend. # - SINGLE_ONLY: Only single-token queries (no spec decode support) # - UNIFORM: Supports uniform multi-token queries (spec decode with uniform lengths) @@ -1677,6 +1789,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # speculative decoding is enabled. query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY + # Whether this builder can flatten a non-causal query block into decode rows. + supports_non_causal_multi_token_decode: ClassVar[bool] = False + # The threshold for reordering the batch into decode and prefill requests. # If > 1, the batch will be reordered such that requests with # query length <= threshold are classified as decode requests. @@ -1777,8 +1892,14 @@ def __init__( self.vllm_config = vllm_config self.device = device self.use_pcp = parallel_config.prefill_context_parallel_size > 1 + self.non_causal_multi_token_decode = getattr( + kv_cache_spec, "non_causal_multi_token_decode", False + ) - self.num_heads = self.model_config.get_num_attention_heads(parallel_config) + # A draft cache group can have a different head count from the target. + self.num_heads = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or self.model_config.get_num_attention_heads(parallel_config) self.mla_dims = get_mla_dims(self.model_config) self.aot_schedule = current_platform.is_cuda() @@ -1904,14 +2025,42 @@ def build( seq_lens = common_attn_metadata.seq_lens dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills( - common_attn_metadata, - decode_threshold=self.reorder_batch_threshold, - require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), - treat_short_extends_as_decodes=not self.use_pcp, + non_causal_decode = common_attn_metadata.causal is False + if non_causal_decode: + if not ( + self.supports_non_causal_multi_token_decode + and self.non_causal_multi_token_decode + ): + raise ValueError( + "Non-causal multi-token MLA requires an explicitly supported " + "attention group." + ) + query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + num_active_reqs = int(torch.count_nonzero(query_lens > 0)) + uniform_active_queries = num_active_reqs > 0 and bool( + torch.all(query_lens[:num_active_reqs] == query_lens[0]) + ) + trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0)) + if not (uniform_active_queries and trailing_graph_padding): + raise ValueError( + "Non-causal MLA requires a uniform query block; got query " + f"lengths {query_lens.tolist()}." + ) + # Use exact GPU sequence lengths instead of the prefill path's CPU + # context-length upper bounds. + num_decodes = num_reqs + num_prefills = 0 + num_decode_tokens = num_tokens + num_prefill_tokens = 0 + else: + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), + treat_short_extends_as_decodes=not self.use_pcp, + ) ) - ) assert num_decodes + num_prefills == num_reqs assert num_decode_tokens + num_prefill_tokens == num_tokens @@ -2002,6 +2151,7 @@ def build( num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, num_prefills=num_prefills, + causal=not non_causal_decode, prefill=prefill_metadata, decode=decode_metadata, ) diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index bb1c995aeb57..85dbd6331b26 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -24,6 +24,7 @@ get_multimodal_config, get_vit_attn_backend, ) +from vllm.platforms import current_platform from vllm.utils.flashinfer import ( is_flashinfer_cudnn_fp8_prefill_attn_supported, ) @@ -32,6 +33,7 @@ from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.vit_attn_wrappers import ( + vit_aiter_fp8_attn_wrapper, vit_flash_attn_wrapper, vit_flashinfer_wrapper, vit_torch_sdpa_wrapper, @@ -380,8 +382,8 @@ def _init_fp8_state(self) -> None: No-op if FP8 is not requested. Raises ``ValueError`` if FP8 is requested but the platform does not support it. """ - # Populate defaults so ``_forward_flashinfer`` can - # check ``self.fp8_enabled`` and others without AttributeError. + # Populate defaults so backend forward methods can check + # ``self.fp8_enabled`` and related state without AttributeError. self.fp8_enabled = False self._fp8_dynamic_scale = False self.fp8_quant: QuantFP8 | None = None @@ -393,13 +395,38 @@ def _init_fp8_state(self) -> None: if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8": return - # FP8 path - if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + if self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA: + if not current_platform.is_rocm(): + raise ValueError("AITER FP8 ViT attention requires ROCm.") + + from vllm.platforms.rocm import on_mi3xx + + if not on_mi3xx(): + raise ValueError( + "AITER FP8 ViT attention requires an MI300-series or " + "MI350-series GPU (gfx942 or gfx950)." + ) + try: + from aiter import flash_attn_varlen_fp8_pertensor_func # noqa: F401 + except ImportError as exc: + raise ValueError( + "mm_encoder_attn_dtype='fp8' with ROCM_AITER_FA requires " + "an AITER build that provides " + "flash_attn_varlen_fp8_pertensor_func." + ) from exc + elif self.attn_backend == AttentionBackendEnum.FLASHINFER: + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + raise ValueError( + "mm_encoder_attn_dtype='fp8' requires the FlashInfer " + "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) " + "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is " + "not available on Hopper (H100/H200) or earlier." + ) + else: raise ValueError( - "mm_encoder_attn_dtype='fp8' requires the FlashInfer " - "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) " - "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is " - "not available on Hopper (H100/H200) or earlier." + "mm_encoder_attn_dtype='fp8' requires either the " + "ROCM_AITER_FA backend on ROCm or the FlashInfer cuDNN " + f"backend on CUDA, got {self.attn_backend}." ) self.fp8_enabled = True @@ -642,6 +669,37 @@ def _record_amax_and_update_scales( buffer_wrapped, ) + def _quantize_qkv_fp8( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert self.fp8_quant is not None + + if self._fp8_dynamic_scale: + self._record_amax_and_update_scales(query, key, value) + + query = quantize_fp8_maybe_pad_head_dim( + query, + self._fp8_q_scale, + skip_scale=self.skip_scale_q, + fp8_quant=self.fp8_quant, + ) + key = quantize_fp8_maybe_pad_head_dim( + key, + self._fp8_k_scale, + skip_scale=self.skip_scale_k, + fp8_quant=self.fp8_quant, + ) + value = quantize_fp8_maybe_pad_head_dim( + value, + self._fp8_v_scale, + skip_scale=self.skip_scale_v, + fp8_quant=self.fp8_quant, + ) + return query, key, value + def _forward_flashinfer( self, query: torch.Tensor, @@ -653,29 +711,7 @@ def _forward_flashinfer( | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: if self.fp8_enabled: - assert self.fp8_quant is not None - - if self._fp8_dynamic_scale: - self._record_amax_and_update_scales(query, key, value) - - query = quantize_fp8_maybe_pad_head_dim( - query, - self._fp8_q_scale, - skip_scale=self.skip_scale_q, - fp8_quant=self.fp8_quant, - ) - key = quantize_fp8_maybe_pad_head_dim( - key, - self._fp8_k_scale, - skip_scale=self.skip_scale_k, - fp8_quant=self.fp8_quant, - ) - value = quantize_fp8_maybe_pad_head_dim( - value, - self._fp8_v_scale, - skip_scale=self.skip_scale_v, - fp8_quant=self.fp8_quant, - ) + query, key, value = self._quantize_qkv_fp8(query, key, value) output = vit_flashinfer_wrapper( q=query, @@ -697,6 +733,44 @@ def _forward_flashinfer( return output + def _forward_aiter_fp8( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, + ) -> torch.Tensor: + assert (cu_seqlens is not None and max_seqlen is not None) or ( + cu_seqlens is None and max_seqlen is None + ), "cu_seqlens and max_seqlen should be both set or both None." + + bsz, q_len = query.size()[:2] + kv_len = key.size(1) + is_reshaped = query.dim() != 4 + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len) + + query, key, value = self._quantize_qkv_fp8(query, key, value) + + output = vit_aiter_fp8_attn_wrapper( + q=query, + k=key, + v=value, + q_descale=self._fp8_q_scale, + k_descale=self._fp8_k_scale, + v_descale=self._fp8_v_scale, + batch_size=bsz, + output_dtype=self.dtype, + scale=self.scale, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + if output.shape[-1] != self.head_size: + output = output[..., : self.head_size].contiguous() + if is_reshaped: + output = output.reshape(bsz, q_len, -1) + return output + def forward_native( self, query: torch.Tensor, @@ -719,7 +793,9 @@ def forward_cuda( sequence_lengths: torch.Tensor | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: - if self.is_flash_attn_backend: + if self.fp8_enabled and self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA: + return self._forward_aiter_fp8(query, key, value, cu_seqlens, max_seqlen) + elif self.is_flash_attn_backend: return self._forward_fa(query, key, value, cu_seqlens, max_seqlen) elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN: return self._forward_triton(query, key, value, cu_seqlens, max_seqlen) diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 19cad7986bf9..90934d2892e5 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -1,29 +1,34 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared forward_mha implementation and metadata builder for sparse MLA backends.""" +"""Shared MHA implementation and metadata builder for sparse MLA backends.""" from shutil import which -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast import numpy as np import torch -from vllm.distributed import get_dcp_group +from vllm import _custom_ops as ops +from vllm.distributed import ( + get_dcp_group, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBaseImpl, + MLACommonMetadata, MLACommonPrefillMetadata, build_mla_chunked_context_metadata, get_mla_dims, ) from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import has_flashinfer -from vllm.utils.torch_utils import np_to_pinned_tensor -from vllm.v1.attention.backend import ( - AttentionMetadata, - AttentionMetadataBuilder, -) +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor +from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.attention.ops.merge_attn_states import merge_attn_states if TYPE_CHECKING: from vllm.config import VllmConfig @@ -35,6 +40,32 @@ T = TypeVar("T", bound=AttentionMetadata) +GLOBAL_TOPK_MASK_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB + + +def _is_masked_mha_available( + num_heads_total: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + kv_cache_dtype: str, +) -> bool: + """Check if masked MHA can ever fire for this model configuration.""" + if not current_platform.is_device_capability_family(100): + return False + if ( + num_heads_total != 128 + or kv_lora_rank != 512 + or qk_nope_head_dim != 128 + or qk_rope_head_dim != 64 + or v_head_dim != 128 + ): + return False + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + fa_version = get_flash_attn_version(head_size=qk_head_dim, head_size_v=v_head_dim) + return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype) + class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): metadata_cls: type[T] @@ -86,6 +117,20 @@ def __init__( dtype=self.model_config.dtype, device=device, ) + self.topk_mask_workspace: torch.Tensor | None = None + if _is_masked_mha_available( + self.model_config.model_arch_config.total_num_attention_heads, + self.mla_dims.kv_lora_rank, + self.mla_dims.qk_nope_head_dim, + self.mla_dims.qk_rope_head_dim, + self.mla_dims.v_head_dim, + vllm_config.cache_config.cache_dtype, + ): + self.topk_mask_workspace = torch.zeros( + GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize, + dtype=torch.int32, + device=device, + ) layer_prefill_backend = vllm_config.compilation_config.static_forward_context[ layer_names[0] ].prefill_backend @@ -194,6 +239,7 @@ def build( block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...], query_start_loc=prefill_query_start_loc, max_query_len=prefill_max_query_len, + query_lens_cpu=prefill_query_lens_cpu, chunked_context=self._build_chunked_context_fields( common_attn_metadata, num_decodes, @@ -203,6 +249,11 @@ def build( q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, + use_dense_mha=( + prefill_max_seq_len <= self.topk_tokens + and not self.vllm_config.attention_config.sparse_mla_force_mqa + ), + topk_mask_workspace=self.topk_mask_workspace, ) self._prefill_backend.prepare_metadata(prefill) @@ -256,9 +307,139 @@ def _build_prefill_fields( ) +@triton.jit +def _scatter_topk_kernel( + mask_ptr, + topk_ptr, + cu_q_lens_ptr, + num_words: tl.constexpr, + mask_row_stride: tl.constexpr, + num_topk: tl.constexpr, + topk_stride: tl.constexpr, + max_q_len: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + BLOCK_WORDS: tl.constexpr, +): + row_idx = tl.program_id(0) + row_ptr = mask_ptr + row_idx * mask_row_stride + word_offsets = tl.arange(0, BLOCK_WORDS) + tl.store( + row_ptr + word_offsets, + tl.zeros([BLOCK_WORDS], dtype=tl.int32), + mask=word_offsets < num_words, + ) + + req_idx = row_idx // max_q_len + q_local = row_idx % max_q_len + q_start = tl.load(cu_q_lens_ptr + req_idx) + q_len = tl.load(cu_q_lens_ptr + req_idx + 1) - q_start + if q_local < q_len: + src_row = q_start + q_local + offsets = tl.arange(0, BLOCK_TOPK) + in_range = offsets < num_topk + indices = tl.load( + topk_ptr + src_row * topk_stride + offsets, + mask=in_range, + other=-1, + ) + valid = in_range & (indices >= 0) + word_indices = indices >> 5 + bits = (1 << (indices & 31)).to(tl.int32) + tl.atomic_or(row_ptr + word_indices, bits, mask=valid) + + +@triton.jit +def _scatter_topk_single_req_kernel( + mask_ptr, + topk_ptr, + num_words: tl.constexpr, + mask_row_stride: tl.constexpr, + num_topk: tl.constexpr, + topk_stride: tl.constexpr, + total_q: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + BLOCK_WORDS: tl.constexpr, +): + row_idx = tl.program_id(0) + row_ptr = mask_ptr + row_idx * mask_row_stride + word_offsets = tl.arange(0, BLOCK_WORDS) + tl.store( + row_ptr + word_offsets, + tl.zeros([BLOCK_WORDS], dtype=tl.int32), + mask=word_offsets < num_words, + ) + if row_idx < total_q: + offsets = tl.arange(0, BLOCK_TOPK) + in_range = offsets < num_topk + indices = tl.load( + topk_ptr + row_idx * topk_stride + offsets, + mask=in_range, + other=-1, + ) + valid = in_range & (indices >= 0) + word_indices = indices >> 5 + bits = (1 << (indices & 31)).to(tl.int32) + tl.atomic_or(row_ptr + word_indices, bits, mask=valid) + + +def _build_topk_mask( + topk_indices_per_req: list[torch.Tensor], + q_lens: list[int], + max_q_len: int, + max_seq_len: int, + out: torch.Tensor, +) -> torch.Tensor: + """Build a bit-packed top-k mask into ``out[:B, :max_Q, :num_words]``.""" + batch_size = len(q_lens) + num_words = (max_seq_len + 31) // 32 + total_rows = batch_size * max_q_len + if total_rows == 0: + return out[:batch_size, :max_q_len, :num_words] + + total_q = sum(q_lens) + mask_row_stride = out.stride(-2) + block_words = triton.next_power_of_2(num_words) + + if batch_size == 1: + topk_packed = topk_indices_per_req[0] + num_topk = topk_packed.shape[1] + _scatter_topk_single_req_kernel[(max_q_len,)]( + out, + topk_packed, + num_words=num_words, + mask_row_stride=mask_row_stride, + num_topk=num_topk, + topk_stride=topk_packed.stride(0), + total_q=total_q, + BLOCK_TOPK=triton.next_power_of_2(num_topk), + BLOCK_WORDS=block_words, + ) + return out[:1, :max_q_len, :num_words] + + topk_packed = torch.cat(topk_indices_per_req, dim=0) + num_topk = topk_packed.shape[1] + q_lens_tensor = np_to_pinned_tensor(np.asarray(q_lens, dtype=np.int32)).to( + out.device, non_blocking=True + ) + cu_q_lens = out.new_zeros(batch_size + 1) + torch.cumsum(q_lens_tensor, dim=0, out=cu_q_lens[1:]) + _scatter_topk_kernel[(total_rows,)]( + out, + topk_packed, + cu_q_lens, + num_words=num_words, + mask_row_stride=mask_row_stride, + num_topk=num_topk, + topk_stride=topk_packed.stride(0), + max_q_len=max_q_len, + BLOCK_TOPK=triton.next_power_of_2(num_topk), + BLOCK_WORDS=block_words, + ) + return out[:batch_size, :max_q_len, :num_words] + + class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]): - """Sparse MLA base: shared dense-MHA prefill (from MLACommonBaseImpl) plus the - sparse top-k MQA decode path. Subclasses implement forward_mqa.""" + """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True @@ -274,7 +455,6 @@ def __init__( logits_soft_cap: float | None, attn_type: str, kv_sharing_target_layer_name: str | None, - # MLA-specific q_lora_rank: int | None, kv_lora_rank: int, qk_nope_head_dim: int, @@ -308,7 +488,6 @@ def __init__( if indexer is not None else topk_indices_buffer ) - self._use_flashinfer_concat_mla_k = ( has_flashinfer() and which("ninja") is not None @@ -316,3 +495,335 @@ def __init__( and (self.qk_nope_head_dim == 128) and (self.qk_rope_head_dim == 64) ) + self.masked_mha_available = _is_masked_mha_available( + num_heads_total=num_heads * get_tensor_model_parallel_world_size(), + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_cache_dtype=kv_cache_dtype, + ) + + @staticmethod + def _slice_topk_per_req( + topk_all: torch.Tensor, + q_lens: list[int], + ) -> list[torch.Tensor]: + topk_per_req = [] + offset = 0 + for q_len in q_lens: + topk_per_req.append(topk_all[offset : offset + q_len]) + offset += q_len + return topk_per_req + + @staticmethod + def _remap_topk_to_ranges( + topk_per_req: list[torch.Tensor], + range_starts: list[int] | torch.Tensor, + range_lens: list[int], + ) -> list[torch.Tensor]: + remapped = [] + for topk, start, length in zip(topk_per_req, range_starts, range_lens): + valid = (topk >= start) & (topk < start + length) + remapped.append(torch.where(valid, topk - start, -1)) + return remapped + + def _project_kv( + self, kv_c_normed: torch.Tensor, k_pe: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + -1, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + return self._concat_k_nope_k_pe(k_nope, k_pe), v + + @staticmethod + def _try_build_global_mask( + topk_per_req: list[torch.Tensor], + q_lens: list[int], + max_query_len: int, + max_seq_len: int, + topk_mask_workspace: torch.Tensor, + ) -> torch.Tensor | None: + """Build a full-sequence top-k mask if it fits within the budget. + + When the mask fits, it is reused across the suffix and all context + chunks, avoiding per-chunk mask rebuilds. Returns None when the + mask is too large, signalling the caller to fall back to per-chunk + index remapping. + """ + batch_size = len(q_lens) + tile_m = 128 if max_query_len <= 128 else 256 + padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m + num_words_padded = (max_seq_len + 31) // 32 + 1 + needed = batch_size * padded_q_len * num_words_padded + if needed * torch.int32.itemsize > GLOBAL_TOPK_MASK_MAX_BYTES: + return None + + mask = topk_mask_workspace[:needed].view( + batch_size, padded_q_len, num_words_padded + ) + _build_topk_mask( + topk_per_req, + q_lens, + padded_q_len, + max_seq_len, + mask, + ) + return mask + + def _run_masked_mha( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + topk_per_req: list[torch.Tensor], + q_lens: list[int], + causal: bool, + return_softmax_lse: bool = False, + dense_mask: torch.Tensor | None = None, + key_starts: torch.Tensor | None = None, + topk_mask_workspace: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from vllm.model_executor.layers.attention.sparse_mla_mask import ( + dense_mask_mod, + offset_dense_mask_mod, + ) + from vllm.vllm_flash_attn import flash_attn_varlen_func + + tile_m = 128 if max_seqlen_q <= 128 else 256 + padded_q_len = triton.cdiv(max_seqlen_q, tile_m) * tile_m + if dense_mask is None: + batch_size = len(q_lens) + num_words = (max_seqlen_k + 31) // 32 + assert topk_mask_workspace is not None + workspace_3d = topk_mask_workspace[ + : batch_size * padded_q_len * num_words + ].view(batch_size, padded_q_len, num_words) + dense_mask = _build_topk_mask( + topk_per_req, + q_lens, + padded_q_len, + max_seqlen_k, + workspace_3d, + ) + if key_starts is not None: + dense_mask[:, 0, -1].copy_(key_starts) + kwargs = { + "q": q, + "k": k, + "v": v, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "softmax_scale": self.scale, + "return_softmax_lse": return_softmax_lse, + "fa_version": 4, + "mask_mod": dense_mask_mod if key_starts is None else offset_dense_mask_mod, + "aux_tensors": [dense_mask], + "aux_tensor_leading_dims": [2], + "causal": causal, + } + + return flash_attn_varlen_func(**kwargs) + + def _compute_context_mha( + self, + q: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + prefill_metadata: MLACommonPrefillMetadata, + k_scale: torch.Tensor, + q_lens: list[int], + topk_per_req: list[torch.Tensor], + dense_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.dcp_world_size > 1: + raise NotImplementedError( + "Masked MHA with context does not yet support decode context " + "parallelism" + ) + + chunked_context = prefill_metadata.chunked_context + assert chunked_context is not None + use_global_mask = dense_mask is not None + output: torch.Tensor | None = None + output_lse: torch.Tensor | None = None + workspace = chunked_context.workspace + + for i, toks in enumerate(chunked_context.seq_tot): + if toks == 0: + continue + ops.gather_and_maybe_dequant_cache( + src_cache=kv_c_and_k_pe_cache, + dst=workspace, + block_table=prefill_metadata.block_table, + cu_seq_lens=chunked_context.cu_seq_lens[i], + token_to_seq=chunked_context.token_to_seq[i], + num_tokens=chunked_context.chunk_total_token[i], + kv_cache_dtype=self.kv_cache_dtype, + scale=k_scale, + seq_starts=chunked_context.starts[i], + ) + + chunk_kv_c = workspace[:toks, : self.kv_lora_rank] + chunk_k_pe = workspace[:toks, self.kv_lora_rank :].unsqueeze(1) + k, v = self._project_kv(chunk_kv_c, chunk_k_pe) + chunk_lens = chunked_context.seq_lens[i].tolist() + chunk_topk = ( + topk_per_req + if use_global_mask + else self._remap_topk_to_ranges( + topk_per_req, + chunked_context.starts[i], + chunk_lens, + ) + ) + attn_out, lse = self._run_masked_mha( + q=q, + k=k, + v=v, + cu_seqlens_q=prefill_metadata.query_start_loc, + cu_seqlens_k=chunked_context.cu_seq_lens[i], + max_seqlen_q=prefill_metadata.max_query_len, + max_seqlen_k=chunked_context.max_seq_lens[i], + topk_per_req=chunk_topk, + q_lens=q_lens, + causal=False, + return_softmax_lse=True, + dense_mask=dense_mask, + key_starts=(chunked_context.starts[i] if use_global_mask else None), + topk_mask_workspace=prefill_metadata.topk_mask_workspace, + ) + + if output is None: + output = attn_out + output_lse = lse + else: + assert output_lse is not None + merge_attn_states( + output=output, + output_lse=output_lse, + prefix_output=output, + prefix_lse=output_lse, + suffix_output=attn_out, + suffix_lse=lse, + ) + + assert output is not None and output_lse is not None + return output, output_lse + + def forward_mha( # type: ignore[override] + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + k_scale: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + ) -> None: + prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] + topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined] + force_dense = getattr(self, "_sparse_mla_force_dense_mha", False) + force_masked = getattr(self, "_sparse_mla_force_masked_mha", False) + if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked): + return super().forward_mha( + q, + kv_c_normed, + k_pe, + kv_c_and_k_pe_cache, + cast(MLACommonMetadata, attn_metadata), + k_scale, + output, + output_scale, + ) + + assert output_scale is None + assert self.masked_mha_available + prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined] + assert prefill_metadata is not None + assert prefill_metadata.query_lens_cpu is not None + assert self.topk_indices_buffer is not None + + q_lens = prefill_metadata.query_lens_cpu.tolist() + num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined] + topk_all = self.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + q.shape[0] + ] + topk_per_req = self._slice_topk_per_req(topk_all, q_lens) + + k, v = self._project_kv(kv_c_normed, k_pe) + chunked_context = prefill_metadata.chunked_context + if chunked_context is None: + attn_out = self._run_masked_mha( + q=q, + k=k, + v=v, + cu_seqlens_q=prefill_metadata.query_start_loc, + cu_seqlens_k=prefill_metadata.query_start_loc, + max_seqlen_q=prefill_metadata.max_query_len, + max_seqlen_k=prefill_metadata.max_query_len, + topk_per_req=topk_per_req, + q_lens=q_lens, + causal=True, + topk_mask_workspace=prefill_metadata.topk_mask_workspace, + ) + assert isinstance(attn_out, torch.Tensor) + output.copy_(attn_out[..., : self.v_head_dim].flatten(start_dim=-2)) + return + + context_lens = chunked_context.seq_lens.sum(dim=0).tolist() + dense_mask = self._try_build_global_mask( + topk_per_req, + q_lens, + prefill_metadata.max_query_len, + prefill_max_seq_len, + prefill_metadata.topk_mask_workspace, + ) + if dense_mask is not None: + suffix_topk = topk_per_req + else: + suffix_topk = self._remap_topk_to_ranges(topk_per_req, context_lens, q_lens) + suffix_output, suffix_lse = self._run_masked_mha( + q=q, + k=k, + v=v, + cu_seqlens_q=prefill_metadata.query_start_loc, + cu_seqlens_k=prefill_metadata.query_start_loc, + max_seqlen_q=prefill_metadata.max_query_len, + max_seqlen_k=prefill_metadata.max_query_len, + topk_per_req=suffix_topk, + q_lens=q_lens, + causal=True, + return_softmax_lse=True, + dense_mask=dense_mask, + key_starts=( + chunked_context.context_lens if dense_mask is not None else None + ), + topk_mask_workspace=prefill_metadata.topk_mask_workspace, + ) + context_output, context_lse = self._compute_context_mha( + q=q, + kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, + prefill_metadata=prefill_metadata, + k_scale=k_scale, + q_lens=q_lens, + topk_per_req=topk_per_req, + dense_mask=dense_mask, + ) + merge_attn_states( + output=output.view(-1, self.num_heads, self.v_head_dim), + prefix_output=context_output[..., : self.v_head_dim], + prefix_lse=context_lse, + suffix_output=suffix_output[..., : self.v_head_dim], + suffix_lse=suffix_lse, + prefill_tokens_with_context=chunked_context.prefill_tokens_with_context, + ) diff --git a/vllm/model_executor/layers/attention/sparse_mla_mask.py b/vllm/model_executor/layers/attention/sparse_mla_mask.py new file mode 100644 index 000000000000..1cd151034c2e --- /dev/null +++ b/vllm/model_executor/layers/attention/sparse_mla_mask.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import cutlass +import cutlass.cute as cute + +from vllm.vllm_flash_attn.cute import utils + + +@cute.jit +def dense_mask_mod( + batch: cute.TensorSSA, + head: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors: list, +) -> cute.TensorSSA: + dense_mask = aux_tensors[0] + batch_idx = utils.ssa_to_scalar(batch) + q_idx = utils.ssa_to_scalar(q_idx) + kv_idx = utils.ssa_to_scalar(kv_idx) + word_idx = kv_idx >> 5 + bit_idx = cutlass.Uint32(kv_idx & 31) + word = dense_mask[batch_idx, q_idx, word_idx] + result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32) + result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx) + return result.load() + + +dense_mask_mod.__vec_size__ = 32 + + +@cute.jit +def offset_dense_mask_mod( + batch: cute.TensorSSA, + head: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors: list, +) -> cute.TensorSSA: + dense_mask = aux_tensors[0] + batch_idx = utils.ssa_to_scalar(batch) + q_idx = utils.ssa_to_scalar(q_idx) + key_start = dense_mask[batch_idx, 0, dense_mask.shape[2] - 1] + kv_idx = utils.ssa_to_scalar(kv_idx) + key_start + word_idx = kv_idx >> 5 + bit_idx = cutlass.Uint32(kv_idx & 31) + word = dense_mask[batch_idx, q_idx, word_idx] + result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32) + result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx) + return result.load() + + +offset_dense_mask_mod.__vec_size__ = 32 diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index 53994114558e..197213dc79fe 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod +import torch + from vllm.config import VllmConfig from vllm.v1.attention.backend import AttentionBackend, AttentionImpl from vllm.v1.kv_cache_interface import KVCacheSpec @@ -21,6 +23,14 @@ class AttentionLayerBase(ABC): impl: "AttentionImpl" supports_dcp: bool = True + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + """Bind the allocated KV cache tensor to this layer. + + The default stores the cache view as-is; subclasses (e.g. Mamba) + override this to unpack the raw buffer into per-state views. + """ + self.kv_cache = kv_cache + @abstractmethod def get_attn_backend(self) -> type[AttentionBackend]: """Get the attention backend class for this layer.""" diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index e49e135b26a9..c28ad88b5e27 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -43,7 +43,7 @@ if flashinfer_comm is not None else None ) -except ImportError: +except (ImportError, AttributeError): flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment] get_fi_ar_workspace = None # type: ignore[assignment] _AR_RESIDUAL_RMS_NORM = None @@ -93,7 +93,7 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool max_token_num=max_token_num, hidden_dim=hidden_size, dtype=hidden_states.dtype, - group=get_tp_group().device_group, + group=get_tp_group().cpu_group, ) if workspace is None: return False, 0 diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 8f434272f561..17f287556b79 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -19,7 +19,7 @@ FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.layer import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( @@ -63,7 +63,7 @@ def get_config() -> dict[str, Any] | None: __all__ = [ - "FusedMoE", + "FusedMoEFactory", "FusedMoERouter", "FusedMoEConfig", "FusedMoEQuantConfig", diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py index d114992504ba..baa6136485e3 100644 --- a/vllm/model_executor/layers/fused_moe/activation.py +++ b/vllm/model_executor/layers/fused_moe/activation.py @@ -22,6 +22,7 @@ class MoEActivation(Enum): # expects the *packed* layout ([all gates; all ups]), as produced by a # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3). SWIGLUOAI = "swigluoai" + SITU = "situ" SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave" SWIGLUSTEP = "swiglustep" @@ -77,6 +78,7 @@ def from_str(cls, s: str) -> "MoEActivation": MoEActivation.SILU: "silu_and_mul", MoEActivation.GELU: "gelu_and_mul", MoEActivation.GELU_TANH: "gelu_tanh_and_mul", + MoEActivation.SITU: "situ_and_mul", MoEActivation.SWIGLUOAI: "swigluoai_and_mul", MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp", MoEActivation.SWIGLUSTEP: "swiglustep_and_mul", @@ -133,6 +135,8 @@ def apply_moe_activation( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """Apply MoE activation function. @@ -163,6 +167,25 @@ def apply_moe_activation( torch.ops._C.gelu_and_mul(output, input) elif activation == MoEActivation.GELU_TANH: torch.ops._C.gelu_tanh_and_mul(output, input) + elif activation == MoEActivation.SITU: + # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries. + # (The pure-torch fallback below upcast both halves to fp32 and + # allocated ~8 temporaries per call, blowing up MoE memory.) + # Both betas come from FusedMoEConfig; a missing beta means the caller + # bypassed the config plumbing, so fail rather than silently use 1.0. + # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel + # (up passed through), matching SituAndMul(linear_beta=None). + assert activation_situ_beta is not None, ( + "SITU requires activation_situ_beta from FusedMoEConfig" + ) + torch.ops._C.situ_and_mul( + output, + input, + activation_situ_beta, + -1.0 + if activation_situ_linear_beta is None + else activation_situ_linear_beta, + ) elif activation == MoEActivation.SWIGLUOAI: torch.ops._C.swigluoai_and_mul(output, input) elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6af93bfde90c..58c9c8d9f6dd 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -19,6 +19,7 @@ FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + BatchedPrepareAndFinalize, make_moe_prepare_and_finalize_naive_dp_ep, make_moe_prepare_and_finalize_no_dp_ep, ) @@ -140,6 +141,16 @@ def maybe_make_prepare_finalize( if not allow_new_interface: return None + # Opt-in XPU batched path: reorganize tokens into E x T x K locally + # (no all-to-all) so BatchedTritonExperts (moe_mmk TD) can run. + if current_platform.is_xpu() and moe.moe_backend == "batched_triton": + return BatchedPrepareAndFinalize( + max_num_tokens=moe.max_num_tokens, + num_local_experts=moe.num_local_experts, + num_dispatchers=1, + rank=moe.moe_parallel_config.ep_rank, + ) + # For DP/TP case, fall back to naive P/F. if moe.moe_parallel_config.dp_size > 1: logger.info_once( diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 02e6c51116d1..f2ee520d09b7 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1134,9 +1134,9 @@ def make( level's of parallelism to use in the fused moe layer. Args: - tp_size_ (int): `tp_size` passed into the FusedMoE constructor. - pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor. - dp_size_ (int): `dp_size` passed into the FusedMoE constructor. + tp_size_ (int): `tp_size` passed into the FusedMoEFactory constructor. + pcp_size_ (int): `pcp_size` passed into the FusedMoEFactory constructor. + dp_size_ (int): `dp_size` passed into the FusedMoEFactory constructor. vllm_parallel_config (ParallelConfig): vLLM's parallel config object which contains the `enable_expert_parallel` flag. @@ -1314,6 +1314,10 @@ class FusedMoEConfig: swiglu_alpha: float | None = None swiglu_beta: float | None = None + # SituGLU parameters used by Kimi sit(u/v2) activations. + activation_situ_beta: float | None = None + activation_situ_linear_beta: float | None = None + max_capture_size: int = 0 # Set by __post_init__ diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 1a0acff058cc..c89e5ddc3044 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -20,6 +20,11 @@ from vllm.utils.torch_utils import direct_register_custom_op _CPU_MOE_LAYER_CACHE = {} +# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert +# intermediate ("N") dimension in blocks of this size and have no tail/ +# remainder handling, so a shard is only eligible for the fast path when +# its per-partition intermediate size is a multiple of it. +_MOE_GROUPED_GEMM_N_TILE = 32 def _swigluoai_forward_native( @@ -230,6 +235,7 @@ class CPUFusedMOE: """CPU-based fused MoE implementation.""" def __init__(self, layer: torch.nn.Module) -> None: + self._pad_moe_intermediate_for_grouped_gemm(layer) use_grouped_gemm, isa = self.check_grouped_gemm(layer) self.isa = isa if use_grouped_gemm: @@ -284,6 +290,69 @@ def __call__( apply_router_weight_on_input, ) + def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None: + """Zero-pad the per-partition MoE intermediate dim up to a multiple + of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels + can be used even when TP sharding (moe_intermediate_size // tp_size) + lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4 + -> 176). Only applies to the x86 (AMX/vec) kernels and half-split + gate/up activations; interleaved layouts (swigluoai) are left + untouched. + """ + if not hasattr(torch.ops._C, "prepack_moe_weight"): + return + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return + + intermediate_size = layer.w2_weight.size(2) + remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE + if remainder == 0: + return + if layer.activation == MoEActivation.SWIGLUOAI: + return + + pad = _MOE_GROUPED_GEMM_N_TILE - remainder + padded_size = intermediate_size + pad + num_experts, _, hidden_size = layer.w13_weight.shape + + new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size) + new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size] + new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[ + :, intermediate_size: + ] + replace_parameter(layer, "w13_weight", new_w13) + + new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size) + new_w2[:, :, :intermediate_size] = layer.w2_weight + replace_parameter(layer, "w2_weight", new_w2) + + if hasattr(layer, "w13_bias"): + new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size) + new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size] + new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[ + :, intermediate_size: + ] + replace_parameter(layer, "w13_bias", new_bias) + + def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str: + # w2's input size is the per-partition MoE intermediate size (the + # dimension TP-sharding splits), and it's what most commonly breaks + # alignment, e.g. moe_intermediate_size=704 at tp=4 gives + # 704 // 4 == 176, which isn't a multiple of 32. + intermediate_size_per_partition = layer.w2_weight.size(2) + return ( + "CPU fused-MoE AMX grouped-gemm kernel cannot be used for a " + f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 " + f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE " + f"intermediate size ({intermediate_size_per_partition}) is not " + f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic " + "zero-padding could not resolve this (typically because the " + "activation uses an interleaved gate/up layout, e.g. " + "swigluoai). vLLM refuses to silently fall back to the much " + "slower per-expert torch loop on AMX-capable CPUs; consider a " + "different --tensor-parallel-size." + ) + def check_grouped_gemm( self, layer: torch.nn.Module, @@ -297,20 +366,25 @@ def check_grouped_gemm( w2_input_size = layer.w2_weight.size(2) w2_output_size = layer.w2_weight.size(1) + supports_amx = torch.cpu._is_amx_tile_supported() + if supports_amx: + if ( + dtype == torch.bfloat16 + and w13_output_size % 32 == 0 + and w2_output_size % 32 == 0 + and w13_input_size % 32 == 0 + and w2_input_size % 32 == 0 + ): + return True, "amx" + raise RuntimeError(self._grouped_gemm_alignment_error(layer)) + if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0): return False, "none" - supports_amx = torch.cpu._is_amx_tile_supported() - if ( - supports_amx - and dtype == torch.bfloat16 - and w13_input_size % 32 == 0 - and w2_input_size % 32 == 0 + layer.activation == MoEActivation.SWIGLUOAI + and w2_input_size % _MOE_GROUPED_GEMM_N_TILE != 0 ): - return True, "amx" - - if supports_amx: return False, "none" supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM @@ -321,8 +395,7 @@ def check_grouped_gemm( and w2_input_size % 4 == 0 ): return True, "neon" - else: - return False, "none" + return False, "none" return True, "vec" diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index f4319b99c816..b25d3e93ae96 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -5,6 +5,8 @@ and updated to fit vllm needs and terminology. """ +import math + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -424,7 +426,7 @@ def ep_gather( num_warps = 2 num_tokens = output_tensor.shape[0] hidden_size = input_tensor.shape[1] - BLOCK_D = min(hidden_size, 1024) + BLOCK_D = math.gcd(hidden_size, 1024) assert hidden_size % BLOCK_D == 0 grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024)) @@ -505,7 +507,7 @@ def deepgemm_moe_permute( dtype=torch.int32, ) else: - aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32) + aq_scale_out = torch.zeros((M_sum, sf_k), device=device, dtype=torch.float32) # DeepGEMM uses negative values in m_indices (here expert_ids) to mark # completely invalid / padded blocks that should be skipped. We always diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 7c3fe5831f3b..53880cb916a2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -20,7 +20,9 @@ __all__ = [ "AiterW4A8ExpertsMonolithic", + "AiterW4A16ExpertsMonolithic", "aiter_triton_kernel_w4a8_moe_forward", + "aiter_triton_kernel_w4a16_moe_forward", ] @@ -46,11 +48,27 @@ def aiter_triton_kernel_w4a8_moe_forward( and quant_config.use_mxfp4_w4a8 and rocm_aiter_ops.is_enabled() ) - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) + + # gfx1250: aiter's in-kernel gather is numerically broken + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -199,12 +217,11 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - # Requires AITER and GFX950 if not rocm_aiter_ops.is_enabled(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return on_gfx950() + return on_gfx950() or on_gfx1250() @staticmethod def _supports_no_act_and_mul() -> bool: @@ -293,3 +310,358 @@ def apply( unpadded_N_w2=self.moe_config.hidden_dim_unpadded, unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, ) + + +def _aiter_raw(t): + if t is None or isinstance(t, torch.Tensor): + return t + return t.storage.data if hasattr(t, "storage") else t + + +def _aiter_w4a16_silu_via_a8w4( + hidden_states: torch.Tensor, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + w1_bias, + w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input: bool, + swiglu_limit: float, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, +) -> torch.Tensor: + """ + MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation. + """ + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant import dynamic_mxfp8_quant + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + should_use_cdna4_mx_scale_swizzle, + ) + + swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + quant_dtype = torch.float8_e4m3fn + + g1_gammas = gammas if apply_router_weight_on_input else None + g2_gammas = None if apply_router_weight_on_input else gammas + + hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype) + raw_gate_up = moe_gemm_a8w4( + hidden_q, + w1_data, + a1_scale, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=g1_gammas, + swizzle_mx_scale=swz, + out_dtype=torch.bfloat16, + apply_swiglu=False, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + if unpadded_N_w1 is not None: + raw_gate_up = raw_gate_up[:, :unpadded_N_w1] + + interim_fp8, a2_scale = fused_clamp_act_mul( + raw_gate_up, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=quant_dtype, + scale_dtype_fmt="ue8m0", + quant_block_size=32, + ) + + out = moe_gemm_a8w4( + interim_fp8, + w2_data, + a2_scale, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=swz, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + return out + + +def aiter_triton_kernel_w4a16_moe_forward( + hidden_states: torch.Tensor, + w1, + w2, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, + num_expert_group: int | None = None, + topk_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + score_mode: str | None = None, +): + assert quant_config is not None and rocm_aiter_ops.is_enabled() + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe_routing import routing as _routing_mod + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing + + if score_mode is not None: + use_grouped_topk = num_expert_group is not None and num_expert_group > 1 + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, + topk, + score_mode=score_mode, + bias=e_score_correction_bias, + renorm=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) + else: + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) + w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + + gammas = routing_data.gate_scal if routing_data else None + + swiglu_alpha = ( + quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults. + if activation == MoEActivation.SILU and on_gfx1250(): + return _aiter_w4a16_silu_via_a8w4( + hidden_states, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + quant_config.w1_bias, + quant_config.w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input, + swiglu_limit, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, + ) + + # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu. + swiglu_add_residual = activation != MoEActivation.SILU + + intermediate = moe_gemm_a16w4( + hidden_states, + w1_data, + None, + w1_wscale, + None, + None, + quant_config.w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=swiglu_add_residual, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + + out = moe_gemm_a16w4( + intermediate, + w2_data, + None, + w2_wscale, + None, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + + return out + + +class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.renormalize = moe_config.routing_method in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + if not rocm_aiter_ops.is_enabled(): + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() or on_gfx1250() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert self.moe_config.intermediate_size_per_partition_unpadded is not None + assert self.moe_config.hidden_dim_unpadded is not None + score_mode = ( + "sqrtsoftplus" + if self.moe_config.routing_method == RoutingMethodType.DeepseekV4 + else None + ) + return aiter_triton_kernel_w4a16_moe_forward( + hidden_states=hidden_states, + w1=w1, + w2=w2, + gating_output=router_logits, + topk=self.topk, + renormalize=self.renormalize, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=self.quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2, + unpadded_K_w1=self.moe_config.hidden_dim_unpadded, + unpadded_N_w2=self.moe_config.hidden_dim_unpadded, + unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, + num_expert_group=num_expert_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + score_mode=score_mode, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c5330f3b4382..59115df13b5e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -6,10 +6,13 @@ ``convert_to_fp8_moe_kernel_format``. """ +import math + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8TritonExpertsBase, ) @@ -17,6 +20,9 @@ logger = init_logger(__name__) +_AITER_SWIGLU_ALPHA = 1.702 +_AITER_SWIGLU_BETA = 1.0 + def is_aiter_mxfp8_moe_available() -> bool: """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` @@ -93,6 +99,27 @@ def is_supported_config( return False, ( "kernel requires the aiter flydsl package, which is not installed" ) + if ( + is_supported + and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE + ): + return False, ( + "kernel hardcodes SwiGLU-OAI activation and requires " + f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; " + f"got activation={moe_config.activation.value}" + ) + if is_supported and ( + moe_config.swiglu_alpha is None + or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA) + or moe_config.swiglu_beta is None + or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA) + ): + return False, ( + "kernel hardcodes SwiGLU-OAI with " + f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; " + f"got swiglu_alpha={moe_config.swiglu_alpha} and " + f"swiglu_beta={moe_config.swiglu_beta}" + ) return is_supported, reason def apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index 90972fc7f0c6..ce30afd5bea2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -204,6 +204,7 @@ def persistent_masked_m_silu_mul_quant( dtype=ys_dtype, device=y.device, ) + y_s.zero_() ceil_ue8m0 = quant_scale_fmt in [ DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 3ed1734cb915..03561ceb8a23 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -9,6 +9,8 @@ CPUQuantAlgo, CPUQuantMethod, convert_weight_packed_scale_zp, + cpu_fused_moe_int8, + cpu_prepack_moe_weight_int8, fused_experts_cpu, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -27,7 +29,7 @@ kInt8StaticChannelSym, kMxfp4Static, ) -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform # =========================================================================== # FP8 W8A16 MoE @@ -552,16 +554,6 @@ def apply( # =========================================================================== -def prepare_int8_moe_layer_for_cpu( - w13: torch.Tensor, - w2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """VNNI-prepack INT8 MoE weights for CPU kernel.""" - packed_w13 = torch.ops._C.convert_weight_packed(w13) - packed_w2 = torch.ops._C.convert_weight_packed(w2) - return packed_w13, packed_w2 - - class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation monolithic MoE experts.""" @@ -586,7 +578,10 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - return current_platform.is_cpu() + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.X86 + ) @staticmethod def _supports_no_act_and_mul() -> bool: @@ -638,7 +633,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """VNNI-prepack INT8 MoE weights for CPU kernel.""" from vllm.model_executor.utils import replace_parameter - w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight) + w13 = torch.ops._C.convert_weight_packed(layer.w13_weight) + w2 = torch.ops._C.convert_weight_packed(layer.w2_weight) replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) @@ -701,3 +697,168 @@ def apply( None, # limit True, # is_vnni ) + + +class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): + """Arm INT8 MoE with per-token activation and channelwise weight quantization.""" + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + 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]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + return supported, reason + if moe_config.in_dtype not in ( + torch.float32, + torch.float16, + torch.bfloat16, + ): + return False, "kernel requires float32, float16, or bfloat16 activations" + if moe_config.hidden_dim % 32 != 0: + return False, "kernel requires hidden dim divisible by 32" + if moe_config.intermediate_size_per_partition % 32 != 0: + return False, "kernel requires intermediate dim divisible by 32" + return True, None + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "cpu_fused_moe_int8") + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + ) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return not moe_parallel_config.use_ep + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == ( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from vllm.model_executor.utils import replace_parameter + + w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "neon") + w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "neon") + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + select_experts, + ) + + topk_weights, topk_ids = select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + use_grouped_topk=num_expert_group is not None, + top_k=self.moe_config.experts_per_token, + renormalize=self.moe_config.routing_method + in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ), + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func="softmax", + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + e_score_correction_bias=e_score_correction_bias, + ) + + if apply_router_weight_on_input: + assert topk_ids.size(1) == 1 + hidden_states.mul_(topk_weights.to(hidden_states.dtype)) + + assert self.w1_scale is not None + assert self.w2_scale is not None + return cpu_fused_moe_int8( + hidden_states, + w1, + w2, + self.w1_scale, + self.w2_scale, + self.w1_bias, + self.w2_bias, + topk_weights, + topk_ids, + activation.value, + "neon", + skip_weighted=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 160cff1eeb54..211364d24bca 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -444,7 +444,14 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP] + # SILU has fused gate+mul+quant kernels; SWIGLUSTEP/SITU take the + # general path (activation applied via self.activation, which forwards + # the situ betas, then FP8 requant). + return activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUSTEP, + MoEActivation.SITU, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -488,16 +495,15 @@ def _act_mul_quant( M_sum, N = input.size() activation_out_dim = self.adjust_N_for_activation(N, activation) - if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: - assert activation == MoEActivation.SILU - return fused_silu_mul_fp8_quant_packed( - input=input, - output_q=output, - group_size=block_k, - clamp_limit=self.gemm1_clamp_limit, - ) - if activation == MoEActivation.SILU: + # Fused gate+mul+quant kernels for the common SILU case. + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return fused_silu_mul_fp8_quant_packed( + input=input, + output_q=output, + group_size=block_k, + clamp_limit=self.gemm1_clamp_limit, + ) use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return silu_mul_per_token_group_quant_fp8_colmajor( input=input, @@ -506,12 +512,24 @@ def _act_mul_quant( clamp_limit=self.gemm1_clamp_limit, ) + # General gated activations (SWIGLUSTEP, SITU): apply the activation + # (self.activation forwards the situ betas from moe_config) then + # FP8-requant into the layout DeepGEMM expects for this scale format. act_out = torch.empty( (M_sum, activation_out_dim), dtype=input.dtype, device=input.device ) self.activation(activation, act_out, input) + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return per_token_group_quant_fp8_packed_for_deepgemm( + act_out, block_k, use_ue8m0=True, out_q=output + ) + use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return per_token_group_quant_fp8( - act_out, block_k, column_major_scales=True, out_q=output + act_out, + block_k, + column_major_scales=True, + out_q=output, + use_ue8m0=use_ue8m0, ) def apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index b512d51c1353..823ae43b5ef3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, @@ -76,7 +79,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( @@ -90,7 +93,7 @@ def _supports_quant_scheme( @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( @@ -163,4 +166,5 @@ def apply( num_local_experts=self.local_num_experts, local_expert_offset=self.local_expert_offset, moe_output=output, + activation_type=activation_to_flashinfer_int(activation), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 21bda8e173fd..5cc15866d5cb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -32,7 +32,15 @@ kFp8StaticTensorSym, ) from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton +from vllm.triton_utils import tl, triton, use_tensor_descriptor +from vllm.triton_utils.allocation import set_triton_allocator + + +def _is_capturing_or_compiling() -> bool: + # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA (XPU) torch. + return torch.compiler.is_compiling() or ( + current_platform.is_cuda_alike() and torch.cuda.is_current_stream_capturing() + ) @triton.jit @@ -71,9 +79,32 @@ def moe_mmk( use_w8a8: tl.constexpr, use_w8a16: tl.constexpr, per_act_token_quant: tl.constexpr, + # TD: a_base_ptr/b_base_ptr are the expert/CTA-offset bases of A[M,K]/B[N,K]. + a_base_ptr=None, + b_base_ptr=None, + M=0, + N=0, + stride_am: tl.int64 = 0, + stride_bn: tl.int64 = 0, + USE_TD: tl.constexpr = False, ): offs_k = tl.arange(0, BLOCK_K) + if USE_TD: + # make_tensor_descriptor requires the last (K) stride to be a + # compile-time 1; the launcher only enables USE_TD for K-contiguous A/B. + a_desc = tl.make_tensor_descriptor( + a_base_ptr, + shape=[M, K], + strides=[stride_am, 1], + block_shape=[BLOCK_M, BLOCK_K], + ) + b_desc = tl.make_tensor_descriptor( + b_base_ptr, + shape=[N, K], + strides=[stride_bn, 1], + block_shape=[BLOCK_N, BLOCK_K], + ) if use_w8a16: b_scale_ptrs = ( b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn @@ -110,12 +141,17 @@ def moe_mmk( for k in range(0, tl.cdiv(K, BLOCK_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. - a = tl.load( - a_ptrs, - mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), - other=0.0, - ) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + if USE_TD: + # B is [N, K]; tile is [BLOCK_N, BLOCK_K], transposed for dot. + a = a_desc.load([0, k * BLOCK_K]) + b = tl.trans(b_desc.load([0, k * BLOCK_K])) + else: + a = tl.load( + a_ptrs, + mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) # We accumulate along the K dimension. if use_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -193,6 +229,7 @@ def expert_triton_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + USE_TD: tl.constexpr = False, ): offs_m = tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) % N @@ -238,6 +275,13 @@ def expert_triton_kernel( use_fp8_w8a8, use_int8_w8a16, per_act_token_quant, + a_ptr, + b_ptr, + M, + N, + stride_am, + stride_bn, + USE_TD, ) # store in C @@ -292,6 +336,7 @@ def batched_triton_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + USE_TD: tl.constexpr = False, ): expert_id = tl.program_id(axis=0) e_num_tokens = tl.load(expert_num_tokens + expert_id) @@ -372,6 +417,7 @@ def batched_triton_kernel( BLOCK_M, BLOCK_N, BLOCK_K, + USE_TD, ) @@ -444,6 +490,18 @@ def invoke_moe_batched_triton_kernel( stride_asm = 0 stride_ask = 0 + use_td = ( + use_tensor_descriptor() + and A.stride(2) == 1 + and B.stride(2) == 1 + and (K * A.element_size()) % 16 == 0 + and (BLOCK_M & (BLOCK_M - 1)) == 0 + and (BLOCK_N & (BLOCK_N - 1)) == 0 + and (BLOCK_K & (BLOCK_K - 1)) == 0 + ) + if use_td: + set_triton_allocator(A.device) + batched_triton_kernel[grid]( A, B, @@ -485,6 +543,7 @@ def invoke_moe_batched_triton_kernel( BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, + USE_TD=use_td, ) @@ -616,10 +675,7 @@ def apply( for expert in range(num_local_experts): # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor - if ( - torch.compiler.is_compiling() - or torch.cuda.is_current_stream_capturing() - ): + if _is_capturing_or_compiling(): num = hidden_states.shape[1] else: num = int(expert_num_tokens[expert].item()) @@ -659,7 +715,7 @@ def batched_moe_kernel_quantize_input( per_act_token_quant: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): + if _is_capturing_or_compiling(): # Note: this does a bunch of extra work because expert_num_tokens is # ignored but it does support torch.compile + cudagraphs. hidden_dim = A.size(-1) @@ -745,7 +801,7 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - return current_platform.is_cuda_alike() + return current_platform.is_cuda_alike() or current_platform.is_xpu() @staticmethod def _supports_no_act_and_mul() -> bool: @@ -758,13 +814,13 @@ def _supports_quant_scheme( ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import get_cdna_version - is_rocm_on_gfx9 = on_gfx9() + _rocm_support_fp8 = get_cdna_version() > 2 else: - is_rocm_on_gfx9 = False + _rocm_support_fp8 = False - device_supports_fp8 = is_rocm_on_gfx9 or ( + device_supports_fp8 = _rocm_support_fp8 or ( p.is_cuda() and p.has_device_capability((8, 9)) ) 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 5f112380cf9b..45836c387492 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 @@ -289,7 +289,14 @@ def moe_problem_size( assert a1.size(0) == num_experts num_tokens = a1.size(1) - return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k + return ( + meta1.num_experts, + num_tokens, + # Logical intermediate width for both gated and non-gated activations + self.layer.intermediate_size_per_partition, + meta1.shape_k, + top_k, + ) def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): from vllm.utils.humming import GemmType as HummingGemmType @@ -327,7 +334,8 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): real_shape_m = M * topk output_shape = (M, K) - down_input_size = N if activation.is_gated else (N * 2) + gate_up_size = N * (2 if activation.is_gated else 1) + down_input_size = N a_dtype = self.layer.humming_metas["w13"].a_dtype c_dtype = self.layer.humming_metas["w13"].c_dtype num_bits = a_dtype.num_bits @@ -347,7 +355,7 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): "dtype": torch_dtype_map[a_dtype], }, "gate_up_output": { - "shape": (real_shape_m, N * 2), + "shape": (real_shape_m, gate_up_size), "dtype": torch_dtype_map[c_dtype], }, "activation_output": { @@ -654,12 +662,9 @@ def apply( topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, - outputs=buffers["output"], + outputs=output, ) - # Note: output is already written to buffers["output"] - # which aliases workspace13/output - class HummingGroupedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -769,16 +774,13 @@ def apply( ) moe_unpermute( - out=buffers["output"], + out=output, permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, inv_permuted_idx=inv_perm, expert_first_token_offset=expert_first_token_offset, ) - # Note: output is already written to buffers["output"] - # which aliases workspace13/output - class BatchedHummingGroupedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -872,13 +874,10 @@ def apply( layer=self.layer, inputs=inputs, input_scale=input_scale, - outputs=buffers["down_output"].view(-1, hidden_states.size(-1)), + outputs=output.view(-1, hidden_states.size(-1)), valid_shape_m=valid_shape_m, expert_layout=expert_num_tokens, compute_config=self.compute_config_str, tuning_config=self.w2_tuning_config_str, sublayer_name="w2", ) - - # Note: output is already written to buffers["down_output"] - # which aliases workspace13/output diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 4b0a0b8ecad3..da2628277d39 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -80,28 +80,15 @@ def _patch_make_bitmatrix_metadata() -> None: import triton.language as tl try: - if current_platform.is_rocm(): - from triton_kernels.tensor_details import bitmatrix as _bm - from triton_kernels.tensor_details.bitmatrix import ( - BitmatrixMetadata, - _keyed_add, - cdiv, - ) - from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 - sum_bitmatrix_rows, - ) - else: - from vllm.third_party.triton_kernels.tensor_details import ( - bitmatrix as _bm, - ) - from vllm.third_party.triton_kernels.tensor_details.bitmatrix import ( - BitmatrixMetadata, - _keyed_add, - cdiv, - ) - from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 - sum_bitmatrix_rows, - ) + from triton_kernels.tensor_details import bitmatrix as _bm + from triton_kernels.tensor_details.bitmatrix import ( + BitmatrixMetadata, + _keyed_add, + cdiv, + ) + from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 + sum_bitmatrix_rows, + ) except ImportError: return diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 20576de5a608..1125130c8904 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -93,6 +93,8 @@ def _fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: assert hidden_states.ndim == 2 M, K = hidden_states.size() @@ -171,6 +173,8 @@ def _fused_marlin_moe( beta=gemm1_beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) if output is None: @@ -256,6 +260,8 @@ def fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function computes a Mixture of Experts (MoE) layer using two sets of @@ -357,6 +363,8 @@ def fused_marlin_moe( num_tokens_post_padded=num_tokens_post_padded, activation=activation, activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, input_global_scale1=input_global_scale1, input_global_scale2=input_global_scale2, global_scale1=global_scale1, @@ -423,6 +431,9 @@ def batched_fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_func: Callable[..., None] = apply_moe_activation, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function massages the inputs so the batched hidden_states can be @@ -530,6 +541,9 @@ def batched_fused_marlin_moe( quant_type=quant_type, apply_router_weight_on_input=apply_router_weight_on_input, activation=activation, + activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, expert_map=expert_map, block_size_m=block_size_m, sorted_token_ids=sorted_token_ids, @@ -645,6 +659,7 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SITU, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, @@ -656,10 +671,10 @@ def _supports_activation(activation: MoEActivation) -> bool: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return not ( - moe_parallel_config.use_fi_nvl_two_sided_kernels - or moe_parallel_config.use_fi_nvl_one_sided_kernels - ) + # One-sided FI-NVL all2all pairs with MarlinExperts fine (the + # compressed-tensors MXFP4 path runs this exact combo); only the + # two-sided kernels are unsupported here. + return not moe_parallel_config.use_fi_nvl_two_sided_kernels @property def quant_type_id(self) -> int: @@ -793,6 +808,10 @@ def apply( global_num_experts=global_num_experts, activation=activation, activation_func=self.activation, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=( + self.moe_config.activation_situ_linear_beta + ), moe_sum=self.moe_sum, expert_map=expert_map, output=output, @@ -833,6 +852,8 @@ def activation_with_lora( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: # act_input = intermediate_cache1 (M*topk, 2N for gated) # act_output = intermediate_cache2 (M*topk, N) @@ -871,6 +892,8 @@ def activation_with_lora( beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) lora_state["cache2"] = act_output @@ -918,6 +941,8 @@ def moe_sum_with_lora( global_num_experts=global_num_experts, activation=activation, activation_func=activation_with_lora, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, moe_sum=moe_sum_with_lora, expert_map=expert_map, output=output, @@ -1021,6 +1046,38 @@ def apply( apply_router_weight_on_input: bool, ): assert expert_tokens_meta is not None, "Num valid tokens per batch is required" + + def activation_func( + act: MoEActivation, + act_output: torch.Tensor, + act_input: torch.Tensor, + *, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + **kwargs, + ) -> None: + if act != MoEActivation.SITU: + self.activation( + act, + act_output, + act_input, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + **kwargs, + ) + return + + num_experts, max_num_tokens = hidden_states.shape[:2] + beta = 1.0 if activation_situ_beta is None else activation_situ_beta + linear_beta = activation_situ_linear_beta + torch.ops._C.masked_situ_and_mul( + act_output.view(num_experts, max_num_tokens, -1), + act_input.view(num_experts, max_num_tokens, -1), + expert_tokens_meta.expert_num_tokens, + beta, + -1.0 if linear_beta is None else linear_beta, + ) + return batched_fused_marlin_moe( hidden_states=hidden_states, expert_num_tokens=expert_tokens_meta.expert_num_tokens, @@ -1051,4 +1108,7 @@ def apply( clamp_limit=self.gemm1_clamp_limit, gemm1_alpha=self.gemm1_alpha, gemm1_beta=self.gemm1_beta, + activation_func=activation_func, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py index 71dd7634a697..ad6083251fb0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py @@ -107,17 +107,13 @@ def activation( limit = self.quant_config.gemm1_clamp_limit if limit is None: raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit") - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) apply_moe_activation( activation, output, input, clamp_limit=float(limit), - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, ) return super().activation(activation, output, input) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index 9839756880ab..e8c7dc969213 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -362,10 +362,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) + # `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`. limit = self.quant_config.gemm1_clamp_limit limit = None if limit is None else float(limit) out = fused_moe_mxfp8_native( @@ -376,8 +373,8 @@ def apply( self.w2_scale_val, topk_weights, topk_ids, - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, limit=limit, global_num_experts=global_num_experts, expert_map=expert_map, diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index 833fa70d9ef6..09fc03028cae 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -69,9 +69,15 @@ def __init__( self.quantization_emulation = True if self.ocp_mx_scheme in { + OCP_MX_Scheme.w_mxfp4, + OCP_MX_Scheme.w_mxfp6_e3m2, + OCP_MX_Scheme.w_mxfp6_e2m3, + }: + # Weight-only schemes leave activations unquantized. + self._quant_dtype = None + elif self.ocp_mx_scheme in { OCP_MX_Scheme.w_mxfp4_a_mxfp4, }: - # Weight has to be dequantized for mxfp4 emulation. self._quant_dtype = "mxfp4" elif self.ocp_mx_scheme in [ OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2, diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 4f191334bb2a..096e7294809d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -5,6 +5,7 @@ import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -262,6 +263,8 @@ def rocm_aiter_fused_experts( elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu") activation_interleave = False + elif activation == MoEActivation.SITU: + activation_method = rocm_aiter_ops.get_aiter_activation_type("situ") else: raise ValueError(f"Unsupported activation: {activation}") @@ -356,10 +359,12 @@ def rocm_aiter_fused_experts( # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099 # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401 - hidden_pad = hidden_pad // 128 * 128 - intermediate_pad = ( - intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) - ) + # SITU's A16W4 FlyDSL kernel pads per gate/up half; pass through unrounded. + if activation != MoEActivation.SITU: + hidden_pad = hidden_pad // 128 * 128 + intermediate_pad = ( + intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) + ) # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs # for interleaved vs separated gate and up weights. @@ -370,7 +375,15 @@ def rocm_aiter_fused_experts( from aiter.ops.flydsl.moe_common import GateMode gate_mode = "" - if quant_config.use_mxfp4_w4a16: + if activation == MoEActivation.SITU: + # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) + # fp8 flydsl kernels; default a16w4 SiTU stays separated. + gate_mode = ( + GateMode.INTERLEAVE.value + if envs.AITER_SITUV2_A8W4 + else GateMode.SEPARATED.value + ) + elif quant_config.use_mxfp4_w4a16: gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: gate_mode = ( @@ -401,6 +414,8 @@ def rocm_aiter_fused_experts( bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + beta=moe_config.activation_situ_beta, + linear_beta=moe_config.activation_situ_linear_beta, ) @@ -457,11 +472,10 @@ def _supports_quant_scheme( ] if (weight_key, activation_key) not in SUPPORTED_W_A: return False - # CK MXFP4 MoE kernels are only supported on gfx950. if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - if not on_gfx950(): + if not on_gfx950() or on_gfx1250(): return False return True diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 3196667b3f7a..d1bab54dda11 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -43,8 +43,13 @@ kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, + kInt8DynamicTensorSym, kInt8DynamicTokenSym, + kInt8Static, kInt8StaticChannelSym, + kInt8StaticTensorSym, ) from vllm.platforms import current_platform from vllm.triton_utils import tl @@ -100,15 +105,25 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - # INT8 requires at least 7.5 (Turing). + # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs + # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and + # the Triton int8_w8a8 fused MoE kernel handles them. device_supports_int8 = ( current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) - ) + ) or current_platform.is_rocm() supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] if device_supports_int8: - supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym)) + # Activations are consumed as float and quantized to int8 + # dynamically inside the kernel, so only dynamic-activation int8 + # schemes are supported (static-activation int8 is not). + supported += [ + # per-channel weight + dynamic per-token activation + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # per-tensor weight + dynamic per-tensor activation + (kInt8StaticTensorSym, kInt8DynamicTensorSym), + ] if current_platform.supports_fp8(): supported += [ (kFp8Static128BlockSym, kFp8Dynamic128Sym), @@ -125,6 +140,7 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SITU, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, @@ -542,40 +558,44 @@ def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: class TritonWNA16Experts(TritonExperts): @staticmethod def _supports_current_device() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return current_platform.is_cuda_alike() or current_platform.is_xpu() @staticmethod def _supports_no_act_and_mul() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return True @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + SUPPORTED_W = [ + kInt4Static, + kInt8Static, + kInt4Static32, + # other group sizes? + ] + return weight_key in SUPPORTED_W @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels ) def apply( @@ -598,7 +618,9 @@ def apply( ): # Check constraints. if self.quant_config.use_int4_w4a16: - assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch" + assert hidden_states.size(-1) // 2 == w1.size(2), ( + f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}" + ) else: assert hidden_states.size(-1) == w1.size(2), ( f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}" diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 143df9e81cd6..e3a85a1365d7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -34,6 +34,9 @@ class TrtLlmBf16ExpertsBase: monolithic interfaces. """ + def supports_routing_replay_capture(self) -> bool: + return True + def __init__( self, moe_config: FusedMoEConfig, @@ -246,7 +249,11 @@ def apply( assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - return flashinfer.fused_moe.trtllm_bf16_moe( + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + out = flashinfer.fused_moe.trtllm_bf16_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -263,4 +270,9 @@ def apply( routing_method_type=self.routing_method_type, activation_type=activation_to_flashinfer_int(activation), tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), + routing_replay_out=routing_replay_out, + ) + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] ) + return out diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 55ddd1a2c964..50ed4ed5c3cc 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -265,6 +265,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. """ + def supports_routing_replay_capture(self) -> bool: + return True + def __init__( self, moe_config: FusedMoEConfig, @@ -403,6 +406,11 @@ def _apply_block_scale( n_group = num_expert_group or 0 selected_topk_group = topk_group or 0 + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + kwargs = dict( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -427,11 +435,16 @@ def _apply_block_scale( use_shuffled_weight=use_shuffled_weight, weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, + routing_replay_out=routing_replay_out, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL: kwargs["activation_type"] = activation_type - return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs) + result = flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs) + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] + ) + return result def _apply_per_tensor( self, @@ -464,6 +477,11 @@ def _apply_per_tensor( else: assert not apply_router_weight_on_input + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -485,6 +503,10 @@ def _apply_per_tensor( routing_method_type=self.routing_method_type, activation_type=activation_type, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), + routing_replay_out=routing_replay_out, + ) + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] ) return out diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 2b8a5529ab15..315ec9f9c056 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -15,6 +15,10 @@ TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, + has_flashinfer_situ_activation, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -79,6 +83,36 @@ def __init__( else: self.gemm1_clamp_limit = None + # SITU (SituGLU) TRTLLM-Gen kernel computes + # left = alpha * tanh(x0 / alpha) * sigmoid(x0) # gate (x0) + # right = beta * tanh(x1 / beta) # up (x1) + # which matches vLLM's situ_and_mul with (beta, linear_beta), so map + # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta -> + # gatedActBeta (gemm1_beta). Both must be > 0. + if moe_config.activation == MoEActivation.SITU: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, ( + "SITU requires activation_situ_beta > 0" + ) + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 " + "(the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = torch.full( + (self.local_num_experts,), + float(situ_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_beta = torch.full( + (self.local_num_experts,), + float(situ_linear_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_clamp_limit = None + self.max_capture_size = moe_config.max_capture_size @staticmethod @@ -103,7 +137,16 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + if activation == MoEActivation.SITU: + return has_flashinfer_situ_activation() + return activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SILU, + ) + + @staticmethod + def _flashinfer_activation_type(activation: MoEActivation) -> int: + return activation_to_flashinfer_int(activation) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -122,6 +165,9 @@ class TrtLlmMxfp4ExpertsMonolithic( Wraps flashinfer.trtllm_fp4_block_scale_moe(). """ + def supports_routing_replay_capture(self) -> bool: + return True + @staticmethod def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, @@ -139,6 +185,7 @@ def _supports_routing_method( activation_key: QuantKey | None, ) -> bool: return routing_method in [ + RoutingMethodType.DeepSeekV3, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, ] @@ -184,9 +231,13 @@ def apply( device=hidden_states.device, ) + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) trtllm_fp4_block_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), - routing_bias=None, + routing_logits=router_logits, + routing_bias=e_score_correction_bias, hidden_states=x_quant, hidden_states_scale=x_scale, gemm1_weights=w1, @@ -203,18 +254,22 @@ def apply( output2_scale_scalar=None, num_experts=global_num_experts, top_k=self.topk, - n_group=None, - topk_group=None, + n_group=(num_expert_group or 0), + topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=None, + routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, do_finalize=True, + activation_type=self._flashinfer_activation_type(activation), tune_max_num_tokens=max(self.max_capture_size, 1), output=output, + routing_replay_out=routing_replay_out, + ) + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] ) - return output @@ -261,6 +316,75 @@ def workspace_shapes( output = (M, self.hidden_dim_unpadded) return (workspace1, workspace2, output) + def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int: + """Max tokens per kernel call before the batched-GEMM grid overflows. + + The TRTLLM-Gen batched GEMM launches a static grid whose batch (Y) + dimension is ``getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts, + tileTokensDim)`` and must stay <= 65535. Solving that for num_tokens + with the smallest tile the kernel may pick (tileTokensDim=8, the runner + default) gives a bound that is safe regardless of the tactic selected. + Without it, large batches (e.g. Kimi-K3 top_k=16, EP16 profiling with + 131072 gathered tokens) overflow the grid and the GEMM launch fails. + """ + MAX_GRID_Y = 65535 + MIN_TILE_TOKENS_DIM = 8 + max_tokens = (MAX_GRID_Y - global_num_experts) * MIN_TILE_TOKENS_DIM // top_k + return max(1, min(300000, max_tokens)) + + def _invoke_kernel( + self, + output: torch.Tensor, + x_quant: torch.Tensor, + x_scale: torch.Tensor | None, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + local_num_experts: int, + local_expert_offset: int, + topk: int, + ) -> None: + from flashinfer import trtllm_fp4_block_scale_routed_moe + + packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=w1, + gemm1_weights_scale=self.w1_scale, + gemm1_bias=self.w1_bias, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, + gemm2_weights=w2, + gemm2_weights_scale=self.w2_scale, + gemm2_bias=self.w2_bias, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=global_num_experts, + top_k=topk, + n_group=None, + topk_group=None, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=None, + # Modular kernel receives pre-routed tokens, so routing is already + # done. Use Renormalize as a safe default the TRTLLM kernel supports. + routing_method_type=RoutingMethodType.Renormalize, + do_finalize=True, + enable_pdl=True, + activation_type=self._flashinfer_activation_type(activation), + output=output, + tune_max_num_tokens=max(self.max_capture_size, 1), + ) + def apply( self, output: torch.Tensor, @@ -281,7 +405,6 @@ def apply( ): topk = topk_ids.size(-1) local_num_experts = w1.size(0) - intermediate_size = self.intermediate_size_per_partition local_expert_offset = self.moe_config.ep_rank * local_num_experts if a1q_scale is not None: @@ -292,48 +415,27 @@ def apply( x_quant = hidden_states x_scale = None - # Pack topk ids and weights into format expected by the kernel. - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) - assert self.w1_scale is not None assert self.w2_scale is not None - kwargs = { - "topk_ids": packed_tensor, - "routing_bias": None, - "hidden_states": x_quant, - "hidden_states_scale": x_scale, - "gemm1_weights": w1, - "gemm1_weights_scale": self.w1_scale, - "gemm1_bias": self.w1_bias, - "gemm1_alpha": self.gemm1_alpha, - "gemm1_beta": self.gemm1_beta, - "gemm1_clamp_limit": self.gemm1_clamp_limit, - "gemm2_weights": w2, - "gemm2_weights_scale": self.w2_scale, - "gemm2_bias": self.w2_bias, - "output1_scale_scalar": None, - "output1_scale_gate_scalar": None, - "output2_scale_scalar": None, - "num_experts": global_num_experts, - "top_k": topk, - "n_group": None, - "topk_group": None, - "intermediate_size": intermediate_size, - "local_expert_offset": local_expert_offset, - "local_num_experts": local_num_experts, - "routed_scaling_factor": None, - # Modular kernel receives pre-routed tokens, so routing - # is already done. Use Renormalize as a safe default that - # the TRTLLM C++ kernel supports. - "routing_method_type": RoutingMethodType.Renormalize, - "do_finalize": True, - "enable_pdl": True, - "output": output, - "tune_max_num_tokens": max(self.max_capture_size, 1), - } - - from flashinfer import trtllm_fp4_block_scale_routed_moe - trtllm_fp4_block_scale_routed_moe(**kwargs) + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. + M = x_quant.size(0) + chunk_size = self._max_supported_tokens(topk, global_num_experts) + for start in range(0, M, chunk_size): + end = min(start + chunk_size, M) + self._invoke_kernel( + output[start:end], + x_quant[start:end], + None if x_scale is None else x_scale[start:end], + topk_ids[start:end], + topk_weights[start:end], + w1, + w2, + activation, + global_num_experts, + local_num_experts, + local_expert_offset, + topk, + ) return output diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index c6e5e70a14ad..a9e725f2965c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -122,6 +122,9 @@ def expects_unquantized_inputs(self) -> bool: # The kernel handles quantization internally. return True + def supports_routing_replay_capture(self) -> bool: + return True + def apply( self, hidden_states: torch.Tensor, @@ -144,7 +147,12 @@ def apply( assert self.w1_scale is not None assert self.w2_scale is not None - return flashinfer_trtllm_mxint4_moe( + + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + result = flashinfer_trtllm_mxint4_moe( x=hidden_states, router_logits=router_logits, w13_weight_packed=w1, @@ -160,4 +168,9 @@ def apply( topk_group=topk_group, e_score_correction_bias=e_score_correction_bias, routing_method_type=self.routing_method, + routing_replay_out=routing_replay_out, + ) + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] ) + return result diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 2a1443417328..a5555dc1323e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -22,6 +22,7 @@ ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, + has_flashinfer_situ_activation, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -110,6 +111,27 @@ def _per_expert(val: float | None) -> torch.Tensor | None: self.gemm1_alpha = None self.gemm1_beta = None + # SITU (Kimi SituGLU) TRTLLM-Gen kernel computes + # left=alpha*tanh(x0/alpha)*sigmoid(x0), right=beta*tanh(x1/beta), + # matching vLLM's situ_and_mul, so map situ beta -> gatedActAlpha + # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta). + # These operate on the dequantized gate/up, so they are NOT folded by + # g1_alphas in process_weights_after_loading. + self.is_situ = moe_config.activation == MoEActivation.SITU + if self.is_situ: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, ( + "SITU requires activation_situ_beta > 0" + ) + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 " + "(the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = _per_expert(situ_beta) + self.gemm1_beta = _per_expert(situ_linear_beta) + self.gemm1_clamp_limit = None + logger.debug_once( "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", moe_config.activation, @@ -142,7 +164,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # (via the in-place mul above) and never changes again, so this is a # static, per-expert constant. Register on the layer so EPLB # rearranges it alongside the other expert tensors. - if self.gemm1_clamp_limit is not None: + # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the + # raw GEMM1 accumulator, so they are registered as-is without the + # g1_alphas fold used by the SwiGLU-OAI clamp/beta below. + if self.gemm1_clamp_limit is not None and not self.is_situ: gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas layer.register_parameter( "gemm1_clamp_limit", @@ -155,7 +180,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # raw. Register both on the layer so EPLB rearranges them with the # other per-expert tensors. if self.gemm1_beta is not None: - gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas + gemm1_beta = ( + self.gemm1_beta + if self.is_situ + else self.gemm1_beta / self.quant_config.g1_alphas + ) layer.register_parameter( "gemm1_beta", torch.nn.Parameter(gemm1_beta, requires_grad=False), @@ -197,7 +226,9 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" + """Supports SITU only when the installed FlashInfer exposes it.""" + if activation == MoEActivation.SITU: + return has_flashinfer_situ_activation() return activation in [ MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, @@ -435,6 +466,9 @@ class TrtLlmNvFp4ExpertsMonolithic( Monolithic version of the kernel (router + experts). """ + def supports_routing_replay_capture(self) -> bool: + return True + @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: """The modular implementation should be used for the Dp/Ep or EPLB case.""" @@ -509,10 +543,14 @@ def apply( output1_scale_gate_scalar = self.quant_config.g1_alphas + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) # Invoke kernel. # NOTE: Activation padding and output # truncation are handled by the MoE runner's - return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -544,4 +582,9 @@ def apply( activation_type=activation_to_flashinfer_int(activation), per_token_scale=per_token_scale, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), + routing_replay_out=routing_replay_out, )[0] + self._maybe_dispatch_routing_replay( + routing_replay_out, num_tokens=hidden_states.shape[0] + ) + return result diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index cde167e5d364..dfc578b44170 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -20,6 +20,7 @@ kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, kMxfp8Static, @@ -64,10 +65,16 @@ def __init__( ) self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit self.fused_moe_impl: XpuFusedMoe | None = None + is_xe2_or_xe3 = torch.ops._xpu_C.is_xe2_arch() or torch.ops._xpu_C.is_xe3_arch() + if not is_xe2_or_xe3: + raise NotImplementedError( + "XPUExperts is only supported on Intel Xe2/Xe3 GPUs" + ) + self._expects_unquantized_inputs = is_xe2_or_xe3 @property def expects_unquantized_inputs(self) -> bool: - return True + return self._expects_unquantized_inputs @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -172,6 +179,7 @@ def apply( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, + a1q_scale=a1q_scale, ) @@ -309,6 +317,24 @@ def __init__( num_dispatchers, ) + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # K = a1q.size(-1). When activations are pre-quantized packed mxfp4, + # K is the packed hidden_size (= logical / 2); the kernel output is at + # logical hidden_size (2 * K). When unquantized (bf16), K is already + # the logical size. + logical_K = K if self.expects_unquantized_inputs else 2 * K + return (0,), (0,), (M, logical_K) + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -316,5 +342,6 @@ def _supports_quant_scheme( ) -> bool: SUPPORTED_W_A = [ (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A diff --git a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py index b1588c4a2a46..27c07815cb09 100644 --- a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py @@ -7,6 +7,7 @@ import os import flydsl.compiler as flyc +import flydsl.expr as fx import torch from aiter.fused_moe import moe_sorting as aiter_moe_sorting from aiter.ops.flydsl.kernels.moe_gemm_2stage import ( @@ -43,6 +44,12 @@ } +def _as_flydsl_ptr(tensor: torch.Tensor | None): + """Adapt a torch tensor to the pointer ABI expected by AITER FlyDSL.""" + data_ptr = None if tensor is None else tensor.data_ptr() + return flyc.from_c_void_p(fx.Uint8, data_ptr) + + def moe_sorting( topk_ids: torch.Tensor, topk_weights: torch.Tensor, @@ -219,6 +226,22 @@ def fused_flydsl_moe_impl( ) compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1) + args1 = ( + _as_flydsl_ptr(out_stage1), + _as_flydsl_ptr(hidden_states), + _as_flydsl_ptr(w1), + _as_flydsl_ptr(scale_x_1d), + _as_flydsl_ptr(w1_scale), + _as_flydsl_ptr(sorted_token_ids), + _as_flydsl_ptr(sorted_expert_ids), + _as_flydsl_ptr(sorted_weights_1d), + _as_flydsl_ptr(num_valid_ids), + tokens, + inter_dim, + model_dim, + int(blocks), + stream, + ) if compiled_exe1 is None: exe1 = compile_moe_gemm1( model_dim=model_dim, @@ -235,41 +258,10 @@ def fused_flydsl_moe_impl( use_cshuffle_epilog=False, scale_is_bf16=scale_is_bf16, ) - compiled_exe1 = flyc.compile( - exe1, - out_stage1, - hidden_states, - w1, - scale_x_1d, - w1_scale, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - num_valid_ids, - tokens, - inter_dim, - model_dim, - int(blocks), - stream, - ) + compiled_exe1 = flyc.compile(exe1, *args1) _FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1 - - compiled_exe1( - out_stage1, - hidden_states, - w1, - scale_x_1d, - w1_scale, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - num_valid_ids, - tokens, - inter_dim, - model_dim, - int(blocks), - stream, - ) + else: + compiled_exe1(*args1) a2_1d = out_stage1.view(-1).contiguous() a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) @@ -291,6 +283,23 @@ def fused_flydsl_moe_impl( ) compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2) + args2 = ( + _as_flydsl_ptr(out_stage2), + _as_flydsl_ptr(a2_1d), + _as_flydsl_ptr(w2), + _as_flydsl_ptr(a2_scale_1d), + _as_flydsl_ptr(w2_scale), + _as_flydsl_ptr(sorted_token_ids), + _as_flydsl_ptr(sorted_expert_ids), + _as_flydsl_ptr(sorted_weights_1d), + _as_flydsl_ptr(num_valid_ids), + tokens, + model_dim, + inter_dim, + int(blocks), + stream, + ) + out_stage2.zero_() if compiled_exe2 is None: exe2 = compile_moe_gemm2( model_dim=model_dim, @@ -306,42 +315,10 @@ def fused_flydsl_moe_impl( doweight_stage2=bool(doweight_stage2), scale_is_bf16=scale_is_bf16, ) - compiled_exe2 = flyc.compile( - exe2, - out_stage2, - a2_1d, - w2, - a2_scale_1d, - w2_scale, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - num_valid_ids, - tokens, - model_dim, - inter_dim, - int(blocks), - stream, - ) + compiled_exe2 = flyc.compile(exe2, *args2) _FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2 - - out_stage2.zero_() - compiled_exe2( - out_stage2, - a2_1d, - w2, - a2_scale_1d, - w2_scale, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - num_valid_ids, - tokens, - model_dim, - inter_dim, - int(blocks), - stream, - ) + else: + compiled_exe2(*args2) return out_stage2 diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 62cf12ea24e2..be4930052a9a 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -28,9 +28,13 @@ from vllm.model_executor.layers.fused_moe.utils import ( enable_swap_ab, moe_kernel_quantize_input, + resolve_moe_use_td, + warn_if_moe_use_td_ineffective, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.triton_utils.allocation import set_triton_allocator +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op @@ -346,6 +350,8 @@ def fused_moe_kernel( per_channel_quant: tl.constexpr, HAS_BIAS: tl.constexpr, SWAP_AB: tl.constexpr, + # Tensor-descriptor path for the A gather and B load in the K-loop. + USE_TD: tl.constexpr = False, ): """ Implements the fused computation for a Mixture of Experts (MOE) using @@ -435,7 +441,25 @@ def fused_moe_kernel( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) - if SWAP_AB: + # TD gather and the SWAP_AB accumulator layout are mutually exclusive. + tl.static_assert(not (USE_TD and SWAP_AB)) + if USE_TD: + # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx. + m_td = num_valid_tokens // top_k + a_desc = tl.make_tensor_descriptor( + base=a_ptr, + shape=(m_td, K), + strides=(stride_am, stride_ak), + block_shape=(1, BLOCK_SIZE_K), + ) + b_desc = tl.make_tensor_descriptor( + base=b_ptr + off_experts * stride_be, + shape=(N, K), + strides=(stride_bn, stride_bk), + block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K), + ) + gather_idx = (offs_token // top_k).to(tl.int32) + elif SWAP_AB: a_ptrs = a_ptr + ( offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am ) @@ -453,7 +477,6 @@ def fused_moe_kernel( + off_experts * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) ) - if use_int8_w8a16: b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn @@ -497,18 +520,21 @@ def fused_moe_kernel( for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. - if SWAP_AB: + if USE_TD: + a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K) + b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T + elif SWAP_AB: a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + a = tl.load(a_ptrs, mask=a_mask, other=0.0) + b = tl.load(b_ptrs, mask=b_mask, other=0.0) else: - a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K) - b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K - a = tl.load( - a_ptrs, - mask=a_mask, - other=0.0, - ) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -535,9 +561,10 @@ def fused_moe_kernel( accumulator += tl.dot(a, b) else: accumulator += tl.dot(a, b) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk + if not USE_TD: + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk if SWAP_AB: accumulator = tl.trans(accumulator, (1, 0)) @@ -764,6 +791,19 @@ def invoke_fused_moe_triton_kernel( else: SWAP_AB = False + # Quantized weights always carry a B_scale (see the asserts below); key off + # that rather than enumerating quant flags, which misses w8a16-fp8/nvfp4/etc. + is_quantized = B_scale is not None + warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized) + + # TD path is unvalidated under quantization; fall back to the pointer path. + use_td = resolve_moe_use_td() and not is_quantized + if use_td: + # The TD path builds a tensor descriptor inside the kernel, which + # requires a PyTorch-backed scratch allocator to be registered + # (Triton raises "no allocator was set" otherwise on CUDA). + set_triton_allocator(A.device) + if use_fp8_w8a8 or use_int8_w8a8: assert B_scale is not None assert block_shape is None or triton.cdiv( @@ -804,6 +844,19 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + if use_td and A.size(1) % BLOCK_SIZE_K != 0: + # TD gather/load feeding tl.dot with a non-block-aligned K + # miscompiles (~74% of output elements wrong) on real HW; + # this is a compiler-codegen issue, not a Python-maskable + # boundary gap. Fall back to the pointer-arith path. + logger.warning_once( + "Disabling VLLM_TRITON_USE_TD for this MoE launch: K=%d is not " + "a multiple of BLOCK_SIZE_K=%d, which triggers a known " + "Triton tensor-descriptor + tl.dot miscompilation.", + A.size(1), + BLOCK_SIZE_K, + ) + use_td = False fused_moe_kernel[grid]( A, B, @@ -846,6 +899,7 @@ def invoke_fused_moe_triton_kernel( HAS_BIAS=HAS_BIAS, BLOCK_SIZE_K=BLOCK_SIZE_K, SWAP_AB=SWAP_AB, + USE_TD=use_td, **config, ) @@ -1298,7 +1352,11 @@ def get_default_config( bit = 4 if dtype == "int4_w4a16" else 8 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit) if use_moe_wna16_cuda: - config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1} + config = { + "BLOCK_SIZE_M": min(16, next_power_of_2(M)), + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } elif M <= 20: config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1} elif M <= 40: diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 9622f9ed886e..cfbadba7aced 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -96,8 +96,7 @@ def determine_expert_counts( return global_num_experts, logical_num_experts, num_fused_shared_experts -# TODO: rename this -def FusedMoE( +def FusedMoEFactory( num_experts: int, # Global number of experts top_k: int, hidden_size: int, @@ -120,6 +119,8 @@ def FusedMoE( swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, e_score_correction_bias: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, activation: str = "silu", @@ -178,6 +179,8 @@ def FusedMoE( scoring_func: Scoring function for routing ("softmax" or others) routed_scaling_factor: Scaling factor applied to topk_weights or output swiglu_limit: SwiGLU activation limit + activation_situ_beta: SituGLU activation beta + activation_situ_linear_beta: SituGLU linear beta e_score_correction_bias: Expert score correction bias tensor apply_router_weight_on_input: Whether to apply router weights on input activation: Activation function name ("silu", "gelu", etc.) @@ -352,6 +355,8 @@ def FusedMoE( swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, skip_final_all_reduce=skip_final_all_reduce, ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 5c188e5660ca..f7dc5dc2c054 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -10,7 +10,6 @@ import torch import vllm.envs as envs -from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -891,6 +890,8 @@ def activation( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: apply_moe_activation( activation, @@ -901,6 +902,16 @@ def activation( beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=( + self.moe_config.activation_situ_beta + if activation_situ_beta is None + else activation_situ_beta + ), + activation_situ_linear_beta=( + self.moe_config.activation_situ_linear_beta + if activation_situ_linear_beta is None + else activation_situ_linear_beta + ), ) @abstractmethod @@ -1000,6 +1011,57 @@ def _supports_router_logits_dtype( def is_monolithic() -> bool: return True + routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None + _routing_replay_buffer: torch.Tensor | None = None + + def supports_routing_replay_capture(self) -> bool: + """Whether this expert supports routing replay capture. + + Subclasses backed by a kernel that exposes routed expert IDs + (e.g. FlashInfer's ``routing_replay_out``) should override. + """ + return False + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.routing_replay_capture_fn = capture_fn + if capture_fn is None: + self._routing_replay_buffer = None + return + self._routing_replay_buffer = torch.empty( + (self.moe_config.max_num_tokens, self.moe_config.experts_per_token), + dtype=torch.int16, + device=self.moe_config.device, + ) + + def _maybe_make_routing_replay_buffer( + self, + num_tokens: int, + device: torch.device, + ) -> torch.Tensor | None: + if self.routing_replay_capture_fn is None: + return None + buf = self._routing_replay_buffer + assert buf is not None + if buf.shape[0] < num_tokens or buf.device != device: + raise ValueError( + "Routing replay buffer was initialized for " + f"{buf.shape[0]} tokens on {buf.device}, but the kernel " + f"received {num_tokens} tokens on {device}." + ) + return buf + + def _maybe_dispatch_routing_replay( + self, + routing_replay_out: torch.Tensor | None, + num_tokens: int, + ) -> None: + if routing_replay_out is None or self.routing_replay_capture_fn is None: + return + self.routing_replay_capture_fn(routing_replay_out[:num_tokens]) + def apply( self, hidden_states: torch.Tensor, @@ -1143,21 +1205,6 @@ def _prepare( The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ - # Skip cudagraph/DP padding tokens uniformly across all a2a backends: - # forcing padded rows' expert ids to -1 makes every prepare_finalize drop - # them (not dispatched / not computed by the experts). The V2 model runner - # marks them in forward_context.is_padding; it is None for runners that do - # not populate it, leaving topk_ids unchanged. - # Gated by VLLM_MOE_SKIP_PADDING (off by default) because this requires the - # experts kernel to treat topk_id == -1 as a skip sentinel, which not all - # MoE backends support yet. - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - n = topk_ids.shape[0] - # TODO: Properly support DBO (padding lives at the batch tail). - topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids) if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't @@ -1272,6 +1319,14 @@ def _fused_experts( activation, ) + use_output_alias = ( + output_alias is not None + and output_alias.shape == fused_out.shape + and output_alias.dtype == fused_out.dtype + and output_alias.device == fused_out.device + and output_alias.is_contiguous() + ) + # If caller's output buffer already matches fused_out shape/dtype, alias # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream. # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the @@ -1279,15 +1334,10 @@ def _fused_experts( if current_platform.is_rocm(): from vllm._aiter_ops import rocm_aiter_ops - if ( - rocm_aiter_ops.is_fused_moe_enabled() - and output_alias is not None - and output_alias.shape == fused_out.shape - and output_alias.dtype == fused_out.dtype - and output_alias.device == fused_out.device - and output_alias.is_contiguous() - ): + if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled(): fused_out = output_alias + elif use_output_alias: + fused_out = output_alias self.fused_experts.apply( output=fused_out, diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 5a2b4c3a75b4..ab7775bf3209 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -24,7 +24,6 @@ kInt8StaticChannelSym, ) from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform logger = init_logger(__name__) @@ -41,20 +40,12 @@ def _get_priority_backends( """ Get available backends in priority order based on platform and config. """ - _AVAILABLE_BACKENDS = [ + return [ Int8MoeBackend.TRITON, Int8MoeBackend.HUMMING, Int8MoeBackend.CPU, ] - def _move_to_front(backends: list[Int8MoeBackend], backend: Int8MoeBackend) -> None: - backends.insert(0, backends.pop(backends.index(backend))) - - if current_platform.is_cpu(): - _move_to_front(_AVAILABLE_BACKENDS, Int8MoeBackend.CPU) - - return _AVAILABLE_BACKENDS - def backend_to_kernel_cls( backend: Int8MoeBackend, @@ -78,14 +69,13 @@ def backend_to_kernel_cls( HummingGroupedExperts, HummingIndexedExperts, ] - elif backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + ArmCPUExpertsInt8, CPUExpertsInt8, ) - return [CPUExpertsInt8] - + return [ArmCPUExpertsInt8, CPUExpertsInt8] else: raise ValueError(f"Unknown Int8 MoE backend: {backend.value}") @@ -176,7 +166,9 @@ def _return_or_raise( logger.debug_once(_make_log_unsupported(backend, reason)) raise NotImplementedError( - "No Int8 MoE backend supports the deployment configuration." + "No Int8 MoE backend supports the deployment configuration " + f"(weight_key={weight_key}, activation_key={activation_key}). " + "Set `VLLM_LOGGING_LEVEL=DEBUG` to see per-backend unsupported reasons." ) @@ -274,13 +266,7 @@ def convert_to_int8_moe_kernel_format( 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, - ) - - w13, w2 = prepare_int8_moe_layer_for_cpu(w13, w2) - elif int8_backend != Int8MoeBackend.TRITON: + elif int8_backend not in (Int8MoeBackend.TRITON, Int8MoeBackend.CPU): raise ValueError(f"Unsupported Int8 MoE backend: {int8_backend.value}") return w13, w2 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 a80013f65baa..2ac53525c9c2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -24,11 +24,15 @@ MarlinExperts, MarlinExpertsBase, ) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( + TritonWNA16Experts, +) from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( TrtLlmMxint4ExpertsMonolithic, ) from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_moe_marlin_supports_config, marlin_act_int8_process_scales, marlin_moe_padded_intermediate, marlin_moe_permute_scales, @@ -50,6 +54,7 @@ class WNA16MoEBackend(Enum): HUMMING = "HUMMING" CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" + TRITON = "TRITON" XPU = "XPU" EMULATION = "EMULATION" @@ -76,6 +81,8 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return [TrtLlmMxint4ExpertsMonolithic] + elif backend == WNA16MoEBackend.TRITON: + return [TritonWNA16Experts] elif backend == WNA16MoEBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, @@ -107,19 +114,75 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] - _AVAILABLE_BACKENDS = [ + return [ WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.TRITON, WNA16MoEBackend.HUMMING, WNA16MoEBackend.EMULATION, ] - return _AVAILABLE_BACKENDS + + +def _backend_incompatibility_reason( + backend: WNA16MoEBackend, + moe_config: FusedMoEConfig, + quant_config: QuantizationConfig | QuantizationArgs, + may_have_zp: bool, + may_have_bias: bool, + allow_tile_padding: bool, +) -> str | None: + if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias): + return "zero points and bias are not supported" + + 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.moe_wna16 import MoeWNA16Config + + if backend == WNA16MoEBackend.TRITON: + if may_have_bias: + return "expert bias is not supported" + if isinstance(quant_config, AutoAWQConfig): + return "the AutoAWQ weight layout is not supported" + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act: + return "GPTQ activation ordering is not supported" + if ( + isinstance(quant_config, QuantizationArgs) + and quant_config.actorder == "group" + ): + return "group activation ordering is not supported" + + # Marlin only supports certain problem/group sizes. + allow_marlin = not isinstance(quant_config, MoeWNA16Config) + + if allow_marlin and backend in ( + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + ): + if isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, QuantizationArgs)): + group_size = quant_config.group_size + else: + return "Marlin not supported for this layer" + + if not check_moe_marlin_supports_config( + moe_config, group_size, allow_tile_padding + ): + return "Marlin not supported for this layer" + + if not allow_marlin and backend in ( + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.EMULATION, + ): + return "the MoeWNA16 checkpoint layout is not supported" + + return None def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: """Map user's MoEBackend to WNA16MoEBackend.""" mapping = { + "triton": WNA16MoEBackend.TRITON, "marlin": WNA16MoEBackend.MARLIN, "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, @@ -136,6 +199,10 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, + quant_config: QuantizationConfig | QuantizationArgs, + may_have_zp: bool, + may_have_bias: bool, + allow_tile_padding: bool = False, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -143,6 +210,9 @@ def select_wna16_moe_backend( config: the shared ``FusedMoEConfig`` for this layer. weight_key: The QuantKey describing the weight quantization. Must have int4 or int8 type. + quant_config: Quantization structure and checkpoint format description. + may_have_zp: Whether the integration can provide weight zero points. + may_have_bias: Whether the integration can provide expert bias. Returns: A tuple of (``WNA16MoEBackend``, experts class or ``None``). @@ -189,6 +259,16 @@ def _return_or_raise( runner_backend = config.moe_backend if runner_backend != "auto": requested_backend = map_wna16_backend(runner_backend) + reason = _backend_incompatibility_reason( + requested_backend, + config, + quant_config, + may_have_zp, + may_have_bias, + allow_tile_padding, + ) + if reason is not None: + raise ValueError(_make_log_unsupported(requested_backend, reason)) return _return_or_raise( requested_backend, config, weight_key, None, activation_format ) @@ -197,6 +277,17 @@ def _return_or_raise( AVAILABLE_BACKENDS = _get_priority_backends() for backend in AVAILABLE_BACKENDS: + reason = _backend_incompatibility_reason( + backend, + config, + quant_config, + may_have_zp, + may_have_bias, + allow_tile_padding, + ) + if reason is not None: + logger.debug_once(_make_log_unsupported(backend, reason), scope="local") + continue activation_key = None # always BF16 activation for WNA16 MoE for k_cls in backend_to_kernel_cls(backend): supported, reason = k_cls.is_supported_config( @@ -294,6 +385,7 @@ def make_wna16_moe_kernel( allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, + TritonWNA16Experts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, @@ -315,6 +407,7 @@ def make_wna16_moe_kernel( assert prepare_finalize is not None logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") + logger.info_once("Using %s", experts_cls.__name__, scope="local") extra_args: dict[str, Any] = {} if backend == WNA16MoEBackend.HUMMING: @@ -329,35 +422,17 @@ def make_wna16_moe_kernel( "is_k_full": is_k_full, } - if experts_cls is XPUExpertsWNA16: - assert ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard - ), ( - "XPUExpertsWNA16 only supports the Standard activation format; " - "xpu_fused_moe(is_int4=True) does not implement BatchedExperts." - ) - experts: mk.FusedMoEExperts = XPUExpertsWNA16( - moe_config=moe_config, - quant_config=moe_quant_config, - ) - elif ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts - ): + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None - experts = experts_cls( - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) - else: - experts = experts_cls( - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) + extra_args["max_num_tokens"] = max_num_tokens + extra_args["num_dispatchers"] = prepare_finalize.num_dispatchers() + + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + **extra_args, + ) return mk.FusedMoEKernel( prepare_finalize, @@ -972,7 +1047,7 @@ def _process_weights_xpu( w2: [E, K, N] int4 (uint8 storage [E, K, N // 2]) w2_scales: [E, K, N // group_size] params_dtype - Input GPTQ layout from FusedMoE.weight_loader: + Input GPTQ layout from MoERunner.weight_loader: w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim) w13_scales: [E, K // group_size, 2*N] params_dtype w2: [E, N // 8, K] int32 @@ -1034,11 +1109,68 @@ def _humming_wna16_weight_schema( "sym": quant_config.is_sym, } raise TypeError( - "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, " + "Humming WNA16 checkpoint schema requires AutoAWQConfig or " + "AutoGPTQConfig, " f"got {type(quant_config).__name__}." ) +def _convert_moe_wna16_humming_tensors( + tensors: dict[str, torch.Tensor], has_zero_point: bool +) -> dict[str, torch.Tensor]: + """Convert MoeWNA16's N-first uint8 packing to Humming's int32 packing.""" + if sys.byteorder != "little": + raise NotImplementedError( + "MoeWNA16 to Humming conversion requires a little-endian host." + ) + + output = { + "weight": tensors["qweight"].contiguous().view(torch.int32), + "weight_scale": tensors["scales"], + } + if has_zero_point: + qzeros = tensors["qzeros"] + output["zero_point"] = ( + qzeros.transpose(-1, -2) + .contiguous() + .view(torch.int32) + .transpose(-1, -2) + .contiguous() + ) + return output + + +class _MoeWNA16HummingWeightSchema: + """Adapter from MoeWNA16's generic packed layout to Humming's layout.""" + + def __init__(self, bits: int, group_size: int, has_zero_point: bool) -> None: + self.bits = bits + self.group_size = group_size + self.has_zero_point = has_zero_point + + def convert_humming( + self, + tensors: dict[str, torch.Tensor], + shape_n_stacks: list[int], + shape_k_stacks: list[int], + param_dtype: torch.dtype, + num_experts: int | None = None, + ) -> tuple[Any, dict[str, torch.Tensor]]: + del shape_n_stacks, shape_k_stacks, num_experts + from vllm.utils.humming import HummingWeightSchema, dtypes + + output = _convert_moe_wna16_humming_tensors( + tensors, has_zero_point=self.has_zero_point + ) + output["weight_scale"] = output["weight_scale"].to(param_dtype) + schema = HummingWeightSchema( + b_dtype=dtypes.DataType.from_str(f"uint{self.bits}"), + weight_scale_group_size=self.group_size, + has_zero_point=self.has_zero_point, + ) + return schema, output + + def _unpack_and_dequant_int4_gptq( w_int32: torch.Tensor, scale: torch.Tensor, @@ -1318,18 +1450,32 @@ def convert_to_wna16_moe_kernel_format( Args: backend: the selected ``WNA16MoEBackend``. - layer: the ``FusedMoE`` layer whose parameters are being prepared. + layer: the ``MoERunner`` layer whose parameters are being prepared. 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.moe_wna16 import MoeWNA16Config 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) - ) + if isinstance(quant_config, MoeWNA16Config): + from vllm.utils.humming import HummingInputSchema + + convert_to_humming_moe_kernel_format( + layer, + weight_schema=_MoeWNA16HummingWeightSchema( + bits=quant_config.weight_bits, + group_size=layer.group_size, + has_zero_point=quant_config.has_zp, + ), + input_schema=HummingInputSchema(), + ) + else: + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_wna16_weight_schema(quant_config) + ) return None if backend in ( @@ -1344,16 +1490,32 @@ def convert_to_wna16_moe_kernel_format( ) if isinstance(quant_config, AutoAWQConfig): - if w13_qzeros is None or w2_qzeros is None: - raise ValueError("AWQ Marlin MoE requires zero-point tensors.") - - weight_bits = quant_config.weight_bits + num_bits = quant_config.weight_bits pack_factor = quant_config.pack_factor group_size = quant_config.group_size + elif isinstance(quant_config, AutoGPTQConfig): + num_bits = quant_config.quant_type.size_bits + pack_factor = quant_config.pack_factor + group_size = quant_config.group_size + actorder = "group" if quant_config.desc_act else None + elif isinstance(quant_config, QuantizationArgs): + num_bits = quant_config.num_bits + pack_factor = 32 // quant_config.num_bits + group_size = quant_config.group_size + actorder = quant_config.actorder + else: + raise TypeError( + "Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or " + f"QuantizationArgs, got {type(quant_config).__name__}." + ) + + if isinstance(quant_config, AutoAWQConfig): + if w13_qzeros is None or w2_qzeros is None: + raise ValueError("AWQ Marlin MoE requires zero-point tensors.") return _process_awq_weights_marlin( layer, - weight_bits, + num_bits, pack_factor, group_size, input_dtype, @@ -1366,41 +1528,28 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - elif isinstance(quant_config, AutoGPTQConfig): - num_bits = quant_config.quant_type.size_bits - pack_factor = quant_config.pack_factor - group_size = quant_config.group_size - actorder = "group" if quant_config.desc_act else None - elif isinstance(quant_config, QuantizationArgs): - num_bits = quant_config.num_bits - pack_factor = 32 // quant_config.num_bits - group_size = quant_config.group_size - actorder = quant_config.actorder else: - raise TypeError( - "Marlin WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig or " - f"QuantizationArgs, got {type(quant_config).__name__}." + if w13_g_idx is None or w2_g_idx is None: + raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") + + return _process_weights_marlin( + layer, + input_dtype, + num_bits, + pack_factor, + group_size, + actorder, + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + w13_qzeros, + w2_qzeros, + w13_bias, + w2_bias, ) - if w13_g_idx is None or w2_g_idx is None: - raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") - return _process_weights_marlin( - layer, - input_dtype, - num_bits, - pack_factor, - group_size, - actorder, - w13, - w2, - w13_scale, - w2_scale, - w13_g_idx, - w2_g_idx, - w13_qzeros, - w2_qzeros, - w13_bias, - w2_bias, - ) elif backend == WNA16MoEBackend.CPU: return _process_weights_cpu( quant_config, @@ -1482,5 +1631,47 @@ def convert_to_wna16_moe_kernel_format( w13_qzeros, w2_qzeros, ) + elif backend == WNA16MoEBackend.TRITON: + # Two possible input layouts depending on the quantization source: + # + # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first + # → just view as uint8 (no-op) + # + # AutoGPTQ/compressed-tensors (int32, K-first): + # (E, K // pack32, N_out) + # → transpose to N-first, then view as uint8 to get + # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] + # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + + if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)): + # These integrations build in K-first format even when the Triton + # backend is selected. Transpose to N-first first. + w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) + w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) + w13_scale = w13_scale.transpose(1, 2).contiguous() + w2_scale = w2_scale.transpose(1, 2).contiguous() + else: + # MoeWNA16 uses N-first uint8 weights and scales. + w13_uint8 = w13.view(torch.uint8) + w2_uint8 = w2.view(torch.uint8) + return ( + w13_uint8, + w2_uint8, + w13_scale, + w2_scale, + None, + None, + None, + None, + w13_qzeros, + w2_qzeros, + None, + None, + w13_bias, + w2_bias, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 921e8f114d5a..1451a0373c63 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -222,11 +222,14 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( + AiterW4A16ExpertsMonolithic, + ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( AiterExperts, ) - return [AiterExperts] + return [AiterExperts, AiterW4A16ExpertsMonolithic] elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( @@ -334,7 +337,10 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: backend-level ``is_supported_config`` check filters by device capability). """ if current_platform.is_rocm(): - return [Mxfp4MoeBackend.AITER_MXFP4_BF16] + return [ + Mxfp4MoeBackend.AITER_MXFP4_BF16, + Mxfp4MoeBackend.EMULATION, + ] if current_platform.is_xpu(): return [Mxfp4MoeBackend.XPU] _AVAILABLE_BACKENDS = [ @@ -502,6 +508,7 @@ def select_mxfp4_moe_backend( _get_priority_backends_for_gpt_oss(), requested_activation_key ) + unsupported_reasons = [] for backend in AVAILABLE_BACKENDS: # Use requested_activation_key if provided, otherwise use backend default act_key = ( @@ -518,6 +525,7 @@ def select_mxfp4_moe_backend( return backend, k_cls else: logger.debug_once(_make_log_unsupported(backend, reason)) + unsupported_reasons.append((backend, reason)) if current_platform.is_xpu(): backend = Mxfp4MoeBackend.XPU @@ -541,26 +549,19 @@ def select_mxfp4_moe_backend( activation_format, ) - if current_platform.is_rocm(): - backend = Mxfp4MoeBackend.TRITON_UNFUSED - logger.info_once(_make_log_backend(backend)) - return _return_or_raise( - Mxfp4MoeBackend.TRITON_UNFUSED, - config, - kMxfp4Static, - None, - activation_format, - ) - - if current_platform.is_cuda(): - raise NotImplementedError( - "No MXFP4 MoE backend supports the deployment configuration. " - f"weight_key=kMxfp4Static, activation_key={activation_key}. " - "Native backends require specific hardware. " - "Set `VLLM_LOGGING_LEVEL=DEBUG` to see detailed unsupported reasons. " - ) - - return Mxfp4MoeBackend.NONE, None + unsupported_log = "; ".join( + [ + f"backend: {backend.value}, reason: {reason}" + for backend, reason in unsupported_reasons + ] + ) + raise NotImplementedError( + "No MXFP4 MoE backend supports the deployment configuration. " + f"weight_key=kMxfp4Static, activation_key={activation_key}. " + f"Candidate backends were: " + f"{[backend.value for backend in AVAILABLE_BACKENDS]}. " + f"Unsupported reasons: {unsupported_log}. " + ) def select_deepseek_v4_mxfp4_moe_backend( @@ -653,7 +654,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( else: hidden_size = round_up(hidden_size, 256) elif backend in TRTLLM_BACKENDS: - intermediate_size = round_up(intermediate_size, 256) + intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 256) elif backend in ( Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, @@ -1268,6 +1269,11 @@ def convert_weight_to_mxfp4_moe_kernel_format( Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends. """ + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( @@ -1440,7 +1446,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250: # Initially introduced for DeepSeekV4 if w13_bias is not None: @@ -1497,7 +1503,9 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend in TRITON_BACKENDS: + elif mxfp4_backend in TRITON_BACKENDS or ( + mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 + ): from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig if mxfp4_backend == Mxfp4MoeBackend.TRITON: @@ -1554,8 +1562,12 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor: w13_bias, w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.XPU: - # No additional transformation needed for XPU backend + elif mxfp4_backend in ( + Mxfp4MoeBackend.XPU, + Mxfp4MoeBackend.EMULATION, + ): + # No additional transformation is needed: XPU consumes the checkpoint + # layout directly, while emulation dequantizes that layout at runtime. return ( w13_weight, w2_weight, @@ -1567,7 +1579,7 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor: else: raise ValueError( f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. " - f"Expected TRTLLM, Triton, AITER, or XPU backend." + f"Expected TRTLLM, Triton, AITER, XPU, or emulation backend." ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index b9086cfa48aa..2a03eacea8e1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -12,10 +12,10 @@ kMxfp8Dynamic, kMxfp8Static, ) -from vllm.platforms import current_platform logger = init_logger(__name__) +# Ordered by priority. _SUPPORTED_BACKENDS = ( Fp8MoeBackend.FLASHINFER_TRTLLM, Fp8MoeBackend.DEEPGEMM, @@ -26,6 +26,8 @@ # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, Fp8MoeBackend.HUMMING, + Fp8MoeBackend.TRITON_MXFP8, + Fp8MoeBackend.EMULATION, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls( return [AiterMxfp8Experts] if backend == Fp8MoeBackend.TRITON_MXFP8: - # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e. - # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise. - # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``. - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + Mxfp8NativeTritonExperts, + ) - return [Mxfp8NativeTritonExperts] + return [Mxfp8NativeTritonExperts] + if backend == Fp8MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, ) @@ -105,35 +104,6 @@ def _select_kernel_cls( ) -def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: - """ROCm fallback when no auto-selected MXFP8 backend is available. - - The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by - ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or - explicitly via ``--moe-backend aiter``; this fallback handles the rest - (native dot_scaled on gfx950, else BF16 emulation). - """ - - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) - - logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.") - return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts - - from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( - Mxfp8EmulationTritonExperts, - ) - - logger.info_once( - "No native MXFP8 MoE backend available on this device; " - "MXFP8 weights will be dequantized to BF16 once at load time and the " - "MoE will run in BF16 (no per-step dequant)." - ) - return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts - - def select_mxfp8_moe_backend( config: FusedMoEConfig, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: @@ -167,8 +137,5 @@ def select_mxfp8_moe_backend( logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) return backend, experts_cls - # simplify the logic for rocm, refactor later when more backends are supported - if current_platform.is_rocm(): - return _select_rocm_mxfp8_backend() - + # TODO: add debug log with reason. raise ValueError("No MXFP8 MoE backends available.") diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 5d4c73363135..bf77a316bb49 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -151,6 +151,7 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend """Map user's MoEBackend to UnquantizedMoeBackend.""" mapping = { "triton": UnquantizedMoeBackend.TRITON, + "batched_triton": UnquantizedMoeBackend.BATCHED_TRITON, "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS, "aiter": UnquantizedMoeBackend.AITER, @@ -233,6 +234,7 @@ def select_unquantized_moe_backend( activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if moe_config.moe_parallel_config.use_batched_activation_format + or moe_config.moe_backend == "batched_triton" else mk.FusedMoEActivationFormat.Standard ) @@ -360,6 +362,13 @@ def make_unquantized_moe_kernel( experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: + from vllm.model_executor.layers.fused_moe.utils import ( + warn_if_moe_use_td_ineffective, + ) + + # Warn against the selected backend, not each probed candidate. + warn_if_moe_use_td_ineffective(backend.value, is_quantized=False) + # Create Prepare/Finalize is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) prepare_finalize = maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py index 943027717bbe..d11beaf98fbd 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py @@ -102,7 +102,7 @@ def prepare( num_local_experts, self.max_num_tokens, hidden_dim ) - b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device) + b_a1_scale = torch.zeros(scale_shape, dtype=torch.float32, device=a1.device) else: assert quant_config.a1_scale is None b_a1_scale = None @@ -112,15 +112,28 @@ def prepare( a1_scale = normalize_scales_shape(quant_config.a1_scale) - for expert_id in range(first_expert, last_expert): - topks = torch.any(topk_ids == expert_id, dim=1).flatten() - rows = torch.count_nonzero(topks.flatten()) - if rows == 0: - continue - idx = expert_id - first_expert - tokens_per_expert[idx] = rows - rhs = a1[: topks.numel()][topks] - if quant_config.quant_dtype is not None: + if quant_config.quant_dtype is None: + # Vectorized dispatch: compute the [E_local, T] hit mask and + # per-expert slot offsets in one shot (single nonzero sync instead + # of a launch-bound per-expert Python loop). + local_ids = torch.arange(first_expert, last_expert, device=a1.device).view( + -1, 1, 1 + ) + hits = (topk_ids.unsqueeze(0) == local_ids).any(dim=2) # [E_local, T] + tokens_per_expert[:num_local_experts] = hits.sum(dim=1).to(torch.int32) + # slot of each token within its expert batch, preserving token order + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T] + e_idx, t_idx = hits.nonzero(as_tuple=True) + b_a1[e_idx, slots[e_idx, t_idx]] = a1[t_idx].to(b_type) + else: + for expert_id in range(first_expert, last_expert): + topks = torch.any(topk_ids == expert_id, dim=1).flatten() + rows = torch.count_nonzero(topks.flatten()) + if rows == 0: + continue + idx = expert_id - first_expert + tokens_per_expert[idx] = rows + rhs = a1[: topks.numel()][topks] if a1_scale is not None: if quant_config.is_per_act_token: rhs_a1_scale = a1_scale[: topks.numel()][topks] @@ -140,8 +153,6 @@ def prepare( b_a1_scale[idx, :rows] = b_s[:rows] else: b_a1_scale[idx, : b_s.shape[0]] = b_s - else: - b_a1[idx, :rows, :] = rhs assert b_a1_scale is None or b_a1_scale.ndim == 3 diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index 89571278c6e6..fd7d45655dd6 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -160,11 +160,6 @@ def _map_global_to_physical_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: return topk_ids return self.global_to_physical[topk_ids] - def _map_local_to_global_ids(self, expert_topk_ids: torch.Tensor) -> torch.Tensor: - if self.local_expert_global_ids is None: - return expert_topk_ids - return self.local_expert_global_ids[expert_topk_ids] - def _do_quant( self, x: torch.Tensor | tuple[torch.Tensor, torch.Tensor], diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 80e4064ea6b6..0cd0ad995078 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -114,6 +114,7 @@ def __init__( self.e_score_correction_bias = e_score_correction_bias self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters + self._loaded_expert_biases: set[str] = set() self.quant_method = self._get_quant_method( self.layer_name, @@ -205,7 +206,6 @@ def _get_quant_method( def _needs_intermediate_size_param(self, quant_method: FusedMoEMethodBase) -> bool: return quant_method.__class__.__name__ in ( "AutoGPTQMoEMethod", - "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", "CompressedTensorsW4A16FlydslMoEMethod", ) @@ -232,17 +232,27 @@ def update_expert_map_info(self): # Update local attributes from ExpertMapManager self.local_num_experts = self.expert_map_manager.local_num_experts self.expert_placement_strategy = self.expert_map_manager.placement_strategy - self.register_buffer("_expert_map", self.expert_map_manager.expert_map) - self.register_buffer("expert_mask", self.expert_map_manager.expert_mask) + self.register_buffer( + "_expert_map", self.expert_map_manager.expert_map, persistent=False + ) + self.register_buffer( + "expert_mask", self.expert_map_manager.expert_mask, persistent=False + ) # Get routing tables from ExpertMapManager routing_tables = self.expert_map_manager.routing_tables if routing_tables is not None: # Register routing tables as buffers for this layer global_to_physical, physical_to_global, local_global = routing_tables - self.register_buffer("expert_global_to_physical", global_to_physical) - self.register_buffer("expert_physical_to_global", physical_to_global) - self.register_buffer("expert_local_to_global", local_global) + self.register_buffer( + "expert_global_to_physical", global_to_physical, persistent=False + ) + self.register_buffer( + "expert_physical_to_global", physical_to_global, persistent=False + ) + self.register_buffer( + "expert_local_to_global", local_global, persistent=False + ) def _expert_routing_tables( self, @@ -411,6 +421,28 @@ def _get_hidden_dim(shard_dim: int, ndim: int) -> int: f"for a {ndim}D tensor (expected {dim_a} or {dim_b})" ) + @staticmethod + def _orient_fused_weight( + fused_weight: torch.Tensor, shard_id: str, unpadded_hidden: int + ) -> torch.Tensor: + """Normalise a fused expert tensor from either checkpoint orientation + to (intermediate, hidden) for w1/w3 and (hidden, intermediate) for w2. + + Only transposes when the hidden dim is definitively on the wrong axis, + leaving tensors that have no hidden dim alone (e.g. a fused per-channel + scale, (2 * intermediate, 1)). + """ + if shard_id == "w2": + hidden_axis, intermediate_axis = -2, -1 + else: + hidden_axis, intermediate_axis = -1, -2 + if ( + fused_weight.shape[hidden_axis] != unpadded_hidden + and fused_weight.shape[intermediate_axis] == unpadded_hidden + ): + return fused_weight.transpose(-1, -2) + return fused_weight + @staticmethod def _narrow_expert_data_for_padding( expert_data: torch.Tensor, @@ -626,7 +658,6 @@ def weight_loader( # TODO (mgoin): check self.quant_method.quant_config.quant_format # against known CompressionFormat enum values that have this quality if quant_method_name in ( - "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", "CompressedTensorsWNA16RDNA3MoEMethod", "CompressedTensorsW4A16FlydslMoEMethod", @@ -691,6 +722,25 @@ def weight_loader( expert_data = param.data if full_load else param.data[expert_id] + if "bias" in weight_name: + self._loaded_expert_biases.add(weight_name.rsplit(".", 1)[-1]) + if shard_id == "w2": + expert_data = self._narrow_expert_data_for_padding( + expert_data, + loaded_weight, + hidden_dim=0, + ) + expert_data.copy_(loaded_weight) + else: + self._load_w13( + shard_id=shard_id, + shard_dim=0, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=self.moe_config.tp_rank, + ) + return True if return_success else None + # Case input scale: input_scale loading is only supported for fp8 if "input_scale" in weight_name: # this is needed for compressed-tensors only @@ -870,7 +920,9 @@ def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[str]: expert_mapping = self.get_expert_mapping(include_fused=True) - unpadded_hidden = self.moe_config.hidden_dim_unpadded + unpadded_hidden = ( + self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim + ) for expert_name, loaded_weight in weights: qual_name = f"{self.layer_name}.{expert_name}" # Fused expert weights can be identified by their 3D tensors @@ -889,17 +941,13 @@ def load_weights( # w1 and w3 share one fused tensor; use a local copy so the # transpose below doesn't mutate loaded_weight across # iterations (else w3 is transposed twice and wrongly chunked) - fused_weight = loaded_weight + fused_weight = self._orient_fused_weight( + loaded_weight, shard_id, unpadded_hidden + ) if shard_id in {"w1", "w3"}: - if fused_weight.shape[-1] != unpadded_hidden: - # [..., hidden, intermediate] -> [..., intermediate, hidden] - fused_weight = fused_weight.transpose(-1, -2) # Repurpose expert_id for deconcatenating w1 and w3 experts_shard = fused_weight.chunk(2, dim=1)[expert_id] else: - if fused_weight.shape[-2] != unpadded_hidden: - # [..., intermediate, hidden] -> [..., hidden, intermediate] - fused_weight = fused_weight.transpose(-1, -2) experts_shard = fused_weight start = 0 else: diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index 115f43a58bfd..c71ee30fb048 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -90,10 +90,20 @@ def __init__( ) -> None: hf_config = vllm_config.model_config.hf_text_config num_experts_per_tok = _get_num_experts_per_tok(hf_config) + num_layers = hf_config.num_hidden_layers + logger.info( + "RoutedExpertsCapturer: allocating buffer with " + "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d " + "(hf_config.model_type=%s)", + max_num_batched_tokens, + num_layers, + num_experts_per_tok, + getattr(hf_config, "model_type", "unknown"), + ) self.device_buffer = torch.zeros( ( max_num_batched_tokens, - hf_config.num_hidden_layers, + num_layers, num_experts_per_tok, ), # Use int32 for the device / host transit buffers: it @@ -170,7 +180,7 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: # (trailing rows are SP ceil-div padding). The TP group # is always initialized on real rollout workers, and # every rank in the group reaches this branch in - # lockstep (bind is per-FusedMoE layer, SP is a global + # lockstep (bind is per-FusedMoEFactory layer, SP is a global # condition), so a bare all_gather here will not # deadlock -- let it raise if the precondition is # violated rather than skip silently. diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index f89eb4910b35..bc019a58e2e2 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -9,6 +9,7 @@ import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -20,6 +21,13 @@ ) +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + return is_padding[:num_tokens] if is_padding is not None else None + return None + + def vllm_topk_softmax( topk_weights: torch.Tensor, topk_indices: torch.Tensor, @@ -35,6 +43,7 @@ def vllm_topk_softmax( gating_output, renormalize, e_score_correction_bias, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -57,6 +66,7 @@ def vllm_topk_sigmoid( renormalize, e_score_correction_bias, routed_scaling_factor, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -144,6 +154,7 @@ def vllm_topk_softplus_sqrt( e_score_correction_bias, input_tokens, hash_indices_table, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index 855fa606565b..28f67b5c6931 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py @@ -5,8 +5,10 @@ import torch import vllm._custom_ops as ops +import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -14,6 +16,13 @@ from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + return is_padding[:num_tokens] if is_padding is not None else None + return None + + def vllm_topk_softmax( topk_weights: torch.Tensor, topk_indices: torch.Tensor, @@ -27,6 +36,7 @@ def vllm_topk_softmax( token_expert_indices, gating_output, renormalize, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -45,6 +55,7 @@ def vllm_topk_sigmoid( token_expert_indices, gating_output, renormalize, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index c7cfccbe64b9..ff4874b20e61 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -9,6 +9,7 @@ from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, + get_routing_method_type, ) from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501 AiterSharedRoutedFusedMoERouter, @@ -67,7 +68,8 @@ def create_fused_moe_router( The selection logic follows this priority order: 1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set 2. ZeroExpertRouter - if zero_expert_type is not None - 3. GroupedTopKRouter - if use_grouped_topk is True + 3. GroupedTopKRouter - if use_grouped_topk is True and the grouping is not + degenerate (at most one group, with topk_group <= 1) 4. CustomRoutingRouter - if custom_routing_function is not None 5. FusedTopKBiasRouter - if e_score_correction_bias is not None 6. AiterSharedRoutedFusedMoERouter - if num_fused_shared_experts > 0 @@ -143,30 +145,48 @@ def create_fused_moe_router( "num_expert_group and topk_group must be provided when " "use_grouped_topk is True" ) - grouped_topk_router = GroupedTopKRouter( - top_k=top_k, - global_num_experts=global_num_experts, - eplb_state=eplb_state, - num_expert_group=num_expert_group, - topk_group=topk_group, - renormalize=renormalize, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=num_fused_shared_experts, + + # For topk_group <= 1, grouped implementation is pure overhead. + degenerate_grouping = num_expert_group <= 1 and topk_group <= 1 + # FusedTopKRouter cannot apply routed_scaling_factor, FusedTopKBiasRouter can. + scaling_handled_downstream = ( + routed_scaling_factor == 1.0 or e_score_correction_bias is not None ) - if ( - grouped_topk_router.routing_method_type != RoutingMethodType.Unspecified - or num_expert_group > 1 - or topk_group > 1 - ): - return grouped_topk_router - # If routing_method for GroupedTopKRouter is Unspecified and there is only - # one group, fallback to standard top-k routing - use_grouped_topk = False - num_expert_group = None - topk_group = None + # Degenerating must not change the advertised routing method, which drives + # kernel selection. num_expert_group only affects it for biased routing. + def advertised_routing_method(groups: int | None) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=scoring_func, + top_k=top_k, + renormalize=renormalize, + num_expert_group=groups, + has_e_score_bias=e_score_correction_bias is not None, + routed_scaling_factor=routed_scaling_factor, + ) + + routing_method_preserved = advertised_routing_method( + num_expert_group + ) == advertised_routing_method(None) + + if not ( + degenerate_grouping + and scaling_handled_downstream + and routing_method_preserved + ): + return GroupedTopKRouter( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=num_fused_shared_experts, + ) + # Otherwise fall through to the non-grouped chain below. if custom_routing_function is not None: return CustomRoutingRouter( diff --git a/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py new file mode 100644 index 000000000000..9daac9fb53d5 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.utils.torch_utils import aux_stream, current_stream + +from .moe_runner import MoERunner, _unpack + +logger = init_logger(__name__) + + +class LatentMoERunner(MoERunner): + """MoE runner for latent MoE with a replicated routed up-projection. + + Fused path (tp>1, un-reduced combine output, shared expert, no SP): + concatenates the un-reduced latent partial (dim d) and the un-reduced + shared partial (dim D) into one contiguous buffer, all-reduces once, then + splits. The latent half is normed and up-projected locally (replicated + up-proj -> full hidden), and the shared add folds into the GEMM epilogue + (``torch.addmm``). One collective total, no post-reduction communication. + + Native path: the replicated up-proj produces the full hidden dim on every + rank, so the base runner combines routed + shared correctly at any TP size + (using two collectives instead of the fused path's one). + """ + + def __init__( + self, + *args, + enable_k3_latent_moe_tail_fusion: bool = False, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion + use_fused_path = self._use_fused_path() + if ( + self.enable_k3_latent_moe_tail_fusion + and use_fused_path + and self.moe_config.tp_size not in (8, 16) + ): + logger.warning_once( + "K3 latent-MoE tail fusion currently supports TP=8 and TP=16, " + "but TP=%d is configured. Falling back to the default path.", + self.moe_config.tp_size, + ) + self.enable_k3_latent_moe_tail_fusion = False + + if self.enable_k3_latent_moe_tail_fusion and use_fused_path: + vllm_config = get_current_vllm_config() + if vllm_config.parallel_config.use_ubatching: + raise ValueError( + "K3 latent-MoE tail fusion does not support DBO or ubatching." + ) + if vllm_config.model_config.enable_sleep_mode: + raise ValueError( + "K3 latent-MoE tail fusion does not support sleep mode." + ) + transform = self.routed_output_transform + assert transform is not None + norm = transform.norm + assert norm is not None + from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( + KimiK3LatentMoETailOp, + ) + + op = KimiK3LatentMoETailOp.initialize( + hidden_size=transform.up_proj.weight.shape[0], + latent_size=norm.weight.shape[0], + dtype=norm.weight.dtype, + device=norm.weight.device, + rms_eps=norm.variance_epsilon, + ) + self._k3_latent_moe_tail_op = op + + def _get_zero_residual( + self, + hidden_states: torch.Tensor, + max_token_num: int, + ) -> torch.Tensor: + """Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel. + + flashinfer requires a residual buffer even when there is no residual to + add. + """ + buf = getattr(self, "_zero_residual", None) + if buf is None: + buf = torch.zeros( + max_token_num * hidden_states.shape[-1], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._zero_residual = buf + + assert buf.dtype == hidden_states.dtype + assert buf.device == hidden_states.device + assert hidden_states.numel() <= buf.numel() + + return buf[: hidden_states.numel()].view_as(hidden_states) + + def _use_fused_path(self) -> bool: + # The fused path merges the latent and shared reductions into one + # all-reduce, so it needs actual TP parallelism, a shared expert (to + # concat), an un-reduced combine output, and no sequence parallelism. + return ( + self.moe_config.tp_size > 1 + and self._shared_experts is not None + and not self._fused_output_is_reduced + and not self.moe_config.is_sequence_parallel + ) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + if self._use_fused_path(): + return self._fused_forward( + hidden_states, router_logits, input_ids, shared_experts_input + ) + return super().forward( + hidden_states, router_logits, input_ids, shared_experts_input + ) + + def _fused_forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + # When the caller pre-applies the routed input transform outside the + # runner (e.g. to overlap it on a separate stream), it passes the + # already-transformed routed input as ``hidden_states`` and the original + # hidden states as ``shared_experts_input``; skip the transform then. + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = ( + self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + ) + + result = self._forward_entry( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + self._encode_layer_name(), + self.moe_config.hidden_dim_unpadded + if self._quant_method.has_unpadded_output + else 0, + ) + + shared_output, fused_output = _unpack(result) + assert shared_output is not None + + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] + + transform = self.routed_output_transform + assert transform is not None + + if self.enable_k3_latent_moe_tail_fusion: + op = self._k3_latent_moe_tail_op + if 0 < fused_output.shape[0] <= op.contract.max_num_tokens: + norm = transform.norm + assert norm is not None + result = op( + fused_output, + shared_output, + norm.weight, + transform.up_proj.weight, + ) + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, output_is_reduced=True + ) + return self._maybe_add_zero_expert_output(result) + + fused_latent = None + if transform.norm is not None: + fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm) + else: + fused_latent = tensor_model_parallel_all_reduce(fused_output) + + shared_expert_stream = ( + aux_stream() + if shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + else None + ) + if shared_expert_stream is not None: + # overlap shared expert allreduce with latent up_proj + main = current_stream() + shared_output.record_stream(shared_expert_stream) + shared_expert_stream.wait_stream(main) + with torch.cuda.stream(shared_expert_stream): + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + main.wait_stream(shared_expert_stream) + else: + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + result.add_(shared_output) + + # Output is already fully reduced; this only strips padding. + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, output_is_reduced=True + ) + + return self._maybe_add_zero_expert_output(result) + + def allreduce_norm_latent_out( + self, + hidden_states: torch.Tensor, + norm: RMSNorm, + ) -> tuple[torch.Tensor, torch.Tensor]: + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.""" + from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, + ) + + if self.moe_config.tp_size == 1: + return norm(hidden_states) + + if flashinfer_trtllm_fused_allreduce_norm is not None: + ok, max_token_num = _can_use_flashinfer( + hidden_states, self.moe_config.tp_size + ) + if ok: + norm_out = torch.empty_like(hidden_states) + # With norm_out provided, the kernel writes the new residual + # (all_reduce(hidden_states) + residual) into the hidden_states + # buffer and the normalized result into norm_out. + flashinfer_trtllm_fused_allreduce_norm( + allreduce_in=hidden_states, + residual=self._get_zero_residual(hidden_states, max_token_num), + rms_gamma=norm.weight, + rms_eps=norm.variance_epsilon, + world_size=self.moe_config.tp_size, + weight_bias=0.0, + launch_with_pdl=True, + fp32_acc=True, + max_token_num=max_token_num, + pattern_code=_AR_RESIDUAL_RMS_NORM, + norm_out=norm_out, + ) + return norm_out + + reduced = tensor_model_parallel_all_reduce(hidden_states) + return norm(reduced) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 7942957f3273..202ecac9e512 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -348,11 +348,12 @@ def _quant_method(self) -> FusedMoEMethodBase: return self.routed_experts.quant_method def apply_routed_input_transform( - self, hidden_states: torch.Tensor + self, + hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply transform for routed experts (e.g., latent projection). - This is called by FusedMoE.forward_native. The original hidden_states + This is called by MoERunner.forward_native. The original hidden_states is saved separately so shared experts get [S, hidden_size] while routed experts get the transformed [S, moe_latent_size]. @@ -416,6 +417,7 @@ def _fused_output_is_reduced(self) -> bool: def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, + fused_output_is_reduced: bool | None = None, ) -> torch.Tensor | None: """All-reduce shared expert output when the combine kernel already reduced fused output. @@ -425,18 +427,43 @@ def _maybe_reduce_shared_expert_output( * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled in the model. """ + if fused_output_is_reduced is None: + fused_output_is_reduced = self._fused_output_is_reduced + if ( shared_output is not None and not self.moe_config.is_sequence_parallel - and self._fused_output_is_reduced + and fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) return shared_output + def _maybe_reduce_routed_output_before_transform( + self, + fused_output: torch.Tensor, + fused_output_is_reduced: bool, + ) -> tuple[torch.Tensor, bool]: + """All-reduce latent routed output before its output transform. + + Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm. + TP partial routed outputs must be summed in latent space before such + transforms are applied. + """ + if ( + self.routed_output_transform is not None + and not self.moe_config.is_sequence_parallel + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not fused_output_is_reduced + ): + fused_output = tensor_model_parallel_all_reduce(fused_output) + fused_output_is_reduced = True + return fused_output, fused_output_is_reduced + def _maybe_reduce_final_output( self, states: torch.Tensor, trunc_size: int | None, + output_is_reduced: bool | None = None, ) -> torch.Tensor: """All-reduce the combined output if needed. @@ -455,11 +482,14 @@ def _maybe_reduce_final_output( # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. + if output_is_reduced is None: + output_is_reduced = self._fused_output_is_reduced + if ( not self.moe_config.is_sequence_parallel and not self.moe_config.skip_final_all_reduce and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - and not self._fused_output_is_reduced + and not output_is_reduced ): states = tensor_model_parallel_all_reduce(states) @@ -616,7 +646,7 @@ def _maybe_sync_shared_experts_stream( ): # If router/gate provided, then apply it here. # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via + # parallel execution of shared experts with the RoutedExperts via # separate cuda stream) if self._shared_experts is not None: assert shared_experts_input is not None @@ -643,6 +673,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: """Invoke the fused moe layer. @@ -664,11 +695,16 @@ def forward( _moe_forward and _moe_forward_shared must be split. """ - # Apply transform for routed experts (e.g., latent projection - # for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) + # Apply transform for routed experts (e.g., latent projection for + # latent MoE). When the caller pre-applies the routed input transform + # outside the runner (e.g. to overlap it on a separate stream), it + # passes the already-transformed routed input as ``hidden_states`` and + # the original hidden states as ``shared_experts_input``; skip the + # transform in that case so shared experts still see the original input. + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) # Record before `_maybe_pad_hidden_states` pads activations to match # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi` @@ -708,9 +744,22 @@ def forward( if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] - # If combine kernel already reduced fused, reduce shared to match. + fused_output_is_reduced = self._fused_output_is_reduced + + # Latent routed output has to be reduced before output transform, + # because the transform may include non-linear normalization. + fused_output, fused_output_is_reduced = ( + self._maybe_reduce_routed_output_before_transform( + fused_output, + fused_output_is_reduced, + ) + ) + + # If routed output is already reduced, reduce shared to match. # See note above re: the two all-reduce points. - shared_output = self._maybe_reduce_shared_expert_output(shared_output) + shared_output = self._maybe_reduce_shared_expert_output( + shared_output, fused_output_is_reduced + ) shared_output, fused_output = self._maybe_apply_routed_scale_to_output( shared_output, fused_output @@ -724,7 +773,9 @@ def forward( else: result = fused_output - result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform) + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, fused_output_is_reduced + ) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py index cc79095ead8a..5129c9f22db3 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py @@ -36,6 +36,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: raise NotImplementedError @@ -61,7 +62,7 @@ def _replace_quant_method(self, quant_method: FusedMoEMethodBase): ######################################################################## # - # FusedMoE layer methods + # FusedMoEFactory layer methods # ######################################################################## diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 837c1498622f..f9543358ecee 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -71,7 +71,7 @@ def apply( assert output.size() == fused_expert_output.size(), ( "output shape is expected to match the fused_expert_output shape. " f"But got output={output.size()}, " - f"used_expert_output={fused_expert_output.size()}" + f"fused_expert_output={fused_expert_output.size()}" ) output.copy_(fused_expert_output, non_blocking=True) return output @@ -164,13 +164,20 @@ def apply( first_expert = num_local_experts * self.rank last_expert = first_expert + num_local_experts - for expert_id in range(first_expert, last_expert): - matching_tokens = topk_ids == expert_id - topks = torch.any(matching_tokens, dim=1).flatten() - rows = torch.count_nonzero(topks) - rhs = fused_expert_output[expert_id - first_expert, :rows, :] - if not apply_router_weight_on_input: - rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1)) - output[topks] = output[topks] + rhs + # Vectorized weighted scatter-add (single nonzero sync instead of a + # per-expert loop; mirrors the dispatch path). + local_ids = torch.arange( + first_expert, last_expert, device=topk_ids.device + ).view(-1, 1, 1) + matching = topk_ids.unsqueeze(0) == local_ids # [E_local, T, topk] + hits = matching.any(dim=2) # [E_local, T] + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T] + e_idx, t_idx = hits.nonzero(as_tuple=True) + gathered = fused_expert_output[e_idx, slots[e_idx, t_idx], :] + if not apply_router_weight_on_input: + # weight of token t at expert e: the topk weight on the matching slot + weights = (matching * topk_weights.unsqueeze(0)).sum(dim=2) # [E_local, T] + gathered = gathered * weights[e_idx, t_idx].unsqueeze(1).to(gathered.dtype) + output.index_add_(0, t_idx, gathered.to(output.dtype)) return output diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 7a2c670a8cea..d913e9ad6765 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -139,10 +139,13 @@ def create_weights( def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor: # Pad the weight tensor. This is an optimization on ROCm platform, which - # can benefit from tensors located far enough from one another in memory + # can benefit from tensors located far enough from one another in memory. + # Skip padding when EPLB is enabled because EPLB requires contiguous + # weights for the view/rearrangement operations. if ( envs.VLLM_ROCM_MOE_PADDING and current_platform.is_rocm() + and not self.moe.moe_parallel_config.enable_eplb and weight.stride(-1) == 1 and (weight.stride(-2) * weight.element_size()) % 512 == 0 ): diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 6aa4eb0cd2aa..cce8ccd073fc 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -7,7 +7,9 @@ import torch import torch.nn.functional as F +import vllm.envs as envs from vllm import _custom_ops as ops +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -17,12 +19,14 @@ ) from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( quant_dequant_mxfp4, + xpu_mxfp4_quantize, ) from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( quant_dequant_mxfp6, ) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, + xpu_mxfp8_quantize, ) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( ref_nvfp4_quant_dequant, @@ -37,6 +41,8 @@ if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +logger = init_logger(__name__) + @triton.jit def _count_expert_num_tokens( @@ -195,6 +201,8 @@ def _mxfp4_quantize( per_act_token_quant: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, None]: + if current_platform.is_xpu(): + return xpu_mxfp4_quantize(A) assert block_shape is None # TODO: native mxfp4 is currently not integrated in vllm, # so simulating even on devices supporting this data type natively. @@ -223,6 +231,8 @@ def _mxfp8_e4m3_quantize( is_sf_swizzled_layout: bool = False, mx_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: + if current_platform.is_xpu(): + return xpu_mxfp8_quantize(A) assert A_scale is None assert not per_act_token_quant assert block_shape is None or block_shape == [1, 32] @@ -309,7 +319,7 @@ def moe_kernel_quantize_input( A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16) return A, None elif quant_dtype == "mxfp4": - if not quantization_emulation: + if not current_platform.is_xpu() and not quantization_emulation: raise NotImplementedError( "moe_kernel_quantize_input should not be used for native" " quant_dtype='mxfp4' MOE. Please open an issue." @@ -318,7 +328,7 @@ def moe_kernel_quantize_input( elif quant_dtype == "mxfp8": # TODO: `quant_dtype == "mxfp8"` is ambiguous, # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`. - if quantization_emulation: + if not current_platform.is_xpu() and quantization_emulation: raise NotImplementedError( "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE " "quantization emulation. Please open an issue." @@ -579,3 +589,80 @@ def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool: and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N >= 64 ) + + +def moe_use_td_hw_supported() -> bool: + """Whether the current device can run the TD (gather) path of + ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override). + + The A-load uses ``tensor_descriptor.gather``, which lowers to the PTX + ``tile::gather4`` instruction. That instruction is part of the + ``tcgen05``/Tensor Memory (TMEM) family introduced with Blackwell and has + no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature + '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike + ``scatter4``, ``gather4`` is supported across the whole sm100+ range + including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498, + which enables ``gather4`` on sm120/sm121 while leaving ``scatter4`` + unsupported there. So this gates on a blanket ``has_device_capability(100)`` + rather than the sm100 *family* check used for the scatter store path. + """ + if current_platform.is_xpu(): + return True + if current_platform.is_cuda(): + return current_platform.has_device_capability(100) + return False + + +def resolve_moe_use_td() -> bool: + """Tri-state resolver for ``VLLM_TRITON_USE_TD``. + + Unset auto-selects the TD path on XPU only, mirroring the attention + dispatcher in ``triton_attn.py``. ``1``/``0`` force it on/off regardless + of hardware; forcing ``1`` where it cannot compile (see + ``moe_use_td_hw_supported``) fails at ptxas. Blackwell CUDA (sm100+) can + compile it but is opt-in only, pending validation. + """ + override = envs.VLLM_TRITON_USE_TD + if override is None: + return current_platform.is_xpu() + return override + + +_warned_moe_use_td_ineffective = False + + +def warn_if_moe_use_td_ineffective( + active_backend: str, is_quantized: bool = False +) -> None: + """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored. + + Fires when the user set the env explicitly and either (a) the active + MoE backend is not the fused Triton kernel, or (b) the model is + quantized (the TD path falls back to the pointer path under any + quantization). + """ + global _warned_moe_use_td_ineffective + if _warned_moe_use_td_ineffective: + return + if envs.VLLM_TRITON_USE_TD is None: + return + is_triton = active_backend.upper() == "TRITON" + if is_triton and not is_quantized: + return + if not is_triton: + reason = ( + f"the active MoE backend is {active_backend!r}; pass " + "`--moe-backend triton` to enable the tensor-descriptor path" + ) + else: + reason = ( + "the model uses quantized MoE weights; the TD path is " + "currently restricted to non-quantized weights and falls " + "back to the pointer path" + ) + logger.warning( + "VLLM_TRITON_USE_TD is set to %s but %s.", + envs.VLLM_TRITON_USE_TD, + reason, + ) + _warned_moe_use_td_ineffective = True diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index c0dbc776acb1..e4662148ff75 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -11,6 +11,7 @@ from typing_extensions import TypeIs import vllm.envs as envs +from vllm.config import get_current_vllm_config from vllm.distributed import ( divide, get_tensor_model_parallel_rank, @@ -49,6 +50,7 @@ "UnquantizedLinearMethod", "CompressedTensorsLinearMethod", "CompressedTensorsLinearTransformMethod", + "QutlassNvFP4LinearMethod", "AutoAWQMarlinLinearMethod", "AutoAWQLinearMethod", "AutoGPTQLinearMethod", @@ -254,6 +256,8 @@ def __init__( *, return_bias: bool = True, disable_tp: bool = False, + tp_rank: int | None = None, + tp_size: int | None = None, ): super().__init__() @@ -277,8 +281,17 @@ def __init__( raise ValueError("All linear layers should support quant method.") self.return_bias = return_bias self.disable_tp = disable_tp - self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0 - self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1 + if disable_tp: + self.tp_rank, self.tp_size = 0, 1 + else: + self.tp_rank = ( + tp_rank if tp_rank is not None else get_tensor_model_parallel_rank() + ) + self.tp_size = ( + tp_size + if tp_size is not None + else get_tensor_model_parallel_world_size() + ) def update_param_tp_status(self): # Single source of truth for a parameter's TP state. BasevLLMParameter @@ -425,6 +438,11 @@ class ColumnParallelLinear(LinearBase): (e.g. model.layers.0.qkv_proj) return_bias: If true, return bias together with outputs in forward pass. disable_tp: If true, weights matrix won't be sharded through tp rank. + tp_rank: Override the tensor-parallel rank used for sharding. Defaults to + the global TP rank. Used to shard at a coarser granularity than one + shard per rank (see ``DCPGroupColumnParallelLinear``). + tp_size: Override the tensor-parallel world size used for sharding. + Defaults to the global TP world size. """ # --8<-- [end:column_parallel_linear] @@ -442,10 +460,21 @@ def __init__( *, return_bias: bool = True, disable_tp: bool = False, + tp_rank: int | None = None, + tp_size: int | None = None, ): # Divide the weight matrix along the last dimension. - self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0 - self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1 + if disable_tp: + self.tp_rank, self.tp_size = 0, 1 + else: + self.tp_rank = ( + tp_rank if tp_rank is not None else get_tensor_model_parallel_rank() + ) + self.tp_size = ( + tp_size + if tp_size is not None + else get_tensor_model_parallel_world_size() + ) self.input_size_per_partition = input_size self.output_size_per_partition = divide(output_size, self.tp_size) self.output_partition_sizes = [self.output_size_per_partition] @@ -465,6 +494,8 @@ def __init__( prefix, return_bias=return_bias, disable_tp=disable_tp, + tp_rank=self.tp_rank, + tp_size=self.tp_size, ) self._maybe_allow_fp8_block_shape_mismatch() @@ -586,6 +617,47 @@ def extra_repr(self) -> str: return s +class DCPGroupColumnParallelLinear(ColumnParallelLinear): + """Column-parallel linear whose weight is sharded across DCP groups. + + With Decode Context Parallelism (DCP) the KV cache is sharded across a DCP + group, so MLA decode must attend the group's full head set. This layer shards + its output across DCP *groups* (effective tp size ``tp_size // + dcp_world_size``) rather than across every rank, so each rank in a group + holds the whole group's heads, letting decode skip the query all-gather. + + :meth:`forward` returns the group's full head set. :meth:`_local_view` + extracts this rank's TP shard for prefill. + """ + + def __init__(self, *args, **kwargs): + dcp_world_size = ( + get_current_vllm_config().parallel_config.decode_context_parallel_size + ) + rank = get_tensor_model_parallel_rank() + world_size = get_tensor_model_parallel_world_size() + self.group_size = max(dcp_world_size, 1) + self.qrep_active = self.group_size > 1 + self.rank_in_group = rank % self.group_size + super().__init__( + *args, + **kwargs, + tp_rank=rank // self.group_size, + tp_size=world_size // self.group_size, + ) + + def _local_view(self, out: torch.Tensor) -> torch.Tensor: + """Slice this rank's tp head shard from a group-heads output. + + ``out`` is head-shaped, i.e. ``(..., group_heads, head_dim)``. + """ + if self.group_size == 1: + return out + n = out.shape[-2] // self.group_size + start = self.rank_in_group * n + return out[..., start : start + n, :].contiguous() + + class MergedColumnParallelLinear(ColumnParallelLinear): """Packed linear layers with column parallelism. diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 496a1dd15304..31298a66beed 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -7,7 +7,6 @@ from vllm.config import get_current_vllm_config from vllm.distributed import ( - get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, tensor_model_parallel_gather, ) @@ -145,7 +144,8 @@ def _get_logits( logits = self._apply_head(lm_head, hidden_states, embedding_bias) # Gather logits for TP - logits = self._gather_logits(logits) + if lm_head.tp_size > 1: + logits = self._gather_logits(logits) # Remove paddings in vocab (if any). if logits is not None: @@ -169,7 +169,7 @@ def get_top_tokens( "The local argmax reduction optimization is not supported for " "non-positive logit scaling factors." ) - tp_size = get_tensor_model_parallel_world_size() + tp_size = lm_head.tp_size logits = self._apply_head(lm_head, hidden_states, embedding_bias) if self.soft_cap is not None: diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 62fd64bd0dea..d06916f697c7 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -2,11 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import abstractmethod from collections.abc import Iterable +from math import prod import torch from vllm.config import VllmConfig from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.selector import get_mamba_attn_backend @@ -24,6 +26,22 @@ class MambaBase(AttentionLayerBase): kv_cache: tuple[torch.Tensor, ...] supports_dcp: bool = False + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views. + + Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm) + packed contiguously; slice them out and reinterpret per dtype/shape. + """ + pages = kv_cache.squeeze(dim=(1, 2)) + states: list[torch.Tensor] = [] + offset = 0 + for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()): + nbytes = prod(shape) * get_dtype_size(dtype) + state = pages[:, offset : offset + nbytes].view(dtype) + states.append(state.view(-1, *shape)) + offset += nbytes + self.kv_cache = tuple(states) + @abstractmethod def get_state_shape(self) -> Iterable[tuple[int, ...]]: """ diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index f95aac54fdfe..bc49226c75b2 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1,32 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + import torch from einops import rearrange from torch import nn +from torch.nn.parameter import Parameter -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import ( - divide, -) -from vllm.forward_context import ForwardContext, get_forward_context -from vllm.logger import init_logger +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig +from vllm.distributed import divide, get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention -from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader -from vllm.model_executor.utils import set_weight_attrs -from vllm.third_party.flash_linear_attention.ops.kda import ( - FusedRMSNormGated, - chunk_kda_with_fused_gate, - fused_kda_gate, - fused_recurrent_kda, +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, ) +from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from ...linear import ( ColumnParallelLinear, + MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) @@ -36,49 +37,118 @@ is_conv_state_dim_first, ) from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from ..ops.gather_initial_states import gather_initial_states -logger = init_logger(__name__) - - -def kda_attention( - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, - g1: torch.Tensor, - beta: torch.Tensor, - core_attn_out: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self._forward( - q_proj_states=q_proj_states, - k_proj_states=k_proj_states, - v_proj_states=v_proj_states, - g1=g1, - beta=beta, - core_attn_out=core_attn_out, - ) - - -def kda_attention_fake( - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, - g1: torch.Tensor, - beta: torch.Tensor, - core_attn_out: torch.Tensor, - layer_name: str, -) -> None: - return - - -direct_register_custom_op( - op_name="kda_attention", - op_func=kda_attention, - mutates_args=["core_attn_out"], - fake_impl=kda_attention_fake, -) +# Empirical lower bound for the KDA gate to avoid numerical underflow. +_KDA_GATE_LOGBOUND_MIN = -5.0 + + +def a_log_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load KDA A_log stored as either old 4D or current 1D weights.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + + if loaded_weight.dim() == 4: + assert loaded_weight.shape[:2] == (1, 1), ( + f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}" + ) + assert loaded_weight.shape[-1] == 1, ( + f"Expected old A_log last dim to be 1, got {loaded_weight.shape}" + ) + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + return default_weight_loader(param, loaded_weight) + + return loader + + +def _make_fused_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, +) -> Callable[..., None]: + sharded_dims = [dim // tp_size for dim in dims] + + def weight_loader( + param: torch.Tensor, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + if loaded_weight.dim() == 2: + loaded_weight = loaded_weight.unsqueeze(1) + shard_size = sharded_dims[loaded_shard_id] + source_start = tp_rank * shard_size + target_start = sum(sharded_dims[:loaded_shard_id]) + loaded_shard = loaded_weight[source_start : source_start + shard_size] + param.data[target_start : target_start + shard_size].copy_(loaded_shard) + + return weight_loader + + +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks. + + The replicated shard is represented as ``size * tp_size`` so the merged + parameter reserves ``size`` local rows on every rank. Loading that shard + from rank zero then gives every rank the complete checkpoint weight. + """ + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_id: int, + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_id = replicated_shard_id + output_sizes = output_sizes.copy() + output_sizes[replicated_shard_id] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank @PluggableLayer.register("kimi_gated_delta_net_attention") @@ -96,7 +166,11 @@ def get_state_shape( self, ) -> tuple[tuple[int, ...], tuple[int, ...]]: return MambaStateShapeCalculator.kda_state_shape( - self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, ) def __init__( @@ -114,122 +188,142 @@ def __init__( assert self.num_heads % self.tp_size == 0 self.local_num_heads = divide(self.num_heads, self.tp_size) - projection_size = self.head_dim * self.num_heads + self.projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(self.projection_size, self.tp_size) self.conv_size = kda_config["short_conv_kernel_size"] - - self.q_proj = ColumnParallelLinear( - self.hidden_size, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.q_proj", - ) - self.k_proj = ColumnParallelLinear( - self.hidden_size, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.k_proj", - ) - self.v_proj = ColumnParallelLinear( - self.hidden_size, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.v_proj", - ) - - self.f_a_proj = ReplicatedLinear( + self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False) + + if self.use_full_rank_gate: + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. The padding also avoids an + # Inductor correctness issue seen with the row-strided G view. + qkvg_output_sizes = [self.projection_size] * 4 + in_proj_output_sizes = qkvg_output_sizes + [ + self.head_dim, + self.num_heads, + ] + local_output_size = ( + 4 * self.local_projection_size + self.head_dim + self.local_num_heads + ) + self.in_proj_padding = -local_output_size % 16 + if self.in_proj_padding: + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) + else: + in_proj_output_sizes = [self.projection_size] * 3 + [ + self.num_heads, + self.head_dim, + ] + self.in_proj_padding = 0 + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( self.hidden_size, - self.head_dim, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, bias=False, quant_config=self.quant_config, - prefix=f"{prefix}.f_a_proj", + prefix=f"{prefix}.in_proj_qkvgfab", ) + if self.in_proj_padding: + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() self.f_b_proj = ColumnParallelLinear( self.head_dim, - projection_size, + self.projection_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.f_b_proj", ) self.dt_bias = nn.Parameter( - torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32) + torch.empty(self.local_projection_size, dtype=torch.float32) ) set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - self.b_proj = ColumnParallelLinear( - self.hidden_size, - self.num_heads, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.b_proj", - ) - - self.q_conv1d = ColumnParallelLinear( + # One packed parameter and cache let decode run a single conv update. + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. + self.conv1d = ColumnParallelLinear( input_size=self.conv_size, - output_size=projection_size, + output_size=3 * self.projection_size, bias=False, params_dtype=torch.float32, - prefix=f"{prefix}.q_conv1d", + prefix=f"{prefix}.conv1d", ) - self.k_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.k_conv1d", + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_fused_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + ) + }, ) - self.v_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.v_conv1d", - ) - # unsqueeze to fit conv1d weights shape into the linear weights shape. - # Can't do this in `weight_loader` since it already exists in - # `ColumnParallelLinear` and `set_weight_attrs` - # doesn't allow to override it - self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1) - self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) - self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) self.A_log = nn.Parameter( - torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) + torch.empty(self.local_num_heads, dtype=torch.float32) ) - set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)}) - - self.g_a_proj = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.g_a_proj", + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) + if self.gate_lower_bound is not None: + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, ( + "KDA gate lower bound must be in " + f"[{_KDA_GATE_LOGBOUND_MIN}, 0). " + f"Got {self.gate_lower_bound}." + ) + self.use_safe_gate = self.gate_lower_bound is not None + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" ) - self.g_b_proj = ColumnParallelLinear( - self.head_dim, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.g_b_proj", + backend = "triton" if backend == "auto" else backend + assert backend == "triton", ( + "The shared Kimi GDN layer only supports the Triton KDA " + f"prefill backend, got {backend!r}." ) + if not self.use_full_rank_gate: + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_b_proj", + ) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") self.o_proj = RowParallelLinear( - projection_size, + self.projection_size, self.hidden_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.o_proj", ) - compilation_config = get_current_vllm_config().compilation_config + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def rearrange_mixed_qkv( + self, mixed_qkv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + seq_len = mixed_qkv.shape[0] + qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim) + # Materialize all three row-strided inputs with one token-major to + # QKV-major permutation. Each unbound tensor is then contiguous. + qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1) + return qkv.unbind(0) + def forward( self, hidden_states: torch.Tensor, @@ -237,42 +331,57 @@ def forward( output: torch.Tensor, ) -> None: num_tokens = hidden_states.size(0) - q = self.q_proj(hidden_states)[0] - k = self.k_proj(hidden_states)[0] - v = self.v_proj(hidden_states)[0] + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + if self.use_full_rank_gate: + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + else: + mixed_qkv, beta, f_a = projected_qkvgfab.split( + [ + 3 * self.local_projection_size, + self.local_num_heads, + self.head_dim, + ], + dim=-1, + ) + g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] - beta = self.b_proj(hidden_states)[0].float().sigmoid() - g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] + g1 = self.f_b_proj(f_a)[0] beta = beta.unsqueeze(0) g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) - g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) - core_attn_out = torch.zeros( + core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) - torch.ops.vllm.kda_attention( - q, - k, - v, - g1, - beta, - core_attn_out, - self.prefix, + + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, ) - core_attn_out = self.o_norm(core_attn_out, g2) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") output[:] = self.o_proj(core_attn_out)[0] + @eager_break_during_capture def _forward( self, - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, + mixed_qkv: torch.Tensor, g1: torch.Tensor, + g2: torch.Tensor, beta: torch.Tensor, core_attn_out: torch.Tensor, ) -> None: @@ -280,165 +389,247 @@ def _forward( attn_metadata_raw = forward_context.attn_metadata if attn_metadata_raw is None: - # # V1 profile run return + # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies + # under kimi_k3/{amd,nvidia}/ops so each can diverge independently. + if current_platform.is_rocm(): + from vllm.models.kimi_k3.amd.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + else: + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + assert isinstance(attn_metadata_raw, dict) attn_metadata_narrowed = attn_metadata_raw[self.prefix] assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) - has_initial_state = attn_metadata_narrowed.has_initial_state - non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc - non_spec_state_indices_tensor = ( - attn_metadata_narrowed.non_spec_state_indices_tensor - ) # noqa: E501 - num_actual_tokens = attn_metadata_narrowed.num_actual_tokens - constant_caches = self.kv_cache - - q_proj_states = q_proj_states[:num_actual_tokens] - k_proj_states = k_proj_states[:num_actual_tokens] - v_proj_states = v_proj_states[:num_actual_tokens] + m = attn_metadata_narrowed + has_initial_state = m.has_initial_state + non_spec_query_start_loc = m.non_spec_query_start_loc + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor + spec_sequence_masks = m.spec_sequence_masks + spec_token_indx = m.spec_token_indx + non_spec_token_indx = m.non_spec_token_indx + spec_state_indices_tensor = m.spec_state_indices_tensor + spec_query_start_loc = m.spec_query_start_loc + num_accepted_tokens = m.num_accepted_tokens + num_actual_tokens = m.num_actual_tokens + mixed_qkv = mixed_qkv[:num_actual_tokens] g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] - (conv_state, recurrent_state) = constant_caches + constant_caches = self.kv_cache + + conv_state, recurrent_state = constant_caches # conv_state must be (..., dim, width-1) for the conv kernels. # DS layout stores it that way directly; SD layout needs a transpose. if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) - conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2) - - q_conv_weights = self.q_conv1d.weight.view( - self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2) + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) ) - k_conv_weights = self.k_conv1d.weight.view( - self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2) + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split( + self.local_projection_size, dim=0 ) - v_conv_weights = self.v_conv1d.weight.view( - self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2) + q_conv_state, k_conv_state, v_conv_state = conv_state.split( + self.local_projection_size, dim=-2 ) - if attn_metadata_narrowed.num_prefills > 0: - q_proj_states = q_proj_states.transpose(0, 1) - k_proj_states = k_proj_states.transpose(0, 1) - v_proj_states = v_proj_states.transpose(0, 1) - q = causal_conv1d_fn( - q_proj_states, - q_conv_weights, - self.q_conv1d.bias, - activation="silu", - conv_states=conv_state_q, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) - k = causal_conv1d_fn( - k_proj_states, - k_conv_weights, - self.k_conv1d.bias, - activation="silu", - conv_states=conv_state_k, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) - v = causal_conv1d_fn( - v_proj_states, - v_conv_weights, - self.v_conv1d.bias, - activation="silu", - conv_states=conv_state_v, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) + + # Split tokens into the multi-query spec-decode part and the remaining + # (prefill / plain decode) part. + if spec_sequence_masks is not None: + if m.num_prefills == 0 and m.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + g1_spec, beta_spec = g1, beta + mixed_qkv_ns = g1_ns = beta_ns = None + else: + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) else: - assert non_spec_state_indices_tensor is not None - decode_conv_indices = non_spec_state_indices_tensor[ - : attn_metadata_narrowed.num_actual_tokens - ] - q = causal_conv1d_update( - q_proj_states, - conv_state_q, - q_conv_weights, - self.q_conv1d.bias, - activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta + + # ---------- spec-decode multi-query path ---------- + core_attn_out_spec = None + if spec_sequence_masks is not None: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] + spec_max_query_len = spec_state_indices_tensor.size(-1) + + # Sibling beta and, for full-rank gates, output-gate views remain + # live, so write the convolution output separately. + spec_conv_out = torch.empty( + mixed_qkv_spec.shape, + dtype=mixed_qkv_spec.dtype, + device=mixed_qkv_spec.device, ) - k = causal_conv1d_update( - k_proj_states, - conv_state_k, - k_conv_weights, - self.k_conv1d.bias, + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, + conv_state_indices=spec_conv_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_max_query_len, + validate_data=False, + out=spec_conv_out, ) - v = causal_conv1d_update( - v_proj_states, - conv_state_v, - v_conv_weights, - self.v_conv1d.bias, - activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1) ) - - q, k, v = map( - lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v) - ) - - if attn_metadata_narrowed.num_prefills > 0: - assert non_spec_state_indices_tensor is not None - assert has_initial_state is not None - zero_idx = non_spec_state_indices_tensor[~has_initial_state] - recurrent_state[zero_idx] = 0 - initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous() - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=g1, - beta=beta, - A_log=self.A_log, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1] + # Spec-only batches write directly into core_attn_out. + spec_out = ( + core_attn_out[:, : q_spec.shape[1]] + if m.num_prefills == 0 and m.num_decodes == 0 + else None ) - # Init cache - recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state - else: - assert non_spec_query_start_loc is not None - g1 = fused_kda_gate( - rearrange(g1, "1 n h d -> n (h d)"), - self.A_log, - self.head_dim, - g_bias=self.dt_bias, - ).unsqueeze(0) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = fused_recurrent_kda( - q=q, - k=k, - v=v, - g=g1, - beta=beta, + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, initial_state=recurrent_state, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc[ - : attn_metadata_narrowed.num_decodes + 1 - ], - ssm_state_indices=non_spec_state_indices_tensor, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, ) - core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ - 0, :num_actual_tokens - ] + + # ---------- non-spec path (prefill or plain decode) ---------- + core_attn_out_non_spec = None + if mixed_qkv_ns is not None: + assert g1_ns is not None and beta_ns is not None + if m.num_prefills > 0: + q_ns, k_ns, v_ns = mixed_qkv_ns.split( + self.local_projection_size, dim=-1 + ) + + # Packed prefill conv would require copying V solely to make + # it dense for KDA. Separate calls accept the strided inputs + # and produce dense Q/K/V without that extra traffic. + # TODO: Use packed conv once every KDA prefill backend accepts + # row-strided Q/K/V directly. + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns, k_ns, v_ns = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in (q_ns, k_ns, v_ns) + ) + + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + initial_state = gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + g_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + ) + # Init cache + recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + + else: + # pure-decode non-spec batch + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : mixed_qkv_ns.size(0) + ] + # Sibling beta and, for full-rank gates, output-gate views + # remain live, so write the conv output separately. + packed_conv_out = torch.empty( + mixed_qkv_ns.shape, + dtype=mixed_qkv_ns.dtype, + device=mixed_qkv_ns.device, + ) + mixed_qkv_ns = causal_conv1d_update( + mixed_qkv_ns, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=decode_conv_indices, + validate_data=True, + out=packed_conv_out, + ) + core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + state_indices=decode_conv_indices, + ) + + # ---------- merge spec and non-spec outputs ---------- + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + # Mixed batches require indexed placement in the original order. + merged = torch.empty( + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), + dtype=core_attn_out_spec.dtype, + device=core_attn_out_spec.device, + ) + merged.index_copy_(1, spec_token_indx, core_attn_out_spec) + merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens] + elif core_attn_out_non_spec is not None: + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ + 0, :num_actual_tokens + ] + else: + assert core_attn_out_spec is not None + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a6524961ea92..7538a2b6b49d 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -32,6 +32,9 @@ causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -494,12 +497,27 @@ def __init__( if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) self.model_config = model_config self.cache_config = cache_config self.prefix = prefix + self.use_replayssm = ( + cache_config.use_replayssm if cache_config is not None else False + ) + self.replayssm_buffer_len = ( + cache_config.replayssm_buffer_len + if cache_config is not None and cache_config.use_replayssm + else None + ) + self.mamba_config = vllm_config.mamba_config + if self.use_replayssm and self.num_heads % self.tp_size != 0: + raise ValueError( + "--use-replayssm requires tensor-parallel heads to divide evenly" + ) + # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode + # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache). + _n_state = 5 if self.use_replayssm else 2 + self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) self.num_spec = vllm_config.num_speculative_tokens if self.num_spec > 0: @@ -591,7 +609,7 @@ def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None: # Triton's autotuner includes tensor dtypes in its cache key, # so state_dtype must match what real inference uses. - _, ssm_state_dtype = self.get_state_dtype() + ssm_state_dtype = self.get_state_dtype()[1] # SSD kernel autotune keys depend on dtype and head dimensions, # not on sequence length or batch size, so a single shape suffices. @@ -702,6 +720,10 @@ def conv_ssm_forward( else self.kv_cache[0].transpose(-1, -2) ) ssm_state = self.kv_cache[1] + if self.use_replayssm: + x_cache, dt_cache, B_cache = self.kv_cache[2:] + else: + x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size @@ -1027,37 +1049,78 @@ def conv_ssm_forward( # - mamba_cache_params.ssm_state's slots will be selected # using state_indices_tensor_d # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt_d, - A_d, - B_d, - C_d, - D_d, - dt_bias, - dt_softplus=True, - state_batch_indices=state_indices_tensor_d_input, - dst_state_batch_indices=state_indices_tensor_d_output, - out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim), - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=query_start_loc_d, - is_blackwell=self.is_blackwell, + preallocated_ssm_out_d = preallocated_ssm_out_d.view( + num_decode_tokens, -1, self.head_dim ) + if self.use_replayssm: + assert self.replayssm_buffer_len is not None + selective_state_update_replayssm_output_only( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=attn_metadata.bc_pre_scratch, + write_pos=attn_metadata.write_pos_d, + is_flush=attn_metadata.is_flush_d, + max_cache_len=self.replayssm_buffer_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + # Stochastic Rounding for the vanilla decode path is read + # from mamba_config inside ssu_dispatch; the replay kernel + # isn't a dispatch backend, so pass it here. + enable_stochastic_rounding=( + self.mamba_config.enable_stochastic_rounding + ), + cache_philox_rounds=( + self.mamba_config.stochastic_rounding_philox_rounds + ), + ) + else: + selective_state_update( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d_input, + dst_state_batch_indices=state_indices_tensor_d_output, + out=preallocated_ssm_out_d, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=query_start_loc_d, + is_blackwell=self.is_blackwell, + ) - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: + def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) + if self.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, self.model_config.dtype + ) + return base_dtype - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + tp_world_size = get_tensor_model_parallel_world_size() + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), + tp_world_size=tp_world_size, n_groups=self.n_groups, num_heads=self.num_heads, head_dim=self.head_dim, @@ -1065,6 +1128,15 @@ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, ) + if self.use_replayssm: + assert self.replayssm_buffer_len is not None + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + self.n_groups, + tp_world_size, + self.replayssm_buffer_len, + ) + return base_shape @property def mamba_type(self) -> MambaAttentionBackendEnum: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 9e78b8222803..e3361ca7dc2a 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -80,6 +80,18 @@ def mamba2_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + @classmethod + def append_replayssm_ring( + cls, + base_dtypes: tuple[torch.dtype, ...], + model_dtype: ModelDType | torch.dtype, + ) -> tuple[torch.dtype, ...]: + """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple: + ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``. + """ + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return (*base_dtypes, activation_dtype, torch.float32, activation_dtype) + @classmethod def _mamba_state_dtype( cls, @@ -186,6 +198,28 @@ def mamba2_state_shape( temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) return conv_state_shape, temporal_state_shape + @classmethod + def append_replayssm_ring( + cls, + base_shapes: tuple[tuple[int, ...], ...], + n_groups: int, + tp_world_size: int, + replayssm_buffer_len: int, + ) -> tuple[tuple[int, ...], ...]: + """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a + base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape + ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended + ``n_groups``. + """ + local_nheads, head_dim, state_size = base_shapes[1] + local_ngroups = divide(n_groups, tp_world_size) + return ( + *base_shapes, + (local_nheads, replayssm_buffer_len, head_dim), + (local_nheads, replayssm_buffer_len), + (local_ngroups, replayssm_buffer_len, state_size), + ) + @classmethod def short_conv_state_shape( cls, @@ -254,7 +288,7 @@ def kda_state_shape( conv_dim = proj_size + 2 * proj_k_size conv_state_shape = cls._orient_conv_shape( - divide(conv_dim, tp_world_size), conv_kernel_size - 1 + divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec ) recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) return (conv_state_shape, recurrent_state_shape) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f7c237ca2db0..8335c849666d 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -8,11 +8,12 @@ import numpy as np import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID -@triton.jit() +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) def _causal_conv1d_fwd_kernel( # continuous batching # Pointers to matrices x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences @@ -33,7 +34,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, - num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + num_cache_lines, # added to support vLLM larger cache lines # Strides stride_x_dim: tl.constexpr, # stride to get to next feature-value, stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) @@ -58,6 +59,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, ): conv_states_ptr = initial_states_ptr conv_state_indices_ptr = cache_indices_ptr @@ -68,6 +70,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching KERNEL_WIDTH - 1 ) # can be passed via argument if it's not the same as this value + if launch_pdl: + tl.extra.cuda.gdc_wait() + # one program handles one chunk in a single sequence # rather than mixing sequences - to make updating initial_states across sequences efficiently @@ -79,6 +84,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) if idx_seq == pad_slot_id: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) @@ -136,6 +143,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching if HAS_NULL_BLOCK: # noqa if conv_states_input_coord == null_block_id: # not processing as this is a null block (padding) + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return conv_states_base = ( conv_states_ptr @@ -408,6 +417,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor w_col3 = tl.load(w_ptrs, mask_w, other=0.0) mask_x_1d = idx_feats < dim + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + for idx_token in range(segment_len): acc = acc_preload @@ -741,11 +754,12 @@ def grid(META): BLOCK_M=BLOCK_M, BLOCK_N=256, num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), ) return out.to(original_x_dtype) -@triton.jit() +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) def _causal_conv1d_update_kernel( # Pointers to matrices x_ptr, # (batch, dim, seqlen) @@ -763,7 +777,7 @@ def _causal_conv1d_update_kernel( dim: tl.constexpr, seqlen: tl.constexpr, state_len: tl.constexpr, - num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + num_cache_lines, # added to support vLLM larger cache lines # Strides stride_x_seq: tl.constexpr, stride_x_dim: tl.constexpr, @@ -789,10 +803,16 @@ def _causal_conv1d_update_kernel( NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, ): + if launch_pdl: + tl.extra.cuda.gdc_wait() + # ruff: noqa: E501 idx_seq = tl.program_id(0) if idx_seq >= batch: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return # [BLOCK_N,] elements along the feature-dimension (channel) @@ -814,6 +834,8 @@ def _causal_conv1d_update_kernel( if HAS_NULL_BLOCK: # noqa if conv_states_input_coord == null_block_id: # not processing as this is not the actual sequence + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return if IS_VARLEN: @@ -831,6 +853,8 @@ def _causal_conv1d_update_kernel( o_offset = idx_seq * stride_o_seq if query_start_index == query_end_index: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return if IS_SPEC_DECODING: @@ -969,6 +993,9 @@ def _causal_conv1d_update_kernel( mask_x_1d = idx_feats < dim # STEP 5: compute each token + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + for idx_token in tl.range(seqlen): acc = acc_preload @@ -1080,6 +1107,7 @@ def causal_conv1d_update( block_idx_last_scheduled_token: torch.Tensor | None = None, initial_state_idx: torch.Tensor | None = None, validate_data=False, + out: torch.Tensor | None = None, ): """ x: Input tensor which can take the following shapes: @@ -1117,7 +1145,8 @@ def causal_conv1d_update( for example: conv_state_indices = [null_block_id, 1, 20, null_block_id] in this case, the kernel will not process entries at indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x` + out: optional output tensor with the same shape as `x`. When omitted, + the input is overwritten. """ if validate_data: assert null_block_id is not None @@ -1129,10 +1158,22 @@ def causal_conv1d_update( original_x_dtype = x.dtype x = x.to(conv_state.dtype) + if out is None: + out = x + else: + if out.shape != x.shape: + raise ValueError( + f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}." + ) + if out.dtype != original_x_dtype or out.device != x.device: + raise ValueError( + "`out` must have the same dtype and device as the input `x`." + ) unsqueeze = query_start_loc is None and x.dim() == 2 if unsqueeze: # make it (batch, dim, seqlen) with seqlen == 1 x = x.unsqueeze(-1) + out = out.unsqueeze(-1) if query_start_loc is None: batch, dim, seqlen = x.shape else: @@ -1159,8 +1200,6 @@ def causal_conv1d_update( assert num_cache_lines >= batch assert weight.stride(1) == 1 # Need this - # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o' - out = x stride_w_dim, stride_w_width = weight.stride() if query_start_loc is None: @@ -1233,7 +1272,18 @@ def grid(META): NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, BLOCK_N=256, + launch_pdl=current_platform.is_arch_support_pdl(), ) if unsqueeze: out = out.squeeze(-1) return out.to(original_x_dtype) + + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu, + causal_conv1d_update_cpu, + ) + + causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore + causal_conv1d_update = causal_conv1d_update_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI325X,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI325X,cache_dtype=float16.json new file mode 100644 index 000000000000..9dfddc603d77 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI325X,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.6.0", + "128": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "131072": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "196608": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "262144": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index b047ca6d6169..c552d93445e6 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -6,18 +6,31 @@ import torch import torch.nn.functional as F +from vllm._custom_ops import causal_conv1d_update_cpu_vec +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID -# for prefill -def causal_conv1d_torch( + +def causal_conv1d_fn_cpu( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, conv_states: torch.Tensor, query_start_loc: torch.Tensor, - cache_indices: torch.Tensor, - has_initial_state: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, + **kwargs, ) -> torch.Tensor: + """CPU implementation for causal_conv1d_fwd.""" + if isinstance(activation, bool) and activation: + activation = "silu" + elif isinstance(activation, bool): + activation = None + + original_x_dtype = x.dtype + x = x.to(conv_states.dtype) + out = torch.empty_like(x) state_len = weight.shape[1] - 1 assert activation in {None, "silu", "swish"} @@ -27,11 +40,21 @@ def causal_conv1d_torch( for idx in range(query_start_loc.shape[0] - 1) ] weight = weight.unsqueeze(1) + for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx): - slot = int(cache_indices[seq_idx].item()) + if bos == eos: + continue + + slot = ( + int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx + ) + + if slot == pad_slot_id: + continue seq_x = x[:, bos:eos].unsqueeze(0) - if bool(has_initial_state[seq_idx].item()): + + if has_initial_state is not None and bool(has_initial_state[seq_idx].item()): initial_state = conv_states[slot, :, :state_len].unsqueeze(0) else: initial_state = torch.zeros( @@ -51,16 +74,48 @@ def causal_conv1d_torch( groups=weight.shape[0], ) seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype) + if activation in ("silu", "swish"): seq_out = F.silu(seq_out) out[:, bos:eos] = seq_out.squeeze(0) conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0)) - return out + return out.to(original_x_dtype) + + +def causal_conv1d_update_cpu( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int | None = None, + **kwargs, +) -> torch.Tensor: + """CPU implementation for causal_conv1d_update.""" + if isinstance(activation, bool): + activation = "silu" if activation else None + + if pad_slot_id is None: + pad_slot_id = kwargs.get("null_block_id", NULL_BLOCK_ID) + if pad_slot_id is None: + pad_slot_id = NULL_BLOCK_ID + + return causal_conv1d_update_cpu_vec( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + query_start_loc, + pad_slot_id, + ) -# for decode def causal_conv1d_update_torch( x: torch.Tensor, conv_state: torch.Tensor, @@ -68,6 +123,11 @@ def causal_conv1d_update_torch( bias: torch.Tensor | None = None, activation: str | None = None, ) -> torch.Tensor: + """ + Pure PyTorch fallback for causal_conv1d_update. + Currently used as a fallback for Arm (aarch64) to leverage + oneDNN/ACL F.conv1d kernels for batched decoding. + """ assert activation in {None, "silu", "swish"} _, dim, seq_len = x.shape diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 9f3aa12c8d17..e84a0c6b723e 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -10,9 +10,13 @@ from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( - causal_conv1d_torch, + causal_conv1d_fn_cpu as causal_conv1d_torch, +) +from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_update_cpu, causal_conv1d_update_torch, ) +from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import ( LayerNameType, _resolve_layer_name, @@ -140,21 +144,30 @@ def _cpu_gdn_attention_nonspec( conv_states=conv_state, weight=layer.conv1d.weight, bias=layer.conv1d.bias, - silu_activation=layer.activation == "silu", + silu_activation=(layer.activation == "silu"), conv_state_indices=decode_state_indices, is_vnni=True, ) else: - decode_conv_state = conv_state[decode_state_indices].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - # [B, dim] -> [B, dim, 1] - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, - weight=conv_weights, - bias=layer.conv1d.bias, - activation=layer.activation, - ).squeeze(-1) - conv_state[decode_state_indices] = decode_conv_state + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + decode_conv_state = conv_state[decode_state_indices].contiguous() + decode_mixed_qkv = causal_conv1d_update_torch( + x=decode_mixed_qkv.unsqueeze(-1), + conv_state=decode_conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + ).squeeze(-1) + conv_state[decode_state_indices] = decode_conv_state + else: + decode_mixed_qkv = causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_state=conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + conv_state_indices=decode_state_indices, + ) query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) @@ -223,22 +236,24 @@ def _cpu_gdn_attention_nonspec( A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias ) - initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 - attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu( + # zero pool slots for sequences without a prior state; the kernel + # gathers/mutates ssm_state in place via prefill_state_indices, so no + # separate scatter-back is needed after the call. + no_initial_state = prefill_state_indices[~prefill_has_initial_state] + if no_initial_state.numel() > 0: + ssm_state[no_initial_state] = 0 + attn_out, _ = ops.chunk_gated_delta_rule_cpu( query=query, key=key, value=value, g=g, beta=beta, - initial_state=initial_state, + initial_state=ssm_state, output_final_state=True, cu_seqlens=prefill_query_start_loc, head_first=False, use_qk_l2norm_in_kernel=True, - ) - ssm_state[prefill_state_indices] = last_recurrent_state.to( - ssm_state.dtype, copy=False + initial_state_indices=prefill_state_indices, ) core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) @@ -293,25 +308,21 @@ def _cpu_gdn_attention_spec_aware( width: int, state_len: int, ) -> None: - mixed_qkv = mixed_qkv.contiguous() - a = a.contiguous() - b = b.contiguous() - spec_sequence_masks = attn_metadata_i.spec_sequence_masks conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len) ssm_state = _ssm_state_view(layer) if spec_sequence_masks is None: # No spec sequences in this batch (e.g. the prompt prefill step while - # speculative decoding is configured). Process as prefill/decode using - # torch conv (which only touches the first ``width-1`` columns of the - # wide buffer, leaving the rolling history untouched). + # speculative decoding is configured). Process as ordinary + # prefill/decode while touching only the first ``width-1`` columns of + # the wide buffer, leaving the rolling history untouched. _spec_aware_nonspec( layer, attn_metadata_i, - mixed_qkv, - b, - a, + mixed_qkv.contiguous(), + b.contiguous(), + a.contiguous(), core_attn_out, conv_buf, ssm_state, @@ -326,9 +337,9 @@ def _cpu_gdn_attention_spec_aware( num_decodes = attn_metadata_i.num_decodes if num_prefills == 0 and num_decodes == 0: - mixed_qkv_spec = mixed_qkv - b_spec = b - a_spec = a + mixed_qkv_spec = mixed_qkv.contiguous() + b_spec = b.contiguous() + a_spec = a.contiguous() spec_out_indx = None else: assert spec_token_indx is not None @@ -392,50 +403,71 @@ def _spec_forward( assert spec_qsl is not None assert num_accepted is not None - spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1].to("cpu", torch.int64) - num_acc_cpu = num_accepted[:num_spec_decodes].to("cpu", torch.int64) + spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1] seq_starts = spec_qsl_cpu[:-1] seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1] # ---- 1. Convolution (per-sequence rolling buffer) ---- - w2d = _unpacked_conv_weight(layer) # (dim, width) - dim = w2d.size(0) - w = w2d.unsqueeze(1) # (dim, 1, width) for F.conv1d depthwise + dim = mixed_qkv_spec.size(-1) bias = layer.conv1d.bias silu = layer.activation == "silu" - conv_out = torch.empty_like(mixed_qkv_spec) - col0 = spec_state_indices[:, 0].to("cpu", torch.int64) - for i in range(num_spec_decodes): - q_i = int(seq_lens[i].item()) - if q_i == 0: - continue - start = int(seq_starts[i].item()) - slot0 = int(col0[i].item()) - a_prev = int(num_acc_cpu[i].item()) - offset = a_prev - 1 - B = conv_buf[slot0] # (dim, state_len) - x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) - prior = B[:, offset : offset + (width - 1)] - conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) # (1, dim, w-1+q) - out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) - if silu: - out = F.silu(out) - conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) - # Roll the buffer: drop ``a_prev`` from the front, append the new - # draft tokens, keep total length == state_len. - keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] - new_B = torch.cat([keep, x_seq], dim=-1) - B.copy_(new_B) + can_use_native_conv = ( + torch.cpu._is_amx_tile_supported() + and not is_conv_state_dim_first() + and width == 4 + and num_spec_decodes > 0 + and bool(torch.all(seq_lens == seq_lens[0]).item()) + and int(seq_lens[0].item()) > 0 + ) + if can_use_native_conv: + q_i = int(seq_lens[0].item()) + conv_out = ops.causal_conv1d_update_cpu( + x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim), + conv_states=conv_buf, + weight=layer.conv1d.weight, + bias=bias, + silu_activation=silu, + conv_state_indices=spec_state_indices[:num_spec_decodes, 0] + .to("cpu", torch.int32) + .contiguous(), + is_vnni=True, + num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32), + ).view_as(mixed_qkv_spec) + else: + w = _unpacked_conv_weight(layer).unsqueeze(1) + col0 = spec_state_indices[:num_spec_decodes, 0] + num_acc_cpu = num_accepted[:num_spec_decodes] + conv_out = torch.empty_like(mixed_qkv_spec) + for i in range(num_spec_decodes): + q_i = int(seq_lens[i].item()) + if q_i == 0: + continue + start = int(seq_starts[i].item()) + slot0 = int(col0[i].item()) + offset = int(num_acc_cpu[i].item()) - 1 + B = conv_buf[slot0] # (dim, state_len) + x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) + prior = B[:, offset : offset + (width - 1)] + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) + out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) + if silu: + out = F.silu(out) + conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) + # Roll the buffer: drop the accepted history from the front, append + # the new draft tokens, keep total length == state_len. + keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] + new_B = torch.cat([keep, x_seq], dim=-1) + B.copy_(new_B) # ---- 2. Recurrent (multi-slot SSM state) ---- # Single fused kernel call: it runs the recurrence over each sequence's # draft tokens internally, resumes from slot ``num_accepted-1`` and stores # the state after token ``t`` into slot ``t`` (rollback for the next step). query, key, value = layer.rearrange_mixed_qkv(conv_out) - query = query.squeeze(0).contiguous() - key = key.squeeze(0).contiguous() - value = value.squeeze(0).contiguous() + query = query.squeeze(0) + key = key.squeeze(0) + value = value.squeeze(0) spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous() num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous() cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous() @@ -445,8 +477,8 @@ def _spec_forward( q=query, k=key, v=value, - a=a_spec.contiguous(), - b=b_spec.contiguous(), + a=a_spec, + b=b_spec, initial_state_source=ssm_state, spec_state_indices=spec_idx, num_accepted_tokens=num_acc, @@ -456,15 +488,6 @@ def _spec_forward( return out_spec -def core_attn_out_like(layer, mixed_qkv_spec: torch.Tensor) -> torch.Tensor: - num_tokens = mixed_qkv_spec.size(0) - return torch.zeros( - (num_tokens, layer.num_v_heads // layer.tp_size, layer.head_v_dim), - dtype=mixed_qkv_spec.dtype, - device=mixed_qkv_spec.device, - ) - - def _spec_aware_nonspec( layer, attn_metadata_i: GDNAttentionMetadata, @@ -476,13 +499,19 @@ def _spec_aware_nonspec( ssm_state: torch.Tensor, width: int, ) -> None: - """Non-spec prefill/decode with a wide conv buffer (torch path).""" + """Non-spec prefill/decode with a wide conv buffer.""" state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor query_start_loc = attn_metadata_i.non_spec_query_start_loc assert state_indices_tensor is not None assert query_start_loc is not None + state_indices_tensor = state_indices_tensor.contiguous() - conv_weights = _unpacked_conv_weight(layer) + is_amx = torch.cpu._is_amx_tile_supported() + if is_amx and is_conv_state_dim_first(): + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + + if not is_amx: + conv_weights = _unpacked_conv_weight(layer) num_decodes = attn_metadata_i.num_decodes num_decode_tokens = attn_metadata_i.num_decode_tokens @@ -494,33 +523,48 @@ def _spec_aware_nonspec( decode_b = b[:num_decode_tokens] decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] - # Only the first ``width-1`` columns hold the real conv state. - decode_conv_state = conv_buf[decode_state_indices][ - :, :, : width - 1 - ].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, - weight=conv_weights, - bias=layer.conv1d.bias, - activation=layer.activation, - ).squeeze(-1) - conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state + if is_amx: + decode_mixed_qkv = ops.causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_states=conv_buf, + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + silu_activation=layer.activation == "silu", + conv_state_indices=decode_state_indices, + is_vnni=True, + ) + else: + # Only the first ``width-1`` columns hold the real conv state. + conv_state_view = conv_buf[:, :, : width - 1] + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + decode_conv_state = conv_state_view[decode_state_indices].contiguous() + decode_mixed_qkv = causal_conv1d_update_torch( + x=decode_mixed_qkv.unsqueeze(-1), + conv_state=decode_conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + ).squeeze(-1) + conv_state_view[decode_state_indices] = decode_conv_state + else: + decode_mixed_qkv = causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_state=conv_state_view, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + conv_state_indices=decode_state_indices, + ) query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) - # rearrange_mixed_qkv can return views whose last dim is not - # contiguous; the fused CPU kernel requires a contiguous last dim. - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( A_log=layer.A_log, dt_bias=layer.dt_bias, q=query, k=key, v=value, - a=decode_a.contiguous(), - b=decode_b.contiguous(), + a=decode_a, + b=decode_b, initial_state_source=ssm_state, initial_state_indices=decode_state_indices, cu_seqlens=query_start_loc[: num_decodes + 1], @@ -546,38 +590,52 @@ def _spec_aware_nonspec( prefill_has_initial_state = has_initial_state[ num_decodes : num_decodes + num_prefills ] - # ``causal_conv1d_torch`` only touches columns [:width-1] of the buffer. - prefill_mixed_qkv = causal_conv1d_torch( - x=prefill_mixed_qkv.transpose(0, 1), - weight=conv_weights, - bias=layer.conv1d.bias, - conv_states=conv_buf, - query_start_loc=prefill_query_start_loc, - cache_indices=prefill_state_indices, - has_initial_state=prefill_has_initial_state, - activation=layer.activation, - ).transpose(0, 1) + if is_amx: + prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( + x=prefill_mixed_qkv.transpose(0, 1), + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + silu_activation=layer.activation == "silu", + is_vnni=True, + ).transpose(0, 1) + else: + prefill_mixed_qkv = causal_conv1d_torch( + x=prefill_mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + activation=layer.activation, + ).transpose(0, 1) query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) g, beta = ops.fused_gdn_gating_cpu( A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias ) - initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 - attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu( + # zero pool slots for sequences without a prior state; the kernel + # gathers/mutates ssm_state in place via prefill_state_indices, so no + # separate scatter-back is needed after the call. + no_initial_state = prefill_state_indices[~prefill_has_initial_state] + if no_initial_state.numel() > 0: + ssm_state[no_initial_state] = 0 + attn_out, _ = ops.chunk_gated_delta_rule_cpu( query=query, key=key, value=value, g=g, beta=beta, - initial_state=initial_state, + initial_state=ssm_state, output_final_state=True, cu_seqlens=prefill_query_start_loc, head_first=False, use_qk_l2norm_in_kernel=True, - ) - ssm_state[prefill_state_indices] = last_recurrent_state.to( - ssm_state.dtype, copy=False + initial_state_indices=prefill_state_indices, ) core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) @@ -596,48 +654,66 @@ def _spec_aware_nonspec_subset( Returns outputs ordered like ``non_spec_token_indx``. """ - out = core_attn_out_like(layer, mixed_qkv) has_initial_state = attn_metadata_i.has_initial_state prefill_state_indices = attn_metadata_i.prefill_state_indices prefill_qsl = attn_metadata_i.prefill_query_start_loc assert prefill_state_indices is not None and prefill_qsl is not None assert has_initial_state is not None + prefill_state_indices = prefill_state_indices.contiguous() - conv_weights = _unpacked_conv_weight(layer) - conv_out = causal_conv1d_torch( - x=mixed_qkv.transpose(0, 1), - weight=conv_weights, - bias=layer.conv1d.bias, - conv_states=conv_buf, - query_start_loc=prefill_qsl, - cache_indices=prefill_state_indices, - has_initial_state=has_initial_state, - activation=layer.activation, - ).transpose(0, 1) + is_amx = torch.cpu._is_amx_tile_supported() + if is_amx and is_conv_state_dim_first(): + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + + if is_amx: + conv_out = ops.causal_conv1d_fwd_cpu( + x=mixed_qkv.transpose(0, 1), + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_qsl, + cache_indices=prefill_state_indices, + has_initial_state=has_initial_state, + silu_activation=layer.activation == "silu", + is_vnni=True, + ).transpose(0, 1) + else: + conv_weights = _unpacked_conv_weight(layer) + conv_out = causal_conv1d_torch( + x=mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_qsl, + cache_indices=prefill_state_indices, + has_initial_state=has_initial_state, + activation=layer.activation, + ).transpose(0, 1) query, key, value = layer.rearrange_mixed_qkv(conv_out) g, beta = ops.fused_gdn_gating_cpu( A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias ) - initial_state = ssm_state[prefill_state_indices] - initial_state[~has_initial_state, ...] = 0 - attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu( + # zero pool slots for sequences without a prior state; the kernel + # gathers/mutates ssm_state in place via prefill_state_indices, so no + # separate scatter-back is needed after the call. + no_initial_state = prefill_state_indices[~has_initial_state] + if no_initial_state.numel() > 0: + ssm_state[no_initial_state] = 0 + attn_out, _ = ops.chunk_gated_delta_rule_cpu( query=query, key=key, value=value, g=g, beta=beta, - initial_state=initial_state, + initial_state=ssm_state, output_final_state=True, cu_seqlens=prefill_qsl, head_first=False, use_qk_l2norm_in_kernel=True, + initial_state_indices=prefill_state_indices, ) - ssm_state[prefill_state_indices] = last_recurrent_state.to( - ssm_state.dtype, copy=False - ) - out[:] = attn_out.squeeze(0) - return out + return attn_out.squeeze(0) def cpu_gdn_attention_core_fake( diff --git a/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py new file mode 100644 index 000000000000..a65793d79245 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm._custom_ops as ops +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _mamba_chunk_scan_combined_fwd_cpu( + x, + dt, + A, + B, + C, + chunk_size, + out, + D=None, + z=None, + dt_bias=None, + initial_states=None, + return_intermediate_states=False, + seq_idx=None, + cu_seqlens=None, + cu_chunk_seqlens=None, + last_chunk_indices=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_dtype=None, + **kwargs, +): + seqlen, nheads, headdim = x.shape + _, ngroups, dstate = B.shape + + assert cu_seqlens is not None + batch = cu_seqlens.size(0) - 1 + + dt_f = dt.float() + if dt_bias is not None: + dt_f = dt_f + dt_bias.float().unsqueeze(0) + if dt_softplus: + dt_f = torch.nn.functional.softplus(dt_f) + if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): + dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) + + all_states = torch.zeros( + batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device + ) + if initial_states is not None: + all_states.copy_(initial_states.float()) + + assert out.is_contiguous(), ( + "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " + "pre-allocated as a contiguous tensor" + ) + + D_1d = None + if D is not None: + d = D.float() + while d.dim() > 1 and d.stride(-1) == 0: + d = d.squeeze(-1) + D_1d = d.contiguous() + + ops.mamba_chunk_scan_fwd_cpu( + out, + all_states, + x, + dt_f, + A, + B, + C, + D_1d, + z, + cu_seqlens.to(torch.int32), + ) + + out_dtype = state_dtype if state_dtype is not None else x.dtype + all_states = all_states.to(out_dtype) + + return all_states + + +def selective_state_update( + state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + dst_state_batch_indices=None, + null_block_id=NULL_BLOCK_ID, + out=None, + num_accepted_tokens=None, + cu_seqlens=None, + is_blackwell=False, + enable_stochastic_rounding=False, + cache_philox_rounds=0, +): + """CPU implementation for selective_state_update.""" + # Ensure out tensor exists + if out is None: + out = torch.empty_like(x if x.dim() == 2 else x) + + _state = state.unsqueeze(1) if state.dim() == 3 else state + _x = x.unsqueeze(1) if x.dim() == 2 else x + _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt + _A = A.unsqueeze(0) if A.dim() == 2 else A + _B = B.unsqueeze(1) if B.dim() == 2 else B + _C = C.unsqueeze(1) if C.dim() == 2 else C + _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D + _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z + _dt_bias = ( + dt_bias.unsqueeze(0) + if (dt_bias is not None and dt_bias.dim() == 1) + else dt_bias + ) + _out = out.unsqueeze(1) if out.dim() == 2 else out + + _sbi = state_batch_indices + _dsbi = dst_state_batch_indices + ops.selective_state_update_cpu( + _state, + _x, + _dt, + _A, + _B, + _C, + _D, + _z, + _dt_bias, + dt_softplus, + _sbi, + _dsbi, + null_block_id, + _out, + num_accepted_tokens, + cu_seqlens, + ) + return _out.squeeze(1) if out.dim() == 2 else _out diff --git a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py new file mode 100644 index 000000000000..b952e3ebbce0 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _gather_initial_states_kernel( + state_ptr, + indices_ptr, + has_initial_state_ptr, + output_ptr, + stride_state_batch, + stride_indices, + stride_has_initial_state, + row_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < row_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + has_initial_state = tl.load( + has_initial_state_ptr + batch_idx * stride_has_initial_state + ).to(tl.int1) + state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) + state_idx = tl.where(has_initial_state, state_idx, 0) + values = tl.load( + state_ptr + state_idx * stride_state_batch + offsets, + mask=mask & has_initial_state, + other=0.0, + ) + tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask) + + +def gather_initial_states( + state: torch.Tensor, + indices: torch.Tensor, + has_initial_state: torch.Tensor, +) -> torch.Tensor: + """Gather dense state rows, replacing uninitialized rows with zeros.""" + assert state.ndim >= 2 + assert state.is_cuda + assert indices.ndim == 1 and has_initial_state.ndim == 1 + assert indices.shape == has_initial_state.shape + assert indices.device == state.device + assert has_initial_state.device == state.device + assert indices.dtype in (torch.int32, torch.int64) + assert has_initial_state.dtype == torch.bool + + row_size = state[0].numel() + # Mamba pages may pad stride(0), but each state row remains dense. + assert state[0].is_contiguous() + output = torch.empty( + (indices.numel(), *state.shape[1:]), + dtype=state.dtype, + device=state.device, + ) + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (triton.cdiv(row_size, block_size), indices.numel()) + _gather_initial_states_kernel[grid]( + state, + indices, + has_initial_state, + output, + state.stride(0), + indices.stride(0), + has_initial_state.stride(0), + row_size=row_size, + BLOCK_SIZE=block_size, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index d348defcc768..af45467886ea 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -845,3 +845,13 @@ def selective_scan_fn( return delta # output written inplace to delta else: return z # output written inplace to z + + +from vllm.platforms import current_platform # noqa: E402 + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + selective_state_update as selective_state_update_cpu, + ) + + selective_state_update = selective_state_update_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py new file mode 100644 index 000000000000..aff4fc1c3f2d --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch-config selection for the ReplaySSM Mamba2 output_only decode kernel. + +Mirrors ``mamba_ssm.py``: a hard-coded heuristic per kernel, plus an +``override`` context manager for benchmarks/tests/config sweeps. Hardware is +auto-detected (Blackwell vs not) so call sites need not thread it through. +""" + +import functools +from contextlib import contextmanager + +from vllm.platforms import current_platform +from vllm.triton_utils import triton + + +@functools.cache +def _is_blackwell() -> bool: + try: + return current_platform.is_device_capability_family(100) + except Exception: + return False + + +# Per-kernel overrides keyed by the kernel name passed to get_replayssm_config. +_overrides: dict[str, tuple] = {} + + +@contextmanager +def override_replayssm_config(kernel: str, config: tuple): + """Pin ``kernel``'s launch config for the duration of the context.""" + prev = _overrides.get(kernel) + _overrides[kernel] = config + try: + yield + finally: + if prev is None: + _overrides.pop(kernel, None) + else: + _overrides[kernel] = prev + + +def _dstate_tile(dstate: int, tile: int) -> int: + return max(16, min(tile, triton.next_power_of_2(dstate))) + + +def _mamba2_output_only(dstate, L, is_blackwell): + # (block_size_m, num_warps, nf_dstate_tile, fl_dstate_tile, num_stages); + # decoupled dstate tiling, serving-batch optimum, dtype-independent per device. + if is_blackwell: + return 64, 1, _dstate_tile(dstate, 32), _dstate_tile(dstate, 64), 2 + return 16, 1, _dstate_tile(dstate, 64), _dstate_tile(dstate, 128), 2 + + +def get_replayssm_config(kernel: str, **shape) -> tuple: + """Return the launch config for ``kernel`` (override > tuned default). + + kernel: "mamba2_output_only". ``shape`` carries the keying dims (dstate; + ``L`` for the buffer length, default 16); hardware is auto-detected. + """ + if kernel in _overrides: + return _overrides[kernel] + bw = _is_blackwell() + if kernel == "mamba2_output_only": + return _mamba2_output_only(shape["dstate"], shape.get("L", 16), bw) + raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}") diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py new file mode 100644 index 000000000000..5b11bb7a08e0 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -0,0 +1,720 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_precompute_kernel( + B_ptr, + C_ptr, + B_cache_ptr, + write_pos_ptr, + is_flush_ptr, + bc_pre_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + ngroups, + dstate, + # Input strides + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + # Cache strides + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_g = tl.program_id(axis=1) + + # On flush steps the main kernel does not read bc_pre, so skip the work. + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + if is_flush: + return + + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load( + state_batch_indices_ptr + + pid_b * stride_state_indices_batch + + 0 * stride_state_indices_T + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + + B_ptr += pid_b * stride_B_batch + pid_g * stride_B_group + C_ptr += pid_b * stride_C_batch + pid_g * stride_C_group + B_cache_ptr += state_batch_idx * stride_B_cache_batch + pid_g * stride_B_cache_group + bc_pre_ptr += pid_b * stride_bc_pre_batch + pid_g * stride_bc_pre_group + + B_cur = tl.load( + B_ptr + offs_n * stride_B_dstate, + mask=offs_n < dstate, + other=0.0, + ) + C = tl.load( + C_ptr + offs_n * stride_C_dstate, + mask=offs_n < dstate, + other=0.0, + ) + B_cache_ptrs = ( + B_cache_ptr + + offs_k[:, None] * stride_B_cache_pos + + offs_n[None, :] * stride_B_cache_dstate + ) + B_cache = tl.load( + B_cache_ptrs, + mask=(offs_k[:, None] < write_pos) & (offs_n[None, :] < dstate), + other=0.0, + ) + B_all = tl.where(offs_k[:, None] == write_pos, B_cur[None, :], B_cache) + bc = tl.sum(B_all.to(tl.float32) * C[None, :].to(tl.float32), axis=1) + + tl.store( + bc_pre_ptr + offs_k * stride_bc_pre_pos, + bc, + mask=(offs_k <= write_pos) & (offs_k < MAX_CACHE_LEN), + ) + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_kernel( + # Pointers to matrices + state_ptr, + rand_seed_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + bc_pre_ptr, + write_pos_ptr, + is_flush_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # State strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # Input strides + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + # Cache strides + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K_CACHE: tl.constexpr, + BLOCK_SIZE_K_DOT: tl.constexpr, + NF_DSTATE_TILE: tl.constexpr, + NF_NDS: tl.constexpr, + FL_DSTATE_TILE: tl.constexpr, + FL_NDS: tl.constexpr, + DOT_INPUT_PRECISION: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + # Resolve the physical state slot for this decode row; skip padded rows. + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load( + state_batch_indices_ptr + + pid_b * stride_state_indices_batch + + 0 * stride_state_indices_T + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n_full = tl.arange(0, BLOCK_SIZE_DSTATE) + + # Buffer cursor (number of cached tokens so far) and the flush flag. + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + + # Advance every pointer to this (row, head, group). + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += ( + state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + ) + B_cache_ptr += ( + state_batch_idx * stride_B_cache_batch + + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + ) + bc_pre_ptr += ( + pid_b * stride_bc_pre_batch + + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + ) + + # Current-token dt (+ bias, softplus), scalar A, and current x. C, the + # checkpoint state S_0, and current-token B are read per dstate tile below. + dt_cur = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32) + if DT_SOFTPLUS: + dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur) + A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32) + x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0) + + if not is_flush: + # Output-only route: read y without materializing the state, using the + # precomputed k^T q products (`bc`): + # y = total_decay * (S_0 q) + sum_j s_j (k_j^T q) v_j. + # Then append the current token to the buffer. + offs_k_cache = tl.arange(0, BLOCK_SIZE_K_CACHE) + # dt over the window (history + current token), then the decay weights. + dt_all_cache = tl.load( + dt_cache_ptr + offs_k_cache * stride_dt_cache_pos, + mask=offs_k_cache < write_pos, + other=0.0, + ).to(tl.float32) + dt_all_cache = tl.where(offs_k_cache == write_pos, dt_cur, dt_all_cache) + dA_cumsum_cache = A * tl.cumsum(dt_all_cache, axis=0) + dA_total_cache = A * tl.sum(dt_all_cache, axis=0) + total_decay_cache = tl.exp(dA_total_cache) + scale_cache = dt_all_cache * tl.exp(dA_total_cache - dA_cumsum_cache) + scale_cache = tl.where(offs_k_cache <= write_pos, scale_cache, 0.0) + + # Gather buffered x over the window (history + current token). + x_all_cache_ptrs = ( + x_cache_ptr + + offs_m[:, None] * stride_x_cache_dim + + offs_k_cache[None, :] * stride_x_cache_pos + ) + x_all_cache = tl.load( + x_all_cache_ptrs, + mask=(offs_m[:, None] < dim) & (offs_k_cache[None, :] < write_pos), + other=0.0, + ) + x_all_cache = tl.where( + offs_k_cache[None, :] == write_pos, x_cur[:, None], x_all_cache + ) + + # Decayed checkpoint readout sum_n S_0(m,n) q(n), streamed over NF dstate + # tiles so the (M, N) state slice is never held whole. + offs_nt = tl.arange(0, NF_DSTATE_TILE) + ck_acc = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(NF_NDS): + offs_n = i * NF_DSTATE_TILE + offs_nt + nmask = offs_n < dstate + st = tl.load( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate, + mask=(offs_m[:, None] < dim) & nmask[None, :], + other=0.0, + ) + c_chunk = tl.load( + C_ptr + offs_n * stride_C_dstate, mask=nmask, other=0.0 + ).to(tl.float32) + ck_acc += tl.sum(st.to(tl.float32) * c_chunk[None, :], axis=1) + checkpoint_out = ck_acc * total_decay_cache + bc_cache = tl.load( + bc_pre_ptr + offs_k_cache * stride_bc_pre_pos, + mask=offs_k_cache <= write_pos, + other=0.0, + ) + cache_out = tl.sum( + x_all_cache.to(tl.float32) * (scale_cache * bc_cache)[None, :], axis=1 + ) + out = checkpoint_out + cache_out + + # Append the current token (x, dt, B) into the buffer at write_pos. + tl.store( + x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos, + x_cur, + mask=offs_m < dim, + ) + if pid_m == 0: + B_cur = tl.load( + B_ptr + offs_n_full * stride_B_dstate, + mask=offs_n_full < dstate, + other=0.0, + ) + tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur) + tl.store( + B_cache_ptr + + write_pos * stride_B_cache_pos + + offs_n_full * stride_B_cache_dstate, + B_cur, + mask=offs_n_full < dstate, + ) + else: + # Flush step: state route. Reconstruct the state from cached inputs, + # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T), persist it as the new + # checkpoint, then read y = S_t q -- streamed over FL dstate tiles. + offs_k_dot = tl.arange(0, BLOCK_SIZE_K_DOT) + dt_all_dot = tl.load( + dt_cache_ptr + offs_k_dot * stride_dt_cache_pos, + mask=offs_k_dot < write_pos, + other=0.0, + ).to(tl.float32) + dt_all_dot = tl.where(offs_k_dot == write_pos, dt_cur, dt_all_dot) + dA_cumsum_dot = A * tl.cumsum(dt_all_dot, axis=0) + dA_total_dot = A * tl.sum(dt_all_dot, axis=0) + total_decay_dot = tl.exp(dA_total_dot) + scale_dot = dt_all_dot * tl.exp(dA_total_dot - dA_cumsum_dot) + scale_dot = tl.where(offs_k_dot <= write_pos, scale_dot, 0.0) + + # Gather buffered x over the window (history + current token). + x_all_dot_ptrs = ( + x_cache_ptr + + offs_m[:, None] * stride_x_cache_dim + + offs_k_dot[None, :] * stride_x_cache_pos + ) + x_all_dot = tl.load( + x_all_dot_ptrs, + mask=(offs_m[:, None] < dim) & (offs_k_dot[None, :] < write_pos), + other=0.0, + ) + x_all_dot = tl.where( + offs_k_dot[None, :] == write_pos, x_cur[:, None], x_all_dot + ) + x_all_ty = x_all_dot.to(x_ptr.dtype.element_ty) + + # Distinct tile locals (_f) from the nf branch: differing tile widths + # would force a shape-mismatched merge at the if/else exit. + offs_nt_f = tl.arange(0, FL_DSTATE_TILE) + out = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(FL_NDS): + offs_n_f = i * FL_DSTATE_TILE + offs_nt_f + nmask_f = offs_n_f < dstate + # Gather buffered B over the window (history + current token). + B_all_dot = tl.load( + B_cache_ptr + + offs_k_dot[:, None] * stride_B_cache_pos + + offs_n_f[None, :] * stride_B_cache_dstate, + mask=(offs_k_dot[:, None] < write_pos) & nmask_f[None, :], + other=0.0, + ) + B_cur_tile = tl.load( + B_ptr + offs_n_f * stride_B_dstate, mask=nmask_f, other=0.0 + ) + B_all_dot = tl.where( + offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot + ) + # CUDA uses tf32x3 to keep fp32 parity with the elementwise + # baseline. Other platforms use their backend default; bf16/fp16 + # inputs are unaffected by this flag. + B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to( + x_ptr.dtype.element_ty + ) + delta_state = tl.dot( + x_all_ty, B_scaled, input_precision=DOT_INPUT_PRECISION + ) + state_ptrs = ( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n_f[None, :] * stride_state_dstate + ) + state_mask = (offs_m[:, None] < dim) & nmask_f[None, :] + st_f = tl.load(state_ptrs, mask=state_mask, other=0.0) + state_new = st_f.to(tl.float32) * total_decay_dot + delta_state.to( + tl.float32 + ) + if USE_RS_ROUNDING: + # Stochastic-round fp32->fp16 (Blackwell cvt.rs), mirroring the + # baseline. Only the flush step stores state, so this runs at 1/L + # the baseline's per-step rate. Absolute per-element offsets seed + # the RNG so each state element draws independently. + rand_seed = tl.load(rand_seed_ptr) + rand_offsets = ( + state_batch_idx * stride_state_batch + + pid_h * stride_state_head + + offs_m[:, None] * stride_state_dim + + offs_n_f[None, :] * stride_state_dstate + ) + if PHILOX_ROUNDS > 0: + rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS) + else: + rand = tl.randint(rand_seed, rand_offsets) + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "stochastic rounding requires an fp16 SSM state cache", + ) + state_store = convert_rs_fp16x2(state_new, rand) + else: + state_store = state_new.to(st_f.dtype) + tl.store(state_ptrs, state_store, mask=state_mask) + c_chunk_f = tl.load( + C_ptr + offs_n_f * stride_C_dstate, mask=nmask_f, other=0.0 + ).to(tl.float32) + out += tl.sum(state_new * c_chunk_f[None, :], axis=1) + + # Skip connection (D) and output gate (z). + if HAS_D: + D_ptr += pid_h * stride_D_head + D = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + out += x_cur.to(tl.float32) * D + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + z = tl.load(z_ptr + offs_m * stride_z_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + out *= z * tl.sigmoid(z) + + tl.store(out_ptr + offs_m * stride_out_dim, out, mask=offs_m < dim) + + +def selective_state_update_replayssm_output_only( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + bc_pre: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + is_flush: torch.Tensor | None = None, + max_cache_len: int = 16, + state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + enable_stochastic_rounding: bool = False, + cache_philox_rounds: int = 0, +) -> torch.Tensor: + """Cached-bc SSM update for vLLM's autoregressive Mamba2 decode path.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out is not None and out.dim() == 2: + out = out.unsqueeze(1) + if state_batch_indices is not None and state_batch_indices.dim() == 1: + state_batch_indices = state_batch_indices.unsqueeze(1) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + assert out is not None and out.shape == x.shape + + assert A.stride(-1) == 0 and A.stride(-2) == 0, ( + "Cached kernel requires TIE_HDIM (A scalar per head)" + ) + assert dt.stride(-1) == 0, "Cached kernel requires TIE_HDIM (dt scalar per head)" + if dt_bias is not None: + assert dt_bias.stride(-1) == 0, ( + "Cached kernel requires TIE_HDIM (dt_bias scalar per head)" + ) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert x_cache.shape[1:] == (nheads, max_cache_len, dim) + assert dt_cache.shape[1:] == (nheads, max_cache_len) + assert B_cache.shape[1:] == (ngroups, max_cache_len, dstate) + assert write_pos is not None and write_pos.shape[0] >= batch + assert write_pos.dtype == torch.int32 + assert is_flush is not None and is_flush.shape[0] >= batch + assert is_flush.dtype in (torch.bool, torch.int8) + assert bc_pre is not None + assert bc_pre.shape[0] >= batch and bc_pre.shape[1] >= ngroups + assert bc_pre.shape[2] == max_cache_len + assert bc_pre.dtype == torch.float32 + if state_batch_indices is not None: + assert state_batch_indices.shape[0] >= batch + assert state_batch_indices.shape[1] >= 1 + + block_size_k_cache = max(1, triton.next_power_of_2(max_cache_len)) + block_size_k_dot = max(16, block_size_k_cache) + block_size_m, num_warps, nf_tile, fl_tile, num_stages = get_replayssm_config( + "mamba2_output_only", dstate=dstate, L=max_cache_len + ) + bs_dstate = triton.next_power_of_2(dstate) + nf_dstate_tile = max(16, min(nf_tile, bs_dstate)) + nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile) + fl_dstate_tile = max(16, min(fl_tile, bs_dstate)) + fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile) + # AMD Triton does not support tf32x3, so use its backend default. CUDA + # retains tf32x3 to preserve fp32 parity with the elementwise baseline. + dot_input_precision = None if current_platform.is_rocm() else "tf32x3" + + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + state_indices_strides = ( + (state_batch_indices.stride(0), state_batch_indices.stride(1)) + if state_batch_indices is not None + else (0, 0) + ) + rand_seed = ( + torch.randint(0, 2**32, (1,), device=state.device) + if enable_stochastic_rounding + else None + ) + + with torch.accelerator.device_index(x.device.index): + # Both kernels always launch: the precompute kernel self-skips flush + # rows per row (the branch can't be hoisted out under CUDA graphs), so + # it is a no-op when every row is flushing. + _replayssm_output_only_precompute_kernel[(batch, ngroups)]( + B, + C, + B_cache, + write_pos, + is_flush, + bc_pre, + state_batch_indices, + null_block_id, + batch, + ngroups, + dstate, + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + max_cache_len, + block_size_k_cache, + num_warps=2, + ) + _replayssm_output_only_kernel[grid]( + state, + rand_seed, + x, + dt, + dt_bias, + A, + B, + C, + D, + z, + out, + x_cache, + dt_cache, + B_cache, + bc_pre, + write_pos, + is_flush, + state_batch_indices, + null_block_id, + batch, + nheads, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + dt.stride(0), + dt.stride(1), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + D.stride(0) if D is not None else 0, + D.stride(1) if D is not None else 0, + z_strides[0], + z_strides[1], + z_strides[2], + out.stride(0), + out.stride(1), + out.stride(2), + x_cache.stride(0), + x_cache.stride(1), + x_cache.stride(3), + x_cache.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + dt_softplus, + max_cache_len, + block_size_m, + block_size_k_cache, + block_size_k_dot, + nf_dstate_tile, + nf_nds, + fl_dstate_tile, + fl_nds, + dot_input_precision, + enable_stochastic_rounding, + cache_philox_rounds, + num_warps=num_warps, + num_stages=num_stages, + ) + + if not has_heads: + out = out.squeeze(1) + return out diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 4c93a768b629..8c645574b9e7 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -225,3 +225,11 @@ def mamba_chunk_scan_combined_varlen( ) return varlen_states + + +from vllm.platforms import current_platform # noqa: E402 + +if current_platform.is_cpu(): + import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm + + _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 92258ef204bd..24bdba688e70 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -4,15 +4,16 @@ Dispatch module for Mamba selective state update (SSU) backends. Provides a unified `selective_state_update` function that dispatches to -either the Triton or FlashInfer backend based on the configured -`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM. +the Triton, FlashInfer, or CPU backend based on the configured +`MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA) +the backend defaults to 'cpu'. """ from abc import ABC, abstractmethod import torch -from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.logger import init_logger from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -125,8 +126,13 @@ def __init__(self, mamba_config: MambaConfig): "Please install flashinfer (>= 0.6.4): " "pip install flashinfer-python" ) from e + logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm) self._kernel = _fi_ssu + @property + def _algorithm(self) -> MambaSSUAlgorithm: + return self._mamba_config.ssu_algorithm or "auto" + @property def name(self) -> str: return "flashinfer" @@ -156,7 +162,6 @@ def __call__( if self._mamba_config.enable_stochastic_rounding else None ) - self._kernel( state, x, @@ -179,12 +184,79 @@ def __call__( out=out, rand_seed=rand_seed, philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10, + algorithm=self._algorithm, + ) + + +class CPUSSUBackend(MambaSSUBackend): + """CPU SSU backend using the compiled C++ VSX/scalar kernel. + + On CPU-only platforms (PowerPC, x86 without CUDA) this dispatches to + the vectorized C++ kernel registered as ``torch.ops._C.selective_state_update_cpu``. + That kernel uses vec_op SIMD intrinsics (VSX on ppc64le, AVX2 on x86, + scalar fallback elsewhere) and is parallelised with OpenMP across heads. + + Falls back to the pure-PyTorch implementation only if the C++ op is + unavailable (e.g. a CPU-less build). + """ + + def __init__(self, mamba_config: MambaConfig): + super().__init__(mamba_config) + from vllm import _custom_ops as ops + + self._cpp_kernel = ops.selective_state_update_cpu + logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.") + + @property + def name(self) -> str: + return "cpu" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, + ) -> None: + # C++ kernel: state shape expected as (nstates, nheads, dim, dstate) + # The kernel writes in-place into `out` and updates `state`. + self._cpp_kernel( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, ) _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { MambaBackendEnum.TRITON: TritonSSUBackend, MambaBackendEnum.FLASHINFER: FlashInferSSUBackend, + MambaBackendEnum.CPU: CPUSSUBackend, } _mamba_ssu_backend: MambaSSUBackend | None = None @@ -210,6 +282,20 @@ def initialize_mamba_ssu_backend( global _mamba_ssu_backend backend = mamba_config.backend + + # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is + # unstable or unavailable. Silently fall back to the CPU + # backend unless the user explicitly chose something other than "triton". + if backend == MambaBackendEnum.TRITON: + from vllm.platforms import current_platform + + if current_platform.is_cpu(): + logger.info( + "CPU platform detected: overriding Mamba SSU backend " + "from 'triton' to 'cpu'." + ) + backend = MambaBackendEnum.CPU + if backend not in _BACKEND_REGISTRY: raise ValueError( f"Unknown Mamba SSU backend: {backend}. " diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index e7e36f2fc538..64fd14c3cfdb 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -23,6 +23,7 @@ causal_conv1d_fn, causal_conv1d_update, ) +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata @@ -42,6 +43,7 @@ def __init__( layer_idx: int, model_config: ModelConfig | None = None, cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() @@ -67,12 +69,14 @@ def __init__( input_size=dim, output_sizes=[dim] * 3, bias=self.bias, + quant_config=quant_config, prefix=f"{prefix}.in_proj", ) self.out_proj = RowParallelLinear( input_size=dim, output_size=dim, bias=self.bias, + quant_config=quant_config, prefix=f"{prefix}.out_proj", ) @@ -94,9 +98,13 @@ def forward_native( # 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_fn_cpu as causal_conv1d_torch, + ) + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_update_cpu, causal_conv1d_update_torch, ) + from vllm.platforms import CpuArchEnum, current_platform forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -164,17 +172,26 @@ def forward_native( 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 + Bx_d = B_d * x_d # (num_decodes, dim) + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + conv_state_view = conv_state[state_indices_d].contiguous() + out_d = causal_conv1d_update_torch( + Bx_d.unsqueeze(-1), + conv_state_view, + conv_weights, + self.conv.bias, + activation=None, + ).squeeze(-1) + conv_state[state_indices_d] = conv_state_view + else: + out_d = causal_conv1d_update_cpu( + Bx_d, + conv_state, + conv_weights, + self.conv.bias, + activation=None, + conv_state_indices=state_indices_d, + ) conv_output_list.insert(0, C_d * out_d) hidden_states_out = torch.vstack(conv_output_list) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 23733f769383..da6bbf2e0084 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -175,7 +175,7 @@ def forward_xpu( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return self.forward_native( + return torch.ops._xpu_C.mhc_pre( residual, fn, hc_scale, @@ -185,9 +185,6 @@ def forward_xpu( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - n_splits, - norm_weight, - norm_eps, ) @@ -260,7 +257,7 @@ def forward_xpu( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - return self.forward_native( + return torch.ops._xpu_C.mhc_post( x, residual, post_layer_mix, @@ -369,16 +366,8 @@ def forward_xpu( out = torch.empty( num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device ) - torch.ops.vllm.hc_head_triton( - hs_flat, - hc_fn, - hc_scale, - hc_base, - out, - hidden_size, - rms_norm_eps, - hc_eps, - hc_mult, + torch.ops._xpu_C.hc_head_fused( + hs_flat, hc_fn, hc_scale, hc_base, out, rms_norm_eps, hc_eps ) return out.view(*outer_shape, hidden_size) @@ -548,7 +537,7 @@ def forward_xpu( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - return self.forward_native( + return torch.ops._xpu_C.mhc_fused_post_pre( x, residual, post_layer_mix, @@ -561,8 +550,4 @@ def forward_xpu( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - n_splits, - tile_n, - norm_weight, - norm_eps, ) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 66a95b43c71f..6c0e8e069fc7 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -8,6 +8,7 @@ from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.platforms import current_platform @dataclass @@ -27,6 +28,7 @@ class MLAModules: is_sparse: bool topk_indices_buffer: torch.Tensor | None indexer_rotary_emb: torch.nn.Module | None = None + g_proj: torch.nn.Module | None = None # --8<-- [start:multi_head_latent_attention] @@ -65,6 +67,8 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", skip_topk: bool = False, + non_causal_multi_token_decode: bool = False, + allow_short_prefill_indexer_scoring_skip: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -87,12 +91,17 @@ def __init__( self.indexer = mla_modules.indexer self.indexer_rope_emb = mla_modules.indexer_rotary_emb self.is_sparse = mla_modules.is_sparse + self.g_proj = mla_modules.g_proj # Whether to skip top-k token selection computation in this layer. # When True, the indexer will not be called, and the layer will reuse # the topk_tokens buffer written by a previous layer in the same pass. # Refer: https://arxiv.org/abs/2603.12201 for more details. self.skip_topk = skip_topk + # qrep is active when the query projection is a DCP-group-sharded layer + # that materializes the full group head set locally. + q_proj_layer = self.q_b_proj if self.q_lora_rank is not None else self.q_proj + self.dcp_q_replicate = getattr(q_proj_layer, "qrep_active", False) if self.indexer is not None: assert hasattr(self.indexer, "topk_tokens") self.topk_tokens = self.indexer.topk_tokens @@ -110,11 +119,32 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.attn", kv_b_proj=self.kv_b_proj, + dcp_q_replicate=self.dcp_q_replicate, use_sparse=self.is_sparse, indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, + non_causal_multi_token_decode=non_causal_multi_token_decode, ) - + indexer_op = getattr(self.indexer, "indexer_op", None) + if indexer_op is not None and hasattr( + indexer_op, "dense_mha_metadata_layer_name" + ): + enable_short_prefill_scoring_skip = ( + allow_short_prefill_indexer_scoring_skip + and not self.skip_topk + and not getattr(indexer_op, "use_pcp", False) + and current_platform.is_cuda() + ) + # The indexer and main MLA use independent decode thresholds and + # may classify the same short extend differently. Bind the main + # MLA layer name so the eager indexer op can check whether the + # batch's top-k indices will be consumed. + # PCP is excluded because indexer cache/scoring ownership differs + # across ranks and the no-consumer invariant has not been + # established there. + indexer_op.dense_mha_metadata_layer_name = ( + self.mla_attn.layer_name if enable_short_prefill_scoring_skip else "" + ) self.prefix = prefix def forward( @@ -143,7 +173,8 @@ def forward( dim=-1, ) q_c = self.q_a_layernorm(q_c) - q = self.q_b_proj(q_c)[0] + q_proj_layer = self.q_b_proj + q_proj_input = q_c else: assert self.kv_a_proj_with_mqa is not None, ( "kv_a_proj_with_mqa is required when q_lora_rank is None" @@ -152,15 +183,20 @@ def forward( "q_proj is required when q_lora_rank is None" ) kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] - q = self.q_proj(hidden_states)[0] + q_proj_layer = self.q_proj + q_proj_input = hidden_states kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) kv_c_normed = self.kv_a_layernorm(kv_c) - - q = q.view(-1, self.num_heads, self.qk_head_dim) # Add head dim of 1 to k_pe k_pe = k_pe.unsqueeze(1) + q = q_proj_layer(q_proj_input)[0] + heads = self.num_heads + if self.dcp_q_replicate: + heads *= q_proj_layer.group_size + q = q.view(-1, heads, self.qk_head_dim) + if self.rotary_emb is not None: q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( positions, q[..., self.qk_nope_head_dim :], k_pe @@ -172,11 +208,19 @@ def forward( if llama_4_scaling is not None: q *= llama_4_scaling + q_dcp_replicated = None + if self.dcp_q_replicate: + q_dcp_replicated, q = q, q_proj_layer._local_view(q) + attn_out = self.mla_attn( q, kv_c_normed, k_pe, output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim), + q_dcp_replicated=q_dcp_replicated, ) + if self.g_proj is not None: + attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid() + return self.o_proj(attn_out)[0] diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 1e49ec3387e4..5e9946ff35b8 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -560,6 +560,10 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, kInt4Static, + quant_config=self.quant_config, + may_have_zp=self.quant_config.zero_point, + may_have_bias=True, + allow_tile_padding=True, ) def create_weights( @@ -779,16 +783,6 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi gemm1_beta=getattr(layer, "swiglu_beta", None), ) - def select_gemm_impl( - self, - prepare_finalize, - layer: RoutedExperts, - ): - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel " - "initialization logic. This function should not be called." - ) - def apply( self, layer: RoutedExperts, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index b056f5222af8..710e947f0bbf 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -489,6 +489,10 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, weight_key, + quant_config=self.quant_config, + may_have_zp=not self.quant_config.is_sym, + may_have_bias=True, + allow_tile_padding=not self.quant_config.desc_act, ) def create_weights( @@ -577,25 +581,25 @@ def create_weights( set_weight_attrs(w2_scales, extra_weight_attrs) # don't shard the w2 scales when running act order set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) - # up_proj scales + # up_proj zero points w13_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size13, 2 * intermediate_size_per_partition // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) layer.register_parameter("w13_qzeros", w13_qzeros) set_weight_attrs(w13_qzeros, extra_weight_attrs) - # down_proj scales + # down_proj zero points w2_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size2, hidden_size // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) @@ -644,6 +648,26 @@ def create_weights( layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + # Some GPTQ checkpoints contain expert biases even when the model + # architecture does not declare them. Zero initialization keeps + # checkpoints without biases equivalent to the bias-free path. + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + if self.experts_cls is not None and issubclass( self.experts_cls, FusedMoEExpertsModular ): @@ -651,12 +675,31 @@ def create_weights( layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + def replace_or_register(name: str, val: torch.Tensor | None): + if val is None: + return + + if hasattr(layer, name): + replace_parameter(layer, name, val) + else: + layer.register_parameter( + name, torch.nn.Parameter(val, requires_grad=False) + ) + is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1 - if is_a_8bit: - assert self.quant_config.quant_type.size_bits == 8, ( - "W8A8-INT8 is not supported by marlin kernel." - ) + assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, ( + "W8A8-INT8 is not supported by marlin kernel." + ) + + w13_bias = getattr(layer, "w13_bias", None) + if "w13_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w13_bias", None) + w13_bias = None + w2_bias = getattr(layer, "w2_bias", None) + if "w2_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w2_bias", None) + w2_bias = None converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, @@ -669,8 +712,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: 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), + w13_bias=w13_bias, + w2_bias=w2_bias, + w13_qzeros=getattr(layer, "w13_qzeros", None), + w2_qzeros=getattr(layer, "w2_qzeros", None), ) if converted is None: @@ -703,42 +748,12 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w2_g_idx", w2_g_idx) replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - if w13_qzeros is not None: - replace_parameter(layer, "w13_qzeros", w13_qzeros) - if w2_qzeros is not None: - replace_parameter(layer, "w2_qzeros", w2_qzeros) - if w13_input_global_scale is not None: - if hasattr(layer, "w13_input_global_scale"): - replace_parameter( - layer, "w13_input_global_scale", w13_input_global_scale - ) - else: - layer.register_parameter( - "w13_input_global_scale", - torch.nn.Parameter(w13_input_global_scale, requires_grad=False), - ) - if w2_input_global_scale is not None: - if hasattr(layer, "w2_input_global_scale"): - replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale) - else: - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - if w13_bias is not None: - if hasattr(layer, "w13_bias"): - replace_parameter(layer, "w13_bias", w13_bias) - else: - layer.register_parameter( - "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False) - ) - if w2_bias is not None: - if hasattr(layer, "w2_bias"): - replace_parameter(layer, "w2_bias", w2_bias) - else: - layer.register_parameter( - "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False) - ) + replace_or_register("w13_input_global_scale", w13_input_global_scale) + replace_or_register("w2_input_global_scale", w2_input_global_scale) + replace_or_register("w13_bias", w13_bias) + replace_or_register("w2_bias", w2_bias) + replace_or_register("w13_qzeros", w13_qzeros) + replace_or_register("w2_qzeros", w2_qzeros) # The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight. layer.w13_weight = layer.w13_qweight @@ -792,16 +807,6 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi w2_bias=getattr(layer, "w2_bias", None), ) - def select_gemm_impl( - self, - prepare_finalize, - layer: RoutedExperts, - ): - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel " - "initialization logic. This function should not be called." - ) - def apply( self, layer: RoutedExperts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 260a8309ff7d..6fea370824f0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -215,7 +215,7 @@ def get_quant_method( def _add_fused_moe_to_target_scheme_map(self): # XXXXXXXXXXXXXXXXXXXXXX """ Helper function to update target_scheme_map - since linear layers get fused into FusedMoE + since linear layers get fused into RoutedExperts targeting 'Linear' needs to also match RoutedExperts modules. """ diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 7c03baf1e005..78d70d492d2b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -5,7 +5,6 @@ import torch from compressed_tensors import CompressionFormat from compressed_tensors.quantization import ( - ActivationOrdering, QuantizationStrategy, QuantizationType, ) @@ -19,9 +18,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa WNA16_SUPPORTED_BITS, ) -from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_moe_marlin_supports_layer, -) from vllm.platforms import current_platform logger = init_logger(__name__) @@ -78,9 +74,6 @@ def get_moe_method( return CompressedTensorsW8A8Mxfp8MoEMethod(layer.moe_config) if quant_config._is_wNa16_group_channel(weight_quant, input_quant): - # group_size=None means channelwise - group_size = weight_quant.group_size or -1 - valid_format_and_bits = ( weight_quant.num_bits in WNA16_SUPPORTED_BITS and format == CompressionFormat.pack_quantized.value @@ -95,87 +88,53 @@ def get_moe_method( f" and bits: {weight_quant.num_bits}", ) - # Prefer to use the MarlinMoE kernel when it is supported. - is_actorder = ( - weight_quant.strategy == QuantizationStrategy.GROUP - and weight_quant.actorder - in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC) - ) - if ( - not check_moe_marlin_supports_layer( - layer, group_size, allow_tile_padding=not is_actorder - ) - or current_platform.is_rocm() - ): - if is_actorder: - raise ValueError( - "WNA16MoE is not supported with actorder=group/dynamic." + # Native ROCm HIP kernels (RDNA3, etc.) + if current_platform.is_rocm(): + from . import rocm_moe_rdna + + if rocm_moe_rdna.is_supported(weight_quant): + return rocm_moe_rdna.make_method( + weight_quant, input_quant, layer.moe_config + ) + from vllm.platforms.rocm import on_gfx950 + + vllm_config = get_current_vllm_config() + is_lora_disabled = vllm_config.lora_config is None + moe_backend = vllm_config.kernel_config.moe_backend + group_size = weight_quant.group_size or -1 + if ( + weight_quant.strategy == QuantizationStrategy.GROUP + and weight_quant.type == QuantizationType.INT + and group_size == 32 + and weight_quant.num_bits == 4 + and is_lora_disabled + and on_gfx950() + and moe_backend == "flydsl" + ): + from .compressed_tensors_moe_w4a16_flydsl import ( + CompressedTensorsW4A16FlydslMoEMethod, ) - # Native ROCm HIP kernels (RDNA3, etc.) - if current_platform.is_rocm(): - from . import rocm_moe_rdna - - if rocm_moe_rdna.is_supported(weight_quant): - return rocm_moe_rdna.make_method( - weight_quant, input_quant, layer.moe_config - ) - from vllm.platforms.rocm import on_gfx950 - - vllm_config = get_current_vllm_config() - is_lora_disabled = vllm_config.lora_config is None - moe_backend = vllm_config.kernel_config.moe_backend - if ( - weight_quant.strategy == QuantizationStrategy.GROUP - and weight_quant.type == QuantizationType.INT - and group_size == 32 - and weight_quant.num_bits == 4 - and is_lora_disabled - and on_gfx950() - and moe_backend == "flydsl" - ): - from .compressed_tensors_moe_w4a16_flydsl import ( - CompressedTensorsW4A16FlydslMoEMethod, - ) - - logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod") - return CompressedTensorsW4A16FlydslMoEMethod( - weight_quant, input_quant, layer.moe_config - ) - elif moe_backend == "emulation": - # Although this is called 'Marlin', actually it selects - # emulation backend by calling select_wna16_moe_backend. - # TODO: we need to update CompressedTensorsWNA16MoeMethod - # to honor "--moe-backend" option - from .compressed_tensors_moe_wna16_marlin import ( - CompressedTensorsWNA16MarlinMoEMethod, - ) - - logger.info_once( - "Using CompressedTensorsWNA16MarlinMoEMethod " - "(emulation backend requested)" - ) - return CompressedTensorsWNA16MarlinMoEMethod( - weight_quant, input_quant, layer.moe_config, layer_name - ) - - from .compressed_tensors_moe_wna16 import ( - CompressedTensorsWNA16MoEMethod, - ) + logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod") + return CompressedTensorsW4A16FlydslMoEMethod( + weight_quant, input_quant, layer.moe_config + ) + elif moe_backend == "emulation": + logger.info_once( + "Using CompressedTensorsWNA16MoEMethod " + "(emulation backend requested)" + ) - logger.info_once("Using CompressedTensorsWNA16MoEMethod") - return CompressedTensorsWNA16MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - else: - from .compressed_tensors_moe_wna16_marlin import ( - CompressedTensorsWNA16MarlinMoEMethod, - ) + from .compressed_tensors_moe_wna16 import ( + CompressedTensorsWNA16MoEMethod, + ) - logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod") - return CompressedTensorsWNA16MarlinMoEMethod( - weight_quant, input_quant, layer.moe_config - ) + logger.info_once("Using CompressedTensorsWNA16MoEMethod") + return CompressedTensorsWNA16MoEMethod( + weight_quant, + input_quant, + layer.moe_config, + ) elif quant_config._is_nvfp4_format(weight_quant): from .compressed_tensors_moe_w4a4_nvfp4 import ( CompressedTensorsW4A4Nvfp4MoEMethod, 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 d304ca56bf17..28100b012e12 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 @@ -162,6 +162,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: routing_tables=layer._expert_routing_tables(), layer=layer, ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def maybe_make_prepare_finalize( self, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 2dabf7a5154b..eaec0fec4540 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -1,28 +1,51 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any import torch from compressed_tensors.quantization import ( + ActivationOrdering, QuantizationArgs, + QuantizationStrategy, ) -import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( + FusedMoEExpertsModular, RoutedExperts, SharedExperts, ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, - int8_w8a16_moe_quant_config, +) +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, + select_wna16_moe_backend, ) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa + WNA16_SUPPORTED_TYPES_MAP, + WNA16_ZP_SUPPORTED_TYPES_MAP, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_moe_marlin_supports_config, + get_marlin_input_dtype, + marlin_make_workspace_new, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static32GroupScale, + kInt4StaticGroupScale, + kInt8StaticGroupScale, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs logger = init_logger(__name__) @@ -39,19 +62,180 @@ def __init__( self.weight_quant = weight_quant self.input_quant = input_quant # Extract properties from weight_quant + self.symmetric = weight_quant.symmetric self.num_bits = weight_quant.num_bits self.packed_factor = 32 // weight_quant.num_bits self.strategy = weight_quant.strategy - # channelwise is not supported by this kernel - assert weight_quant.strategy == "group" self.group_size = weight_quant.group_size - # grouped actorder isn't supported by this kernel - assert weight_quant.actorder != "group" - assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE. " - "Try --moe-backend emulation." + self.actorder = weight_quant.actorder + + # Extract quant_type and create weight key for oracle selection + self.quant_type = ( + WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + if self.symmetric + else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits] + ) + + if self.num_bits == 4: + if self.group_size == 32: + scale = kInt4Static32GroupScale + else: + scale = kInt4StaticGroupScale + elif self.num_bits == 8: + assert self.group_size == -1 + scale = kInt8StaticGroupScale + else: + raise ValueError( + "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now." + ) + + weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) + + is_actorder = self.strategy == QuantizationStrategy.GROUP and self.actorder in ( + ActivationOrdering.GROUP, + ActivationOrdering.DYNAMIC, ) + # Select WNA16 MoE backend via oracle. + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + weight_key=weight_key, + quant_config=self.weight_quant, + may_have_zp=not self.symmetric, + may_have_bias=False, + allow_tile_padding=not is_actorder, + ) + + self.is_marlin = self.wna16_backend in [ + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + ] + self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM + + if self.is_marlin: + assert check_moe_marlin_supports_config( + self.moe, self.group_size, allow_tile_padding=not is_actorder + ) + self.input_dtype = get_marlin_input_dtype(layer_name) + else: + # channelwise is not supported by this kernel + assert weight_quant.strategy == "group" + # grouped actorder isn't supported by this kernel + assert weight_quant.actorder != "group" + + assert self.symmetric, "Only symmetric quantization is supported for MoE" + + # Non-Marlin WNA16 always uses bf16/fp16 inputs + self.input_dtype = torch.bfloat16 + + def get_weight_shape( + self, + weight_name: str, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + num_groups_w2: int | None = None, + num_groups_w13: int | None = None, + ) -> tuple[int, int, int]: + """ + Get the shape of the weight based on the weight name, number of experts + hidden size, intermediate size per partition, number of groups for w2, + and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 + for weight scales/zero_points. + """ + if weight_name in ("w13_scale", "w13_zp"): + assert num_groups_w13 is not None, ( + "num_groups_w13 must be provided for weight scales/zero_points" + ) + if weight_name in ("w2_scale", "w2_zp"): + assert num_groups_w2 is not None, ( + "num_groups_w2 must be provided for weight scales/zero_points" + ) + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + shape_map = { + "w13_weight": { + "Flashinfer": ( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size // self.packed_factor, + ), + "Marlin": ( + num_experts, + hidden_size // self.packed_factor, + w13_num_shards * intermediate_size_per_partition, + ), + }, + "w13_scale": { + "Flashinfer": ( + num_experts, + w13_num_shards * intermediate_size_per_partition, + num_groups_w13, + ), + "Marlin": ( + num_experts, + num_groups_w13, + w13_num_shards * intermediate_size_per_partition, + ), + }, + "w13_zp": { + "Marlin": ( + num_experts, + num_groups_w13, + w13_num_shards + * intermediate_size_per_partition + // self.packed_factor, + ), + }, + "w2_weight": { + "Flashinfer": ( + num_experts, + hidden_size, + intermediate_size_per_partition // self.packed_factor, + ), + "Marlin": ( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + ), + }, + "w2_scale": { + "Flashinfer": (num_experts, hidden_size, num_groups_w2), + "Marlin": (num_experts, num_groups_w2, hidden_size), + }, + "w2_zp": { + "Marlin": ( + num_experts, + num_groups_w2, + hidden_size // self.packed_factor, + ), + }, + } + backend_key = "Marlin" if self.is_transposed else "Flashinfer" + return shape_map[weight_name][backend_key] + + @staticmethod + def _w2_scale_sharding( + actorder, + group_size: int, + intermediate_size_per_partition: int, + intermediate_size_full: int, + ) -> tuple[bool, int, bool]: + """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE. + + Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime + and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``. + ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at + quantization time, so scales shard normally per TP rank. + """ + load_full_w2 = (actorder == "group") and group_size != -1 + w2_scales_size = ( + intermediate_size_full if load_full_w2 else intermediate_size_per_partition + ) + is_k_full = (actorder != "group") or ( + intermediate_size_per_partition == intermediate_size_full + ) + return load_full_w2, w2_scales_size, is_k_full + def create_weights( self, layer: torch.nn.Module, @@ -61,18 +245,23 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): + intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") + # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims extra_weight_attrs.update( - {"is_transposed": True, "quant_method": self.strategy} + {"is_transposed": self.is_transposed, "quant_method": self.strategy} ) - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + w13_weight = torch.nn.Parameter( torch.empty( - num_experts, - hidden_size // self.packed_factor, - w13_num_shards * intermediate_size_per_partition, + *self.get_weight_shape( + "w13_weight", + num_experts, + hidden_size, + intermediate_size_per_partition, + ), dtype=torch.int32, ), requires_grad=False, @@ -82,9 +271,12 @@ def create_weights( w2_weight = torch.nn.Parameter( torch.empty( - num_experts, - intermediate_size_per_partition // self.packed_factor, - hidden_size, + *self.get_weight_shape( + "w2_weight", + num_experts, + hidden_size, + intermediate_size_per_partition, + ), dtype=torch.int32, ), requires_grad=False, @@ -92,7 +284,12 @@ def create_weights( layer.register_parameter("w2_weight_packed", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) - w2_scales_size = intermediate_size_per_partition + load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding( + self.actorder, + self.group_size, + intermediate_size_per_partition, + intermediate_size_full, + ) if self.strategy == "channel": num_groups_w2 = num_groups_w13 = 1 @@ -104,23 +301,34 @@ def create_weights( f"({hidden_size}) to be divisible by group_size " f"({self.group_size})." ) - if intermediate_size_per_partition % self.group_size != 0: + if ( + not load_full_w2 + and intermediate_size_per_partition % self.group_size != 0 + ): raise ValueError( - "CompressedTensors WNA16 MoE with static group scales " - "requires the MoE intermediate size per tensor-parallel " - f"partition ({intermediate_size_per_partition}) to be " - f"divisible by group_size ({self.group_size}). Scale " - "groups would otherwise cross TP shard boundaries; use a " - "compatible TP size or enable expert parallelism." + "CompressedTensors WNA16 MoE with static group " + "scales requires the MoE intermediate size per " + "tensor-parallel partition " + f"({intermediate_size_per_partition}) to be divisible by " + f"group_size ({self.group_size}). Scale groups would " + "otherwise cross TP shard boundaries; use a compatible TP " + "size or enable expert parallelism." ) num_groups_w2 = w2_scales_size // self.group_size num_groups_w13 = hidden_size // self.group_size + layer.num_groups_w13 = num_groups_w13 + layer.num_groups_w2 = num_groups_w2 + w13_scale = torch.nn.Parameter( torch.ones( - num_experts, - num_groups_w13, - w13_num_shards * intermediate_size_per_partition, + *self.get_weight_shape( + "w13_scale", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w13=num_groups_w13, + ), dtype=params_dtype, ), requires_grad=False, @@ -129,12 +337,54 @@ def create_weights( set_weight_attrs(w13_scale, extra_weight_attrs) w2_scale = torch.nn.Parameter( - torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + torch.ones( + *self.get_weight_shape( + "w2_scale", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w2=num_groups_w2, + ), + dtype=params_dtype, + ), requires_grad=False, ) layer.register_parameter("w2_weight_scale", w2_scale) set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": False}) + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + + if not self.symmetric: + w13_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w13_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w13=num_groups_w13, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_zero_point", w13_zp) + set_weight_attrs(w13_zp, extra_weight_attrs) + + w2_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w2_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w2=num_groups_w2, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_zero_point", w2_zp) + set_weight_attrs(w2_zp, extra_weight_attrs) w2_weight_shape = torch.nn.Parameter( torch.empty(num_experts, 2), requires_grad=False @@ -195,65 +445,161 @@ def create_weights( layer.a13_scale = None layer.a2_scale = None - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Reconfigure packed weights and scales to match moe_wna16 format - layer.w13_weight_packed = torch.nn.Parameter( - layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), - requires_grad=False, - ) - layer.w2_weight_packed = torch.nn.Parameter( - layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), - requires_grad=False, - ) - layer.w13_weight_scale = torch.nn.Parameter( - layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False + def _setup_kernel(self, layer: RoutedExperts): + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + + # Add Marlin-specific arguments + marlin_args: dict[str, Any] = {} + if self.is_marlin: + marlin_args = { + "w13_g_idx": layer.w13_weight_g_idx, + "w2_g_idx": layer.w2_weight_g_idx, + "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices, + "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices, + "is_k_full": self.is_k_full, + } + + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + **marlin_args, ) - layer.w2_weight_scale = torch.nn.Parameter( - layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Process weights using the shared oracle infrastructure + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.weight_quant, + input_dtype=self.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: + self._setup_kernel(layer) + return + + ( + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, + w13_g_idx_processed, + w2_g_idx_processed, + w13_g_idx_sort_indices, + w2_g_idx_sort_indices, + w13_qzeros, + w2_qzeros, + w13_input_global_scale, + w2_input_global_scale, + _, # w13_bias + _, # w2_bias + ) = converted + + # Replace common parameters + replace_parameter(layer, "w13_weight_packed", w13_qweight) + replace_parameter(layer, "w2_weight_packed", w2_qweight) + replace_parameter(layer, "w13_weight_scale", w13_scales) + replace_parameter(layer, "w2_weight_scale", w2_scales) + + # CPU fused_experts_cpu requires zero points even for symmetric quant + if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU: + assert w13_qzeros is not None and w2_qzeros is not None + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) + + # Marlin-specific parameters (not needed for Flashinfer) + if self.is_marlin: + if w13_g_idx_processed is not None: + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) + if w2_g_idx_processed is not None: + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) + if w13_g_idx_sort_indices is not None: + replace_parameter( + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices + ) + if w2_g_idx_sort_indices is not None: + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + + # Register input global scales if present + if w13_input_global_scale is not None: + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + if w2_input_global_scale is not None: + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + + # Marlin workspace — only needed for Marlin-family backends, not emulation. + if ( + self.experts_cls is not None + and issubclass(self.experts_cls, FusedMoEExpertsModular) + and self.wna16_backend != WNA16MoEBackend.EMULATION + ): + layer.workspace = marlin_make_workspace_new( + layer.w13_weight_g_idx.device, + 4, + existing=getattr(layer, "workspace", None), + ) + + # Alias packed weights to w13_weight/w2_weight for the modular kernel interface + layer.w13_weight = layer.w13_weight_packed + layer.w2_weight = layer.w2_weight_packed + + self._setup_kernel(layer) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - assert self.num_bits == 4 or self.num_bits == 8 - config_builder = ( - int4_w4a16_moe_quant_config - if self.num_bits == 4 - else int8_w8a16_moe_quant_config - ) - - return config_builder( + return make_wna16_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, self.group_size], + group_size=self.group_size, + num_bits=self.num_bits, + w1_zp=getattr(layer, "w13_weight_zero_point", None), + w2_zp=getattr(layer, "w2_weight_zero_point", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) - def select_gemm_impl( + def apply_monolithic( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - if self.moe.is_lora_enabled: - assert self.moe_quant_config is not None - from vllm.triton_utils import HAS_TRITON - - if HAS_TRITON: - from vllm.model_executor.layers.fused_moe import TritonWNA16Experts - - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - return TritonWNA16Experts( - moe_config=self.moe, quant_config=self.moe_quant_config - ) - else: - raise NotImplementedError( - "TritonExperts requires Triton. " - "Install triton or disable LoRA for MoE." - ) - - raise NotImplementedError + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) def apply( self, @@ -264,21 +610,22 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, - layer.w13_weight_packed, - layer.w2_weight_packed, + layer.w13_weight, + layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - quant_config=self.moe_quant_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, ) @property def supports_eplb(self) -> bool: - return True + return self.wna16_backend == WNA16MoEBackend.TRITON 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 deleted file mode 100644 index 2b3317d00f3c..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ /dev/null @@ -1,592 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import torch -from compressed_tensors.quantization import ( - QuantizationArgs, -) - -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import ( - FusedMoEExpertsModular, - RoutedExperts, - SharedExperts, -) -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEQuantConfig, -) -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, - select_wna16_moe_backend, -) -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 - CompressedTensorsMoEMethod, -) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa - WNA16_SUPPORTED_TYPES_MAP, - WNA16_ZP_SUPPORTED_TYPES_MAP, -) -from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - get_marlin_input_dtype, - marlin_make_workspace_new, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kInt4Static32GroupScale, - kInt4StaticGroupScale, - kInt8StaticGroupScale, -) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs - -logger = init_logger(__name__) - - -class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs | None, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - self.symmetric = weight_quant.symmetric - # Extract properties from weight_quant - self.num_bits = weight_quant.num_bits - self.packed_factor = 32 // weight_quant.num_bits - self.strategy = weight_quant.strategy - self.group_size = weight_quant.group_size - self.actorder = weight_quant.actorder - - self.quant_type = ( - WNA16_SUPPORTED_TYPES_MAP[self.num_bits] - if self.symmetric - else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits] - ) - - self.marlin_input_dtype = get_marlin_input_dtype(layer_name) - - if self.num_bits == 4: - if self.group_size == 32: - scale = kInt4Static32GroupScale - else: - scale = kInt4StaticGroupScale - elif self.num_bits == 8: - scale = kInt8StaticGroupScale - else: - raise ValueError( - "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now." - ) - - weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) - - # Select WNA16 MoE backend via oracle. - self.wna16_backend, self.experts_cls = select_wna16_moe_backend( - config=self.moe, - weight_key=weight_key, - ) - - def get_weight_shape( - self, - weight_name: str, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - num_groups_w2: int | None = None, - num_groups_w13: int | None = None, - ) -> tuple[int, int, int]: - """ - Get the shape of the weight based on the weight name, number of experts - hidden size, intermediate size per partition, number of groups for w2, - and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 - for weight scales/zero_points. - """ - if weight_name in ("w13_scale", "w13_zp"): - assert num_groups_w13 is not None, ( - "num_groups_w13 must be provided for weight scales/zero_points" - ) - if weight_name in ("w2_scale", "w2_zp"): - assert num_groups_w2 is not None, ( - "num_groups_w2 must be provided for weight scales/zero_points" - ) - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM - shape_map = { - "w13_weight": { - "Flashinfer": ( - num_experts, - w13_num_shards * intermediate_size_per_partition, - hidden_size // self.packed_factor, - ), - "Marlin": ( - num_experts, - hidden_size // self.packed_factor, - w13_num_shards * intermediate_size_per_partition, - ), - }, - "w13_scale": { - "Flashinfer": ( - num_experts, - w13_num_shards * intermediate_size_per_partition, - num_groups_w13, - ), - "Marlin": ( - num_experts, - num_groups_w13, - w13_num_shards * intermediate_size_per_partition, - ), - }, - "w13_zp": { - "Marlin": ( - num_experts, - num_groups_w13, - w13_num_shards - * intermediate_size_per_partition - // self.packed_factor, - ), - }, - "w2_weight": { - "Flashinfer": ( - num_experts, - hidden_size, - intermediate_size_per_partition // self.packed_factor, - ), - "Marlin": ( - num_experts, - intermediate_size_per_partition // self.packed_factor, - hidden_size, - ), - }, - "w2_scale": { - "Flashinfer": (num_experts, hidden_size, num_groups_w2), - "Marlin": (num_experts, num_groups_w2, hidden_size), - }, - "w2_zp": { - "Marlin": ( - num_experts, - num_groups_w2, - hidden_size // self.packed_factor, - ), - }, - } - backend_key = "Flashinfer" if is_flashinfer else "Marlin" - return shape_map[weight_name][backend_key] - - @staticmethod - def _w2_scale_sharding( - actorder, - group_size: int, - intermediate_size_per_partition: int, - intermediate_size_full: int, - ) -> tuple[bool, int, bool]: - """Decide how to shard w2 group scales across TP for WNA16 Marlin MoE. - - Only ``actorder="group"`` permutes activations by ``g_idx`` at runtime - and therefore needs the full-K (unsharded) w2 scales plus ``is_k_full``. - ``actorder="weight"``/``"static"`` (and ``None``) reorder weights at - quantization time, so scales shard normally per TP rank. - """ - load_full_w2 = (actorder == "group") and group_size != -1 - w2_scales_size = ( - intermediate_size_full if load_full_w2 else intermediate_size_per_partition - ) - is_k_full = (actorder != "group") or ( - intermediate_size_per_partition == intermediate_size_full - ) - return load_full_w2, w2_scales_size, is_k_full - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") - - # Will transpose the loaded weight along the - # intermediate and hidden dim sizes. Will - # shard for TP along the transposed dims - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM - extra_weight_attrs.update( - {"is_transposed": is_transposed, "quant_method": self.strategy} - ) - - w13_weight = torch.nn.Parameter( - torch.empty( - *self.get_weight_shape( - "w13_weight", - num_experts, - hidden_size, - intermediate_size_per_partition, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - *self.get_weight_shape( - "w2_weight", - num_experts, - hidden_size, - intermediate_size_per_partition, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - load_full_w2, w2_scales_size, self.is_k_full = self._w2_scale_sharding( - self.actorder, - self.group_size, - intermediate_size_per_partition, - intermediate_size_full, - ) - - if self.strategy == "channel": - num_groups_w2 = num_groups_w13 = 1 - self.group_size = -1 - else: - if hidden_size % self.group_size != 0: - raise ValueError( - "CompressedTensors WNA16 Marlin MoE requires hidden_size " - f"({hidden_size}) to be divisible by group_size " - f"({self.group_size})." - ) - if ( - not load_full_w2 - and intermediate_size_per_partition % self.group_size != 0 - ): - raise ValueError( - "CompressedTensors WNA16 Marlin MoE with static group " - "scales requires the MoE intermediate size per " - "tensor-parallel partition " - f"({intermediate_size_per_partition}) to be divisible by " - f"group_size ({self.group_size}). Scale groups would " - "otherwise cross TP shard boundaries; use a compatible TP " - "size or enable expert parallelism." - ) - num_groups_w2 = w2_scales_size // self.group_size - num_groups_w13 = hidden_size // self.group_size - - layer.num_groups_w13 = num_groups_w13 - layer.num_groups_w2 = num_groups_w2 - - w13_scale = torch.nn.Parameter( - torch.ones( - *self.get_weight_shape( - "w13_scale", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w13=num_groups_w13, - ), - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_scale) - set_weight_attrs(w13_scale, extra_weight_attrs) - - w2_scale = torch.nn.Parameter( - torch.ones( - *self.get_weight_shape( - "w2_scale", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w2=num_groups_w2, - ), - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_scale) - set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) - - if not self.symmetric: - w13_zp = torch.nn.Parameter( - torch.zeros( - *self.get_weight_shape( - "w13_zp", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w13=num_groups_w13, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_zero_point", w13_zp) - set_weight_attrs(w13_zp, extra_weight_attrs) - - w2_zp = torch.nn.Parameter( - torch.zeros( - *self.get_weight_shape( - "w2_zp", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w2=num_groups_w2, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_zero_point", w2_zp) - set_weight_attrs(w2_zp, extra_weight_attrs) - - w2_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - layer.register_parameter("w2_weight_shape", w2_weight_shape) - set_weight_attrs(w2_weight_shape, extra_weight_attrs) - w13_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - - layer.register_parameter("w13_weight_shape", w13_weight_shape) - set_weight_attrs(w13_weight_shape, extra_weight_attrs) - - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - - layer.a13_scale = None - layer.a2_scale = None - - 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, - w13_scales, - w2_scales, - w13_g_idx_processed, - w2_g_idx_processed, - w13_g_idx_sort_indices, - w2_g_idx_sort_indices, - w13_qzeros, - w2_qzeros, - w13_input_global_scale, - w2_input_global_scale, - _, # w13_bias - _, # w2_bias - ) = converted - - # Replace common parameters - replace_parameter(layer, "w13_weight_packed", w13_qweight) - replace_parameter(layer, "w2_weight_packed", w2_qweight) - replace_parameter(layer, "w13_weight_scale", w13_scales) - replace_parameter(layer, "w2_weight_scale", w2_scales) - - # CPU fused_experts_cpu requires zero points even for symmetric quant - if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU: - replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) - replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) - - # Marlin-specific parameters (not needed for Flashinfer) - if not is_flashinfer: - if w13_g_idx_processed is not None: - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) - if w2_g_idx_processed is not None: - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) - if w13_g_idx_sort_indices is not None: - replace_parameter( - layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices - ) - if w2_g_idx_sort_indices is not None: - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - - # Register input global scales if present - if w13_input_global_scale is not None: - layer.register_parameter( - "w13_input_global_scale", - torch.nn.Parameter(w13_input_global_scale, requires_grad=False), - ) - if w2_input_global_scale is not None: - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - - # Marlin workspace — only needed for Marlin-family backends, not emulation. - if ( - self.experts_cls is not None - and issubclass(self.experts_cls, FusedMoEExpertsModular) - and self.wna16_backend != WNA16MoEBackend.EMULATION - ): - layer.workspace = marlin_make_workspace_new( - layer.w13_weight_g_idx.device, 4 - ) - - # Alias packed weights to w13_weight/w2_weight for the modular kernel interface - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - - assert self.experts_cls is not None - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - assert self.moe_quant_config is not None - - # Add Marlin-specific arguments - marlin_args: dict[str, Any] = {} - if not is_flashinfer: - marlin_args = { - "w13_g_idx": layer.w13_weight_g_idx, - "w2_g_idx": layer.w2_weight_g_idx, - "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices, - "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices, - "is_k_full": self.is_k_full, - } - - self.moe_kernel = make_wna16_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - routing_tables=layer._expert_routing_tables(), - **marlin_args, - ) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return make_wna16_moe_quant_config( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - group_size=self.group_size, - num_bits=self.num_bits, - w1_zp=getattr(layer, "w13_weight_zero_point", None), - w2_zp=getattr(layer, "w2_weight_zero_point", None), - gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), - gemm1_alpha=getattr(layer, "swiglu_alpha", None), - gemm1_beta=getattr(layer, "swiglu_beta", None), - ) - - def apply_monolithic( - self, - layer: RoutedExperts, - x: torch.Tensor, - router_logits: torch.Tensor, - input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: - assert self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routed_scaling_factor=layer.routed_scaling_factor, - ) - - def apply( - self, - layer: RoutedExperts, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts: SharedExperts | None, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - assert not self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts=shared_experts, - shared_experts_input=shared_experts_input, - ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index 7b994e10e448..2eb8ba717bde 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py @@ -9,6 +9,9 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.model_executor.parameter import ( GroupQuantScaleParameter, ModelWeightParameter, @@ -35,7 +38,9 @@ class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme): def __init__(self): self.group_size = 32 - self.kernel = init_mxfp4_linear_kernel() + self.kernel = init_mxfp4_linear_kernel( + activation_quant_key=kMxfp4Dynamic, + ) @classmethod def get_min_capability(cls) -> int: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py index bd1964e667d9..0cde020627d2 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501 TransformTuple, ) +from vllm.platforms import current_platform class CompressedTensorsLinearTransformMethod(LinearMethodBase): @@ -48,11 +49,12 @@ def from_schemes( assert input_tfms or output_tfms - if is_qutlass_fp4_scheme(quant_scheme, input_tfms): + if is_qutlass_fp4_scheme( + quant_scheme, input_tfms + ) and current_platform.has_device_capability(100): return QutlassNvFP4LinearMethod(quant_method, input_tfms, output_tfms) # hadacore or dense gemm is selected by Transform module - return cls(quant_method, input_tfms, output_tfms) def __init__( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py index f5589c8c07fa..755059349382 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py @@ -32,7 +32,7 @@ class HadamardTransform(torch.nn.Module): transforms: dict[int, TransformTuple] # info parsed from transforms config weight: SharedWeightParameter # container for shared tensors - scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0)) + scaled_data_ptrs: set[int] = set() def __init__( self, @@ -44,7 +44,6 @@ def __init__( ): super().__init__() self.transforms = transforms - self.scales = {} if get_tensor_model_parallel_world_size() > 1: raise NotImplementedError( @@ -64,26 +63,26 @@ def __init__( ) data_key = self._get_data_key(scheme, weight_size) + # load up in model's default precision, rather than using scheme.precision self.weight.add_partition( part_index, data_key, size=(weight_size, weight_size), - dtype=scheme.precision, ) # validate that shared tensors and schemes are correct self._validate_input_transforms() def process_weights_after_loading(self): - for part_id in self.weight.partitions: - data = self.weight.partitions[part_id].data - + for part_id, partition in self.weight.partitions.items(): # required by torch.compile self.weight.process_weights_after_loading() - # precompute scale as a runtime multiply, not division - # do not fold into weight in order to utilize FWHT - self.scales[part_id] = 1 / math.sqrt(data.size(0)) + # Merge normalization scale directly into weight, must be done only once + data_ptr = partition.data.data_ptr() + if data_ptr not in HadamardTransform.scaled_data_ptrs: + partition.data.div_(math.sqrt(partition.data.size(0))) + HadamardTransform.scaled_data_ptrs.add(data_ptr) # FUTURE: avoid runtime transpose by processing weights # prior to apply @@ -111,26 +110,19 @@ def forward(self, value: Tensor, part_id: int = 0) -> Tensor: weight = ( weight if self.transforms[part_id].args.inverse else weight.T ) # linear := x(W.T) - scale = self.scales[part_id] if self.transforms[part_id].scheme.head_dim is not None: value = value.unflatten(-1, (-1, weight.size(0))) - value = ( - dispatch_unquantized_gemm()( - self, value.to(weight.dtype), weight, None - ).to(value.dtype) - * scale - ) + value = dispatch_unquantized_gemm()( + self, value.to(weight.dtype), weight, None + ).to(value.dtype) value = value.flatten(-2, -1) return value - return ( - dispatch_unquantized_gemm()( - self, value.to(weight.dtype), weight, None - ).to(value.dtype) - * scale - ) + return dispatch_unquantized_gemm()( + self, value.to(weight.dtype), weight, None + ).to(value.dtype) def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable: return (id(scheme), weight_size) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py index f0bb47a728ad..b8b771e3c7a9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py @@ -2,7 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from torch.nn.parameter import Parameter +from vllm._custom_ops import fusedQuantizeNv +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + NvFp4LinearKernel, +) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 CompressedTensorsScheme, CompressedTensorsW4A4Fp4, @@ -11,18 +17,32 @@ CompressedTensorsLinearTransformMethod, TransformTuple, ) +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + slice_nvfp4_output, +) +from vllm.utils.flashinfer import ( + flashinfer_scaled_fp4_mm, +) __all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"] +NVFP4_MAX = 6.0 + +# QUTLASS supports transform block sizes (16, 32, 64, 128) for NVFP4 +# https://github.com/IST-DASLab/qutlass/blob/v0.2.0/qutlass/csrc/bindings.cpp#L413-L414 def is_qutlass_fp4_scheme( quant_scheme: CompressedTensorsScheme | None, input_tfms: dict[int, TransformTuple], ) -> bool: return ( - isinstance(quant_scheme, (CompressedTensorsW4A4Fp4,)) - and len(input_tfms) == 1 - and input_tfms[0].scheme.head_dim == quant_scheme.group_size + isinstance(quant_scheme, CompressedTensorsW4A4Fp4) + and len(input_tfms) >= 1 + and all( + input_tfm.scheme.head_dim in (16, 32, 64, 128) + for input_tfm in input_tfms.values() + ) ) @@ -50,15 +70,90 @@ def create_weights( ) assert self.input_transform is not None - assert len(self.input_transform.weight) == 1 - assert self.input_transform.weight[0].size(0) == layer.scheme.group_size + assert len(self.input_transform.weight.partitions) >= 1 return ret + @staticmethod + def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str: + """ + Given a kernel, find the string that is needed to be passed into + `flashinfer_scaled_fp4_mm`, using + vllm.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth + """ + kernel_type = type(kernel) + for key, kernels in _LINEAR_BACKEND_KERNEL_MAP.items(): + if not key.startswith("flashinfer_") or kernel_type not in kernels: + continue + backend = key.removeprefix("flashinfer_") + # flashinfer GEMM backend uses "cute-dsl", not "cutedsl" + return backend.replace("cutedsl", "cute-dsl") + raise ValueError( + f"QutlassNvFP4 transform requires a FlashInfer kernel, " + f"got {kernel_type.__name__}" + ) + + def process_weights_after_loading(self, layer): + super().process_weights_after_loading(layer) + + assert self.input_transform is not None + layer.hadamard_matrix = self.input_transform.weight.partitions[0].data + + # fusedQuantizeNv stores raw absmax as block scales (sf = absmax), + # while CT weights use sf = absmax * SFScaleVal / 6.0. The GEMM + # computes alpha * sum(fp4_a * sf_a * fp4_w * sf_w), so alpha must + # compensate: alpha = weight_global_scale / 6.0 + layer.fused_alpha = Parameter( + layer.weight_global_scale / NVFP4_MAX, requires_grad=False + ) + + layer.fused_global_scale = Parameter( + torch.tensor( + [NVFP4_MAX], + dtype=torch.float32, + device=layer.weight_global_scale.device, + ), + requires_grad=False, + ) + + layer.flashinfer_gemm_backend = self._get_flashinfer_gemm_backend( + layer.scheme.kernel + ) + def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - raise NotImplementedError() + assert bias is None + output_size = layer.output_size_per_partition + output_shape = [*x.shape[:-1], output_size] + + x_flat = x.contiguous().flatten(end_dim=-2) + + x_fp4, x_scales = fusedQuantizeNv( + x_flat, layer.hadamard_matrix, layer.fused_global_scale + ) + + x_scales_blocked = to_blocked(x_scales, backend="triton").view(x_scales.shape) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_scales_blocked, + layer.weight_scale, + layer.fused_alpha, + x.dtype, + backend=layer.flashinfer_gemm_backend, + ) + + out = slice_nvfp4_output(out, output_size) + + if self.output_transform is not None: + for part_id, (start, length) in enumerate(self.partition_ranges): + out[:, start : start + length] = self.output_transform( + out[:, start : start + length].clone(), part_id=part_id + ) + + return out.view(*output_shape) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py index afb899cd6d7e..872771af9abc 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py @@ -145,8 +145,8 @@ def find_matched_target( matched_target = ( _find_first_match(layer_name, targets) - or _find_first_match(module.__class__.__name__, targets, True) or _match_fused_layer(layer_name, targets, fused_mapping) + or _find_first_match(module.__class__.__name__, targets, True) ) return matched_target diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 626fc83cdff5..49dd3a6d2249 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -254,16 +254,6 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): return out -def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None: - """Copies any attrs present in `old` but not in `new` to `new`""" - new_attrs = set(dir(new)) - attrs_to_set = {} - for attr in dir(old): - if attr not in new_attrs: - attrs_to_set[attr] = getattr(old, attr) - set_weight_attrs(new, attrs_to_set) - - class Fp8LinearMethod(LinearMethodBase): """Linear method for FP8. Supports loading FP8 checkpoints with static weight scale and diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index 603b80b7cd00..7a0d8496b872 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -38,11 +38,11 @@ def is_wna16_int(self) -> bool: @property def is_mxfp4(self) -> bool: - return self.data_type == "mx_fp" and self.bits == 4 + return "mx_fp" in self.data_type and self.bits == 4 @property def is_mxfp8(self) -> bool: - return self.data_type == "mx_fp" and self.bits == 8 + return "mx_fp" in self.data_type and self.bits == 8 class INCConfigParser: @@ -142,6 +142,14 @@ def get_config(name: str, quantized: bool = True) -> tuple[int, int, bool]: if self._config.extra_config and layer_name in self._config.extra_config: return get_config(layer_name) + # Suffix match: handle cases where extra_config keys use short names + # (e.g. "lm_head") but the layer_name is fully qualified + # (e.g. "model.language_model.lm_head") due to model nesting. + if self._config.extra_config: + for cfg_key in self._config.extra_config: + if layer_name.endswith(f".{cfg_key}"): + return get_config(cfg_key) + quantized = not isinstance(layer, ParallelLMHead) if self._config.block_name_to_quantize: quantized = any( diff --git a/vllm/model_executor/layers/quantization/inc/inc.py b/vllm/model_executor/layers/quantization/inc/inc.py index 86fa7cefcfce..219fc2b04692 100644 --- a/vllm/model_executor/layers/quantization/inc/inc.py +++ b/vllm/model_executor/layers/quantization/inc/inc.py @@ -35,8 +35,12 @@ class INCConfig(QuantizationConfig): """ SUPPORTED_BITS = {2, 3, 4, 8} - SUPPORTED_DTYPES = {"int"} - SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq"} + SUPPORTED_DTYPES = {"int", "mx_fp"} + SUPPORTED_FORMATS = { + "auto_round:auto_gptq", + "auto_round:auto_awq", + "auto_round:llm_compressor", + } SUPPORTED_BACKENDS = { "auto", "gptq", @@ -45,6 +49,11 @@ class INCConfig(QuantizationConfig): "awq:marlin", "marlin", } + MXFP8_BITS = 8 + MXFP8_GROUP_SIZE = 32 + MXFP8_DATA_TYPE = "mx_fp" + MXFP8_PACKING_FORMAT = "auto_round:llm_compressor" + MXFP8_SUPPORTED_ACT_DTYPES = {"mx_fp", "mx_fp_rceil"} def __init__( self, @@ -63,10 +72,13 @@ def __init__( f"Unsupported weight_bits: {weight_bits}, " f"currently only support {self.SUPPORTED_BITS}." ) - if data_type not in self.SUPPORTED_DTYPES: + # auto-round mxfp data_type is e.g. "mx_fp4" / "mx_fp4e2m1"; match the + # "mx_fp" family by substring like auto_round.compressors.is_mx_fp. + is_mxfp = "mx_fp" in data_type + if data_type not in self.SUPPORTED_DTYPES and not is_mxfp: raise ValueError( - f"Unsupported data_type: {data_type}," - f" currently only support {self.SUPPORTED_DTYPES}." + f"Unsupported data_type: {data_type}, " + f"currently only support {self.SUPPORTED_DTYPES}." ) if packing_format not in self.SUPPORTED_FORMATS: raise ValueError( @@ -75,7 +87,7 @@ def __init__( ) if backend not in self.SUPPORTED_BACKENDS: raise ValueError( - f"Unsupported backend: {backend}, " + f"Unsupported backend: {backend}, " f"currently only support {self.SUPPORTED_BACKENDS}." ) @@ -94,12 +106,67 @@ def __init__( self.pack_factor = Fraction(32, weight_bits) self.config_parser = INCConfigParser(self) + self._validate_supported_quantization() + def __repr__(self) -> str: return ( f"INCConfig(weight_bits={self.weight_bits}, " f"group_size={self.group_size}, sym={self.sym})" ) + @property + def is_mxfp(self) -> bool: + # auto-round mxfp data_type is the "mx_fp" family (e.g. "mx_fp", + # "mx_fp4e2m1"); match by substring like auto_round.compressors.is_mx_fp. + return "mx_fp" in self.data_type + + @property + def is_mxfp8(self) -> bool: + # MXFP4 and MXFP8 share data_type "mx_fp" and differ only by bit width. + return self.is_mxfp and self.weight_bits == self.MXFP8_BITS + + def _validate_supported_quantization(self) -> None: + if self.is_mxfp8: + assert self.group_size == self.MXFP8_GROUP_SIZE, ( + "INC MXFP8 only supports group_size=32, " + f"but found group_size={self.group_size}." + ) + assert self.sym, "INC MXFP8 only supports symmetric weights." + assert self.packing_format == self.MXFP8_PACKING_FORMAT, ( + "INC MXFP8 only supports " + f"packing_format={self.MXFP8_PACKING_FORMAT!r}, " + f"but found {self.packing_format!r}." + ) + assert self.backend == "auto", ( + "INC MXFP8 only supports backend='auto', " + f"but found backend={self.backend!r}." + ) + elif self.packing_format == self.MXFP8_PACKING_FORMAT and not self.is_mxfp: + raise ValueError( + f"packing_format={self.MXFP8_PACKING_FORMAT!r} requires " + f"an {self.MXFP8_DATA_TYPE!r} data_type." + ) + + def _validate_raw_config(self, config: dict[str, Any]) -> None: + if not self.is_mxfp8: + return + + expected_fields = { + "act_bits": self.MXFP8_BITS, + "act_data_type": self.MXFP8_DATA_TYPE, + "act_group_size": self.MXFP8_GROUP_SIZE, + "act_sym": True, + "act_dynamic": True, + "enable_quanted_input": False, + } + for field_name, expected_value in expected_fields.items(): + actual_value = self.get_from_keys_or(config, [field_name], expected_value) + assert actual_value == expected_value, ( + "INC MXFP8 only supports " + f"{field_name}={expected_value!r}, " + f"but found {field_name}={actual_value!r}." + ) + @classmethod def get_name(cls) -> QuantizationMethods: return "inc" @@ -118,7 +185,7 @@ def get_config_filenames(cls) -> list[str]: @classmethod def from_config(cls, config: dict[str, Any]) -> "INCConfig": - return cls( + quant_config = cls( weight_bits=cls.get_from_keys(config, ["bits"]), group_size=cls.get_from_keys(config, ["group_size"]), sym=cls.get_from_keys(config, ["sym"]), @@ -132,6 +199,8 @@ def from_config(cls, config: dict[str, Any]) -> "INCConfig": data_type=cls.get_from_keys_or(config, ["data_type"], "int"), backend=cls.get_from_keys_or(config, ["backend", "vllm_backend"], "auto"), ) + quant_config._validate_raw_config(config) + return quant_config def get_layer_config(self, layer, layer_name: str): return self.config_parser.get_layer_config(layer, layer_name) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py index ea6c0a00d86d..1f2d5f4401b2 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py @@ -2,12 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .factory import resolve_scheme +from .inc_mxfp4_scheme import INCMxfp4Scheme +from .inc_mxfp8_scheme import INCMxfp8Scheme from .inc_scheme import INCLinearScheme, INCScheme from .inc_wna16_scheme import INCWna16Scheme __all__ = [ "INCScheme", "INCLinearScheme", + "INCMxfp8Scheme", "INCWna16Scheme", + "INCMxfp4Scheme", "resolve_scheme", ] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/factory.py b/vllm/model_executor/layers/quantization/inc/schemes/factory.py index 4ae85ed9a831..61335e602fff 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/factory.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/factory.py @@ -9,10 +9,14 @@ def resolve_scheme(layer_config: "INCLayerConfig") -> "INCScheme": + from .inc_mxfp4_scheme import INCMxfp4Scheme + from .inc_mxfp8_scheme import INCMxfp8Scheme from .inc_wna16_scheme import INCWna16Scheme scheme_list: list[type[INCScheme]] = [ + INCMxfp8Scheme, INCWna16Scheme, + INCMxfp4Scheme, ] for scheme_cls in scheme_list: diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py new file mode 100644 index 000000000000..04b3cd65cf26 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, Any + +import torch +from torch.nn.parameter import Parameter + +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + ModelWeightParameter, +) + +from .inc_scheme import INCLinearScheme + +if TYPE_CHECKING: + from ..config_parser import INCLayerConfig + + +class INCMxfp4LinearMethod(INCLinearScheme): + """MXFP4 (W4A4) linear method for AutoRound checkpoints. + + E2M1 weights packed two per byte with per-group E8M0 scales + (group_size=32, no global scale). The platform kernel is selected by + ``init_mxfp4_linear_kernel`` (FlashInfer / Marlin on CUDA, ``fp4_gemm`` + on XPU). + """ + + def __init__(self, layer_config: "INCLayerConfig") -> None: + self.group_size = layer_config.group_size or 32 + self.kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs: Any, + ) -> None: + del input_size, output_size + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + weight_loader = extra_weight_attrs.get("weight_loader") + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_packed", weight) + + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.group_size, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight = Parameter(layer.weight_packed.data, requires_grad=False) + del layer.weight_packed + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py new file mode 100644 index 000000000000..7cb3adce090d --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoeWeightScaleSupported, + RoutedExperts, + SharedExperts, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + mxfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, +) +from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( + MarlinExperts, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, + select_mxfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + prepare_moe_fp4_layer_for_marlin, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class INCMxfp4MoEMethod(FusedMoEMethodBase): + """W4A4 MXFP4 group MoE for AutoRound ``auto_round:llm_compressor`` exports. + + Registers the packed MXFP4 layout (uint8 ``weight_packed`` + uint8 E8M0 + ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to the + best backend for the current device: CUTLASS (true W4A4 on supported + GPUs), the native XPU kernel, or Marlin weight-only as a fallback. The + per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are folded + into the stacked ``w13`` / ``w2`` parameters by ``make_expert_params_mapping``. + """ + + def __init__(self, moe) -> None: + super().__init__(moe) + self.group_size = 32 + # Backend selection must stay consistent with the weight preparation in + # process_weights_after_loading / get_fused_moe_quant_config, which only + # implement three layouts: CUTLASS swizzle (true W4A4), the native XPU + # kernel (packed passthrough), and Marlin weight-only. XPU dispatch is + # deferred to the shared oracle; every other non-CUTLASS device falls + # back to Marlin (mirroring CompressedTensorsW4A4Mxfp4MoEMethod). + self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() + self.mxfp4_backend = Mxfp4MoeBackend.MARLIN + self.experts_cls: type[mk.FusedMoEExperts] | None = None + if self.use_cutlass_mxfp4: + self.experts_cls = CutlassExpertsMxfp4 + logger.info_once("Using CutlassExpertsMxfp4 for AutoRound MXFP4 MoE") + elif current_platform.is_xpu(): + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + else: + self.experts_cls = MarlinExperts + logger.info_once( + "Using MarlinExperts (weight-only FP4) for AutoRound MXFP4 MoE" + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + layer.num_experts = num_experts + layer.params_dtype = params_dtype + + # gate + up fused on the output dim; two FP4 packed per input byte. + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # Per-group E8M0 block scales (group_size=32), stored as uint8. + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + if self.use_cutlass_mxfp4: + # W4A4: both weights and activations quantized to MXFP4. + return mxfp4_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + ) + # Weight-only (Marlin) or native XPU kernel. + return make_mxfp4_moe_quant_config( + mxfp4_backend=self.mxfp4_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + layer.w13_weight = torch.nn.Parameter( + layer.w13_weight_packed.data, requires_grad=False + ) + delattr(layer, "w13_weight_packed") + layer.w2_weight = torch.nn.Parameter( + layer.w2_weight_packed.data, requires_grad=False + ) + delattr(layer, "w2_weight_packed") + + if self.use_cutlass_mxfp4: + # Swizzle weight scales from flat checkpoint layout [E, N, K//32] + # to the CUTLASS tiled layout. + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + swizzle_mxfp4_scales, + ) + + E = layer.w13_weight_scale.shape[0] + w13_N = layer.w13_weight_scale.shape[1] + w13_scale_K = layer.w13_weight_scale.shape[2] + w13_K = w13_scale_K * 32 + + w2_M = layer.w2_weight_scale.shape[1] + w2_scale_N = layer.w2_weight_scale.shape[2] + w2_N = w2_scale_N * 32 + + swizzled_w13 = [] + swizzled_w2 = [] + for e_idx in range(E): + s13 = layer.w13_weight_scale[e_idx] + sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K) + swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K)) + s2 = layer.w2_weight_scale[e_idx] + sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N) + swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N)) + layer.w13_weight_scale = torch.nn.Parameter( + torch.stack(swizzled_w13), requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + torch.stack(swizzled_w2), requires_grad=False + ) + elif current_platform.is_xpu(): + # The XPU fused-MoE kernel consumes the packed layout directly; no + # swizzle / repack / transpose is required. + pass + else: + logger.warning_once( + "This device lacks native FP4 compute; using weight-only FP4 " + "via the Marlin kernel, which may reduce performance for " + "compute-heavy workloads." + ) + prepare_moe_fp4_layer_for_marlin(layer) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None: + assert self.experts_cls is not None + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + mxfp4_backend=self.mxfp4_backend, + routing_tables=layer._expert_routing_tables(), + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py new file mode 100644 index 000000000000..e22bb900cea3 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +from vllm.logger import init_logger + +from ..inc_linear import INCLinearMethod +from .inc_scheme import INCScheme + +if TYPE_CHECKING: + import torch + + from ..config_parser import INCLayerConfig + from ..inc import INCConfig + +logger = init_logger(__name__) + + +class INCMxfp4Scheme(INCScheme): + """MXFP4 (W4A4) scheme for AutoRound checkpoints. + + Dispatches to :class:`INCMxfp4LinearMethod` for linear layers and + :class:`INCMxfp4MoEMethod` for fused MoE layers; see those classes for the + per-module weight layout and kernel-selection details. + """ + + @staticmethod + def can_handle(layer_config: "INCLayerConfig") -> bool: + return layer_config.is_mxfp4 + + def get_linear_method( + self, + config: "INCConfig", + layer: "torch.nn.Module", + prefix: str, + layer_config: "INCLayerConfig", + ): + del config, layer, prefix + from .inc_mxfp4_linear import INCMxfp4LinearMethod + + return INCLinearMethod(INCMxfp4LinearMethod(layer_config)) + + def get_moe_method( + self, + config: "INCConfig", + layer: "torch.nn.Module", + prefix: str, + layer_config: "INCLayerConfig", + ): + del config, prefix, layer_config + from .inc_mxfp4_moe import INCMxfp4MoEMethod + + return INCMxfp4MoEMethod(layer.moe_config) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py new file mode 100644 index 000000000000..775b5f5bba65 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + ModelWeightParameter, +) + +from .inc_scheme import INCLinearScheme + + +class INCMxfp8LinearScheme(INCLinearScheme): + def __init__(self) -> None: + self.kernel = init_mxfp8_linear_kernel() + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + if input_size_per_partition % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + "INC MXFP8 requires input_size_per_partition " + f"({input_size_per_partition}) to be divisible by " + f"{MXFP8_BLOCK_SIZE}." + ) + + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=MXFP8_VALUE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=extra_weight_attrs.get("weight_loader"), + ) + layer.register_parameter("weight", weight) + + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // MXFP8_BLOCK_SIZE, + dtype=MXFP8_SCALE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=extra_weight_attrs.get("weight_loader"), + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py new file mode 100644 index 000000000000..2e251a8a5428 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +from ..inc_linear import INCLinearMethod +from .inc_scheme import INCScheme + +if TYPE_CHECKING: + import torch + + from ..config_parser import INCLayerConfig + from ..inc import INCConfig + + +class INCMxfp8Scheme(INCScheme): + @staticmethod + def can_handle(layer_config: "INCLayerConfig") -> bool: + return layer_config.is_mxfp8 + + def get_linear_method( + self, + config: "INCConfig", + layer: "torch.nn.Module", + prefix: str, + layer_config: "INCLayerConfig", + ): + del config, layer, prefix, layer_config + from .inc_mxfp8_linear import INCMxfp8LinearScheme + + return INCLinearMethod(INCMxfp8LinearScheme()) 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 dd6c2fa2eaae..5c99fd98b54d 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 @@ -166,6 +166,7 @@ class INCXPULinearBase(INCLinearScheme): def __init__(self, layer_config: "INCLayerConfig") -> None: self.weight_bits = layer_config.bits self.group_size = layer_config.group_size + self.sym = layer_config.sym self.pack_factor = 32 // self.weight_bits self.is_awq_packed = layer_config.is_awq 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 46ad24ea5b47..f83cc22e8e08 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 @@ -7,7 +7,6 @@ from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.platforms import current_platform -from vllm.scalar_type import scalar_types from ..inc_linear import INCLinearMethod from .inc_scheme import INCScheme @@ -97,8 +96,8 @@ def get_moe_method( layer_config: "INCLayerConfig", ): del config, prefix - # XPU and CPU do not support MoE quantization yet - if current_platform.is_xpu() or current_platform.is_cpu(): + # CPU does not support quantized MoE yet. + if current_platform.is_cpu(): from vllm.model_executor.layers.fused_moe import ( UnquantizedFusedMoEMethod, ) @@ -120,21 +119,17 @@ def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"): MoeWNA16Method, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_marlin_supported, check_moe_marlin_supports_layer, ) - gptq_type_map = { - (4, True): scalar_types.uint4b8, - (8, True): scalar_types.uint8b128, - } - use_marlin = (layer_config.bits, layer_config.sym) in gptq_type_map - if use_marlin: - use_marlin = check_marlin_supported( - gptq_type_map[(layer_config.bits, layer_config.sym)], - layer_config.group_size, - has_zp=not layer_config.sym, - ) and check_moe_marlin_supports_layer(layer, layer_config.group_size) + # AutoGPTQMoEMethod selects its fused-MoE backend through the WNA16 oracle + # (Marlin on CUDA, XPUExpertsWNA16 on XPU). Gate only on the layer-shape + # check like compressed-tensors does; the capability-based + # check_marlin_supported is skipped so the XPU path is reachable. + use_marlin = (layer_config.bits, layer_config.sym) in { + (4, True), + (8, True), + } and check_moe_marlin_supports_layer(layer, layer_config.group_size) if use_marlin: return AutoGPTQMoEMethod( @@ -169,21 +164,12 @@ def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"): MoeWNA16Method, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_marlin_supported, check_moe_marlin_supports_layer, ) - awq_type_map = { - 4: scalar_types.uint4, - 8: scalar_types.uint8, - } - use_marlin = layer_config.bits in awq_type_map - if use_marlin: - use_marlin = check_marlin_supported( - awq_type_map[layer_config.bits], - layer_config.group_size, - not layer_config.sym, - ) and check_moe_marlin_supports_layer(layer, layer_config.group_size) + use_marlin = layer_config.bits in (4, 8) and check_moe_marlin_supports_layer( + layer, layer_config.group_size + ) if use_marlin: return AutoAWQMoEMethod( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index e8810919c204..2eb34630aa61 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -139,11 +139,6 @@ def forward_hip( scale_ub: torch.Tensor | None = None, use_triton: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if self.is_group_quant and use_triton: - assert scale is None, "Dynamic group quantization does not use scale" - - return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size) - use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous() use_aiter_per_tensor_quant = ( use_aiter_quant and self.group_shape.is_per_tensor() diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 7df8178ca71c..e3edb88ae2e0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,8 +12,6 @@ from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( - MarlinNvFp4LinearKernel, - NvFp4LinearLayerConfig, init_fp8_linear_kernel, init_mxfp8_linear_kernel, init_nvfp4_linear_kernel, @@ -58,9 +56,6 @@ QuantizeMethodBase, ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - swap_w13_to_w31, -) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( process_fp8_input_tensor_strategy_moe, process_fp8_weight_channel_strategy, @@ -410,7 +405,7 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: - return 89 + return 80 @classmethod def override_quantization_method( @@ -524,6 +519,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight, layer.weight_scale, layer.logical_widths ) layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) self.fp8_linear.process_weights_after_loading(layer) @@ -783,16 +780,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: RoutedExperts, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def create_weights( self, layer: RoutedExperts, @@ -1250,16 +1237,16 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): """Linear method for ModelOpt NVFP4 W4A16. 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names - directly (no on-disk conversion) and dispatches to the FP4 Marlin GEMM: + directly (no on-disk conversion) and dispatches to a W4A16 GEMM: weight uint8 packed NVFP4 (2 nibbles/byte along input dim) weight_scale fp8-e4m3 per 16-elem group along input dim weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0) - No activation quantization. Marlin expects the global scale in the same - form ModelOpt stores (amax/2688), so we rename weight_scale_2 -> - weight_global_scale **without reciprocation** -- the CT W4A16 path - reciprocates only because CT stores the inverse on disk. + No activation quantization. ModelOpt stores the global scale as + amax/2688, so we rename weight_scale_2 -> weight_global_scale without + reciprocation. The selected kernel converts it to its runtime format. + The CT W4A16 path reciprocates because CT stores the inverse on disk. We also register a placeholder input_scale parameter so that W4A4-shaped checkpoints (which contain *_proj.input_scale tensors) can be loaded @@ -1270,17 +1257,12 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): def __init__(self, quant_config: ModelOptNvFp4Config) -> None: self.quant_config = quant_config - # Vestigial slot mirrored from ModelOptNvFp4LinearMethod: the parent - # config's get_quant_method only fills marlin_input_dtype when - # backend == "marlin"; we don't set that since we pin the kernel - # below, but we keep the attribute for shape parity. self.marlin_input_dtype = None - # Direct-instantiate the Marlin NVFP4 adapter rather than going through - # init_nvfp4_linear_kernel(): the latter's priority list returns a - # cutlass W4A4 kernel as first-pick on this hardware, which would - # silently try to quantize activations (we have no input_scale). For - # W4A16 there is exactly one valid kernel, so we pin it. - self.kernel = MarlinNvFp4LinearKernel(NvFp4LinearLayerConfig()) + # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds: + # 1. `use_a16=True` forces `Marlin`: https://github.com/vllm-project/vllm/commit/e68988a#diff-7135ab92aa94dfacb1ad3c77fc13f9c4ffe0b977f8eac5d86c2afe243e5f92a6R842-R889 + # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale. + # 2. Specifying e.g. `--linear-backend=humming` will override. + self.kernel = init_nvfp4_linear_kernel(use_a16=True) def create_weights( self, @@ -1303,6 +1285,7 @@ def create_weights( layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition + layer.output_partition_sizes = output_partition_sizes if input_size_per_partition % 16 != 0: raise ValueError( @@ -1358,6 +1341,9 @@ def create_weights( layer.register_parameter("input_scale", input_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if not hasattr(layer, "has_bias"): + layer.has_bias = getattr(layer, "bias", None) is not None + # Discard the input_scale placeholder. Whether it carries values # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16 # ckpt), W4A16 mode does not quantize activations, so this is unused. @@ -2014,90 +2000,6 @@ def _check_weight_dtypes(layer: torch.nn.Module) -> None: f"Expected {name} dtype {expected_dtype}, got {actual}." ) - def _shuffle_weights_for_trtllm(self, layer: torch.nn.Module) -> None: - """Shuffle weights and scales into FlashInfer TRTLLM MXFP8 layout.""" - from flashinfer import ( - reorder_rows_for_gated_act_gemm, - shuffle_matrix_a, - shuffle_matrix_sf_a, - ) - - epilogue_tile_m = 128 - num_experts = layer.w13_weight.shape[0] - is_gated = self.moe.is_act_and_mul - intermediate_size_factor = 2 if is_gated else 1 - - w13_weight = layer.w13_weight.data - w13_scale = layer.w13_weight_scale.data - if is_gated: - # FI TRTLLM gated kernels use W31 ordering. Model checkpoints store - # gated projection as W13, so convert once before shuffling. - w13_weight = swap_w13_to_w31(w13_weight) - w13_scale = swap_w13_to_w31(w13_scale) - - w13_weight_shuffled = [] - w2_weight_shuffled = [] - w13_scale_shuffled = [] - w2_scale_shuffled = [] - for i in range(num_experts): - w13_i = w13_weight[i].reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, -1 - ) - w13_sf_i = w13_scale[i].reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, -1 - ) - if is_gated: - # Reorder rows for gated activation layout expected by TRTLLM. - w13_i = reorder_rows_for_gated_act_gemm(w13_i.clone()) - w13_sf_i = reorder_rows_for_gated_act_gemm(w13_sf_i.clone()) - - w13_shuffled_i = shuffle_matrix_a(w13_i.view(torch.uint8), epilogue_tile_m) - w2_shuffled_i = shuffle_matrix_a( - layer.w2_weight.data[i].view(torch.uint8), epilogue_tile_m - ) - w13_weight_shuffled.append( - w13_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE) - ) - w2_weight_shuffled.append( - w2_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE) - ) - w13_sf_shuffled_i = shuffle_matrix_sf_a( - w13_sf_i.view(torch.uint8).reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, - -1, - ), - epilogue_tile_m, - ) - w2_sf_shuffled_i = shuffle_matrix_sf_a( - layer.w2_weight_scale.data[i] - .view(torch.uint8) - .reshape(layer.hidden_size, -1), - epilogue_tile_m, - ) - w13_scale_shuffled.append( - w13_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE) - ) - w2_scale_shuffled.append( - w2_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE) - ) - - replace_parameter( - layer, "w13_weight", torch.stack(w13_weight_shuffled).contiguous() - ) - replace_parameter( - layer, "w2_weight", torch.stack(w2_weight_shuffled).contiguous() - ) - replace_parameter( - layer, - "w13_weight_scale", - torch.stack(w13_scale_shuffled).contiguous(), - ) - replace_parameter( - layer, - "w2_weight_scale", - torch.stack(w2_scale_shuffled).contiguous(), - ) - def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None: """One-time MXFP8->BF16 weight dequant for the emulation path. @@ -2194,16 +2096,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: RoutedExperts, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: @@ -2453,7 +2345,7 @@ def _resolve_quant_algo(self, prefix: str) -> str | None: if key.startswith(prefix_dot): return info["quant_algo"].upper() - # FusedMoE expert prefix is e.g. "...moe.experts", while ModelOpt's + # RoutedExperts expert prefix is e.g. "...moe.experts", while ModelOpt's # quantized_layers entries use "...moe.gate_proj" / "...moe.up_proj". if prefix.endswith(".experts"): parent_dot = prefix.rsplit(".experts", 1)[0] + "." diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 23e175fd6249..76f2d740bc73 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -15,8 +15,13 @@ ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, - int8_w8a16_moe_quant_config, +) +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, + select_wna16_moe_backend, ) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, @@ -27,7 +32,15 @@ QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + INT4_DTYPE, + INT8_DTYPE, + QuantKey, + kInt4Static32GroupScale, + kInt4StaticGroupScale, + kInt8StaticGroupScale, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform @@ -199,6 +212,34 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: super().__init__(moe) self.quant_config = quant_config + num_bits = self.quant_config.weight_bits + group_size = self.quant_config.group_size + + if num_bits == 4: + quant_type = INT4_DTYPE + if group_size == 32: + scale = kInt4Static32GroupScale + else: + scale = kInt4StaticGroupScale + elif num_bits == 8: + assert group_size == -1 + quant_type = INT8_DTYPE + scale = kInt8StaticGroupScale + else: + raise ValueError("MoeWNA16Method only supports int4 and int8 now.") + + weight_key = QuantKey(quant_type, scale) + + # Select WNA16 MoE backend via oracle. + # handle ZP? + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + weight_key=weight_key, + quant_config=self.quant_config, + may_have_zp=self.quant_config.has_zp, + may_have_bias=False, + ) + def create_weights( self, layer: RoutedExperts, @@ -322,23 +363,104 @@ def create_weights( def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: - weight_bits = self.quant_config.weight_bits - has_zp = self.quant_config.has_zp - assert weight_bits == 4 or weight_bits == 8 - config_builder = ( - int4_w4a16_moe_quant_config - if weight_bits == 4 - else int8_w8a16_moe_quant_config - ) + if self.wna16_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 config_builder( + has_zp = self.quant_config.has_zp + return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, w1_zp=layer.w13_qzeros if has_zp else None, w2_zp=layer.w2_qzeros if has_zp else None, - block_shape=[0, layer.group_size], + group_size=layer.group_size, + num_bits=self.quant_config.weight_bits, ) + def _setup_kernel(self, layer: RoutedExperts): + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + backend=self.wna16_backend, + layer=layer, + routing_tables=layer._expert_routing_tables(), + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + has_zp = self.quant_config.has_zp + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=None, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_g_idx=None, + w2_g_idx=None, + w13_qzeros=layer.w13_qzeros if has_zp else None, + w2_qzeros=layer.w2_qzeros if has_zp else None, + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + + ( + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, + _, + _, + _, + _, + w13_qzeros, + w2_qzeros, + w13_input_global_scale, + w2_input_global_scale, + _, # w13_bias + _, # w2_bias + ) = converted + + # Replace common parameters + replace_parameter(layer, "w13_qweight", w13_qweight) + replace_parameter(layer, "w2_qweight", w2_qweight) + replace_parameter(layer, "w13_scales", w13_scales) + replace_parameter(layer, "w2_scales", w2_scales) + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight + + if has_zp: + assert w13_qzeros is not None and w2_qzeros is not None + replace_parameter(layer, "w13_qzeros", w13_qzeros) + replace_parameter(layer, "w2_qzeros", w2_qzeros) + + # Marlin-specific parameters (not needed for Flashinfer) + if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM: + # Register input global scales if present + if w13_input_global_scale is not None: + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + if w2_input_global_scale is not None: + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + + self._setup_kernel(layer) + def apply( self, layer: RoutedExperts, @@ -348,19 +470,44 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, - layer.w13_qweight, - layer.w2_qweight, + layer.w13_weight, + layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - quant_config=self.moe_quant_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) @staticmethod diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 5ef5fd40d5eb..e9f3d0a05376 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -3,6 +3,7 @@ import torch +import vllm.envs as envs from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( @@ -14,9 +15,13 @@ SharedExperts, ) from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import ( + mxfp4_w4a16_moe_quant_config, +) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, + backend_to_kernel_cls, convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, @@ -33,6 +38,7 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -68,8 +74,11 @@ def get_supported_act_dtypes(cls) -> list[torch.dtype]: def get_config_filenames(cls) -> list[str]: return [] - # TODO (zyongye) This is only temporaty fallback. - # We should have `Mxfp4MoEMethod` after this migration is complete. + def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: + """MoE method for RoutedExperts. Subclasses override to pick a + checkpoint-specific kernel family.""" + return Mxfp4MoEMethod(moe) + def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": @@ -86,7 +95,7 @@ def get_quant_method( ) return UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): - return GptOssMxfp4MoEMethod(layer.moe_config) + return self._make_moe_method(layer.moe_config) elif isinstance(layer, Attention): logger.debug_once( "MXFP4 attention layer is not implemented. " @@ -131,6 +140,9 @@ def override_quantization_method( return None return "gpt_oss_mxfp4" + def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: + return GptOssMxfp4MoEMethod(moe) + class GptOssMxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" @@ -419,16 +431,6 @@ def get_fused_moe_quant_config( layer=layer, ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, - layer: RoutedExperts, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel " - "initialization logic. This function should not be called." - ) - def apply( self, layer: RoutedExperts, @@ -472,16 +474,50 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) +def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool: + """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel. + + K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic + MXFP4 backend selector does not cover; route it to AITER on gfx950. + """ + if not current_platform.is_rocm(): + return False + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.platforms.rocm import on_gfx950 + + return ( + rocm_aiter_ops.is_fused_moe_enabled() + and on_gfx950() + and moe.activation == MoEActivation.SITU + and moe.activation_situ_linear_beta is not None + and rocm_aiter_ops.get_aiter_activation_type("situ") is not None + ) + + class Mxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" def __init__(self, moe: FusedMoEConfig): super().__init__(moe) self.weight_dtype = "mxfp4" - self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) + self.is_k3_situ_aiter = _use_k3_situ_aiter(moe) + self.experts_cls: type[mk.FusedMoEExperts] | None + if self.is_k3_situ_aiter: + self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] + logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.") + else: + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend( + moe + ) self.max_capture_size = moe.max_capture_size @@ -523,6 +559,11 @@ def maybe_roundup_sizes( act_dtype=act_dtype, moe_parallel_config=moe_parallel_config, ) + if self.is_k3_situ_aiter: + # K3's AITER A16W4 kernel handles K3's native intermediate size + # (moe_intermediate 3072; e.g. 384/partition at TP8); the generic + # 256 round-up would inflate weights and OOM. + return hidden_size, intermediate_size_per_partition return mxfp4_round_up_hidden_size_and_intermediate_size( self.mxfp4_backend, hidden_size, intermediate_size_per_partition ) @@ -690,7 +731,16 @@ def _setup_kernel( # For TRITON backends, weights are wrapped tensors from triton_kernels # that don't support .detach(). Manually assign parameters. - if self.mxfp4_backend not in TRITON_BACKENDS: + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() + + uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 + ) + if not uses_triton_weight_format: replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w13_weight_scale", w13_scale) @@ -702,7 +752,10 @@ def _setup_kernel( self.w2_precision_config = w2_scale # AITER backend requires weights to be marked as shuffled. - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + if ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 + and not uses_triton_weight_format + ): layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True @@ -724,7 +777,58 @@ def _setup_kernel( layer=layer, ) + def _setup_kernel_k3_situ(self, layer: RoutedExperts) -> None: + # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all]) + # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in + # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here. + from aiter.utility.fp4_utils import e8m0_shuffle + + from vllm._aiter_ops import rocm_aiter_ops + + fp4_dtype = torch.float4_e2m1fn_x2 + e8m0_dtype = torch.float8_e8m0fnu + num_experts = layer.w13_weight.shape[0] + + # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) fp8 + # flydsl kernels, which need w13 weight+scale in interleave layout. + # Default a16w4 keeps the separated layout. + guinterleave = envs.AITER_SITUV2_A8W4 + w13 = rocm_aiter_ops.shuffle_weight_a16w4( + layer.w13_weight.data.view(fp4_dtype), 16, guinterleave + ) + w2 = rocm_aiter_ops.shuffle_weight_a16w4( + layer.w2_weight.data.view(fp4_dtype), 16, False + ) + w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype) + w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype) + w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( + w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave + ) + w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None and self.experts_cls is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + mxfp4_backend=self.mxfp4_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + def process_weights_after_loading(self, layer): + if self.is_k3_situ_aiter: + self._setup_kernel_k3_situ(layer) + return + w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale @@ -745,7 +849,15 @@ def get_fused_moe_quant_config( w2_bias = getattr(layer, "w2_bias", None) swiglu_limit = getattr(layer, "swiglu_limit", None) - if self.mxfp4_backend in TRITON_BACKENDS: + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() + + if self.mxfp4_backend in TRITON_BACKENDS or ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 + ): # TRITON backends free w13/w2_weight_scale after swizzling; the # swizzled scales live inside the precision configs instead. assert self.w13_precision_config is not None @@ -756,6 +868,18 @@ def get_fused_moe_quant_config( w1_scale = layer.w13_weight_scale w2_scale = layer.w2_weight_scale + if self.mxfp4_backend == Mxfp4MoeBackend.EMULATION: + # Canonical ``mxfp4`` checkpoints are weight-only W4A16. The + # generic EMULATION config is W4A4, so preserve BF16 activations + # while the fallback dequantizes only the weights. + return mxfp4_w4a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + gemm1_clamp_limit=swiglu_limit, + ) + return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=w1_scale, @@ -766,16 +890,6 @@ def get_fused_moe_quant_config( layer=layer, ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, - layer: RoutedExperts, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel " - "initialization logic. This function should not be called." - ) - def apply( self, layer: RoutedExperts, @@ -819,4 +933,8 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index ef76594164a8..870e4d2a0f26 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -114,6 +114,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: routing_tables=layer._expert_routing_tables(), layer=layer, ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def get_fused_moe_quant_config( self, layer: torch.nn.Module diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 15023d7ca39a..c5c9c6988f88 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -33,6 +33,13 @@ make_fp8_moe_quant_config, select_fp8_moe_backend, ) +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + Int8MoeBackend, + convert_to_int8_moe_kernel_format, + make_int8_moe_kernel, + make_int8_moe_quant_config, + select_int8_moe_backend, +) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, @@ -59,6 +66,10 @@ kFp8DynamicTokenSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt8DynamicTensorSym, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, + kInt8StaticTensorSym, kMxfp4Dynamic, kNvfp4Dynamic, kNvfp4Static, @@ -502,6 +513,35 @@ def __init__( self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor") self.static_input_scales = not self.input_quant.get("is_dynamic", False) + self.moe_quant_config: FusedMoEQuantConfig | None = None + self.moe_kernel: mk.FusedMoEKernel | None = None + self.int8_backend: Int8MoeBackend | None = None + self.experts_cls: type[mk.FusedMoEExperts] | None = None + + # Dynamic-activation INT8 MoE goes through the oracle + modular kernel. + # The modular TritonExperts kernel consumes float activations and + # quantizes them to int8 itself, so it cannot apply a loaded static + # activation scale (this matches CompressedTensorsW8A8Int8MoEMethod). + # TODO: Static-activation INT8 therefore stays on the legacy fused_experts + # path (see apply()) for now, preserving pre-refactor behavior. + # Needs to be migrated to expert backend. + if not self.static_input_scales: + # Map the Quark weight scheme to oracle quant keys. Per-channel + # weights pair with dynamic per-token activations; per-tensor + # weights with dynamic per-tensor activations. + if self.weight_qscheme == "per_channel": + weight_key = kInt8StaticChannelSym + activation_key = kInt8DynamicTokenSym + else: + weight_key = kInt8StaticTensorSym + activation_key = kInt8DynamicTensorSym + + self.int8_backend, self.experts_cls = select_int8_moe_backend( + config=moe, + weight_key=weight_key, + activation_key=activation_key, + ) + def create_weights( self, layer: torch.nn.Module, @@ -563,7 +603,7 @@ def create_weights( set_weight_attrs(w13_weight_scale, extra_weight_attrs) set_weight_attrs(w2_weight_scale, extra_weight_attrs) else: - # per-tensor: one scalar per expert + # per-tensor: one scalar per expert (two for the fused w1/w3) w13_weight_scale = torch.nn.Parameter( torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False, @@ -582,6 +622,8 @@ def create_weights( # INPUT_SCALES if self.static_input_scales: + # Static activations: the per-expert scales are loaded from the + # checkpoint (used by the legacy fused_experts path). w13_input_scale = torch.nn.Parameter( torch.ones(num_experts, dtype=torch.float32), requires_grad=False, @@ -596,6 +638,7 @@ def create_weights( layer.register_parameter("w2_input_scale", w2_input_scale) set_weight_attrs(w2_input_scale, extra_weight_attrs) else: + # Dynamic activations are quantized in-kernel (no stored scale). layer.w13_input_scale = None layer.w2_input_scale = None @@ -673,7 +716,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if hasattr(layer, attr): delattr(layer, attr) - # For static input scales, collapse per-expert scales to single max + # For static input scales, collapse the per-expert scales to a single + # value (the legacy fused_experts path expects one scale per layer). if self.static_input_scales: if layer.w13_input_scale is None or layer.w2_input_scale is None: raise ValueError( @@ -709,7 +753,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ), ) - # For per-tensor weights, merge w1/w3 scales into single per-expert + # For per-tensor weights, merge the w1/w3 scales into a single + # per-expert scale (dequant -> requant at the max scale). if self.weight_qscheme == "per_tensor": assert layer.w13_weight_scale is not None shard_size = layer.intermediate_size_per_partition @@ -734,10 +779,42 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: max_w13_scales, requires_grad=False ) + # Dynamic activations run through the oracle's modular kernel; static + # activations use the legacy fused_experts path in apply(). + if not self.static_input_scales: + assert self.int8_backend is not None + assert self.experts_cls is not 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_weight_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 + + if not self.static_input_scales: + assert self.int8_backend 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: - if self.weight_qscheme == "per_channel" and not self.static_input_scales: + # Static-activation INT8 has no oracle backend (it uses the legacy + # fused_experts path); build its config directly. + if self.int8_backend is None: return int8_w8a8_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, @@ -745,21 +822,18 @@ def get_fused_moe_quant_config( a2_scale=layer.w2_input_scale, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - per_act_token_quant=True, + per_act_token_quant=False, ) - is_dynamic = not self.static_input_scales - is_per_channel = self.weight_qscheme == "per_channel" - return FusedMoEQuantConfig.make( - torch.int8, + 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, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - per_act_token_quant=is_dynamic, - per_out_ch_quant=is_per_channel, - block_shape=None, + per_act_token_quant=(self.weight_qscheme == "per_channel"), + layer=layer, ) def apply( @@ -771,6 +845,22 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: + if self.moe_kernel is not None: + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts_input=shared_experts_input, + ) + + # Static-activation INT8 MoE: legacy monolithic path (the modular kernel + # quantizes activations dynamically and cannot apply a loaded scale). from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index ea63d1bcf122..8b6edd6cae6d 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -3,25 +3,28 @@ from collections.abc import Callable from fractions import Fraction -from functools import partial from typing import Any import torch -import torch.nn.functional as F -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( - dequant_mxfp4, - quant_dequant_mxfp4, -) -from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( - dequant_mxfp6, - quant_dequant_mxfp6, +from vllm.model_executor.kernels.linear import ( + MxFp4LinearKernel, + MxFp6LinearKernel, + init_mxfp4_linear_kernel, + init_mxfp6_linear_kernel, ) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, - OCP_MX_Scheme, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp4Static, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, ) from vllm.model_executor.parameter import ( GroupQuantScaleParameter, @@ -35,165 +38,54 @@ logger = init_logger(__name__) +_WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = { + "mxfp4": kMxfp4Static, + "mxfp6_e3m2": kMxfp6E3M2Static, + "mxfp6_e2m3": kMxfp6E2M3Static, +} -# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP -# which can force the engine core to spawn instead of fork. -# is_aiter_found_and_supported() checks platform + arch + library availability via -# find_spec/amdsmi, so it stays HIP-free. -# Actual aiter imports are deferred to the functions/methods that need them, -# where HIP initialization is expected. -if is_aiter_found_and_supported(): - from vllm.utils.torch_utils import direct_register_custom_op - - def gemm_with_dynamic_quant( - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - rocm_use_aiter_fp4_asm_gemm: bool = False, - out_dtype: torch.dtype | None = torch.bfloat16, - x_scales: torch.Tensor | None = None, - ) -> torch.Tensor: - from aiter.ops.triton.gemm_afp4wfp4 import ( - gemm_afp4wfp4, - gemm_afp4wfp4_preshuffled_weight_scales, - ) - from aiter.ops.triton.quant import dynamic_mxfp4_quant - - if rocm_use_aiter_fp4_asm_gemm: - from aiter import gemm_a4w4, per_1x32_f4_quant_hip - - M = x.shape[0] - N = weight.shape[0] - K = weight.shape[1] - if rocm_use_aiter_fp4_asm_gemm: - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): - if x_scales is None: - # use hip quant kernel for performance - if M >= 32: - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) - else: - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False) - else: - x_q = x - x_s = x_scales - - if M >= 32: - x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1) - else: - x_s = x_s[:M, ...].view(torch.uint8) - - y = torch.empty(M, N, device=x_q.device, dtype=out_dtype) - gemm_afp4wfp4_preshuffled_weight_scales( - x_q.view(torch.uint8), - weight.view(torch.uint8).view(weight.shape[0] // 16, -1), - x_s, - weight_scale.view(torch.uint8).view( - weight_scale.shape[0] // 32, -1 - ), - out_dtype, - y, - ) - else: - if x_scales is None: - # use hip quant kernel for performance - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) - else: - x_q = x - x_s = x_scales - - y = gemm_a4w4( - x_q, - weight.view(x_q.dtype), - x_s, - weight_scale.view(x_s.dtype), - dtype=out_dtype, - bpreshuffle=True, - ) - return y[:M] - else: - if x_scales is None: - x_q, x_s = dynamic_mxfp4_quant(x) - else: - x_q = x - x_s = x_scales - y = torch.empty( - x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype - ) - - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) - return y - - def gemm_with_dynamic_quant_fake( - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - x_scales: torch.Tensor = None, - rocm_use_aiter_fp4_asm_gemm: bool = False, - out_dtype: torch.dtype | None = torch.bfloat16, - ) -> torch.Tensor: - return torch.empty( - (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device - ) - - direct_register_custom_op( - op_name="gemm_with_dynamic_quant", - op_func=gemm_with_dynamic_quant, - mutates_args=[], - fake_impl=gemm_with_dynamic_quant_fake, - dispatch_key=current_platform.dispatch_key, - ) -elif current_platform.is_rocm(): - logger.warning( - "AITER is not found or not supported on the current platform, " - "QuarkOCP_MX will fall back to emulation." - "Native MXFP4/MXFP6 acceleration will not be available." - ) +_ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = { + "mxfp4": kMxfp4Dynamic, + "mxfp6_e3m2": kMxfp6E3M2Dynamic, + "mxfp6_e2m3": kMxfp6E2M3Dynamic, +} class QuarkOCP_MX(QuarkScheme): + ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel + def __init__( self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any] | None, dynamic_mxfp4_quant: bool = False, ): - self.out_dtype = torch.get_default_dtype() - self.qscheme = "per_group" self.weight_quant_spec = weight_quant_spec self.input_quant_spec = input_quant_spec self.dynamic_mxfp4_quant = dynamic_mxfp4_quant self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp") self.input_dtype: str | None = None if input_quant_spec is not None: - input_quant = input_quant_spec["dtype"] - if input_quant == "fp8_e4m3": - self.input_dtype = "fp8" - else: - self.input_dtype = input_quant.replace("fp", "mxfp") + self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp") - self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype( - self.input_dtype, self.weight_dtype + if self.input_dtype not in [None, *_ACTIVATION_QUANT_KEY_MAP]: + raise ValueError( + f"Unsupported input_dtype={self.input_dtype} for QuarkOCP_MX. " + f"Supported activation dtypes are {_ACTIVATION_QUANT_KEY_MAP.keys()}, " + "or None for weight-only quantization." + ) + + self.weight_quant_key = _WEIGHT_QUANT_KEY_MAP[self.weight_dtype] + self.activation_quant_key = ( + _ACTIVATION_QUANT_KEY_MAP[self.input_dtype] + if self.input_dtype is not None + else None ) if self.weight_dtype == "mxfp4": self.packed_factor: int | Fraction = 2 - self.dequant_func = dequant_mxfp4 else: self.packed_factor = Fraction(numerator=8, denominator=6) - self.dequant_func = partial( - dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "") - ) - - if self.input_dtype is None: - self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( - lambda x: x - ) # no input Q/DQ for weight-only - elif self.input_dtype == "mxfp4": - self.quant_dequant_func = quant_dequant_mxfp4 - else: - self.quant_dequant_func = partial( - quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "") - ) if input_quant_spec is None: self.static_input_scales = False @@ -206,23 +98,6 @@ def __init__( "implemented. Please open an issue." ) - # TODO: integrate (or test) mixed-precision kernel. - self.emulate = not current_platform.supports_mx() or ( - self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4" - ) - - self.rocm_use_aiter_fp4_asm_gemm = ( - rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() - ) - - if not self.emulate and not is_aiter_found_and_supported(): - # Currently need AITER kernels if not emulating - raise NotImplementedError( - f"{self.__class__.__name__} requires AITER to be installed " - "for non-emulation mode! Please refer to " - "https://github.com/ROCm/aiter for installation details." - ) - if not current_platform.supports_mx(): logger.warning_once( "The current platform does not support native MXFP4/MXFP6 " @@ -268,47 +143,16 @@ def process_dynamic_mxfp4_weights_after_loading( from aiter.ops.triton.quant import dynamic_mxfp4_quant w_q, w_s = dynamic_mxfp4_quant(layer.weight) - layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_s, requires_grad=False) layer.weight = torch.nn.Parameter(w_q, requires_grad=False) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) - if self.emulate: - if self.dynamic_mxfp4_quant: - self.process_dynamic_mxfp4_weights_after_loading(layer) - else: - layer.weight_scale = torch.nn.Parameter( - layer.weight_scale.data, requires_grad=False - ) - else: - if self.dynamic_mxfp4_quant: - self.process_dynamic_mxfp4_weights_after_loading(layer) - elif self.rocm_use_aiter_fp4_asm_gemm: - from aiter.ops.shuffle import shuffle_weight - - # shuffle weight scale - weight_scale_shuffle = layer.weight_scale.data - sm, sn = weight_scale_shuffle.shape - weight_scale_shuffle = weight_scale_shuffle.view( - sm // 32, 2, 16, sn // 8, 2, 4, 1 - ) - weight_scale_shuffle = weight_scale_shuffle.permute( - 0, 3, 5, 2, 4, 1, 6 - ).contiguous() - weight_scale_shuffle = weight_scale_shuffle.view(sm, sn) - layer.weight_scale = torch.nn.Parameter( - weight_scale_shuffle, requires_grad=False - ) + if self.dynamic_mxfp4_quant: + self.process_dynamic_mxfp4_weights_after_loading(layer) - # shuffle weight - weight_shuffle = layer.weight.data - weight_shuffle = shuffle_weight(weight_shuffle, layout=(16, 16)) - layer.weight = torch.nn.Parameter(weight_shuffle, requires_grad=False) - else: - layer.weight_scale = torch.nn.Parameter( - layer.weight_scale.data.T.contiguous(), requires_grad=False - ) + self.ocp_mx_linear.process_weights_after_loading(layer) def create_weights( self, @@ -365,25 +209,20 @@ def create_weights( ) layer.register_parameter("weight_scale", weight_scale) + if self.weight_quant_key == kMxfp4Static: + self.ocp_mx_linear = init_mxfp4_linear_kernel( + activation_quant_key=self.activation_quant_key, + ) + elif self.weight_quant_key in [kMxfp6E2M3Static, kMxfp6E3M2Static]: + self.ocp_mx_linear = init_mxfp6_linear_kernel( + weight_quant_key=self.weight_quant_key, + activation_quant_key=self.activation_quant_key, + ) + def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if self.emulate: - dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) - qdq_x = self.quant_dequant_func(x) - return F.linear(qdq_x, dq_w, bias) - y = torch.ops.vllm.gemm_with_dynamic_quant( - x, - layer.weight, - layer.weight_scale, - self.rocm_use_aiter_fp4_asm_gemm, - self.out_dtype, - ) - # gemm_with_dynamic_quant has no bias argument; add it here so the - # native path matches F.linear (e.g. qkv_proj with qkv_bias=True). - if bias is not None: - y = y + bias - return y + return self.ocp_mx_linear.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 29283c7bbda4..e92ebf250827 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -64,9 +64,7 @@ def __init__( kernel_supported_gpu = False if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950 - - kernel_supported_gpu = on_gfx950() + kernel_supported_gpu = current_platform.supports_mx() self.use_aiter_kernel = ( is_aiter_found_and_supported() diff --git a/vllm/model_executor/layers/quantization/qutlass_utils.py b/vllm/model_executor/layers/quantization/qutlass_utils.py index 315ecd0c009d..86b0548307af 100644 --- a/vllm/model_executor/layers/quantization/qutlass_utils.py +++ b/vllm/model_executor/layers/quantization/qutlass_utils.py @@ -84,6 +84,7 @@ def triton_scale_swizzle( ) +@torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=()) def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: """ Rearranges an E8M0 tensor scale from row-major format to @@ -142,6 +143,14 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: return out +@triton_mx_block_rearrange.register_fake +def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor: + rows, cols = scale_tensor.shape + padded_rows = cdiv(rows, 128) * 128 + padded_cols = cdiv(cols, 4) * 4 + return scale_tensor.new_empty((padded_rows, padded_cols)) + + def to_blocked( input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton" ) -> torch.Tensor: @@ -157,7 +166,7 @@ def to_blocked( backend: "torch" (PyTorch path) or "triton" (Triton kernel) Returns: - Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4)) + Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4)) """ if backend == "triton": return triton_mx_block_rearrange(input_matrix).flatten() diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 6f0e237785e0..1b513e9c0ce0 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -103,22 +103,24 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend. Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper, - and interleaves w13 gate/linear rows. + and interleaves w13 gate/linear rows for gated activations. Non-gated + activations use a single w13 projection and keep its row order unchanged. """ from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout # Global scaling factors (same as other FlashInfer backends). num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) - a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) + a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) - half = w13.shape[1] // 2 - w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) - w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) + if layer.activation.is_gated: + half = w13.shape[1] // 2 + w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) + w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) - # Interleave up/gate rows for w13 weights and scales. - w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) - w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) + # Interleave up/gate rows for w13 weights and scales. + w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) + w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) # Convert w13 scale factors: linear → swizzled → MMA layout. w13_scale = swizzle_blockscale(w13_scale) @@ -338,8 +340,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # For some FI kernels, the input scales are shared by all experts. if is_global_sf_supported_for_nvfp4_backend(backend): num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) - a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) + a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) else: a13_scale = a13_scale.max(dim=1).values.to(torch.float32) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py index 4e08a73a69dc..1b5320615f4a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py @@ -190,6 +190,7 @@ def flashinfer_trtllm_mxint4_moe( topk_group: int | None = None, e_score_correction_bias: torch.Tensor | None = None, routing_method_type: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: """ Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel. @@ -262,6 +263,7 @@ def flashinfer_trtllm_mxint4_moe( do_finalize=True, output=None, tune_max_num_tokens=8192, + routing_replay_out=routing_replay_out, ) if isinstance(out, (tuple, list)): out = out[0] diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 632cca1fda29..28c3bbae4c73 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -18,20 +18,34 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int: return activation_to_flashinfer_type(activation).value +def has_flashinfer_situ_activation() -> bool: + try: + from flashinfer.fused_moe.core import ActivationType + except ImportError: + return False + return hasattr(ActivationType, "Situ") + + def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType": from flashinfer.fused_moe.core import ActivationType + if activation == MoEActivation.SITU: + situ = getattr(ActivationType, "Situ", None) + if situ is None: + raise ValueError("The installed FlashInfer does not support SITU") + return situ + # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively ACTIVATION_TO_FI_ACTIVATION = { MoEActivation.SILU_NO_MUL: ActivationType.Silu, MoEActivation.GELU_NO_MUL: ActivationType.Gelu, MoEActivation.SILU: ActivationType.Swiglu, # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args. + MoEActivation.SWIGLUOAI: ActivationType.Swiglu, MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, MoEActivation.GELU: ActivationType.Geglu, MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, - MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, } return ACTIVATION_TO_FI_ACTIVATION[activation] diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 83e56a4567b0..1fce36824710 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -34,7 +34,6 @@ transform_sf_into_required_layout, ) from vllm.utils.platform_utils import get_device_name_as_file_name -from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -45,39 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool: return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz -def _triton_per_token_group_quant_fp8_impl( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return per_token_group_quant_fp8( - x, group_size, column_major_scales=False, use_ue8m0=False - ) - - -def _triton_per_token_group_quant_fp8_fake( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - M, N = x.shape - x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device) - out_bs = torch.empty( - ( - M, - (N + group_size - 1) // group_size, - ), - dtype=torch.float32, - device=x.device, - ) - return x_fp8, out_bs - - -direct_register_custom_op( - "triton_per_token_group_quant_fp8", - _triton_per_token_group_quant_fp8_impl, - fake_impl=_triton_per_token_group_quant_fp8_fake, -) - - def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -909,10 +875,40 @@ def w8a8_triton_block_scaled_mm( Bs: The per-block quantization scale for `B`. block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. - output_dtype: The dtype of the returned tensor. + output_dytpe: The dtype of the returned tensor. Returns: torch.Tensor: The result of matmul. """ + + from vllm.platforms.rocm import on_gfx1250 + + if on_gfx1250(): + # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32. + # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow. + _bn, _bk = block_size[0], block_size[1] + _As = ( + _upcast_e8m0_to_fp32(As) + if As.dtype == torch.float8_e8m0fnu + else As.to(torch.float32) + ) + _Bs = ( + _upcast_e8m0_to_fp32(Bs) + if Bs.dtype == torch.float8_e8m0fnu + else Bs.to(torch.float32) + ) + _K = A.shape[-1] + _N = B.shape[0] + _Af = A.to(torch.float32).reshape(-1, _K) + _Asf = ( + _As.to(torch.float32) + .reshape(-1, _As.shape[-1]) + .repeat_interleave(_bk, dim=1)[:, :_K] + ) + _Bf = B.to(torch.float32) + _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K] + _out = (_Af * _Asf) @ (_Bf * _Bsf).t() + return _out.to(output_dtype).reshape(*A.shape[:-1], _N) + assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index 691d80b0b747..c73c42f4ff02 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from copy import deepcopy from types import MappingProxyType -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import regex as re import torch @@ -69,6 +69,20 @@ def get_dynamic_override( return default_value +def flatten_list(lst: list[Any]) -> list[Any]: + output = [] + + def _flatten(lst: list[Any]): + for i in lst: + if isinstance(i, list): + _flatten(i) + else: + output.append(i) + + _flatten(lst) + return output + + def is_layer_gptq_quantized( prefix: str, quantized_layers: list[str], @@ -83,6 +97,8 @@ def is_layer_gptq_quantized( proj_name = prefix.split(".")[-1] + quantized_layers = flatten_list(quantized_layers) + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index e08e80ec4bb7..15c6a178b05e 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -788,7 +788,7 @@ def _convert_sublayer_to_humming( shape_k_stacks = [shape_k] shape_n_stacks = [shape_n] - if sublayer_name == "w13": + if sublayer_name == "w13" and layer.moe_config.activation.is_gated: shape_n_stacks = [shape_n // 2] * 2 converted_weight_schema, converted_tensors = weight_schema.convert_humming( @@ -991,15 +991,15 @@ def convert_to_humming_moe_kernel_format( # Build sublayer configs from layer properties if not provided if sublayer_configs is None: is_gated = layer.moe_config.activation.is_gated + intermediate_size = layer.moe_config.intermediate_size_per_partition sublayer_configs = { "w13": { - "shape_n": layer.moe_config.intermediate_size_per_partition * 2, + "shape_n": intermediate_size * (2 if is_gated else 1), "shape_k": layer.moe_config.hidden_dim, }, "w2": { "shape_n": layer.moe_config.hidden_dim, - "shape_k": layer.moe_config.intermediate_size_per_partition - * (1 if is_gated else 2), + "shape_k": intermediate_size, }, } diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index f6d96f574d8a..3658761d6f85 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -12,6 +12,7 @@ from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.int8_utils import ( @@ -351,10 +352,12 @@ def marlin_moe_padded_intermediate(intermediate_size: int, group_size: int = -1) return padded -def check_moe_marlin_supports_layer( - layer: RoutedExperts, group_size: int, allow_tile_padding: bool = False +def check_moe_marlin_supports_config( + config: FusedMoEConfig, + group_size: int, + allow_tile_padding: bool = False, ) -> bool: - """Whether the fused MoE Marlin kernel supports ``layer``. + """Whether the fused MoE Marlin kernel supports ``config``. Callers without act-order may pass ``allow_tile_padding=True``: a tile-misaligned intermediate size is then zero-padded to a valid thread @@ -364,15 +367,11 @@ def check_moe_marlin_supports_layer( """ if current_platform.is_rocm(): return False - hidden_size = layer.hidden_size + hidden_size = config.hidden_dim # The layer has not rounded intermediate_size yet; use the stable unpadded # size. gate-up needs n=2*intermediate % 128, down needs k=intermediate % 64. - intermediate_size_per_partition = ( - layer.moe_config.intermediate_size_per_partition_unpadded - ) + intermediate_size_per_partition = config.intermediate_size_per_partition_unpadded assert intermediate_size_per_partition is not None - # apply_router_weight_on_input is not supported for moe marlin - supports_router_weight = not layer.apply_router_weight_on_input if allow_tile_padding: supports_shape = hidden_size % 128 == 0 and ( @@ -384,7 +383,17 @@ def check_moe_marlin_supports_layer( and intermediate_size_per_partition % max(64, group_size) == 0 ) supports_group_size = group_size in [-1, 32, 64, 128] - return supports_shape and supports_group_size and supports_router_weight + return supports_shape and supports_group_size + + +def check_moe_marlin_supports_layer( + layer: RoutedExperts, + group_size: int, + allow_tile_padding: bool = False, +) -> bool: + return check_moe_marlin_supports_config( + layer.moe_config, group_size, allow_tile_padding + ) def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tensor): @@ -397,14 +406,31 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso def marlin_make_workspace_new( - device: torch.device, max_blocks_per_sm: int = 1 + device: torch.device, + max_blocks_per_sm: int = 1, + existing: torch.Tensor | None = None, ) -> torch.Tensor: # In the new marlin kernel, we use the num of threadblocks as workspace # size. The num of threadblocks is sms_count * max_blocks_per_sm. sms = num_compute_units(device.index) - return torch.zeros( - sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False - ) + size = sms * max_blocks_per_sm + # On weight reload, reuse the existing storage so the workspace address + # captured by CUDA graphs stays valid. + if existing is not None: + if ( + existing.device != device + or existing.dtype != torch.int + or existing.numel() != size + ): + raise ValueError( + f"Existing Marlin workspace is incompatible " + f"(device={existing.device}, dtype={existing.dtype}, " + f"numel={existing.numel()}; expected device={device}, " + f"dtype={torch.int}, numel={size}). Reload must reuse the " + f"workspace storage captured by CUDA graphs." + ) + return existing.zero_() + return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False) def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool: diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index f7174ab8f1ff..01d28e5e91ae 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -240,7 +240,9 @@ def prepare_fp4_layer_for_marlin( device = layer.weight.device # WORKSPACE - layer.workspace = marlin_make_workspace_new(device) + layer.workspace = marlin_make_workspace_new( + device, existing=getattr(layer, "workspace", None) + ) # WEIGHT # Repack weights to marlin format @@ -396,7 +398,9 @@ def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 # WORKSPACE - layer.workspace = marlin_make_workspace_new(device, 4) + layer.workspace = marlin_make_workspace_new( + device, 4, existing=getattr(layer, "workspace", None) + ) perm = torch.empty(0, dtype=torch.int, device=device) # WEIGHT @@ -474,14 +478,20 @@ def prepare_moe_fp4_layer_for_marlin( group_size = 16 if is_nvfp4 else 32 - e = layer.moe_config.num_experts + # Use the per-rank (local) expert count: under expert parallelism the + # w13_weight/w2_weight tensors only hold this rank's experts (created with + # local_num_experts), whereas moe_config.num_experts is the global count. + # With no EP the two are equal, so the non-EP path is unchanged. + e = layer.moe_config.num_local_experts k = layer.moe_config.hidden_dim n = layer.moe_config.intermediate_size_per_partition # WORKSPACE device = layer.w13_weight.device param_dtype = layer.params_dtype - layer.workspace = marlin_make_workspace_new(device, 4) + layer.workspace = marlin_make_workspace_new( + device, 4, existing=getattr(layer, "workspace", None) + ) perm = torch.empty(0, dtype=torch.int, device=device) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 739f76659cd7..aaa1d820905d 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -132,7 +132,9 @@ def prepare_fp8_layer_for_marlin( device = layer.weight.device # WORKSPACE - layer.workspace = marlin_make_workspace_new(device) + layer.workspace = marlin_make_workspace_new( + device, existing=getattr(layer, "workspace", None) + ) # WEIGHT # Repack weights to marlin format @@ -278,9 +280,9 @@ def prepare_fp8_moe_layer_for_marlin( # WORKSPACE device = layer.w13_weight.device - # NOTE(rob): we do not need to register the workspace as a param - # because it is not used as part of the weight reloading process. - layer.workspace = marlin_make_workspace_new(device, 4) + layer.workspace = marlin_make_workspace_new( + device, 4, existing=getattr(layer, "workspace", None) + ) perm = torch.empty(0, dtype=torch.int, device=device) # WEIGHT @@ -463,7 +465,9 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: device = layer.weight.device # WORKSPACE - layer.workspace = marlin_make_workspace_new(device) + layer.workspace = marlin_make_workspace_new( + device, existing=getattr(layer, "workspace", None) + ) # WEIGHT - repack FP8 weights to Marlin format perm = torch.empty(0, dtype=torch.int, device=device) @@ -549,7 +553,9 @@ def prepare_mxfp8_moe_layer_for_marlin( param_dtype = torch.get_default_dtype() perm = torch.empty(0, dtype=torch.int, device=device) - layer.workspace = marlin_make_workspace_new(device, 4) + layer.workspace = marlin_make_workspace_new( + device, 4, existing=getattr(layer, "workspace", None) + ) def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: if "w13" in name: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index d076dce758ba..932bf9235289 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -41,8 +41,8 @@ def _mxfp8_e4m3_quantize_torch( ) -> tuple[torch.Tensor, torch.Tensor]: """Naive MXFP8 quantization. For each block of 32 elements along the last dimension, compute a - shared e8m0 scale (the biased exponent of the block-wise amax) - and quantize each element to float8_e4m3fn. + shared e8m0 scale that fits the block-wise amax into the finite + float8_e4m3fn range, and quantize each element to float8_e4m3fn. Returns (quantized_values [same shape, fp8], scales uint8). Scale shape depends on is_sf_swizzled_layout: @@ -58,7 +58,8 @@ def _mxfp8_e4m3_quantize_torch( amax = x_blocked.abs().amax(dim=-1) amax = amax.clamp(min=torch.finfo(torch.float32).tiny) - scale_biased = torch.floor(torch.log2(amax)) + 127.0 + fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max + scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0 scale_biased = scale_biased.clamp(0, 254) scales_uint8 = scale_biased.to(torch.uint8) diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index ad6b272371e4..abb2043d7c45 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -116,19 +116,6 @@ def _dequantize_nvfp4_kernel( tl.store(output_ptr + out_indices, result, mask=block_mask[:, None]) -@triton.jit -def _e2m1_lookup(magnitude): - """Lookup E2M1 float value from 3-bit magnitude.""" - result = tl.where(magnitude == 1, 0.5, 0.0) - result = tl.where(magnitude == 2, 1.0, result) - result = tl.where(magnitude == 3, 1.5, result) - result = tl.where(magnitude == 4, 2.0, result) - result = tl.where(magnitude == 5, 3.0, result) - result = tl.where(magnitude == 6, 4.0, result) - result = tl.where(magnitude == 7, 6.0, result) - return result - - @triton.jit def _round_to_fp4(x): """Round float values to the nearest E2M1 representable value. diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 705da43c373f..6678aa15e2c2 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -90,6 +90,7 @@ def __str__(self): GroupShape.PER_CHANNEL: "per_channel", } group_shape = d.get(self.group_shape, str(self.group_shape)) + return ( f"{fx.graph.dtype_abbrs[self.dtype]}," f"{'static' if self.static else 'dynamic'},{group_shape}" @@ -106,15 +107,24 @@ class QuantKey: symmetric: symmetric if True, asymmetric if False """ - dtype: torch.dtype + # TODO: QuantKey.dtype is assumed to be `torch.dtype` in matcher_utils.py, + # but #37990 introduced e.g. `kInt4Static` that uses a `ScalarType` dtype, + # same for kMxfp6 that does not have a native torch representation. + # Logical dtype and storage (torch) dtype should be separated (see #48949). + dtype: torch.dtype | ScalarType scale: ScaleDesc scale2: ScaleDesc | None = None symmetric: bool = True def __str__(self): scale2_str = f"scale2({self.scale2})," if self.scale2 else "" + dtype_description = ( + fx.graph.dtype_abbrs[self.dtype] + if isinstance(self.dtype, torch.dtype) + else self.dtype + ) return ( - f"QuantKey({fx.graph.dtype_abbrs[self.dtype]}," + f"QuantKey({dtype_description}," f"scale({self.scale}),{scale2_str}" f"{'a' if not self.symmetric else ''}symmetric)" ) @@ -172,6 +182,26 @@ def __str__(self): kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True) +kMxfp6E3M2StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp6E3M2Static = QuantKey( + scalar_types.float6_e3m2f, scale=kMxfp6E3M2StaticGroupScale, symmetric=True +) + +kMxfp6E3M2DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp6E3M2Dynamic = QuantKey( + scalar_types.float6_e3m2f, scale=kMxfp6E3M2DynamicGroupScale, symmetric=True +) + +kMxfp6E2M3StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp6E2M3Static = QuantKey( + scalar_types.float6_e2m3f, scale=kMxfp6E2M3StaticGroupScale, symmetric=True +) + +kMxfp6E2M3DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp6E2M3Dynamic = QuantKey( + scalar_types.float6_e2m3f, scale=kMxfp6E2M3DynamicGroupScale, symmetric=True +) + # TODO: convert this to use SCALAR_TYPE. This is not right. kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True) @@ -190,6 +220,8 @@ def __str__(self): kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True) kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True) +kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True) +kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True) # INT4 W4A8 quantization keys diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 3c946dd130cc..29ce9e5000d9 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -5,6 +5,7 @@ import numpy as np import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .base import RotaryEmbeddingBase @@ -24,16 +25,17 @@ def _triton_mrope_forward( rd: tl.constexpr, pad_n_qh: tl.constexpr, pad_n_kh: tl.constexpr, - pad_hd: tl.constexpr, + pad_rd: tl.constexpr, mrope_section_t: tl.constexpr, mrope_section_h: tl.constexpr, mrope_section_w: tl.constexpr, is_interleaved: tl.constexpr, + is_neox_style: tl.constexpr, ): # Adapted from # https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py # This version supports flatten input tensors from vllm - # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2) + # and supports cos and sin cache with shape (3, num_tokens, rotary_dim // 2) # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary pid = tl.program_id(0) # locate start address @@ -44,9 +46,9 @@ def _triton_mrope_forward( # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position # m of this program instance # #################################################################### - # Note: cos and sin now have shape (3, num_tokens, head_dim // 2) + # Note: cos and sin now have shape (3, num_tokens, rotary_dim // 2) - # Updated stride calculation for half head_dim + # Updated stride calculation for half rotary_dim half_rd = rd // 2 t_cos = cos + pid * half_rd h_cos = t_cos + num_tokens * half_rd @@ -55,12 +57,17 @@ def _triton_mrope_forward( h_sin = t_sin + num_tokens * half_rd w_sin = h_sin + num_tokens * half_rd - # Updated offsets for half head_dim - cos_offsets = tl.arange(0, pad_hd // 2) + # Updated offsets for half rotary_dim + cos_offsets = tl.arange(0, pad_rd // 2) if is_interleaved: - h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) - w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) - t_mask = ~(h_mask | w_mask) + valid_mask = cos_offsets < half_rd + h_mask = ( + valid_mask & ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) + ) + w_mask = ( + valid_mask & ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) + ) + t_mask = valid_mask & ~(h_mask | w_mask) else: t_end = mrope_section_t h_end = t_end + mrope_section_h @@ -79,55 +86,74 @@ def _triton_mrope_forward( sin_row = t_sin_row + h_sin_row + w_sin_row # #################################################################### - # Load the left and right half of q and k for the current - # program instance (i.e. for the current token) separately + # Load the two values in each rotary pair for the current token. + # NeoX pairs the first and second halves, while GPT-J pairs + # adjacent values. # #################################################################### - # left half of the head - first_half_q_offsets = ( - tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_half_k_offsets = ( - tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) - first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) + if is_neox_style: + rotary_offsets = tl.arange(0, pad_rd // 2) + first_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + first_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd // 2 + ) + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd // 2 + ) - q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to( - sin_row.dtype - ) + q_tile_1 = tl.load(q_ptr + first_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_1 = tl.load(k_ptr + first_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - # right half of the head - second_half_q_offsets = first_half_q_offsets + (rd // 2) - second_half_k_offsets = first_half_k_offsets + (rd // 2) - second_q_mask = first_q_mask - second_k_mask = first_k_mask + second_q_offsets = first_q_offsets + (rd // 2) + second_k_offsets = first_k_offsets + (rd // 2) + q_tile_2 = tl.load(q_ptr + second_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_2 = tl.load(k_ptr + second_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to( - sin_row.dtype - ) + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + tl.store(q_ptr + first_q_offsets, new_q_tile_1, mask=first_q_mask) + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + tl.store(q_ptr + second_q_offsets, new_q_tile_2, mask=first_q_mask) + + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + tl.store(k_ptr + first_k_offsets, new_k_tile_1, mask=first_k_mask) + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + tl.store(k_ptr + second_k_offsets, new_k_tile_2, mask=first_k_mask) + else: + # Load and store adjacent rotary pairs contiguously. Using stride-two + # even/odd offsets makes Triton emit scalar 16-bit memory operations on + # AMD, while split/interleave only rearranges values in registers. + rotary_offsets = tl.arange(0, pad_rd) + q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd + ) + k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd + ) + + q_tile = tl.load(q_ptr + q_offsets, mask=q_mask, other=0).to(sin_row.dtype) + k_tile = tl.load(k_ptr + k_offsets, mask=k_mask, other=0).to(sin_row.dtype) + q_tile_1, q_tile_2 = tl.split(tl.reshape(q_tile, (pad_n_qh, pad_rd // 2, 2))) + k_tile_1, k_tile_2 = tl.split(tl.reshape(k_tile, (pad_n_kh, pad_rd // 2, 2))) - # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin] - # Since cos and sin are now half-size, - # we use the same cos_row and sin_row for both halves - new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row - tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask) - new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row - tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask) + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + new_q_tile = tl.interleave(new_q_tile_1, new_q_tile_2) + tl.store(q_ptr + q_offsets, new_q_tile, mask=q_mask) - new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row - tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask) - new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row - tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask) + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + new_k_tile = tl.interleave(new_k_tile_1, new_k_tile_2) + tl.store(k_ptr + k_offsets, new_k_tile, mask=k_mask) def triton_mrope( @@ -139,23 +165,26 @@ def triton_mrope( head_size: int, rotary_dim: int, mrope_interleaved: bool, + is_neox_style: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Qwen2VL mrope kernel. Args: q: [num_tokens, num_heads * head_size] k: [num_tokens, num_kv_heads * head_size] - cos: [3, num_tokens, head_size //2 ] + cos: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) - sin: [3, num_tokens, head_size //2 ] + sin: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) mrope_section: [t, h, w] head_size: int + is_neox_style: Whether rotary pairs use split-half (NeoX) or + adjacent (GPT-J) layout. """ n_row, n_q_head_head_dim = q.shape n_q_head = n_q_head_head_dim // head_size n_kv_head = k.shape[1] // head_size - pad_hd = triton.next_power_of_2(head_size) + pad_rd = triton.next_power_of_2(rotary_dim) pad_n_q_head = triton.next_power_of_2(n_q_head) pad_n_kv_head = triton.next_power_of_2(n_kv_head) @@ -166,6 +195,11 @@ def triton_mrope( cos = cos.contiguous() sin = sin.contiguous() + # Small adjacent-pair tiles perform best with one wave per program on + # ROCm. Keep the existing launch shape for larger rotary dimensions, + # NeoX, and other backends. + use_single_wave = current_platform.is_rocm() and not is_neox_style and pad_rd <= 64 + num_warps = 1 if use_single_wave else 4 _triton_mrope_forward[(n_row,)]( q, k, @@ -178,11 +212,13 @@ def triton_mrope( rotary_dim, pad_n_q_head, pad_n_kv_head, - pad_hd, + pad_rd, mrope_section[0], mrope_section[1], mrope_section[2], mrope_interleaved, + is_neox_style, + num_warps=num_warps, ) return q, k @@ -349,6 +385,7 @@ def forward_cuda( self.head_size, self.rotary_dim, self.mrope_interleaved, + self.is_neox_style, ) return q.reshape(query_shape), k.reshape(key_shape) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 5b8e2bf008e9..9a671be55639 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -8,7 +8,7 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import get_current_vllm_config +from vllm.config import CUDAGraphMode, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -310,6 +310,7 @@ def sparse_attn_indexer( topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -317,7 +318,8 @@ def sparse_attn_indexer( skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run - attn_metadata = get_forward_context().attn_metadata + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) @@ -357,6 +359,7 @@ def sparse_attn_indexer( topk_indices_buffer, skip_k_cache_insert, use_pcp, + dense_mha_metadata_layer_name, use_fp4_cache, ) attn_metadata_narrowed = attn_metadata[k_cache_prefix] @@ -402,6 +405,24 @@ def sparse_attn_indexer( scale_fmt, ) + # The indexer and main MLA may classify the same short extend differently + # because they use independent decode thresholds. Only the main MLA route + # can determine whether the top-k indices will be consumed. + if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name) + if dense_mha_layer: + mla_metadata = attn_metadata.get(dense_mha_layer) + prefill_metadata = getattr(mla_metadata, "prefill", None) + if ( + getattr(prefill_metadata, "use_dense_mha", False) + and getattr(mla_metadata, "num_decode_tokens", -1) == 0 + and not torch.cuda.is_current_stream_capturing() + ): + # Deliberately leave the buffer untouched. Dense MHA does not + # consume top-k indices for this batch; clearing it would be + # unnecessary work. + return topk_indices_buffer + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the # top-k kernels scatter valid indices into it. On the fused deepseek_v32 # nvidia path, _fused_norm_rope_kernel already cleared the same @@ -684,6 +705,7 @@ def sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -739,6 +761,7 @@ def __init__( self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache + self.dense_mha_metadata_layer_name = "" # DCP scalars are constant for the run; resolve them here (config is set # during model construction) and pass them into the custom op, rather # than threading them through per-step metadata. @@ -800,6 +823,7 @@ def forward_cuda( self.topk_indices_buffer, self.skip_k_cache_insert, self.use_pcp, + _encode_layer_name(self.dense_mha_metadata_layer_name), self.use_fp4_cache, self.dcp_rank, self.dcp_world_size, diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 4d2e50420f42..f7e0b6510547 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -164,7 +164,11 @@ def rocm_unquantized_gemm_impl( if use_skinny_reduce_counting: return ops.wvSplitKrc(x, weight, cu_count, bias) - if use_aiter_triton_gemm(n, m, k, x.dtype): + # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires + # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a + # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b + # hidden), so skip aiter there and fall back to the torch GEMM path below. + if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) @@ -172,6 +176,8 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) + # build (gfx9/gfx11 ISA); fall back to torch GEMM there. + # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) @@ -234,10 +240,14 @@ def dispatch_cpu_unquantized_gemm( layer.cpu_linear = torch.nn.functional.linear return + # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights). + # These layers are handled by their own specialized methods. if layer.weight.ndim != 2: # this is not a linear layer - # For now it should be a causal_conv1d op - if torch.cpu._is_amx_tile_supported(): + # For now it should be a causal_conv1d op or MoE 3D expert weights + if torch.cpu._is_amx_tile_supported() and hasattr( + ops, "causal_conv1d_weight_pack" + ): # prepack conv weight unpacked = ( layer.weight.view( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 8d9a7ccbacae..6eed2b5249b6 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer): padding_size: padding size for the vocabulary. quant_config: quant config for the layer prefix: full name of the layer in the state dict + disable_tp: If true, tensor parallelism will be disabled for this layer. """ # noqa: E501 # --8<-- [end:vocab_parallel_embedding] @@ -245,12 +246,19 @@ def __init__( padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + disable_tp: bool = False, ): super().__init__() # Keep the input dimensions. - tp_rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + self.disable_tp = disable_tp + if disable_tp: + tp_rank, self.tp_size = 0, 1 + else: + tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = tp_rank self.num_embeddings = num_embeddings self.padding_size = padding_size self.org_vocab_size = org_num_embeddings or num_embeddings @@ -323,6 +331,13 @@ def __init__( params_dtype=params_dtype, weight_loader=self.weight_loader, ) + self.update_param_tp_status() + + def update_param_tp_status(self): + for param in self.parameters(): + if isinstance(param, BasevLLMParameter): + param.tp_rank = self.tp_rank + param.tp_size = self.tp_size @classmethod def _get_indices( @@ -487,9 +502,9 @@ def forward(self, input_): # Mask the output embedding. if self.tp_size > 1: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) - # Reduce across all the model parallel GPUs. - output = tensor_model_parallel_all_reduce(output_parallel) - return output + # Reduce across all the model parallel GPUs. + return tensor_model_parallel_all_reduce(output_parallel) + return output_parallel def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" @@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding): params_dtype: type of the parameters. org_num_embeddings: original vocabulary size (without LoRA). padding_size: padding size for the vocabulary. + disable_tp: If true, tensor parallelism will be disabled for this layer. """ # --8<-- [end:parallel_lm_head] @@ -530,6 +546,8 @@ def __init__( padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + disable_tp: bool = False, ): super().__init__( num_embeddings, @@ -539,6 +557,7 @@ def __init__( padding_size, quant_config, prefix, + disable_tp=disable_tp, ) self.quant_config = quant_config if bias: diff --git a/vllm/model_executor/model_loader/bitsandbytes_loader.py b/vllm/model_executor/model_loader/bitsandbytes_loader.py index cc9af05af726..525de8d6a8f9 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -467,7 +467,7 @@ def _get_bnb_target_modules(self, model: nn.Module) -> None: elif isinstance(module, RoutedExperts) and hasattr( module.quant_method, "quant_config" ): - # TODO: support FusedMoE with prequant and 8bit. + # TODO: support RoutedExperts with prequant and 8bit. if self.pre_quant and self.load_8bit: raise ValueError( "Prequant BitsAndBytes 8bit models with RoutedExperts " diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 3ea76f4d9b3a..039a577f80d1 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -220,7 +220,7 @@ def _prepare_weights( # safetensors file. Using both breaks. # Here, we download the `model.safetensors.index.json` and filter # any files not found in the index. - if not is_local: + if not is_local and len(hf_weights_files) > 1: download_safetensors_index_file_from_hf( model_name_or_path, index_file, diff --git a/vllm/model_executor/model_loader/mtp_validation.py b/vllm/model_executor/model_loader/mtp_validation.py new file mode 100644 index 000000000000..3f20756abce2 --- /dev/null +++ b/vllm/model_executor/model_loader/mtp_validation.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Scoped controls for MTP checkpoint completeness validation.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar + +_mtp_completeness_check_enabled: ContextVar[bool] = ContextVar( + "mtp_completeness_check_enabled", default=True +) + + +def is_mtp_completeness_check_enabled() -> bool: + """Return whether MTP completeness validation is enabled in this scope.""" + return _mtp_completeness_check_enabled.get() + + +@contextmanager +def disable_mtp_completeness_check() -> Iterator[None]: + """Temporarily disable MTP completeness validation for one weight load.""" + token = _mtp_completeness_check_enabled.set(False) + try: + yield + finally: + _mtp_completeness_check_enabled.reset(token) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 92f454f9a5f4..d6d4eb4b9308 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -14,7 +14,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from .meta import ( - SKIP_TENSORS, + SKIP_LOAD_TENSORS, capture_layer_to_meta, get_numel_loaded, materialize_layer, @@ -140,7 +140,7 @@ def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: """Wrap each parameter's weight loader.""" # Note that nested wrapping will occur for shared tensors for name, tensor in get_layer_tensors(layer).items(): - if name in SKIP_TENSORS: + if name in SKIP_LOAD_TENSORS: continue if _get_weight_loader(tensor).__name__ != "online_process_loader": tensor.weight_loader = make_online_process_loader(layer, name) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 33af4ab71f22..213a9aaeaf8e 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -20,9 +20,11 @@ "get_numel_loaded", ] +# Modules whose tensors are never moved to, or materialized from, the meta device. SKIP_MODULES: set[str] = {"HadamardTransform"} -SKIP_TENSORS: set[str] = { +# Tensors never loaded by a weight loader, so the layerwise trigger ignores them. +SKIP_LOAD_TENSORS: set[str] = { "_expert_map", "expert_mask", "expert_global_to_physical", @@ -31,6 +33,10 @@ "e_score_correction_bias", } +# Tensors which are never moved to, or materialized from, the meta device. +# `bias` is built after create_weights(), so it is never on meta to begin with. +SKIP_TENSORS: set[str] = SKIP_LOAD_TENSORS | {"bias"} + def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor: """Convert a tensor to a meta tensor while preserving class and attributes.""" diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py index f0078d0f9d85..b6d129a79403 100644 --- a/vllm/model_executor/model_loader/reload/utils.py +++ b/vllm/model_executor/model_loader/reload/utils.py @@ -33,15 +33,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors: def get_layer_size(layer: torch.nn.Module) -> int: """Calculate total number of elements across loadable tensors in a layer. - Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta - device and never loaded via weight_loader during layerwise reload. + Excludes SKIP_LOAD_TENSORS (e.g. _expert_map) which are never loaded via + weight_loader during layerwise reload. """ - from .meta import SKIP_TENSORS + from .meta import SKIP_LOAD_TENSORS return sum( tensor.numel() for name, tensor in get_layer_tensors(layer).items() - if name not in SKIP_TENSORS + if name not in SKIP_LOAD_TENSORS ) diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index 008abb6fdfe1..befe560929d1 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -66,6 +66,7 @@ ] logger = init_logger(__name__) +_TENSORIZER_ENGINE_CLEANUP_GRACE_S = 10.0 def is_valid_deserialization_uri(uri: str | None) -> bool: @@ -730,10 +731,38 @@ def tensorize_vllm_model( from vllm.v1.engine.llm_engine import LLMEngine engine = LLMEngine.from_vllm_config(engine_config) - engine.collective_rpc( - "save_tensorized_model", - kwargs={"tensorizer_config": tensorizer_config.to_serializable()}, - ) + error: BaseException | None = None + try: + engine.collective_rpc( + "save_tensorized_model", + kwargs={"tensorizer_config": tensorizer_config.to_serializable()}, + ) + except BaseException as operation_error: + error = operation_error + + def shutdown_engine_core() -> None: + engine.engine_core.shutdown( + timeout=( + envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS + + _TENSORIZER_ENGINE_CLEANUP_GRACE_S + ) + ) + + for name, callback in ( + ("renderer", engine.renderer.shutdown), + ("engine core", shutdown_engine_core), + ): + try: + callback() + except BaseException as shutdown_error: + logger.exception("Failed to shut down tensorization %s", name) + if error is None: + error = shutdown_error + elif hasattr(error, "add_note"): + error.add_note(f"{name} shutdown also failed: {shutdown_error!r}") + + if error is not None: + raise error def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig): diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 3367f4833e6c..36683225d170 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -16,10 +16,9 @@ from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import ( - Attention, - MLAAttention, MMEncoderAttention, ) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -122,12 +121,13 @@ def process_weights_after_loading( # the caching allocator, which starves the OS on UMA devices. release_device_memory_under_pressure(target_device) - # Initialize post-load attention weights for Attention, MLA, and MM encoder. - # NOTE: Happens after other modules so we can easily decompress weights. + # Initialize post-load attention weights for any attention layer and MM + # encoder. NOTE: Happens after other modules so we can easily decompress + # weights. for _, module in model.named_modules(): - if isinstance( - module, (Attention, MLAAttention, MMEncoderAttention) - ) and hasattr(module, "process_weights_after_loading"): + if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr( + module, "process_weights_after_loading" + ): # TODO(lucas): see if there is a way to unify the signatures # of process_weights_after_loading with device_loading_context(module, target_device): @@ -142,6 +142,10 @@ def process_weights_after_loading( if isinstance(module, HpcModule): module.process_weights_after_loading(model) + # Model-level post-load hook, after the per-layer quant finalize. + if hasattr(model, "process_weights_after_loading"): + model.process_weights_after_loading() + # Needed for torchao model reloading via model.reload_weights # @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights` if model_config.quantization == "torchao": @@ -246,9 +250,9 @@ def get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], if key in _MODEL_ARCH_BY_HASH: return _MODEL_ARCH_BY_HASH[key] - model_arch = _get_model_architecture(model_config) - _MODEL_ARCH_BY_HASH[key] = model_arch - return model_arch + model_cls_and_arch = _get_model_architecture(model_config) + _MODEL_ARCH_BY_HASH[key] = model_cls_and_arch + return model_cls_and_arch def get_model_cls(model_config: ModelConfig) -> type[nn.Module]: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index db161a589885..e1897a77f102 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -694,12 +694,10 @@ def _get_checkpoints_size_bytes(files: list[str]) -> int: def _get_available_ram_bytes() -> int: - """Return available RAM, honoring cgroup limits on ROCm.""" + """Return available RAM, honoring cgroup limits.""" import psutil host_available = psutil.virtual_memory().available - if not current_platform.is_rocm(): - return host_available from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 7818200dff5d..862604921ef1 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -43,7 +43,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -168,7 +168,7 @@ def __init__( prefix=f"{prefix}.shared_experts", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=config.n_routed_experts, diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 0122d019588a..6ed67861e9a7 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -18,7 +18,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, MoERunner, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -122,8 +122,8 @@ def __init__( prefix=f"{prefix}.shared_experts", ) - # Routed experts using FusedMoE - self.experts = FusedMoE( + # Routed experts using FusedMoEFactory + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.num_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index b997e153a994..74b17129c989 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -211,8 +211,8 @@ def forward( ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q) - k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k) + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view_as(q) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view_as(k) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 5117f541cfd0..d4423e83cf39 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -13,7 +13,9 @@ from vllm.config.multimodal import BaseDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -222,7 +224,7 @@ def __init__( bias=config.mlp_bias, ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.moe_num_experts, top_k=config.moe_topk, diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 642d07ee659d..926f4b829efa 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -41,7 +41,7 @@ ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -289,7 +289,7 @@ def __init__( else: self.shared_experts = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=self.num_experts, top_k=self.top_k, diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 5cc057c3b458..87e1d4dba082 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -18,7 +18,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -329,8 +329,8 @@ def __init__( else: self.shared_experts = None - # Routed experts using FusedMoE - self.experts = FusedMoE( + # Routed experts using FusedMoEFactory + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=self.num_experts, top_k=self.top_k, @@ -602,7 +602,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: (".gate_up_proj", ".up_proj", 1), ] - # Expert parameter mappings from FusedMoE + # Expert parameter mappings from FusedMoEFactory expert_mappings = list(self.get_expert_mapping()) def load_param(name: str, tensor: torch.Tensor, shard_id=None) -> bool: diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py index da6b1ddb8b66..53679edf89b9 100644 --- a/vllm/model_executor/models/bailing_moe_mtp.py +++ b/vllm/model_executor/models/bailing_moe_mtp.py @@ -20,6 +20,9 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -370,7 +373,10 @@ def load_lm_head(loaded_weight: torch.Tensor) -> None: 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: + if ( + layer_idx not in loaded_mtp_layers + and is_mtp_completeness_check_enabled() + ): raise ValueError( f"Bailing MTP speculative decoding layer {layer_idx} " "weights are missing from checkpoint. Use a checkpoint " diff --git a/vllm/model_executor/models/bert_with_rope.py b/vllm/model_executor/models/bert_with_rope.py index 020ef993474a..02ba333ca603 100644 --- a/vllm/model_executor/models/bert_with_rope.py +++ b/vllm/model_executor/models/bert_with_rope.py @@ -320,7 +320,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits, _ = self.router(hidden_states) # FIXME(Isotr0py): This implementation is too tricky, - # we should use FusedMoE instead in the future + # we should use FusedMoEFactory instead in the future # after supporting ungated activation for it. topk_weights, topk_ids, _ = fused_topk( hidden_states, router_logits, self.top_k, renormalize=False diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index b993247ec6f3..3365ab06fc56 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -16,7 +16,7 @@ ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -300,7 +300,7 @@ def __init__( self.shared_experts = None self.shared_expert_combination_strategy = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.num_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -318,7 +318,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: orig_shape = hidden_states.shape hidden_states = hidden_states.view(-1, self.hidden_size) router_logits, _ = self.gate(hidden_states) - # FusedMoE handles shared expert overlap internally and returns + # FusedMoEFactory handles shared expert overlap internally and returns # shared_output + routed_output when shared_experts is set. final_hidden_states = self.experts(hidden_states, router_logits) if self.shared_expert_combination_strategy == "average": diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 96232237a06e..66427dfcdd48 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -339,6 +339,39 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: ) +class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig): + """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface. + + Kimi-K3 ships its routed experts as compressed-tensors + ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which + lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel + selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config`` + (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4 + checkpoint is left untouched. Covers both the main model and the MTP draft. + + ``model_arch_config.quantization_config`` is a separate dict, snapshotted in + ``ModelConfig.__init__`` before this hook runs, and it is what + ``_verify_quantization`` reads when resolving the quant method. Patch it + alongside the hf configs so the rewrite lands before resolution; otherwise + the main model still resolves to compressed-tensors. + """ + + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + for cfg in ( + model_config.hf_config, + model_config.hf_text_config, + model_config.model_arch_config, + ): + quant_config = getattr(cfg, "quantization_config", None) + if ( + isinstance(quant_config, dict) + and quant_config.get("quant_method") == "compressed-tensors" + and quant_config.get("format") == "mxfp4-pack-quantized" + ): + quant_config["quant_method"] = "mxfp4" + + class GptOssForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -441,6 +474,28 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: pooler_config.use_activation = False +class JinaEmbeddingsV5ModelConfig(VerifyAndUpdateConfig): + """Config handler for Jina Embeddings V5 embedding models.""" + + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + """Enable the bidirectional encoder backbone for -nano checkpoints. + + The V5 family ships more than one backbone under a single + `architectures` entry: the `-small` variants are Qwen3 decoders, while + `-nano` is a bidirectional EuroBERT encoder. Upstream ships a separate + `configuration_*.py` per repository, so the config carries no backbone + field and the encoder variants are only identifiable by + `is_decoder=False`. For encoder checkpoints, set `is_causal=False` so the + Llama backbone uses EncoderOnlyAttention; `JinaEmbeddingsV5Model` then + dispatches to its encoder implementation. + """ + if getattr(model_config.hf_config, "is_decoder", True): + return + + model_config.hf_config.is_causal = False + + class JinaForRankingConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -768,18 +823,56 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +class Qwen3_5ForCausalLMConfig(Qwen3_5ForConditionalGenerationConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + # Text-only Qwen3.5 models use one-dimensional positions. Remove the + # M-RoPE fields inherited from the multimodal configuration. + hf_text_config = vllm_config.model_config.hf_text_config + rope_parameters = getattr(hf_text_config, "rope_parameters", None) + if rope_parameters is not None: + rope_parameters.pop("mrope_section", None) + rope_parameters.pop("mrope_interleaved", None) + + class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig): - """ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache - handling and additionally serves BIDIRECTIONAL attention: ColPali-style - document/query encoding attends over the whole sequence, not causally. Set - is_causal=False so Qwen3NextAttention builds its full_attention layers with - AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are - unaffected). Generation arches keep the parent (causal) and are untouched. - """ + """Apply the attention contract declared by a ColQwen3.5 checkpoint.""" @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: - model_config.hf_config.is_causal = False + configs = { + id(config): config + for config in ( + model_config.hf_config, + model_config.hf_text_config, + ) + } + declarations = [ + contract + for config in configs.values() + if (contract := getattr(config, "retrieval_attention_contract", None)) + is not None + ] + supported = {"causal", "bidirectional"} + if not declarations: + raise ValueError( + "ColQwen3.5 checkpoints must declare " + "retrieval_attention_contract as 'causal' or 'bidirectional'" + ) + if ( + any(not isinstance(contract, str) for contract in declarations) + or any(contract not in supported for contract in declarations) + or len(set(declarations)) != 1 + ): + raise ValueError( + "unsupported or conflicting ColQwen3.5 " + f"retrieval_attention_contract declarations: {declarations!r}" + ) + is_causal = declarations[0] == "causal" + for config in configs.values(): + config.is_causal = is_causal class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): @@ -844,8 +937,11 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "GteNewForSequenceClassification": GteNewModelConfig, "GteNewModel": GteNewModelConfig, "JambaForSequenceClassification": JambaForSequenceClassificationConfig, + "JinaEmbeddingsV5Model": JinaEmbeddingsV5ModelConfig, "JinaForRanking": JinaForRankingConfig, "JinaVLForRanking": JinaVLForSequenceClassificationConfig, + "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig, + "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig, "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, "LlamaBidirectionalModel": LlamaBidirectionalConfig, "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig, @@ -860,7 +956,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig, "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, + "Qwen3_5ForCausalLM": Qwen3_5ForCausalLMConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "Qwen3_5MoeForCausalLM": Qwen3_5ForCausalLMConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, diff --git a/vllm/model_executor/models/cosmos3.py b/vllm/model_executor/models/cosmos3.py index 71629196a78e..7e4e82745772 100644 --- a/vllm/model_executor/models/cosmos3.py +++ b/vllm/model_executor/models/cosmos3.py @@ -37,6 +37,11 @@ class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration): ".to_out.": ".o_proj.", ".norm_q.": ".q_norm.", ".norm_k.": ".k_norm.", + # ModelOpt-native dialect (diffusers/transformers read these; vLLM reads + # weight_scale/input_scale instead), drop so AutoWeightsLoader passes + ".input_quantizer.": None, + ".weight_quantizer.": None, + ".output_quantizer.": None, }, orig_to_new_prefix={ "proj_in.": None, diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py index be61d564ccf2..62ab7922c6f0 100644 --- a/vllm/model_executor/models/cosmos3_edge.py +++ b/vllm/model_executor/models/cosmos3_edge.py @@ -276,7 +276,13 @@ def get_hf_processor(self, **kwargs: object) -> ProcessorMixin: class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor): - pass + @staticmethod + def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool: + # Cosmos renders each video as a full + # <|vision_start|><|video_pad|><|vision_end|> placeholder, and its + # reference processor replaces that entire triplet with timestamped, + # per-frame vision sequences. + return False class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder): @@ -517,6 +523,7 @@ class Cosmos3EdgeForConditionalGeneration( }, orig_to_new_substr={ "_moe_gen": None, + "k_norm_und_for_gen": None, ".add_q_proj.": None, ".add_k_proj.": None, ".add_v_proj.": None, diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py index c28cf939241f..7f1ddc442ae6 100644 --- a/vllm/model_executor/models/dbrx.py +++ b/vllm/model_executor/models/dbrx.py @@ -16,7 +16,7 @@ ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, RoutedExperts, ) from vllm.model_executor.layers.linear import ( @@ -159,7 +159,7 @@ def __init__( self.router = DbrxRouter(config, self.params_dtype) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.ffn_config.moe_num_experts, top_k=config.ffn_config.moe_top_k, hidden_size=config.d_model, diff --git a/vllm/model_executor/models/deepencoder2.py b/vllm/model_executor/models/deepencoder2.py index fdec155d5345..9e251c39b292 100644 --- a/vllm/model_executor/models/deepencoder2.py +++ b/vllm/model_executor/models/deepencoder2.py @@ -10,6 +10,8 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. +from functools import lru_cache + import torch import torch.nn as nn import transformers @@ -90,8 +92,6 @@ def forward( return_dict=None, cache_position=None, ): - # token_type_ids - self._current_token_type_ids = token_type_ids causal_mask_mapping = { "full_attention": self._update_causal_mask( attention_mask, @@ -131,15 +131,12 @@ def _update_causal_mask( input_tensor.shape[1], ) - token_type_ids = self._current_token_type_ids - # attention mask causal_mask = self._create_custom_4d_mask( sequence_length=sequence_length, dtype=dtype, device=device, batch_size=batch_size, - token_type_ids=token_type_ids, ) # padding mask @@ -150,44 +147,43 @@ def _update_causal_mask( return causal_mask + @classmethod + @lru_cache(maxsize=8) + def compute_mask_base(cls, sequence_length, dtype, device): + # token_type_ids is the fixed pattern [0]*n_query + [1]*n_query, + # identical across the batch, so the mask depends only on + # sequence_length: img tokens (first half) attend to + # everything, txt tokens (second half) attend causally among + # themselves. lru_cache keeps one batch-invariant [1, 1, S, S] + # mask per (S, dtype, device). + min_dtype = torch.finfo(dtype).min + n_query = sequence_length // 2 + img = torch.arange(sequence_length, device=device) < n_query + txt = ~img + causal = torch.tril( + torch.ones( + sequence_length, + sequence_length, + dtype=torch.bool, + device=device, + ) + ) + allow = img[None, :] | (txt[:, None] & txt[None, :] & causal) + return torch.where( + allow, + torch.zeros((), dtype=dtype, device=device), + torch.full((), min_dtype, dtype=dtype, device=device), + )[None, None] + def _create_custom_4d_mask( self, sequence_length, dtype, device, batch_size, - token_type_ids, ): - min_dtype = torch.finfo(dtype).min - - masks = [] - for b in range(batch_size): - mask = torch.full( - (sequence_length, sequence_length), - fill_value=min_dtype, - dtype=dtype, - device=device, - ) - - type_ids = token_type_ids[b] - - image_positions = (type_ids == 0).nonzero(as_tuple=True)[0] - text_positions = (type_ids == 1).nonzero(as_tuple=True)[0] - - # non-casual - if len(image_positions) > 0: - mask[image_positions[:, None], image_positions] = 0.0 - - # causal - for i, text_pos in enumerate(text_positions): - if len(image_positions) > 0: - mask[text_pos, image_positions] = 0.0 - mask[text_pos, text_positions[: i + 1]] = 0.0 - - masks.append(mask) - - mask = torch.stack(masks, dim=0).unsqueeze(1) - return mask + base = self.compute_mask_base(sequence_length, dtype, device) + return base.expand(batch_size, -1, -1, -1) return CustomQwen2ModelInner(config) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 3a0c21fe7d29..35cfce3cb98d 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -21,6 +21,9 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -41,25 +44,6 @@ ) -def _restore_full_token_layout_if_needed( - hidden_states: torch.Tensor, - residual: torch.Tensor, - num_tokens: int, - is_sequence_parallel: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """Restore full token rows for the MTP proposer after SP MoE layers.""" - if not is_sequence_parallel and hidden_states.shape[0] == num_tokens: - return hidden_states, residual - - combined_states = torch.cat([hidden_states, residual], dim=-1) - combined_states = tensor_model_parallel_all_gather(combined_states, 0) - combined_states = combined_states[:num_tokens] - hidden_states, residual = combined_states.split( - [hidden_states.shape[-1], residual.shape[-1]], dim=-1 - ) - return hidden_states, residual - - class SharedHead(nn.Module): def __init__( self, @@ -139,13 +123,10 @@ def forward( hidden_states=hidden_states, residual=None, ) - hidden_states, residual = _restore_full_token_layout_if_needed( - hidden_states, - residual, - positions.shape[0], - is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, - ) hidden_states = residual + hidden_states # pre-final-norm (logits hidden) + if self.mtp_block.use_sequence_parallel_moe: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[0]] # Recycle the post-final-norm hidden into the next draft step. # compute_logits applies shared_head (== final norm) to the pre-norm # element, so logits and the recycle each get exactly one final-norm. @@ -516,7 +497,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled(): raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint. The checkpoint may have " diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index b811afafb0e8..58cedbb0f9bc 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -65,6 +65,7 @@ from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, EncoderCudaGraphConfig, + EncoderCudaGraphPathConfig, EncoderCudaGraphReplayBuffers, EncoderItemSpec, ) @@ -697,9 +698,15 @@ def get_encoder_cudagraph_config(self): modalities=["image"], buffer_keys=["pixel_values"], out_hidden_size=self.projector_config.n_embed, - enable_dual_path_graph=True, - global_token_per_image=self.global_image_output_token, - local_token_per_patch=self.single_patch_output_token, + paths={ + "global": EncoderCudaGraphPathConfig( + min_token_budget=self.global_image_output_token + ), + "local": EncoderCudaGraphPathConfig( + min_token_budget=self.single_patch_output_token, + allow_zero_tokens=True, + ), + }, ) def get_encoder_cudagraph_budget_range( @@ -730,8 +737,10 @@ def get_encoder_cudagraph_item_specs( EncoderItemSpec( input_size=num_input_tokens, output_tokens=num_output_tokens, - global_output_tokens=global_output_token, - local_output_tokens=local_output_token, + path_output_tokens={ + "global": global_output_token, + "local": local_output_token, + }, ) ) return item_specs @@ -917,32 +926,34 @@ def encoder_eager_forward( def postprocess_encoder_output( self, - output: torch.Tensor, + outputs: dict[str, torch.Tensor], indices: list[int], per_item_out_tokens: list[int], dest: dict[int, torch.Tensor] | list[torch.Tensor | None], clone: bool = False, batch_mm_kwargs: dict[str, Any] | None = None, - local_output: torch.Tensor | None = None, ) -> None: """ Assemble per-image embeddings from global and local encoder outputs. - ``output`` contains global-image features with newlines already + ``output['global']`` contains global-image features with newlines already inserted (from CUDA graph replay or eager fallback): ``[B * 272, n_embed]``. - ``local_output`` contains local-patch features without + ``output['local']`` contains local-patch features without newlines (from CUDA graph replay or eager fallback): ``[P * 100, n_embed]``. May be ``None`` if no patches in batch. This method: - 1. Splits ``output`` into per-image global portions. - 2. Splits ``local_output`` into per-image patch groups. + 1. Splits ``output['global']`` into per-image global portions. + 2. Splits ``output['local']`` into per-image patch groups. 3. For each image: assembles patch grid with newlines via ``_assemble_patch_grid``, then concatenates ``[local_tiled, global, view_seperator]``. """ + output = outputs["global"] + local_output = outputs.get("local") + assert batch_mm_kwargs is not None bsz = len(indices) n_embed = output.shape[-1] diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index daf4ea3de9f3..f076b4fb3fc1 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -33,6 +33,7 @@ from transformers import DeepseekV2Config, DeepseekV3Config import vllm._custom_ops as ops +import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config @@ -49,13 +50,14 @@ from vllm.model_executor.layers.attention import Attention, RSWAAttention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + DCPGroupColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, @@ -359,7 +361,7 @@ def __init__( prefix=f"{prefix}.shared_experts", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=config.n_routed_experts, @@ -849,11 +851,16 @@ def _try_load_fp8_indexer_wk( # We have both weight and scale: dequantize FP8 to BF16. weight_fp8, scale_inv = entry["weight"], entry["scale"] del buf[layer_prefix] - block_size = weight_fp8.shape[1] // scale_inv.shape[1] + if scale_inv.ndim == 1: + # Per-channel scale: one scale per row of [out, in] + group_shape = GroupShape(1, weight_fp8.shape[1]) + else: + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + group_shape = GroupShape(block_size, block_size) weight_bf16 = scaled_dequantize( weight_fp8, scale_inv, - group_shape=GroupShape(block_size, block_size), + group_shape=group_shape, out_dtype=torch.bfloat16, ) @@ -976,6 +983,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None = None, input_size: int | None = None, reduce_results: bool = True, + non_causal_multi_token_decode: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -1015,9 +1023,17 @@ def __init__( prefix=f"{prefix}.kv_a_proj_with_mqa", ) + qrep_enabled = ( + envs.VLLM_DCP_Q_REPLICATE + and vllm_config.parallel_config.decode_context_parallel_size > 1 + and vllm_config.parallel_config.prefill_context_parallel_size <= 1 + ) + q_proj_cls = ( + DCPGroupColumnParallelLinear if qrep_enabled else ColumnParallelLinear + ) if self.q_lora_rank is not None: self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) - self.q_b_proj = ColumnParallelLinear( + self.q_b_proj = q_proj_cls( self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, @@ -1025,7 +1041,7 @@ def __init__( prefix=f"{prefix}.q_b_proj", ) else: - self.q_proj = ColumnParallelLinear( + self.q_proj = q_proj_cls( proj_input_size, self.num_heads * self.qk_head_dim, bias=False, @@ -1163,6 +1179,10 @@ def __init__( # the V1 proposer. A frozen True would leave the draft reading a # never-written topk buffer. skip_topk=_skip_topk and not is_mtp_layer, + non_causal_multi_token_decode=non_causal_multi_token_decode, + # Do not skip scoring for MTP layers: their top-k buffer may be + # reused by later draft iterations through index sharing. + allow_short_prefill_indexer_scoring_skip=not is_mtp_layer, ) def forward( @@ -1463,8 +1483,6 @@ def forward( hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) - # fused_add_rms_norm requires a contiguous residual - residual = residual.contiguous() if idx in self.aux_hidden_state_layers: aux_hidden_state = hidden_states + residual if aux_hidden_state.shape[0] != positions.shape[0]: @@ -1489,8 +1507,6 @@ def forward( hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) - # fused_add_rms_norm requires a contiguous residual - residual = residual.contiguous() if self.end_layer in self.aux_hidden_state_layers: aux_hidden_states.append(hidden_states + residual) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 70566871e094..1d457dd39587 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -49,6 +49,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.v1.outputs import LogprobsTensors +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch @@ -60,7 +61,6 @@ from .interfaces import ( SupportsMultiModal, - SupportsPP, SupportsQuant, ) @@ -142,7 +142,6 @@ class DiffusionGemmaForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsQuant, - SupportsPP, ): """DiffusionGemma for vLLM. @@ -265,10 +264,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): eps=getattr(text_config, "rms_norm_eps", 1e-6), ) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - def compute_self_conditioning( self, inputs_embeds: torch.Tensor, @@ -1262,16 +1257,31 @@ def __call__( valid_canvas_len_np.astype(np.int64), device=device ) + # Per-request top_k/top_p, mirroring the AR sampler. Masked tokens + # become -inf and survive the temperature scaling in the compiled + # step, so Gumbel sampling, probs, and entropy all see the filtered + # distribution. The committed argmax (always the top-1 token) is + # unaffected; only the canvas exploration is constrained. Applied + # before canvas padding so phantom positions stay uniform. + if num_decode > 0: + top_k, top_p = self.sampling_states.get_top_k_top_p( + decode_slots.repeat_interleave(valid_canvas_len), decode_slots_np + ) + if top_k is not None or top_p is not None: + logits = apply_top_k_top_p(logits.float(), top_k, top_p) + # Pad any truncated canvas back to CL so the uniform-CL sampler math # holds. Phantom (padded) positions are zeroed → uniform logits → high # entropy (no premature convergence) and argmax 0 (stable); they are - # never committed (num_sampled == real length). + # never committed (num_sampled == real length). masked_fill (not + # multiply) so -inf entries from top_k/top_p filtering above don't + # turn phantom rows into NaN. if num_decode > 0 and valid_canvas_len_np.min() < CL: ar = torch.arange(CL, device=device) starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] 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) + logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0) # Clear once: the tiled loop below only scatters its own decode slots, # so it must not re-clear earlier tiles' writes. diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index fea390ca21cc..6cc6a056fda7 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -41,7 +41,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -183,7 +183,7 @@ def __init__( else: self.shared_experts = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.moe_num_experts, top_k=config.moe_k, diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py index 0bdb567c0aaf..5d36ca6d3756 100644 --- a/vllm/model_executor/models/ernie45_vl_moe.py +++ b/vllm/model_executor/models/ernie45_vl_moe.py @@ -37,7 +37,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -260,7 +260,7 @@ def __init__( prefix=f"{prefix}.text_experts_gate", ) - self.text_experts = FusedMoE( + self.text_experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.moe_num_experts[0], top_k=config.moe_k, @@ -297,7 +297,7 @@ def __init__( prefix=f"{prefix}.vision_experts_gate", ) - self.vision_experts = FusedMoE( + self.vision_experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.moe_num_experts[1], top_k=config.moe_k, diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index dc88c15fc019..8e8252f45922 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -31,7 +31,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -70,6 +70,7 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, bias: bool = False, + swiglu_limit: float | None = None, prefix: str = "", use_data_parallel: bool = False, ) -> None: @@ -95,7 +96,9 @@ def __init__( raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) - self.act_fn = SiluAndMul() + self.act_fn = ( + SiluAndMul() if swiglu_limit is None else SiluAndMulWithClamp(swiglu_limit) + ) def forward(self, x): gate_up, _ = self.gate_up_proj(x) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index 086040e2eaf7..0d7da9232d2e 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -29,21 +29,25 @@ get_pp_group, get_tensor_model_parallel_world_size, ) -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.config import set_default_rope_theta -from .exaone4 import Exaone4Attention as ExaoneMoeAttention from .exaone4 import Exaone4GatedMLP as ExaoneMoeGatedMLP from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -61,6 +65,7 @@ class ExaoneMoe(nn.Module): def __init__( self, config: PretrainedConfig, + swiglu_limit: float | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", enable_eplb: bool = False, @@ -121,12 +126,13 @@ def __init__( hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, + swiglu_limit=swiglu_limit, prefix=f"{prefix}.shared_experts", ) else: self.shared_experts = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=self.n_routed_experts, @@ -135,6 +141,8 @@ def __init__( intermediate_size=config.moe_intermediate_size, renormalize=config.norm_topk_prob, quant_config=quant_config, + activation="silu", + swiglu_limit=swiglu_limit, use_grouped_topk=True, num_expert_group=config.n_group, topk_group=config.topk_group, @@ -159,13 +167,138 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return final_hidden_states.view(orig_shape) +class ExaoneMoeAttention(nn.Module): + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + is_mtp: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr(config, "head_dim", None) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + is_neox_style = True + if quant_config is not None and quant_config.get_name() == "gguf": + is_neox_style = False + + layer_idx = extract_layer_index(prefix) + + if config.sliding_windows is not None: + self.sliding_window_size = config.sliding_windows[layer_idx] + + if config.layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + + if is_mtp: + self.sliding_window = ( + config.mtp_layer_types[layer_idx] == "sliding_attention" + ) + if config.mtp_sliding_windows is not None: + self.sliding_window_size = config.mtp_sliding_windows[layer_idx] + + if config.mtp_layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + + # apply rotary embeddings to every layer in full attention models + self.apply_rope_all_layers = "sliding_attention" not in config.layer_types + + set_default_rope_theta(config, default_theta=1000000) + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=is_neox_style, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=self.sliding_window_size, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + q = self.q_norm(q) + q = q.flatten(-2, -1) + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + k = self.k_norm(k) + k = k.flatten(-2, -1) + + if self.sliding_window_size or self.apply_rope_all_layers: + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + class ExaoneMoeDecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, - mtp_layer: bool = None, + is_mtp: bool = None, prefix: str = "", ) -> None: super().__init__() @@ -189,17 +322,29 @@ def __init__( bias=attention_bias, cache_config=cache_config, prefix=f"{prefix}.self_attn", + is_mtp=is_mtp, ) - if config.is_moe_layer[layer_idx] and not mtp_layer: + swiglu_limits = getattr(config, "swiglu_limits", None) + if swiglu_limits is not None: + swiglu_limit = swiglu_limits[layer_idx] + swiglu_limit = swiglu_limit if swiglu_limit > 0 else None + else: + swiglu_limit = None + + if config.mlp_layer_types[layer_idx] == "sparse" and not is_mtp: self.mlp = ExaoneMoe( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" + config=config, + swiglu_limit=swiglu_limit, + quant_config=quant_config, + prefix=f"{prefix}.mlp", ) else: self.mlp = ExaoneMoeGatedMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, + swiglu_limit=swiglu_limit, quant_config=quant_config, bias=getattr(config, "mlp_bias", False), prefix=f"{prefix}.mlp", diff --git a/vllm/model_executor/models/exaone_moe_mtp.py b/vllm/model_executor/models/exaone_moe_mtp.py index a37da487dba8..ddd570d604e7 100644 --- a/vllm/model_executor/models/exaone_moe_mtp.py +++ b/vllm/model_executor/models/exaone_moe_mtp.py @@ -81,7 +81,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config.model_config.hf_config, quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", - mtp_layer=True, + is_mtp=True, ) for idx in range(self.num_mtp_layers) ) diff --git a/vllm/model_executor/models/flex_olmo.py b/vllm/model_executor/models/flex_olmo.py index 2ff9d860567d..a0abf9e0fa77 100644 --- a/vllm/model_executor/models/flex_olmo.py +++ b/vllm/model_executor/models/flex_olmo.py @@ -21,7 +21,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear @@ -73,7 +73,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.gate", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=hf_config.num_experts, top_k=hf_config.num_experts_per_tok, hidden_size=hf_config.hidden_size, diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 72b12e26b5cf..2d3334f27513 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -257,25 +257,6 @@ def __init__(self, config: Any): def dtype(self) -> torch.dtype: return self.conv1.weight.dtype - def _prepare_attention_mask( - self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor - ) -> torch.Tensor | None: - if getattr(self.config, "_attn_implementation", "eager") == "flash_attention_2": - return None - - seq_length = inputs_tensor.shape[0] - attention_mask = torch.full( - (1, 1, seq_length, seq_length), - torch.finfo(inputs_tensor.dtype).min, - device=inputs_tensor.device, - dtype=inputs_tensor.dtype, - ) - for i in range(1, len(cu_seqlens)): - start = int(cu_seqlens[i - 1].item()) - end = int(cu_seqlens[i].item()) - attention_mask[..., start:end, start:end] = 0 - return attention_mask - def forward( self, input_features: torch.Tensor, diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 63d97fbd1b50..0fe539886b9f 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -38,7 +38,7 @@ from vllm.model_executor.layers.activation import get_act_and_mul_fn from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -84,6 +84,12 @@ logger = init_logger(__name__) +_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper( + orig_to_new_regex={ + re.compile(r"(? str: return re.sub(r"(? torch.Tensor: class Gemma4MoE(nn.Module): - """Mixture of Experts for Gemma4 using vLLM's FusedMoE. + """Mixture of Experts for Gemma4 using vLLM's MoERunner. - Wraps FusedMoE with custom routing. The router projection is + Wraps MoERunner with custom routing. The router projection is external (Gemma4Router) — this class only handles expert dispatch. Gemma4 routing: softmax over ALL experts → top-k → renormalize. per_expert_scale is folded into routing weights for mathematical - correctness with FusedMoE's fused kernel. + correctness with MoERunner's fused kernel. """ def __init__( @@ -320,11 +326,11 @@ def __init__( self.num_experts = config.num_experts # Per-expert output scale folded into routing weights so that - # FusedMoE's fused kernel computes: Σ_e (expert_e * w_e * scale_e) + # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e) self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) # Gemma4 routing: softmax over ALL experts → top-k → renormalize. - # FusedMoE's built-in fused_topk scopes softmax differently, so + # MoERunner's built-in fused_topk scopes softmax differently, so # a custom routing function is needed for numerical correctness. # NOTE: self.per_expert_scale is read at call time (not captured into # a local) so that torch.func.functional_call parameter substitution @@ -344,8 +350,8 @@ def routing_function( gating_output, topk, self.per_expert_scale ) - # FusedMoE experts with custom Gemma4 routing - self.experts = FusedMoE( + # MoERunner experts with custom Gemma4 routing + self.experts = FusedMoEFactory( num_experts=config.num_experts, top_k=config.top_k_experts, hidden_size=config.hidden_size, @@ -1376,10 +1382,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # MoE expert weight mapping: checkpoint can have either: # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D) # 2. Already per-expert 2D weights (if quantized) - # Map to FusedMoE parameters: - # moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13) - # moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13) - # moe.experts.{id}.down_proj → FusedMoE w2 + # Map to MoERunner parameters: + # moe.experts.{id}.gate_proj → MoERunner w1 (shard of w13) + # moe.experts.{id}.up_proj → MoERunner w3 (shard of w13) + # moe.experts.{id}.down_proj → MoERunner w2 num_experts = getattr(self.config, "num_experts", None) or 0 # Strategy A: dot-separated suffix # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight) @@ -1469,7 +1475,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue param = params_dict[moe_name] # Expert weights are already in the correct - # orientation for FusedMoE after _weight_iterator: + # orientation for MoERunner after _weight_iterator: # gate/up: [I, H] → w1/w3 expects [I, H] # down: [H, I] → w2 expects [H, I] # Scales and other quantization params may be 1D or scalar. @@ -1508,7 +1514,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Gemma4ForCausalLM( nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3 ): - hf_to_vllm_mapper = WeightsMapper( + hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper( orig_to_new_prefix={ # Gemma4ForConditionalGeneration already loads the text stack # from `model.language_model.*`. We reuse that same checkpoint @@ -1662,19 +1668,19 @@ def _weight_iterator(): # MoE expert weights: checkpoint stores as 3D packed # tensors. Explode into per-expert 2D weights for - # FusedMoE weight_loader. + # MoERunner weight_loader. # # Checkpoint format: # moe.gate_up_proj: [E, 2*I, H] (fused gate + up) # moe.down_proj: [E, H, I] # - # FusedMoE expects per-expert: + # MoERunner expects per-expert: # w1 (gate): [I, H] — first half of gate_up # w3 (up): [I, H] — second half of gate_up # w2 (down): [H, I] — as-is from checkpoint # # No transpose needed: checkpoint orientation already - # matches FusedMoE's expected layout. + # matches MoERunner's expected layout. if "moe.gate_up_proj" in name and weight.dim() == 3: num_experts = weight.size(0) intermediate_size = weight.size(1) // 2 diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py index 792f10884296..a77448da506a 100644 --- a/vllm/model_executor/models/gemma4_dspark.py +++ b/vllm/model_executor/models/gemma4_dspark.py @@ -22,7 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer -from .qwen3_dflash import DFlashQwen3Model +from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM from .utils import extract_layer_index, maybe_prefix @@ -37,7 +37,8 @@ def __init__( quant_config: QuantizationConfig | None, prefix: str, ) -> None: - is_full = config.layer_types[extract_layer_index(prefix)] == "full_attention" + layer_idx = extract_layer_index(prefix) + is_full = config.layer_types[layer_idx] == "full_attention" head_dim = ( getattr(config, "global_head_dim", config.head_dim) if is_full @@ -62,6 +63,7 @@ def __init__( prefix=prefix, ) self.is_kv_shared_layer = False + self.causal = _dflash_layer_causal(config, layer_idx) self.use_k_eq_v = use_k_eq_v self.kv_size = self.num_kv_heads * self.head_dim attn_bias = getattr(config, "attention_bias", False) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index d733a181a004..e158a7850164 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -16,7 +16,7 @@ import math from collections.abc import Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal import numpy as np import torch @@ -40,7 +40,10 @@ from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM +from vllm.model_executor.models.gemma4 import ( + _GEMMA4_EXPERT_PARENT_MAPPER, + Gemma4ForCausalLM, +) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.transformers.utils import recursive_replace_linear from vllm.multimodal import MULTIMODAL_REGISTRY @@ -69,6 +72,7 @@ from .interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -83,6 +87,12 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, + ) logger = init_logger(__name__) @@ -977,7 +987,9 @@ class Gemma4ForConditionalGeneration( SupportsPP, SupportsLoRA, SupportsEagle3, + SupportsEncoderCudaGraph, ): + supports_encoder_cudagraph: ClassVar[Literal[True]] = True # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model # runner reads this to keep image bidirectional ranges that exceed the @@ -998,7 +1010,7 @@ class Gemma4ForConditionalGeneration( } # Maps checkpoint prefixes to vLLM module paths. - hf_to_vllm_mapper = WeightsMapper( + hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper( orig_to_new_prefix={ # vision tower "model.vision_tower": "vision_tower", @@ -1010,7 +1022,7 @@ class Gemma4ForConditionalGeneration( "model.language_model.": "language_model.model.", "lm_head.": "language_model.lm_head.", "model": "language_model.model", - } + }, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -1022,6 +1034,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.quant_config = quant_config self.multimodal_config = multimodal_config self.model_dtype = vllm_config.model_config.dtype + self.vllm_config = vllm_config # Only quantize towers when the quant method supports their # dimensions. BNB/torchao handle arbitrary sizes; other methods @@ -1522,6 +1535,416 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: return multimodal_embeddings + # ------------------------------------------------------------------ # + # EncoderCudaGraph protocol methods + # ------------------------------------------------------------------ # + + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.zero_() + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.fill_(-1) + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + return EncoderCudaGraphConfig( + modalities=["image", "video"], + buffer_keys=[ + "pixel_values", + "pixel_position_ids", + "gather_indices", + ], + out_hidden_size=self.config.text_config.hidden_size, + max_frames_per_video=_VIDEO_MAX_FRAMES, + padding_logics={ + "pixel_values": pad_pixel_values, + "pixel_position_ids": pad_pixel_position_ids, + }, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = _SUPPORTED_SOFT_TOKENS[0] + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str: + if "pixel_values" in mm_kwargs: + return "image" + elif "pixel_values_videos" in mm_kwargs: + return "video" + raise ValueError("Unsupported modality in mm_kwargs") + + def get_max_frames_per_video(self) -> int: + return _VIDEO_MAX_FRAMES + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ) -> list["EncoderItemSpec"]: + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + return [ + EncoderItemSpec( + input_size=pv.shape[0], + output_tokens=pv.shape[0] // pool_ratio, + ) + for pv in pixel_values + ] + else: + return [ + EncoderItemSpec( + input_size=pixel_values.shape[1], + output_tokens=pixel_values.shape[1] // pool_ratio, + ) + for _ in range(pixel_values.shape[0]) + ] + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + np_patches = pixel_values_videos.shape[1] + return [ + EncoderItemSpec( + input_size=fc * np_patches, + output_tokens=fc * (np_patches // pool_ratio), + ) + for fc in fc_list + ] + raise ValueError(f"Unknown modality: {modality}") + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + if len(indices) == 0: + is_pv_list = isinstance(pixel_values, list) + is_pp_list = isinstance(pixel_position_ids, list) + return { + "pixel_values": ([] if is_pv_list else pixel_values[:0]), + "pixel_position_ids": ( + [] if is_pp_list else pixel_position_ids[:0] + ), + } + if isinstance(pixel_values, list): + return { + "pixel_values": [pixel_values[i] for i in indices], + "pixel_position_ids": [pixel_position_ids[i] for i in indices], + } + return { + "pixel_values": pixel_values[indices], + "pixel_position_ids": pixel_position_ids[indices], + } + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + + if len(indices) == 0: + is_fc_tensor = isinstance(video_frame_counts, torch.Tensor) + return { + "pixel_values_videos": pixel_values_videos[:0], + "pixel_position_ids_videos": pixel_position_ids_videos[:0], + "video_frame_counts": ( + video_frame_counts[:0] if is_fc_tensor else [] + ), + } + + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + cum_frames = [0] + for fc in fc_list: + cum_frames.append(cum_frames[-1] + fc) + + selected_pv = torch.cat( + [ + pixel_values_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_pp = torch.cat( + [ + pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_fc = ( + video_frame_counts[indices] + if isinstance(video_frame_counts, torch.Tensor) + else [video_frame_counts[i] for i in indices] + ) + return { + "pixel_values_videos": selected_pv, + "pixel_position_ids_videos": selected_pp, + "video_frame_counts": selected_fc, + } + raise ValueError(f"Unknown modality: {modality}") + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int = 256, + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphCaptureInputs": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + dtype = dtype or torch.float32 + + max_size = max(max_batch_size, max_frames_per_batch) + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + # Retrieve the model's actual configured maximum tokens: + configured_max_tokens = getattr( + self.config.vision_config, + "num_soft_tokens", + _SUPPORTED_SOFT_TOKENS[2], + ) + # Dynamically compute the slot capacity per item bounded by both the + # current graph budget and the user's maximum config: + per_item_output = min(token_budget, configured_max_tokens) + # Satisfy k^2 * per_item_output = per_item_patches + per_item_patches = per_item_output * pool_ratio + + patch_size = self.vision_tower.config.patch_size + num_channels = getattr(self.vision_tower.config, "num_channels", 3) + patch_pixels = (patch_size**2) * num_channels + + dummy_pixel_values = torch.zeros( + (max_size, per_item_patches, patch_pixels), + device=device, + dtype=dtype, + ) + dummy_pixel_position_ids = torch.full( + (max_size, per_item_patches, 2), + -1, + device=device, + dtype=torch.long, + ) + dummy_gather_indices = torch.zeros( + (token_budget,), + device=device, + dtype=torch.long, + ) + + return EncoderCudaGraphCaptureInputs( + values={ + "pixel_values": dummy_pixel_values, + "pixel_position_ids": dummy_pixel_position_ids, + "gather_indices": dummy_gather_indices, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphReplayBuffers": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + elif modality == "video": + pixel_values = mm_kwargs["pixel_values_videos"] + pixel_position_ids = mm_kwargs["pixel_position_ids_videos"] + else: + raise ValueError(f"Unsupported modality: {modality}") + + if isinstance(pixel_values, list): + max_patches = max(pv.shape[0] for pv in pixel_values) + batch_size = len(pixel_values) + pv_tensor = torch.zeros( + (batch_size, max_patches, pixel_values[0].shape[1]), + dtype=pixel_values[0].dtype, + device=pixel_values[0].device, + ) + pp_tensor = torch.full( + (batch_size, max_patches, 2), + -1, + dtype=pixel_position_ids[0].dtype, + device=pixel_position_ids[0].device, + ) + for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)): + pv_tensor[i, : pv.shape[0]].copy_(pv) + pp_tensor[i, : pp.shape[0]].copy_(pp) + pixel_values = pv_tensor + pixel_position_ids = pp_tensor + + item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs) + per_item_out_tokens = [spec.output_tokens for spec in item_specs] + total_tokens = sum(per_item_out_tokens) + + device = pixel_values.device + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + # ONLY allocate an array of exact size `total_tokens`. + # DO NOT pad it. The upstream Graph Manager handles the padding securely. + gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device) + + if modality == "image": + dst_offset = 0 + for i, n_tok in enumerate(per_item_out_tokens): + safe_n_tok = min(n_tok, per_item_output) + src_start = i * per_item_output + src_end = src_start + safe_n_tok + gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_n_tok + elif modality == "video": + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + np_patches = pixel_values.shape[1] + frame_output_tokens = np_patches // pool_ratio + safe_frame_output_tokens = min(frame_output_tokens, per_item_output) + + dst_offset = 0 + frame_idx = 0 + for fc in fc_list: + for f in range(fc): + src_start = (frame_idx + f) * per_item_output + src_end = src_start + safe_frame_output_tokens + gather_indices[ + dst_offset : dst_offset + safe_frame_output_tokens + ] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_frame_output_tokens + frame_idx += fc + + return EncoderCudaGraphReplayBuffers( + values={ + "pixel_values": pixel_values, + "pixel_position_ids": pixel_position_ids, + "gather_indices": gather_indices, + } + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + pixel_values = inputs["pixel_values"] + pixel_position_ids = inputs["pixel_position_ids"] + gather_indices = inputs["gather_indices"] + + pad_tensor = (pixel_position_ids == -1).all(dim=-1) + + vt = self.vision_tower + inputs_embeds = vt.patch_embedder( + pixel_values, + pixel_position_ids, + pad_tensor, + ).to(self.model_dtype) + + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_tensor, + pixel_position_ids=pixel_position_ids, + ) + hidden_states = encoder_outputs.last_hidden_state + + pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + pooled_states, _ = vt.pooler( + hidden_states=hidden_states, + pixel_position_ids=pixel_position_ids, + padding_positions=pad_tensor, + output_length=per_item_output, + ) + + if getattr(vt.config, "standardize", False): + pooled_states = (pooled_states - vt.std_bias) * vt.std_scale + + flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1]) + gathered_states = flat_pooled[gather_indices] + + # Cast to the projection layer's dtype to resolve mixed-precision crash + target_dtype = self.embed_vision.embedding_projection.weight.dtype + gathered_states = gathered_states.to(target_dtype) + + flat_proj_embs = self.embed_vision( + inputs_embeds=gathered_states.unsqueeze(0) + ).squeeze(0) + + return flat_proj_embs + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + image_input = self._parse_and_validate_image_input(**mm_kwargs) + assert image_input is not None + embeddings = self._process_image_input(image_input) + elif modality == "video": + video_input = self._parse_and_validate_video_input(**mm_kwargs) + assert video_input is not None + embeddings = self._process_video_input(video_input) + else: + raise ValueError(f"Unsupported modality: {modality}") + + return torch.cat(embeddings, dim=0) + def embed_input_ids( self, input_ids: torch.Tensor, diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 810d9de87b47..19389d844631 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1082,8 +1082,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size @@ -1403,6 +1403,11 @@ def _get_video_second_idx_glmga( timestamps_list = full_second_idxs[::2] return list(timestamps_list) + def _get_video_frame_embed_token_id(self, hf_processor: object) -> int: + if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA: + return hf_processor.image_token_id + return hf_processor.video_token_id + def _construct_video_placeholder( self, video_array: np.ndarray, @@ -1440,13 +1445,7 @@ def _construct_video_placeholder( num_tokens_per_frame = int(H * W) // merge_length placeholder = [] placeholder.append(bov_token_id) - # Glm46VProcessor uses image_token_id for video frame embeddings; - # Glm4vProcessor uses video_token_id. - frame_embed_token_id = ( - hf_processor.video_token_id - if isinstance(hf_processor, Glm4vProcessor) or not TRANSFORMERS_WITH_GA - else hf_processor.image_token_id - ) + frame_embed_token_id = self._get_video_frame_embed_token_id(hf_processor) for frame_idx in frames_idx_token: placeholder.append(boi_token_id) placeholder.extend([frame_embed_token_id] * num_tokens_per_frame) @@ -1604,11 +1603,6 @@ def _call_hf_processor( processor = self.info.get_hf_processor(**mm_kwargs) - # Glm46VProcessor and GLMGA handle image/video placeholders together - # via the direct path. Only Glm4vProcessor (GLM-4.1V) needs the - # split-video path because it uses image_token_id as the video - # placeholder. The direct path requires transformers >= 5.5.0 - # (Glm46VProcessor / GlmgaVideoProcessor support). use_direct_path = ( not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA ) @@ -1630,6 +1624,8 @@ def _call_hf_processor( ): video_grid_thw_lst = [] pixel_values_videos_lst = [] + frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor) + swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id for item in mm_data.pop("videos", []): video_array, metadata = item @@ -1650,9 +1646,10 @@ def _call_hf_processor( tok_kwargs=tok_kwargs, ) input_ids = video_outputs.pop("input_ids") - input_ids[input_ids == processor.image_token_id] = ( - processor.video_token_id - ) + if swap_video_frame_tokens: + input_ids[input_ids == processor.image_token_id] = ( + processor.video_token_id + ) video_placeholder = processor.tokenizer.batch_decode(input_ids)[0] prompt = prompt.replace( "<|begin_of_video|><|video|><|end_of_video|>", @@ -1668,6 +1665,7 @@ def _call_hf_processor( ) else: video_outputs = dict() + swap_video_frame_tokens = False processed_outputs = super()._call_hf_processor( prompt=prompt, @@ -1675,6 +1673,10 @@ def _call_hf_processor( mm_kwargs=mm_kwargs, tok_kwargs=tok_kwargs, ) + if swap_video_frame_tokens: + input_ids = processed_outputs["input_ids"] + input_ids[input_ids == processor.video_token_id] = processor.image_token_id + processed_outputs["input_ids"] = input_ids combined_outputs = dict( processed_outputs, **video_outputs, @@ -1701,7 +1703,7 @@ def _get_prompt_updates( merge_length = image_processor.merge_size**2 - def get_image_replacement_glm4v(item_idx: int): + def get_image_replacement(item_idx: int): out_item = out_mm_kwargs["image"][item_idx] grid_thw = out_item["image_grid_thw"].data assert isinstance(grid_thw, torch.Tensor) @@ -1709,21 +1711,7 @@ def get_image_replacement_glm4v(item_idx: int): num_tokens = int(grid_thw.prod()) // merge_length return [hf_processor.image_token_id] * num_tokens - def get_video_replacement_glm4v(item_idx: int): - out_item = out_mm_kwargs["video"][item_idx] - grid_thw = out_item["video_grid_thw"].data - assert isinstance(grid_thw, torch.Tensor) - - video, metadata = mm_items["video"][item_idx] - placeholder = self.info._construct_video_placeholder( - video, metadata, grid_thw - ) - return PromptUpdateDetails.select_token_id( - placeholder, - embed_token_id=hf_processor.video_token_id, - ) - - def get_video_replacement_glm46v(item_idx: int): + def get_video_replacement(item_idx: int): out_item = out_mm_kwargs["video"][item_idx] grid_thw = out_item["video_grid_thw"].data assert isinstance(grid_thw, torch.Tensor) @@ -1734,25 +1722,19 @@ def get_video_replacement_glm46v(item_idx: int): ) return PromptUpdateDetails.select_token_id( placeholder, - embed_token_id=hf_processor.image_token_id, + embed_token_id=self.info._get_video_frame_embed_token_id(hf_processor), ) - is_glm46v = not isinstance(hf_processor, Glm4vProcessor) - return [ PromptReplacement( modality="image", target=hf_processor.image_token, - replacement=get_image_replacement_glm4v, + replacement=get_image_replacement, ), PromptReplacement( modality="video", target="<|begin_of_video|><|video|><|end_of_video|>", - replacement=( - get_video_replacement_glm46v - if is_glm46v and TRANSFORMERS_WITH_GA - else get_video_replacement_glm4v - ), + replacement=get_video_replacement, ), ] diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index abb9970c403e..615df7cef386 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -42,7 +42,9 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -164,7 +166,7 @@ def __init__( ) # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py - # pattern (see Glm4MoE / FusedMoE wiring there). + # pattern (see Glm4MoE / MoERunner wiring there). self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() self.is_fusion_moe_shared_experts_enabled = ( rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() @@ -183,7 +185,7 @@ def __init__( prefix=f"{prefix}.shared_experts", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.n_routed_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 5fc7e62640f5..27982eb1f3b1 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -21,7 +21,7 @@ ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig @@ -211,7 +211,7 @@ def __init__( return_bias=False, ) assert config.intermediate_size % self.world_size == 0 - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -1010,7 +1010,7 @@ def _load_weights_other( tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size) # Use centralized weight remapping for MoE expert parameters. - # The FusedMoE refactor moved expert params under + # The MoERunner refactor moved expert params under # `mlp.experts.routed_experts.*`; this remaps checkpoint names so # MoE weight/bias keys resolve against params_dict. for name, weight in remap_moe_expert_weights(weights, params_dict): diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index 5909604bd543..328ea9ce6326 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -40,7 +40,7 @@ ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -101,7 +101,7 @@ def __init__( prefix=f"{prefix}.gate", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index 4f70a9662896..af26b39ee05c 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -43,7 +43,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -238,10 +238,10 @@ def forward( ori_k = k if self.use_qk_norm: q = self.query_layernorm( - q.view(-1, self.num_heads, self.head_dim).contiguous() + q.view(-1, self.num_heads, self.head_dim), ) k = self.key_layernorm( - k.view(-1, self.num_kv_heads, self.head_dim).contiguous() + k.view(-1, self.num_kv_heads, self.head_dim), ) attn_output = self.attn(q, k, v) @@ -346,10 +346,10 @@ def forward( q, _ = self.rotary_emb(positions, q, k_tmp) if self.use_qk_norm: q = self.query_layernorm( - q.view(-1, self.num_heads, self.head_dim).contiguous() + q.view(-1, self.num_heads, self.head_dim), ) k = self.key_layernorm( - k.view(-1, self.num_kv_heads, self.head_dim).contiguous() + k.view(-1, self.num_kv_heads, self.head_dim), ) attn_output = self.attn(q, k, v) @@ -441,7 +441,7 @@ def __init__( else: self.shared_mlp = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_mlp, num_experts=self.n_routed_experts, top_k=top_k, diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index b980e8ae46a7..fc05b17d3baa 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/vllm/model_executor/models/hunyuan_vision.py @@ -698,9 +698,12 @@ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) hf_processor = self.info.get_hf_processor(typ=HunYuanVLProcessor) - image_token: str = hf_processor.image_token + image_placeholder = ( + f"{hf_processor.image_start_token}{hf_processor.image_token}" + f"{hf_processor.image_end_token}" + ) - return image_token * num_images + return image_placeholder * num_images def get_dummy_mm_data( self, @@ -779,7 +782,11 @@ def get_replacement_hunyuan_vl(item_idx: int, modality: str): return [ PromptReplacement( modality=modality, - target=[token_ids[modality]], + target=[ + token_ids[f"{modality}_start"], + token_ids[modality], + token_ids[f"{modality}_end"], + ], replacement=partial(get_replacement_hunyuan_vl, modality=modality), ) for modality in ("image",) diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 68c6f2382798..f75fbb0041d5 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -44,7 +44,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -184,7 +184,7 @@ def __init__( scoring_func = "sigmoid" e_score_correction_bias = self.expert_bias - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=self.n_routed_experts, top_k=top_k, hidden_size=config.hidden_size, @@ -271,7 +271,7 @@ def __init__( self.total_num_heads, self.total_num_kv_heads, quant_config=quant_config, - bias=None, + bias=False, prefix=f"{prefix}.qkv_proj", ) self.o_proj = RowParallelLinear( diff --git a/vllm/model_executor/models/hyperclovax_vision.py b/vllm/model_executor/models/hyperclovax_vision.py index 53923d88438a..593be5d88f6b 100644 --- a/vllm/model_executor/models/hyperclovax_vision.py +++ b/vllm/model_executor/models/hyperclovax_vision.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # copied from : https://github.com/huggingface/transformers import ast -from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence from functools import partial from itertools import accumulate @@ -891,37 +890,6 @@ def forward_videos( for i in range(len(feats_per_video)) ) - def _prepare_multimodal_kwargs(self, **kwargs: object): - output = defaultdict(list) - for k, v in kwargs.items(): - if len(v) < 1 or len(v[0]) < 1: - continue # if empty batch of empty sample - - new_k, is_video = k, False - if not k.endswith("_images") and not k.endswith("_videos"): - pass - else: - new_k, is_video = k.split("_")[:-1], k.split("_")[-1] - new_k = "_".join(new_k) - is_video = is_video == "videos" - - for _sample_idx, _v in enumerate(v): # batch -> sample - if new_k not in ["pixel_values"]: - if len(output[new_k]) < _sample_idx + 1: - output[new_k].append(list()) - _v = _v.detach().cpu().numpy().tolist() - output[new_k][_sample_idx] += _v - elif isinstance(_v, torch.Tensor): - if len(output[new_k]) < _sample_idx + 1: - output[new_k].append(list()) - output["is_videos"].append(list()) - _v = list(torch.unbind(_v, dim=0)) - output[new_k][_sample_idx] += _v - output["is_videos"][_sample_idx] += [ - is_video, - ] * len(_v) - return dict(output) - def compute_logits( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index ad94719241c5..4b3cbcabd2ff 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -138,30 +138,6 @@ def _resize_output_size( return height, width - def _get_resize_output_image_size( - self, - *, - image_width: int, - image_height: int, - resolution_max_side: int, - ) -> tuple[int, int]: - hf_processor = self.get_hf_processor() - image_processor: Idefics3ImageProcessor = hf_processor.image_processor - max_image_size = image_processor.size["longest_edge"] - if resolution_max_side > max_image_size: - raise ValueError( - "`resolution_max_side` cannot be larger than `max_image_size`" - ) - - height, width = image_height, image_width - - # Find the output size, when rescaling the longest edge to max_len and - # preserving the aspect ratio - height, width = self._resize_output_size( - height=height, width=width, max_len=resolution_max_side - ) - return height, width - def _get_image_feature_grid_size( self, *, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 32534f8fbb79..efdab8c4d105 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -11,6 +11,7 @@ Sequence, ) from contextlib import ExitStack, contextmanager, nullcontext +from dataclasses import dataclass from typing import ( TYPE_CHECKING, Any, @@ -41,6 +42,7 @@ SpeechToTextParams, VllmConfig, ) + from vllm.config.multimodal import VideoPruningMethod from vllm.inputs import PromptType, TokensPrompt from vllm.lora.model_manager import LoRAModelManager from vllm.model_executor.layers.fused_moe import MoERunner @@ -70,6 +72,16 @@ """ +@dataclass(frozen=True, slots=True) +class DiarizedTranscriptionSegment: + """A timestamped, speaker-attributed segment produced by an ASR model.""" + + start: float + end: float + speaker: str + text: str + + class StreamingTranscriptionPostProcessor: """Stateful streaming post-processor for transcription deltas.""" @@ -424,6 +436,13 @@ class SupportsMultiModalPruning(Protocol): supports_multimodal_pruning: ClassVar[Literal[True]] = True + supported_video_pruning_methods: ClassVar[tuple["VideoPruningMethod", ...]] = ( + "evs", + ) + """Video pruning methods (as reported by + `MultiModalConfig.get_video_pruning_spec`) implemented by this model. + Models supporting methods beyond EVS should override this.""" + def recompute_mrope_positions( self, input_ids: list[int] | torch.Tensor, @@ -984,6 +1003,31 @@ def supports_mamba_prefix_caching( return getattr(model, "supports_mamba_prefix_caching", False) +@runtime_checkable +class SupportsReplaySSM(Protocol): + """The interface for models whose Mamba2 layers support ReplaySSM cached + standard decode. + + This is currently experimental. + """ + + supports_replayssm: ClassVar[Literal[True]] = True + + +@overload +def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ... + + +@overload +def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ... + + +def supports_replayssm( + model: type[object] | object, +) -> TypeIs[type[SupportsReplaySSM]] | TypeIs[SupportsReplaySSM]: + return getattr(model, "supports_replayssm", False) + + @runtime_checkable class SupportsCrossEncoding(Protocol): """The interface required for all models that support cross encoding.""" @@ -1100,6 +1144,9 @@ class SupportsTranscription(Protocol): Enables the segment timestamp option for supported models by setting this to `True`. """ + supports_diarized_transcription: ClassVar[bool] = False + """Enables the ``diarized_json`` response format for the model.""" + supports_explicit_language_detection: ClassVar[bool] = False """ Transcription models that require an explicit language detection step @@ -1205,6 +1252,15 @@ def post_process_output(cls, text: str) -> str: """ return text + @classmethod + def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]: + """Parse the model-specific diarized transcript format. + + Only models that set ``supports_diarized_transcription`` must override + this method. + """ + raise NotImplementedError + @classmethod def get_streaming_post_processor_cls( cls, @@ -1351,7 +1407,7 @@ def _maybe_add_hidden_state( aux_hidden_states: list[torch.Tensor], layer_idx: int, hidden_states: torch.Tensor, - residual: torch.Tensor, + residual: torch.Tensor | None, ) -> list[torch.Tensor]: if layer_idx in self.aux_hidden_state_layers: value = hidden_states + residual if residual is not None else hidden_states @@ -1640,13 +1696,12 @@ def select_encoder_cudagraph_items( def postprocess_encoder_output( self, - output: torch.Tensor, + outputs: dict[str, torch.Tensor], indices: list[int], per_item_out_tokens: list[int], dest: dict[int, torch.Tensor] | list[torch.Tensor | None], clone: bool = False, batch_mm_kwargs: dict[str, Any] | None = None, - local_output: torch.Tensor | None = None, ) -> None: """ Post-process encoder output, directly call scatter_output_slices by default. @@ -1658,7 +1713,9 @@ def postprocess_encoder_output( """ from vllm.model_executor.models.utils import scatter_output_slices - scatter_output_slices(output, indices, per_item_out_tokens, dest, clone) + scatter_output_slices( + outputs["default"], indices, per_item_out_tokens, dest, clone + ) def prepare_encoder_cudagraph_capture_inputs( self, diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index c04b47294541..935abb04ed95 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -42,7 +42,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -173,7 +173,7 @@ def __init__( # For custom routing function self.n_groups = getattr(config, "router_n_groups", -1) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -570,10 +570,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index 33a8c6364176..baee41ae1bd3 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -14,7 +14,9 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import get_pp_group from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, @@ -81,7 +83,7 @@ def __init__( prefix=f"{prefix}.router", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( self.num_total_experts, self.top_k, self.hidden_size, diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 06f5ce282c6c..63d690bb3d99 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -24,6 +24,7 @@ ) from .interfaces import SupportsLateInteraction from .interfaces_base import VllmModelForPooling +from .llama import LlamaForCausalLM from .qwen3 import Qwen3ForCausalLM, Qwen3Model from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -185,11 +186,79 @@ def _build_lora_pairs(adapter_weights: dict) -> dict: return dict(lora_pairs) -class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling): - """Jina Embeddings V5 with task-specific LoRA adapters merged at load time. +def _setup_jina_v5_task_and_pooler(model: nn.Module, vllm_config: VllmConfig) -> None: + """Shared init for jina-embeddings-v5 wrappers: select task + build pooler.""" + model._model_name = vllm_config.model_config.model + model._revision = vllm_config.model_config.revision - Extends Qwen3ForCausalLM (the underlying architecture) and declares itself - as a pooling model so that as_embedding_model() does not wrap it. + model._task = getattr( + vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK + ) + if model._task not in _SUPPORTED_TASKS: + logger.warning( + "Unknown jina_task=%r. Falling back to %r.", + model._task, + _DEFAULT_TASK, + ) + model._task = _DEFAULT_TASK + + pooler_config = vllm_config.model_config.pooler_config + assert pooler_config is not None + model.pooler = DispatchPooler.for_embedding(pooler_config) + + +def _load_jina_v5_weights( + model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]] +) -> set[str]: + """Shared loader: merge the selected task LoRA adapter into the base weights.""" + lora_pairs: dict = {} + scaling = 1.0 + + result = _load_adapter(model._model_name, model._task, model._revision) + if result is None: + logger.warning( + "No adapter found for task %r in %r. Loading raw base weights.", + model._task, + model._model_name, + ) + else: + adapter_config, adapter_weights = result + scaling = adapter_config["lora_alpha"] / adapter_config["r"] + lora_pairs = _build_lora_pairs(adapter_weights) + logger.info( + "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)", + len(adapter_weights), + model._task, + scaling, + len(lora_pairs), + ) + + def _merge_weights( + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + clean_name = name + if clean_name.startswith("model."): + clean_name = clean_name[len("model.") :] + + if clean_name in lora_pairs: + pair = lora_pairs[clean_name] + if "A" in pair and "B" in pair: + lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype) + lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype) + tensor = tensor + (lora_B @ lora_A) * scaling + yield name, tensor + + loader = AutoWeightsLoader(model, ignore_unexpected_prefixes=["lm_head."]) + weights = _merge_weights(weights) + return loader.load_weights(weights, mapper=model.hf_to_vllm_mapper) + + +class JinaEmbeddingsV5DecoderModel(Qwen3ForCausalLM, VllmModelForPooling): + """jina-embeddings-v5 with a Qwen3 decoder backbone (e.g. -small). + + Task-specific LoRA adapters are merged into the base weights at load time. + Declares itself a pooling model so that as_embedding_model() does not wrap it. """ is_pooling_model = True @@ -199,64 +268,46 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) + _setup_jina_v5_task_and_pooler(self, vllm_config) - self._model_name = vllm_config.model_config.model - self._revision = vllm_config.model_config.revision + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_jina_v5_weights(self, weights) - self._task = getattr( - vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK - ) - if self._task not in _SUPPORTED_TASKS: - logger.warning( - "Unknown jina_task=%r. Falling back to %r.", - self._task, - _DEFAULT_TASK, - ) - self._task = _DEFAULT_TASK - - pooler_config = vllm_config.model_config.pooler_config - assert pooler_config is not None - self.pooler = DispatchPooler.for_embedding(pooler_config) + +class JinaEmbeddingsV5EncoderModel(LlamaForCausalLM, VllmModelForPooling): + """jina-embeddings-v5 with a bidirectional EuroBERT (Llama) encoder backbone. + + Used by encoder checkpoints such as jina-embeddings-v5-text-nano + (``is_decoder=False``). EuroBERT is architecturally a bidirectional Llama, so + the LlamaModel backbone switches to EncoderOnlyAttention when the config + carries ``is_causal=False`` (set by ``JinaEmbeddingsV5ModelConfig``). + """ + + is_pooling_model = True + hf_to_vllm_mapper = LlamaForCausalLM.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"": "model."} + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + _setup_jina_v5_task_and_pooler(self, vllm_config) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - lora_pairs: dict = {} - scaling = 1.0 - - result = _load_adapter(self._model_name, self._task, self._revision) - if result is None: - logger.warning( - "No adapter found for task %r in %r. Loading raw base weights.", - self._task, - self._model_name, - ) - else: - adapter_config, adapter_weights = result - scaling = adapter_config["lora_alpha"] / adapter_config["r"] - lora_pairs = _build_lora_pairs(adapter_weights) - logger.info( - "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)", - len(adapter_weights), - self._task, - scaling, - len(lora_pairs), - ) - - def _merge_weights( - weights: Iterable[tuple[str, torch.Tensor]], - ) -> Iterable[tuple[str, torch.Tensor]]: - for name, tensor in weights: - clean_name = name - if clean_name.startswith("model."): - clean_name = clean_name[len("model.") :] - - if clean_name in lora_pairs: - pair = lora_pairs[clean_name] - if "A" in pair and "B" in pair: - lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype) - lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype) - tensor = tensor + (lora_B @ lora_A) * scaling - yield name, tensor - - loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["lm_head."]) - weights = _merge_weights(weights) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return _load_jina_v5_weights(self, weights) + + +class JinaEmbeddingsV5Model(JinaEmbeddingsV5DecoderModel): + """Dispatcher for the jina-embeddings-v5 family. + + The family ships two backbones under one ``architectures`` entry: Qwen3 + decoders (-small) and bidirectional EuroBERT encoders (-nano), told apart by + ``is_decoder``. Inherits the decoder implementation so registry introspection + still sees a valid pooling model, and ``__new__`` swaps in the encoder + variant for encoder checkpoints. + """ + + def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""): + is_decoder = getattr(vllm_config.model_config.hf_config, "is_decoder", True) + if not is_decoder: + return JinaEmbeddingsV5EncoderModel(vllm_config=vllm_config, prefix=prefix) + return super().__new__(cls) diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index 125d7e71c7b5..b1a5f78b1b33 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -409,8 +409,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index dd1fb892ad19..c3d69836a790 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -983,7 +983,7 @@ def _get_vision_info( else: preprocessed_size = ImageSize(width=image_width, height=image_height) - padded_num_frames = num_frames + num_frames % temporal_patch_size + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index bcb8dd32de31..dca62d596b42 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -34,6 +34,7 @@ is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model, ) +from vllm.platforms import current_platform from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig logger = init_logger(__name__) @@ -60,7 +61,7 @@ def wrapper(org, interpolation_mode, shape): @get_rope_shape_decorate -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, disable=current_platform.simple_compile_backend == "tpu") def get_rope_shape(org, interpolation_mode, shape): return ( F.interpolate( @@ -154,9 +155,7 @@ def __init__( def reset_parameters(self): nn.init.normal_(self.weight) - def forward( - self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] - ) -> torch.Tensor: + def get_pos_embeds(self, grid_thws: torch.Tensor | list[list[int]]) -> torch.Tensor: pos_embs = [] grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() for t, h, w in grid_thw_list: @@ -179,8 +178,12 @@ def forward( pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) - out = x + torch.cat(pos_embs) - return out + return torch.cat(pos_embs) + + def forward( + self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] + ) -> torch.Tensor: + return x + self.get_pos_embeds(grid_thws) class MoonVision3dPatchEmbed(nn.Module): @@ -195,6 +198,8 @@ def __init__( pos_emb_width: int = 14, pos_emb_time: int = 4, pos_emb_type: str = "divided_fixed", + patch_embed_proj_bias: bool = True, + pos_emb_interpolation_mode: str = "bicubic", ): super().__init__() assert isinstance(patch_size, int | Sequence), ( @@ -208,7 +213,11 @@ def __init__( self.patch_size = patch_size self.proj = nn.Conv2d( - in_dim, out_dim, kernel_size=patch_size, stride=patch_size + in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias, ) if pos_emb_type == "divided_fixed": @@ -217,17 +226,37 @@ def __init__( width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode, ) else: raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") def forward( - self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] + self, + x: torch.Tensor, + grid_thws: torch.Tensor | list[list[int]] | None, + *, + pos_embeds: torch.Tensor | None = None, ) -> torch.Tensor: - x = self.proj(x).view(x.size(0), -1) - # apply positional embedding - x = self.pos_emb(x, grid_thws) - return x + x = self._proj(x).view(x.size(0), -1) + if pos_embeds is not None: + return x + pos_embeds + assert grid_thws is not None + return self.pos_emb(x, grid_thws) + + def _proj(self, x: torch.Tensor) -> torch.Tensor: + # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton. + if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16): + from aiter.ops.triton.conv.conv2d import conv2d + + return conv2d( + x, + self.proj.weight, + self.proj.bias, + stride=self.patch_size, + layout="nchw", + ) + return self.proj(x) class Rope2DPosEmbRepeated(nn.Module): @@ -342,6 +371,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +def _make_vision_norm(norm_type: str, hidden_dim: int) -> nn.Module: + if norm_type == "layernorm": + return nn.LayerNorm(hidden_dim) + if norm_type == "rmsnorm": + return nn.RMSNorm(hidden_dim) + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + class MoonViTEncoderLayer(nn.Module): """Single encoder layer for MoonViT with TP/DP support.""" @@ -355,23 +392,33 @@ def __init__( *, activation=F.gelu, attn_bias: bool = False, + qkv_hidden_size: int | None = None, + norm_type: str = "layernorm", + mlp_type: str = "mlp2", + linear_bias: bool = True, ): super().__init__() self.use_data_parallel = is_vit_use_data_parallel(num_heads) self.num_heads = num_heads self.hidden_dim = hidden_dim - self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads + self.qkv_hidden_size = ( + hidden_dim if qkv_hidden_size is None else qkv_hidden_size + ) + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads self.tp_size = ( 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() ) self.num_attention_heads_per_partition = divide(num_heads, self.tp_size) - self.norm0 = nn.LayerNorm(hidden_dim) - self.norm1 = nn.LayerNorm(hidden_dim) + self.norm0 = _make_vision_norm(norm_type, hidden_dim) + self.norm1 = _make_vision_norm(norm_type, hidden_dim) + if mlp_type != "mlp2": + raise NotImplementedError(f"Not support mlp_type: {mlp_type}") self.mlp = MLP2( [hidden_dim, mlp_dim, hidden_dim], activation, + bias=linear_bias, quant_config=quant_config, prefix=f"{prefix}.mlp", use_data_parallel=self.use_data_parallel, @@ -387,7 +434,7 @@ def __init__( disable_tp=self.use_data_parallel, ) self.wo = RowParallelLinear( - hidden_dim, + self.qkv_hidden_size, hidden_dim, bias=attn_bias, quant_config=quant_config, @@ -493,8 +540,9 @@ def __init__( f'video_attn_type must be "spatial_temporal", got {video_attn_type}' ) self.video_attn_type = video_attn_type + qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"] self.rope_2d = Rope2DPosEmbRepeated( - block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 + qkv_hidden_size // block_cfg["num_heads"], 512, 512 ) self.blocks = nn.ModuleList( [ @@ -506,13 +554,17 @@ def __init__( for layer_idx in range(num_layers) ] ) - self.final_layernorm = nn.LayerNorm(hidden_dim) + self.final_layernorm = _make_vision_norm( + block_cfg.get("norm_type", "layernorm"), hidden_dim + ) def prepare_encoder_metadata( self, grid_thw_list: list[list[int]], *, device: torch.device, + max_batch_size: int | None = None, + max_seqlen_override: int | None = None, ) -> dict[str, torch.Tensor | None]: metadata: dict[str, torch.Tensor | None] = {} metadata["rope_freqs_cis"] = self.rope_2d.get_freqs_cis( @@ -524,15 +576,30 @@ def prepare_encoder_metadata( cu_seqlens = np.concatenate( [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)] ) + if max_batch_size is not None: + num_seqs = len(cu_seqlens) - 1 + if num_seqs < max_batch_size: + cu_seqlens = np.concatenate( + [ + cu_seqlens, + np.full( + max_batch_size - num_seqs, + cu_seqlens[-1], + dtype=np.int32, + ), + ] + ) attn_backend = self.blocks[0].attn.attn_backend metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens( attn_backend, cu_seqlens, device ) - metadata["max_seqlen"] = torch.tensor( - MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens), - dtype=torch.int32, + max_seqlen = ( + max_seqlen_override + if max_seqlen_override is not None + else MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens) ) + metadata["max_seqlen"] = torch.tensor(max_seqlen, dtype=torch.int32) metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens( attn_backend, cu_seqlens, @@ -545,11 +612,12 @@ def prepare_encoder_metadata( def forward( self, hidden_states: torch.Tensor, - grid_thws: torch.Tensor | list[list[int]], + grid_thws: torch.Tensor | list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: if encoder_metadata is None: + assert grid_thws is not None grid_thw_list = ( grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() ) @@ -604,6 +672,34 @@ def tpool_patch_merger( return outputs +def build_image_merge_gather_idx( + grid_thws: list[list[int]] | list[tuple[int, int, int]], + merge_kernel_size: tuple[int, int], +) -> np.ndarray: + """Build packed spatial-merge indices for image-only CUDA graphs.""" + kh, kw = merge_kernel_size + parts: list[np.ndarray] = [] + offset = 0 + for t, h, w in grid_thws: + if t != 1: + raise ValueError("Image encoder CUDA graphs require grid T == 1") + idx = np.arange(h * w, dtype=np.int64).reshape(h, w) + idx = idx.reshape(h // kh, kh, w // kw, kw) + parts.append(idx.transpose(0, 2, 1, 3).reshape(-1, kh * kw) + offset) + offset += h * w + if not parts: + return np.empty((0, kh * kw), dtype=np.int64) + return np.concatenate(parts) + + +def tpool_patch_merger_packed( + x: torch.Tensor, + merge_gather_idx: torch.Tensor, +) -> torch.Tensor: + """Apply the image-only spatial merge using precomputed tensor indices.""" + return x[merge_gather_idx] + + class MoonViT3dPretrainedModel(nn.Module): """Main vision tower model. @@ -630,6 +726,10 @@ def __init__( pos_emb_width=config.init_pos_emb_width, pos_emb_time=config.init_pos_emb_time, pos_emb_type=config.pos_emb_type, + patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), + pos_emb_interpolation_mode=getattr( + config, "pos_emb_interpolation_mode", "bicubic" + ), ) self.encoder = MoonViT3dEncoder( @@ -638,9 +738,15 @@ def __init__( block_cfg={ "num_heads": config.num_attention_heads, "hidden_dim": config.hidden_size, + "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), "mlp_dim": config.intermediate_size, - "activation": get_act_fn("gelu_pytorch_tanh"), - "attn_bias": True, + "activation": get_act_fn( + getattr(config, "activation_func", "gelu_pytorch_tanh") + ), + "attn_bias": getattr(config, "attn_bias", True), + "norm_type": getattr(config, "norm_type", "layernorm"), + "mlp_type": getattr(config, "mlp_type", "mlp2"), + "linear_bias": getattr(config, "linear_bias", True), }, video_attn_type=config.video_attn_type, quant_config=quant_config, @@ -650,7 +756,7 @@ def __init__( def forward( self, pixel_values: torch.Tensor, - grid_thws: torch.Tensor | list[list[int]], + grid_thws: torch.Tensor | list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: @@ -662,6 +768,22 @@ def forward( Returns: torch.Tensor: The output tokens. """ + if encoder_metadata is not None and "pos_embeds" in encoder_metadata: + hidden_states = self.patch_embed( + pixel_values, + None, + pos_embeds=encoder_metadata["pos_embeds"], + ) + hidden_states = self.encoder( + hidden_states, + None, + encoder_metadata=encoder_metadata, + ) + merge_gather_idx = encoder_metadata["merge_gather_idx"] + assert merge_gather_idx is not None + return tpool_patch_merger_packed(hidden_states, merge_gather_idx) + + assert grid_thws is not None grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() if encoder_metadata is None: encoder_metadata = self.encoder.prepare_encoder_metadata( @@ -685,13 +807,40 @@ def forward( return hidden_states + def prepare_encoder_cudagraph_metadata( + self, + grid_thw_list: list[list[int]], + *, + max_batch_size: int, + max_seqlen_override: int | None = None, + device: torch.device, + ) -> dict[str, torch.Tensor | None]: + """Precompute fixed-buffer metadata for image encoder CUDA graphs.""" + grid_thw_list = [list(map(int, grid)) for grid in grid_thw_list] + metadata = self.encoder.prepare_encoder_metadata( + grid_thw_list, + device=device, + max_batch_size=max_batch_size, + max_seqlen_override=max_seqlen_override, + ) + metadata["pos_embeds"] = self.patch_embed.pos_emb.get_pos_embeds( + grid_thw_list + ).to(device=device) + merge_gather_idx = build_image_merge_gather_idx( + grid_thw_list, self.merge_kernel_size + ) + metadata["merge_gather_idx"] = torch.from_numpy(merge_gather_idx).to( + device=device, non_blocking=True + ) + return metadata + @torch.inference_mode() def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]): """Apply MM projector to vision tower outputs.""" num_embedding_list = [x.shape[0] for x in vt_output] batched = torch.cat(vt_output, dim=0) - projector_dtype = mm_projector.pre_norm.weight.dtype + projector_dtype = next(mm_projector.parameters()).dtype if batched.dtype != projector_dtype: batched = batched.to(projector_dtype) proj_out = mm_projector(batched) @@ -747,11 +896,34 @@ def __init__( ): super().__init__() self.use_data_parallel = use_data_parallel + self.mm_projector_type = getattr(config, "mm_projector_type", "patchmerger") # Hidden size after patch merging merge_h, merge_w = config.merge_kernel_size self.hidden_size = config.hidden_size * merge_h * merge_w + if self.mm_projector_type == "patchmergerv2": + self.linear_1 = ReplicatedLinear( + self.hidden_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.linear_1", + ) + self.linear_2 = ReplicatedLinear( + self.hidden_size, + getattr(config, "text_hidden_size", config.mm_hidden_size), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.linear_2", + ) + self.post_norm = torch.nn.RMSNorm( + getattr(config, "text_hidden_size", config.mm_hidden_size), + eps=config.projector_ln_eps, + ) + self.act = GELUActivation() + return + self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-5) self.linear_1 = ReplicatedLinear( self.hidden_size, @@ -770,6 +942,13 @@ def __init__( self.act = GELUActivation() def forward(self, image_features: torch.Tensor) -> torch.Tensor: + if self.mm_projector_type == "patchmergerv2": + hidden_states = image_features.view(image_features.shape[0], -1) + hidden_states, _ = self.linear_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return self.post_norm(hidden_states) + hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size) hidden_states, _ = self.linear_1(hidden_states) hidden_states = self.act(hidden_states) diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py deleted file mode 100644 index 057d4d01cb05..000000000000 --- a/vllm/model_executor/models/kimi_linear.py +++ /dev/null @@ -1,646 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Iterable - -import torch -from torch import nn - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import ( - get_pp_group, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - fused_moe_make_expert_params_mapping, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - ReplicatedLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( - KimiGatedDeltaNetAttention, -) -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, -) -from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -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.sequence import IntermediateTensors -from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig - -from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - get_spec_layer_idx_from_weight_name, - is_pp_missing_parameter, - make_layers, - maybe_prefix, -) - -logger = init_logger(__name__) - - -class KimiMLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: QuantizationConfig | None = None, - reduce_results: bool = True, - prefix: str = "", - ) -> None: - super().__init__() - - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - reduce_results=reduce_results, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class KimiMoE(nn.Module): - def __init__( - self, - config: KimiLinearConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - layer_idx: int = 0, - ): - super().__init__() - hidden_size = config.hidden_size - intermediate_size = config.intermediate_size - moe_intermediate_size = config.moe_intermediate_size - num_experts = config.num_experts - moe_renormalize = config.moe_renormalize - self.tp_size = get_tensor_model_parallel_world_size() - self.routed_scaling_factor = config.routed_scaling_factor - self.num_shared_experts = config.num_shared_experts - self.layer_idx = layer_idx - - if config.hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {config.hidden_act}. " - "Only silu is supported for now." - ) - - # Gate always runs at half / full precision for now. - self.gate = ReplicatedLinear( - hidden_size, - num_experts, - bias=False, - quant_config=None, - prefix=f"{prefix}.gate", - ) - - self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) - - if self.num_shared_experts is not None: - intermediate_size = moe_intermediate_size * self.num_shared_experts - self.shared_experts = KimiMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=False, - prefix=f"{prefix}.shared_experts", - ) - else: - self.shared_experts = None - - self.experts = FusedMoE( - shared_experts=self.shared_experts, - num_experts=num_experts, - top_k=config.num_experts_per_token, - hidden_size=hidden_size, - intermediate_size=moe_intermediate_size, - renormalize=moe_renormalize, - quant_config=quant_config, - use_grouped_topk=config.use_grouped_topk, - num_expert_group=config.num_expert_group, - topk_group=config.topk_group, - prefix=f"{prefix}.experts", - scoring_func=config.moe_router_activation_func, - e_score_correction_bias=self.gate.e_score_correction_bias, - routed_scaling_factor=self.routed_scaling_factor, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - num_tokens, hidden_size = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_size) - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, router_logits=router_logits - ) - return final_hidden_states.view(num_tokens, hidden_size) - - -class KimiMLAAttention(nn.Module): - """ - Main reference: DeepseekV2 vllm Implementation - """ - - def __init__( - self, - config: KimiLinearConfig, - hidden_size: int, - num_heads: int, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - v_head_dim: int, - q_lora_rank: int | None, - kv_lora_rank: int, - use_nope: bool = False, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - **kwargs, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.qk_nope_head_dim = qk_nope_head_dim - self.qk_rope_head_dim = qk_rope_head_dim - self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - self.v_head_dim = v_head_dim - self.q_lora_rank = q_lora_rank - self.kv_lora_rank = kv_lora_rank - self.num_heads = num_heads - tp_size = get_tensor_model_parallel_world_size() - self.num_local_heads = num_heads // tp_size - self.scaling = self.qk_head_dim**-0.5 - self.use_nope = use_nope - assert self.use_nope is True - assert self.q_lora_rank is None - assert num_heads % tp_size == 0 - self.kv_a_proj_with_mqa = ReplicatedLinear( - self.hidden_size, - self.kv_lora_rank + self.qk_rope_head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.kv_a_proj_with_mqa", - ) - self.q_proj = ColumnParallelLinear( - self.hidden_size, - self.num_heads * self.qk_head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.q_proj", - ) - self.kv_a_layernorm = RMSNorm( - self.kv_lora_rank, - eps=config.rms_norm_eps, - ) - self.kv_b_proj = ColumnParallelLinear( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.kv_b_proj", - ) - self.o_proj = RowParallelLinear( - self.num_heads * self.v_head_dim, - self.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - mla_modules = MLAModules( - kv_a_layernorm=self.kv_a_layernorm, - kv_b_proj=self.kv_b_proj, - rotary_emb=None, - o_proj=self.o_proj, - fused_qkv_a_proj=None, - kv_a_proj_with_mqa=self.kv_a_proj_with_mqa, - q_a_layernorm=None, - q_b_proj=None, - q_proj=self.q_proj, - indexer=None, - is_sparse=False, - topk_indices_buffer=None, - ) - self.mla_attn = MultiHeadLatentAttentionWrapper( - self.hidden_size, - self.num_local_heads, - self.scaling, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - self.v_head_dim, - self.q_lora_rank, - self.kv_lora_rank, - mla_modules, - cache_config, - quant_config, - prefix, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - output: torch.Tensor, - ) -> None: - output[:] = self.mla_attn(positions, hidden_states) - - -class KimiDecoderLayer(nn.Module): - def __init__( - self, - config: KimiLinearConfig, - vllm_config: VllmConfig, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - - self.is_moe = config.is_moe - layer_idx = int(prefix.rsplit(".", 1)[1]) - model_config = vllm_config.model_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - if config.is_kda_layer(layer_idx): - self.self_attn = KimiGatedDeltaNetAttention( - config, - vllm_config, - prefix=f"{prefix}.self_attn", - ) - else: - self.self_attn = KimiMLAAttention( - layer_idx=layer_idx, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - quant_config=quant_config, - cache_config=cache_config, - model_config=model_config, - prefix=f"{prefix}.self_attn", - config=config, - qk_nope_head_dim=config.qk_nope_head_dim, - qk_rope_head_dim=config.qk_rope_head_dim, - v_head_dim=config.v_head_dim, - q_lora_rank=config.q_lora_rank, - kv_lora_rank=config.kv_lora_rank, - use_nope=config.mla_use_nope, - ) - - if ( - self.is_moe - and config.num_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % config.moe_layer_freq == 0 - ): - self.block_sparse_moe = KimiMoE( - config=config, - quant_config=quant_config, - prefix=f"{prefix}.block_sparse_moe", - ) - self.mlp = self.block_sparse_moe - else: - self.mlp = KimiMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - - attn_output = torch.empty_like(hidden_states) - self.self_attn( - hidden_states=hidden_states, - positions=positions, - output=attn_output, - ) - hidden_states = attn_output - - # Fully Connected - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -@support_torch_compile -class KimiLinearModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_text_config - self.config = config - - self.vocab_size = config.vocab_size - - if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = PPMissingLayer() - - def get_layer(prefix: str): - return KimiDecoderLayer( - config, - vllm_config, - prefix, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - get_layer, - prefix=f"{prefix}.layers", - ) - - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - world_size = get_tensor_model_parallel_world_size() - assert config.num_attention_heads % world_size == 0, ( - "num_attention_heads must be divisible by world_size" - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - for _, layer in enumerate(self.layers[self.start_layer : self.end_layer]): - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), - ] - if self.config.is_moe: - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = fused_moe_make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.num_experts, - ) - else: - expert_params_mapping = [] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for args in weights: - name, loaded_weight = args[:2] - kwargs = args[2] if len(args) > 2 else {} - if "rotary_emb.inv_freq" in name: - continue - - spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) - if spec_layer is not None: - continue # skip spec decode layers for main model - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - # We have mlp.experts[0].gate_proj in the checkpoint. - # Since we handle the experts below in expert_params_mapping, - # we need to skip here BEFORE we update the name, otherwise - # name will be updated to mlp.experts[0].gate_up_proj, which - # will then be updated below in expert_params_mapping - # for mlp.experts[0].gate_gate_up_proj, which breaks load. - if ("mlp.experts." in name) and name not in params_dict: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - for idx, (param_name, weight_name, expert_id, shard_id) in enumerate( - expert_params_mapping - ): - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader( - param, - loaded_weight, - name, - expert_id=expert_id, - shard_id=shard_id, - ) - break - else: - # Skip loading extra bias for GPTQ models. - if ( - name.endswith(".bias") - and name not in params_dict - and not self.config.is_linear_attn - ): # noqa: E501 - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight, **kwargs) - loaded_params.add(name) - return loaded_params - - -class KimiLinearForCausalLM( - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid -): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - self.model_config = vllm_config.model_config - self.vllm_config = vllm_config - self.config = self.model_config.hf_config - quant_config = vllm_config.quant_config - self.quant_config = quant_config - self.model = KimiLinearModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - self.config.vocab_size, - self.config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - else: - self.lm_head = PPMissingLayer() - logit_scale = getattr(self.config, "logit_scale", 1.0) - self.logits_processor = LogitsProcessor( - self.config.vocab_size, scale=logit_scale - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs - ) - return hidden_states - - @classmethod - def get_mamba_state_dtype_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype - ) - - @classmethod - def get_mamba_state_shape_from_config( - cls, vllm_config: "VllmConfig" - ) -> tuple[tuple[int, ...], tuple[int, ...]]: - parallel_config = vllm_config.parallel_config - hf_config = vllm_config.model_config.hf_config - tp_size = parallel_config.tensor_parallel_size - num_spec = ( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config - else 0 - ) - return MambaStateShapeCalculator.kda_state_shape( - tp_size, - hf_config.linear_attn_config["num_heads"], - hf_config.linear_attn_config["head_dim"], - conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], - num_spec=num_spec, - ) - - @classmethod - def get_mamba_state_copy_func( - cls, - ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: - return MambaStateCopyFuncCalculator.kda_state_copy_func() - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - return self.logits_processor(self.lm_head, hidden_states) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py index e71054f4da39..3fcf03b7e666 100644 --- a/vllm/model_executor/models/laguna.py +++ b/vllm/model_executor/models/laguna.py @@ -19,7 +19,7 @@ ) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -174,7 +174,7 @@ def __init__( prefix=f"{prefix}.gate", ) - # Shared expert (optional) - passed to FusedMoE for overlap optimization + # Shared expert (optional) - passed to FusedMoEFactory for overlap optimization self.shared_expert: LagunaMLP | None if config.shared_expert_intermediate_size > 0: self.shared_expert = LagunaMLP( @@ -191,20 +191,22 @@ def __init__( # Auxiliary-loss-free load-balancing bias (arXiv:2408.15664). The # checkpoint stores one [num_experts] tensor per MoE layer at # `mlp.experts.e_score_correction_bias`; registering it as a Parameter - # on the FusedMoE lets the weight loader pick it up and the router + # on the MoERunner lets the weight loader pick it up and the router # add it during top-k selection. The fused top-k bias router requires # float32 regardless of model dtype. e_score_correction_bias = torch.nn.Parameter( torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False, ) - - # FusedMoE with SIGMOID routing. Passing `shared_experts=` lets the + # MoERunner with SIGMOID routing. Passing `shared_experts=` lets the + # layer overlap the shared-expert compute with the all2all dispatch. + # `apply_routed_scale_to_output=True` makes MoERunner handle the + # FusedMoEFactory with SIGMOID routing. Passing `shared_experts=` lets the # layer overlap the shared-expert compute with the all2all dispatch. - # `apply_routed_scale_to_output=True` makes FusedMoE handle the + # `apply_routed_scale_to_output=True` makes FusedMoEFactory handle the # routed_scaling_factor, shared+routed combine, and TP all-reduce # internally, so forward() just returns the final hidden states. - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_expert, num_experts=config.num_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index 601bd22b2fac..e18edb0b36e3 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/vllm/model_executor/models/lfm2.py @@ -169,8 +169,8 @@ def forward( n_tokens, _ = hidden_states.shape qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous() - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous() + q = q.view(n_tokens, self.num_heads, self.head_dim) + k = k.view(n_tokens, self.num_kv_heads, self.head_dim) q = self.q_layernorm(q) k = self.k_layernorm(k) q, k = self.rotary_emb(positions, q, k) @@ -257,6 +257,7 @@ def __init__( layer_idx=layer_idx, model_config=model_config, cache_config=cache_config, + quant_config=quant_config, prefix=f"{prefix}.conv", ) diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 698f1fb72ef2..0b41cb0f0fe7 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -15,7 +15,9 @@ ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -143,7 +145,7 @@ def __init__( else: self.gate.e_score_correction_bias = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -255,8 +257,8 @@ def forward( n_tokens, _ = hidden_states.shape qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous() - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous() + q = q.view(n_tokens, self.num_heads, self.head_dim) + k = k.view(n_tokens, self.num_kv_heads, self.head_dim) q = self.q_layernorm(q) k = self.k_layernorm(k) q, k = self.rotary_emb(positions, q, k) @@ -351,6 +353,7 @@ def __init__( layer_idx=layer_idx, model_config=model_config, cache_config=cache_config, + quant_config=quant_config, prefix=f"{prefix}.conv", ) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 71df54a42417..9577c27f5494 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -37,7 +37,7 @@ ChunkedLocalAttention, ) from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -131,7 +131,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): self.n_physical_experts = self.n_local_experts + self.n_redundant_experts self.n_local_physical_experts = self.n_physical_experts // self.ep_size - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_expert, num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 552e2a9e8f8d..9dc11d493a0b 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -24,8 +24,6 @@ import hashlib import importlib -import json -import os from collections.abc import Callable, Iterable, Mapping, Sequence from functools import lru_cache from typing import ( @@ -39,7 +37,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from huggingface_hub import hf_hub_download from PIL import Image from transformers import AutoProcessor, AutoTokenizer, BatchFeature from transformers.dynamic_module_utils import get_class_from_dynamic_module @@ -108,6 +105,7 @@ ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import _merge_mm_kwargs +from vllm.transformers_utils.repo_utils import get_hf_file_to_dict from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -165,13 +163,10 @@ def _load_ov2_processor( # codec video backend keeps its configured defaults. codec_config: dict = {} try: - config_file = os.path.join(path, "preprocessor_config.json") - if not os.path.isfile(config_file): - config_file = hf_hub_download( - path, "preprocessor_config.json", revision=revision - ) - with open(config_file, encoding="utf-8") as f: - codec_config = json.load(f).get("codec", {}) or {} + preprocessor_config = get_hf_file_to_dict( + "preprocessor_config.json", path, revision + ) + codec_config = (preprocessor_config or {}).get("codec", {}) or {} except Exception: logger.debug("OV2: no codec defaults found in preprocessor_config.json") @@ -608,10 +603,6 @@ def _expand_video_markers_in_prompt( _OV2_FPS_MIN_FRAMES = 4 -def _round_by_factor(n: float, factor: int) -> int: - return round(n / factor) * factor - - def _ceil_by_factor(n: float, factor: int) -> int: import math as _math @@ -1354,7 +1345,7 @@ def _get_vision_info( preprocessed = ImageSize(width=rw, height=rh) else: preprocessed = ImageSize(width=image_width, height=image_height) - padded_frames = num_frames + num_frames % temporal_patch_size + padded_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_frames // temporal_patch_size, 1) grid_h = preprocessed.height // patch_size grid_w = preprocessed.width // patch_size diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 18628a64cef7..2c6dedae3588 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -47,7 +47,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -297,7 +297,7 @@ def __init__( ) assert config.zero_expert_type is not None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( zero_expert_type=config.zero_expert_type, e_score_correction_bias=self.router.e_score_correction_bias, num_experts=num_experts, @@ -317,7 +317,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) - # Align to FusedMoE padded hidden size to avoid dim mismatch + # Align to MoERunner padded hidden size to avoid dim mismatch padded_hidden = self.experts.moe_config.hidden_dim if hidden_dim < padded_hidden: hidden_states_padded = torch.nn.functional.pad( @@ -333,7 +333,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states_padded.to(self.router_params_dtype) ) - # FusedMoE handles routing memoization and zero expert computation + # MoERunner handles routing memoization and zero expert computation # internally. Pass full router_logits (including zero experts) so that # zero experts can be properly identified in routing. final_hidden_states = self.experts( diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 4c2ebd958b58..30ddd5eb1211 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -171,7 +171,7 @@ def __init__( torch.empty(config.n_routed_experts, dtype=self.gate_dtype) ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index d0d9589ae1da..747cb0e88b2b 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -715,7 +715,7 @@ def _get_vision_info( effective_frames = num_frames * tokens_per_second else: effective_frames = num_frames - padded_num_frames = effective_frames + effective_frames % temporal_patch_size + padded_num_frames = effective_frames + (-effective_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size grid_w = preprocessed_size.width // patch_size diff --git a/vllm/model_executor/models/minicpm3.py b/vllm/model_executor/models/minicpm3.py index e61e9d06103d..30a9cd34ed0a 100644 --- a/vllm/model_executor/models/minicpm3.py +++ b/vllm/model_executor/models/minicpm3.py @@ -146,7 +146,7 @@ def forward( latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states) kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) latent_cache = latent_cache.unsqueeze(1) - kv_a = self.kv_a_layernorm(kv_a.contiguous()) + kv_a = self.kv_a_layernorm(kv_a) kv, _ = self.kv_b_proj(kv_a) kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index dc75928eb6c1..ac9efabaa3ab 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -651,9 +651,9 @@ def get_num_image_tokens( class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ - "q_proj": ("qkv_proj", "q"), - "k_proj": ("qkv_proj", "k"), - "v_proj": ("qkv_proj", "v"), + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), } ) diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 79bd8f439e3f..b94d6fef0ac3 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -39,9 +39,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -95,7 +93,7 @@ def __init__( else: self.e_score_correction_bias = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, scoring_func=config.scoring_func, diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 8305ecc330f9..991729354bbb 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -39,9 +39,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, @@ -123,7 +121,7 @@ def __init__( prefix=f"{prefix}.gate", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 5236f5c61be9..819881e0f213 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -14,6 +14,7 @@ from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Any, Literal, TypeAlias +import regex as re import torch from torch import nn from transformers import BatchFeature @@ -23,6 +24,7 @@ from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt from vllm.model_executor.models.interfaces import ( + DiarizedTranscriptionSegment, MultiModalEmbeddings, SupportsMultiModal, SupportsPP, @@ -35,6 +37,8 @@ _merge_multimodal_embeddings, init_vllm_registered_model, maybe_prefix, + parse_diarized_speaker, + parse_diarized_timestamp, ) from vllm.model_executor.models.whisper import ( WhisperEncoder, @@ -67,6 +71,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape WHISPER_ENCODER_STRIDE = 2 +MAX_AUDIO_DURATION_S = 90 * 60 AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" @@ -76,6 +81,11 @@ "并在段末标注结束时间戳，以清晰标明该段语音范围。" ) +_MOSS_DIARIZED_HEADER_RE = re.compile( + r"\[(?P[0-9.]{1,32})\]\s*\[(?PS[0-9]{1,15})\]" +) +_MOSS_DIARIZED_END_RE = re.compile(r"\[(?P[0-9.]{1,32})\]\s*\Z") + class MossTranscribeDiarizeAudioInputs(TensorSchema): """ @@ -145,9 +155,7 @@ def _compute_total_audio_tokens( def _get_max_audio_samples(feature_extractor: Any) -> int: - if hasattr(feature_extractor, "chunk_length"): - return int(feature_extractor.chunk_length * feature_extractor.sampling_rate) - return int(feature_extractor.n_samples) + return int(MAX_AUDIO_DURATION_S * feature_extractor.sampling_rate) def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]: @@ -564,6 +572,7 @@ class MossTranscribeDiarizeForConditionalGeneration( supports_transcription = True supports_transcription_only = True supports_segment_timestamp = False + supports_diarized_transcription = True supported_languages = ISO639_1_SUPPORTED_LANGS hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -632,6 +641,46 @@ def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: def post_process_output(cls, text: str) -> str: return text.strip() + @classmethod + def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]: + """Parse MOSS's canonical ``[start][Sxx]text[end]`` transcript.""" + headers: list[tuple[re.Match[str], float, str]] = [] + for match in _MOSS_DIARIZED_HEADER_RE.finditer(text): + start = parse_diarized_timestamp(match["start"]) + speaker = parse_diarized_speaker(match["speaker"]) + if start is not None and speaker is not None: + headers.append((match, start, speaker)) + + if not headers: + return [] + + segments: list[DiarizedTranscriptionSegment] = [] + for index, (header, start, speaker) in enumerate(headers): + next_header_start = ( + headers[index + 1][0].start() if index + 1 < len(headers) else len(text) + ) + body = text[header.end() : next_header_start] + end_match = _MOSS_DIARIZED_END_RE.search(body) + if end_match is None: + return [] + + end = parse_diarized_timestamp(end_match["end"]) + if end is None or end < start: + return [] + + segment_text = body[: end_match.start()].strip() + if segment_text: + segments.append( + DiarizedTranscriptionSegment( + start=start, + end=end, + speaker=speaker, + text=segment_text, + ) + ) + + return segments + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() self.config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 64667503d578..5b4233b0781c 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -42,10 +42,6 @@ maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_retained_tokens_count, - compute_retention_mask, -) from vllm.multimodal.inputs import ( AudioItem, BatchedTensorInputs, @@ -74,6 +70,10 @@ PromptReplacement, PromptUpdate, ) +from vllm.multimodal.video_prune.evs import ( + compute_retained_tokens_count, + compute_retention_mask, +) from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 84f24094a3d8..d68a619eb287 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -33,7 +33,7 @@ from vllm.model_executor.layers.activation import ReLUSquaredActivation from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, activation_without_mul, fused_moe_make_expert_params_mapping, @@ -69,6 +69,7 @@ SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, + SupportsReplaySSM, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -208,7 +209,7 @@ def __init__( self.fc1_latent_proj = None self.fc2_latent_proj = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.n_routed_experts, top_k=config.num_experts_per_tok, @@ -707,6 +708,7 @@ class NemotronHForCausalLM( SupportsQuant, MixtureOfExperts, SupportsMambaPrefixCaching, + SupportsReplaySSM, ): # Relevant only if self.has_moe is True is_non_gated_moe: bool = True @@ -742,18 +744,24 @@ class NemotronHForCausalLM( def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( + ) -> tuple[torch.dtype, ...]: + cache_config = vllm_config.cache_config + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, + cache_config.mamba_cache_dtype, + cache_config.mamba_ssm_cache_dtype, ) + if cache_config.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, vllm_config.model_config.dtype + ) + return base_dtype @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> tuple[tuple[int, ...], ...]: """Calculate shapes for Mamba's convolutional and state caches. Args: @@ -763,12 +771,14 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + - x_cache/dt_cache/B_cache ring-buffer shapes (use_replayssm only) """ parallel_config = vllm_config.parallel_config + cache_config = vllm_config.cache_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim - return MambaStateShapeCalculator.mamba2_state_shape( + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -778,6 +788,14 @@ def get_mamba_state_shape_from_config( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) + if cache_config.use_replayssm: + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + hf_config.n_groups, + parallel_config.tensor_parallel_size, + cache_config.replayssm_buffer_len, + ) + return base_shape @classmethod def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py index a57ef9cf4307..7867caa97c57 100644 --- a/vllm/model_executor/models/olmoe.py +++ b/vllm/model_executor/models/olmoe.py @@ -32,9 +32,7 @@ from vllm.distributed.utils import split_tensor_along_last_dim from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, @@ -94,7 +92,7 @@ def __init__( prefix=f"{prefix}.gate", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index fd40033e7154..6528b8488c10 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -43,7 +43,7 @@ Attention, StaticSinkAttention, ) -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -195,7 +195,7 @@ def __init__( else: self.shared_experts = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=config.n_routed_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/ouro.py b/vllm/model_executor/models/ouro.py deleted file mode 100644 index 527eeaa13bc6..000000000000 --- a/vllm/model_executor/models/ouro.py +++ /dev/null @@ -1,448 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py -# Copyright 2024 The Qwen team. -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only Ouro model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from typing import Any - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors -from vllm.v1.attention.backend import AttentionType - -from .interfaces import SupportsLoRA -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - extract_layer_index, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class OuroMLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class OuroAttention(nn.Module): - def __init__( - self, - config: PretrainedConfig, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - max_position: int = 4096 * 32, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - attn_type: str = AttentionType.DECODER, - dual_chunk_attention_config: dict[str, Any] | None = None, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = hidden_size // self.total_num_heads - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.dual_chunk_attention_config = dual_chunk_attention_config - - # Get total_ut_steps from config, default to 4 if not specified - total_ut_steps = getattr(config, "total_ut_steps", 4) - - # Use total number of hidden layers instead of hardcoded 24 - total_layers = config.num_hidden_layers - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=config.rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - self.attn = nn.ModuleList() - for ut_step in range(total_ut_steps): - base_layer_idx = extract_layer_index(prefix) - unique_layer_idx = ut_step * total_layers + base_layer_idx - - unique_prefix = prefix.replace( - f"layers.{base_layer_idx}", f"layers.{unique_layer_idx}" - ) - - self.attn.append( - Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - attn_type=attn_type, - prefix=f"{unique_prefix}.attn", - **{ - "layer_idx": unique_layer_idx, - "dual_chunk_attention_config": dual_chunk_attention_config, - } - if dual_chunk_attention_config - else {}, - ) - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - current_ut: int, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn[current_ut](q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class OuroDecoderLayer(nn.Module): - def __init__( - self, - config: PretrainedConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - dual_chunk_attention_config = getattr( - config, "dual_chunk_attention_config", None - ) - - if getattr(config, "is_causal", True): - attn_type = AttentionType.DECODER - else: - attn_type = AttentionType.ENCODER_ONLY - - self.self_attn = OuroAttention( - config=config, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - max_position=config.max_position_embeddings, - num_kv_heads=config.num_key_value_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.self_attn", - attn_type=attn_type, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - self.mlp = OuroMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - self.post_attention_layernorm_2 = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - current_ut: int, - residual: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn( - positions=positions, hidden_states=hidden_states, current_ut=current_ut - ) - hidden_states = self.input_layernorm_2(hidden_states) - - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - hidden_states = self.post_attention_layernorm_2(hidden_states) - - return hidden_states, residual - - -@support_torch_compile( - dynamic_arg_dims={ - "input_ids": 0, - "positions": -1, - "intermediate_tensors": 0, - "inputs_embeds": 0, - } -) -class OuroModel(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - decoder_layer_type: type[nn.Module] = OuroDecoderLayer, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - # TODO (@robertgshaw2): see if this can be moved out - if cache_config.sliding_window is not None and hasattr( - config, "max_window_layers" - ): - assert config.max_window_layers == config.num_hidden_layers, ( - "Sliding window for some but all layers is not supported. " - "This model uses sliding window but `max_window_layers` = {} " - "is less than `num_hidden_layers` = {}. Please open an issue " - "to discuss this feature.".format( - config.max_window_layers, - config.num_hidden_layers, - ) - ) - - self.config = config - self.quant_config = quant_config - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - - # Use the provided decoder layer type or default to OuroDecoderLayer - decoder_layer_type = decoder_layer_type or OuroDecoderLayer - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: decoder_layer_type( - config=config, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ), - prefix=f"{prefix}.layers", - ) - - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.early_exit_gate = RowParallelLinear(config.hidden_size, 1, bias=True) - - self.total_ut_steps = getattr(self.config, "total_ut_steps", 4) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - - for current_ut in range(self.total_ut_steps): - residual = None - for layer in self.layers[self.start_layer : self.end_layer]: - hidden_states, residual = layer( - positions, hidden_states, current_ut, residual - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - -class OuroForCausalLM(nn.Module, SupportsLoRA): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - ".gate_proj": (".gate_up_proj", 0), - ".up_proj": (".gate_up_proj", 1), - } - ) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - - self.quant_config = quant_config - self.model = OuroModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - - self.logits_processor = LogitsProcessor(config.vocab_size) - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py index 6dbed78a6fc6..19dd04740760 100644 --- a/vllm/model_executor/models/ovis2_5.py +++ b/vllm/model_executor/models/ovis2_5.py @@ -402,12 +402,11 @@ def _get_prompt_updates( hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs: MultiModalKwargsItems, ) -> list[PromptReplacement]: - tokenizer = self.info.get_tokenizer() - vocab = tokenizer.get_vocab() + hf_processor = self.info.get_hf_processor() placeholder = { - "image": vocab[IMAGE_TOKEN], - "video": vocab[VIDEO_TOKEN], + "image": hf_processor.get_token_value("image_token"), + "video": hf_processor.get_token_value("video_token"), } def get_replacement_ovis(item_idx, modality: str): diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 48c0c714ed8d..7830272c324e 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -32,7 +32,9 @@ ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -325,7 +327,7 @@ def __init__( else: self.shared_experts = None # type: ignore[assignment] - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=self.num_experts, top_k=self.top_k, @@ -342,7 +344,7 @@ def __init__( routed_scaling_factor=self.routed_scaling_factor, ) - def maybe_get_fused_moe(self) -> FusedMoE: + def maybe_get_fused_moe(self) -> FusedMoEFactory: return self.experts def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py index dc2af1e8d6fa..b3b7b995ee3e 100644 --- a/vllm/model_executor/models/phimoe.py +++ b/vllm/model_executor/models/phimoe.py @@ -35,9 +35,7 @@ from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.linear import ( QKVParallelLinear, ReplicatedLinear, @@ -273,7 +271,7 @@ def __init__( prefix=f"{prefix}.gate", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, diff --git a/vllm/model_executor/models/plamo2.py b/vllm/model_executor/models/plamo2.py deleted file mode 100644 index 5fd925cf0bee..000000000000 --- a/vllm/model_executor/models/plamo2.py +++ /dev/null @@ -1,992 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only PLaMo2 model.""" - -from collections.abc import Iterable -from itertools import islice -from typing import TYPE_CHECKING - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import divide, get_tensor_model_parallel_world_size -from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import ForwardContext, get_forward_context -from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.abstract import MambaBase -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, - is_conv_state_dim_first, -) -from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from vllm.model_executor.layers.mamba.ops.ssd_combined import ( - mamba_chunk_scan_combined_varlen, -) -from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import ( - composed_weight_loader, - default_weight_loader, - sharded_weight_loader, -) -from vllm.model_executor.models.interfaces import ( - HasInnerState, - IsHybrid, - SupportsLoRA, - SupportsPP, -) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.sequence import IntermediateTensors -from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backend import AttentionMetadata -from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum - -# Only used for type hinting. -if TYPE_CHECKING: - - class Plamo2Config(PretrainedConfig): # type: ignore - model_type: str = "plamo2" - - hidden_size: int - num_hidden_layers: int - rms_norm_eps: float - # Attention - num_attention_heads: int - hidden_size_per_head: int - num_key_value_heads: int - # Mamba - mamba_d_state: int - mamba_d_conv: int - mamba_num_heads: int - mamba_step: int - # MLP - intermediate_size: int - # Tokenizer - vocab_size: int - - -def is_mamba(config: "Plamo2Config", i: int) -> bool: - assert config.mamba_step > 1 - - if config.num_hidden_layers <= (config.mamba_step // 2): - # use attention in last layer - return i != config.num_hidden_layers - 1 - return (i % config.mamba_step) != (config.mamba_step // 2) - - -# Adapted from: -# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2 -# transformers.models.mamba.modeling_mamba.MambaMixer -# --8<-- [start:plamo2_mamba_mixer] -@PluggableLayer.register("plamo2_mamba_mixer") -class Plamo2MambaMixer(MambaBase, PluggableLayer): - # --8<-- [end:plamo2_mamba_mixer] - - def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None: - super().__init__() - self.config = vllm_config.model_config.hf_config - self.cache_config = vllm_config.cache_config - self.model_config = vllm_config.model_config - self.quant_config = vllm_config.quant_config - self.is_lora_enabled = bool(vllm_config.lora_config) - self.hidden_size = self.config.hidden_size - self.ssm_state_size = self.config.mamba_d_state - self.conv_kernel_size = self.config.mamba_d_conv - self.intermediate_size = ( - self.config.mamba_num_heads * self.config.hidden_size_per_head - ) - self.tp_size = get_tensor_model_parallel_world_size() - self.head_dim = self.config.hidden_size_per_head - self.num_heads = self.config.mamba_num_heads - self.time_step_rank = max(64, self.hidden_size // 16) - self.conv1d = ColumnParallelLinear( - input_size=self.conv_kernel_size, - output_size=self.intermediate_size, - bias=False, - prefix=f"{prefix}.conv1d", - return_bias=False, - ) - # unsqueeze to fit conv1d weights shape into the linear weights shape. - # Can't do this in `weight_loader` since it already exists in - # `ColumnParallelLinear` and `set_weight_attrs` - # doesn't allow to override it - self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) - - self.in_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.in_proj", - return_bias=False, - ) - # selective projection used to make dt, B and C input dependent - self.bcdt_proj = RowParallelLinear( - self.intermediate_size, - self.time_step_rank + self.ssm_state_size * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.bcdt_proj", - return_bias=False, - ) - # time step projection (discretization) - - # In the forward we need to apply dt_proj without the bias, - # as the bias is added in the selective scan kernel. - self.dt_proj = ColumnParallelLinear( - self.time_step_rank, - self.num_heads, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.dt_proj", - return_bias=False, - ) - - self.A = nn.Parameter( - torch.empty( - divide(self.num_heads, self.tp_size), - dtype=torch.float32, - ) - ) - self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - - set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)}) - a_weight_loader = composed_weight_loader( - sharded_weight_loader(0), lambda x: -torch.exp(x.float()) - ) - set_weight_attrs(self.A, {"weight_loader": a_weight_loader}) - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - - self.out_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - input_is_parallel=True, - quant_config=self.quant_config, - prefix=f"{prefix}.out_proj", - return_bias=False, - ) - # The activation function is fixed to SiLU. - self.activation = "silu" - - self.dt_norm = RMSNorm(self.time_step_rank, eps=self.config.rms_norm_eps) - self.B_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - self.C_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - - self.chunk_size = self.config.mamba_chunk_size - - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) - assert self.chunk_size != -1, "chunk_size must be set for v1" - - self.prefix = prefix - - def _project_ssm_parameters(self, hidden_states): - if self.is_lora_enabled: - # Lora kernel requires contiguous tensor. - ssm_parameters = self.bcdt_proj(hidden_states.contiguous()) - else: - ssm_parameters = self.bcdt_proj(hidden_states) - B, C, time_step = torch.split( - ssm_parameters, - [self.ssm_state_size, self.ssm_state_size, self.time_step_rank], - dim=-1, - ) - - # vllm._custom_ops.rms_norm requires contiguous input tensors. - time_step = self.dt_norm(time_step.contiguous()) - B = self.B_norm(B.contiguous()) - C = self.C_norm(C.contiguous()) - dt = self.dt_proj(time_step) - return B, C, dt - - def forward( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - torch.ops.vllm.plamo2_mamba_mixer( - hidden_states, - output, - self.prefix, - ) - - def forward_impl( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - forward_context = get_forward_context() - # attn_metadata contains metadata necessary for the mamba2 triton - # kernels to operate in continuous batching and in chunked prefill - # modes; they are computed at top-level model forward since they - # stay the same and reused for all mamba layers in the same iteration - attn_metadata: AttentionMetadata = forward_context.attn_metadata - - if attn_metadata is not None: - assert isinstance(attn_metadata, dict) - attn_metadata = attn_metadata[self.prefix] - assert isinstance(attn_metadata, Mamba2AttentionMetadata) - self_kv_cache = self.kv_cache - # conv_state = (..., dim, width-1) yet contiguous along 'dim' - # conv_state must be (..., dim, width-1) for the conv kernels. - # DS layout stores it that way directly; SD layout needs a transpose. - conv_state = ( - self_kv_cache[0] - if is_conv_state_dim_first() - else self_kv_cache[0].transpose(-1, -2) - ) - ssm_state = self_kv_cache[1] - state_indices_tensor_p = attn_metadata.state_indices_tensor_p - state_indices_tensor_d = attn_metadata.state_indices_tensor_d - has_initial_states_p = attn_metadata.has_initial_states_p - prep_initial_states = attn_metadata.prep_initial_states - chunk_size = attn_metadata.chunk_size - seq_idx_p = attn_metadata.seq_idx_p - query_start_loc_p = attn_metadata.query_start_loc_p - cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p - last_chunk_indices_p = attn_metadata.last_chunk_indices_p - - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states) - gate, hidden_states = projected_states.chunk(2, dim=-1) - - # 2. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) - - if attn_metadata is None: - # profile run - hidden_states = ( - hidden_states.transpose(0, 1).clone().transpose(0, 1) - ).contiguous() - output[:] = self.out_proj(hidden_states) - return - - num_prefills = attn_metadata.num_prefills # request count - num_decodes = attn_metadata.num_decode_tokens # token count (=request) - num_prefill_tokens = attn_metadata.num_prefill_tokens # token count - has_prefill = num_prefills > 0 - has_decode = num_decodes > 0 - num_actual_tokens = num_prefill_tokens + num_decodes - - # Separate prefill and decode by splitting varlen input - # Split along token dimension - hidden_states_d, hidden_states_p = torch.split( - hidden_states[:num_actual_tokens], - [num_decodes, num_prefill_tokens], - dim=0, - ) - gate_d, gate_p = torch.split( - gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 - ) - # Preallocate output tensor to avoid memcpy cost for merging prefill - # and decode outputs - preallocated_ssm_out = torch.empty( - [ - num_prefill_tokens + num_decodes, - (self.num_heads // self.tp_size) * self.head_dim, - ], - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split( - preallocated_ssm_out, - [num_decodes, num_prefill_tokens], - dim=0, - ) - - # Process prefill requests - if has_prefill: - # 2. Convolution sequence transformation - # - "cache_indices" updates the conv_state cache in positions - # pointed to by "state_indices_tensor_p" - x = hidden_states_p.transpose(0, 1) # this is the form that causal-conv see - hidden_states_p = causal_conv1d_fn( - x, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_states_p, - cache_indices=state_indices_tensor_p, - metadata=attn_metadata, - query_start_loc=query_start_loc_p, - ) - hidden_states_p = hidden_states_p.transpose(0, 1) - hidden_states_p = hidden_states_p[:num_prefill_tokens] - # In some instances, the following `bcdt_proj` op - # requires contiguous inputs - # (e.g. if the Marlin kernel is used). - hidden_states_p = hidden_states_p.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_p) - - # 3. State Space Model sequence transformation - initial_states = None - if has_initial_states_p is not None and prep_initial_states: - # making a copy of the states - initial_states = torch.where( - has_initial_states_p[:, None, None, None], - ssm_state[state_indices_tensor_p], - 0, - ) - - varlen_state = mamba_chunk_scan_combined_varlen( - hidden_states_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt, - self.A, - B.view(num_prefill_tokens, 1, -1), - C.view(num_prefill_tokens, 1, -1), - chunk_size=chunk_size, - D=self.D, - z=gate_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt_bias=self.dt_bias, - seq_idx=seq_idx_p, - cu_seqlens=query_start_loc_p, - cu_chunk_seqlens=cu_chunk_seqlen_p, - last_chunk_indices=last_chunk_indices_p, - initial_states=initial_states, - dt_softplus=True, - dt_limit=(0.0, float("inf")), - out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim), - state_dtype=ssm_state.dtype, - ) - - # update ssm states - # - varlen state is a (batch, nheads, headdim, dstate) tensor - ssm_state[state_indices_tensor_p] = varlen_state - - # Process decode requests - if has_decode: - # 2. Convolution sequence transformation - hidden_states_d = causal_conv1d_update( - hidden_states_d, - conv_state, - conv_weights, - self.conv1d.bias, - self.activation, - conv_state_indices=state_indices_tensor_d, - ) - - # ROCm: Ensure contiguous tensor for bcdt_proj linear layer. - # causal_conv1d_update returns a non-contiguous view (stride 8192 - # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM - # results when batch > 1 on ROCm. - if current_platform.is_rocm(): - hidden_states_d = hidden_states_d.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_d) - - # 3. State Space Model sequence transformation - A = self.A[:, None, ...][:, :, None].expand( - -1, self.head_dim, self.config.mamba_d_state - ) - dt = dt[:, :, None].expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.unsqueeze(1) - C = C.unsqueeze(1) - hidden_states_d = hidden_states_d.view( - -1, self.num_heads // self.tp_size, self.head_dim - ) - - # - the hidden is reshaped into (bs, num_heads, head_dim) - # - ssm_state's slots will be selected - # using state_indices_tensor_d - - # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt, - A, - B, - C, - D, - dt_bias, - z=gate_d.reshape(num_decodes, -1, self.head_dim), - dt_softplus=True, - state_batch_indices=state_indices_tensor_d, - out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim), - ) - - # 4. Final linear projection - output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out) - - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: - assert self.model_config is not None - assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( - self.model_config.dtype, - self.cache_config.mamba_cache_dtype, - self.cache_config.mamba_ssm_cache_dtype, - ) - - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), - n_groups=0, - num_heads=self.num_heads, - head_dim=self.head_dim, - state_size=self.ssm_state_size, - conv_kernel=self.conv_kernel_size, - ) - - @property - def mamba_type(self) -> MambaAttentionBackendEnum: - return MambaAttentionBackendEnum.MAMBA2 - - -def plamo2_mamba_mixer( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self.forward_impl(hidden_states=hidden_states, output=output) - - -def plamo2_mamba_mixer_fake( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - return - - -direct_register_custom_op( - op_name="plamo2_mamba_mixer", - op_func=plamo2_mamba_mixer, - mutates_args=["output"], - fake_impl=plamo2_mamba_mixer_fake, -) - - -class DenseMLP(nn.Module): - def __init__( - self, - config: "Plamo2Config", - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_up_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - prefix=f"{prefix}.gate_up_proj", - quant_config=quant_config, - return_bias=False, - ) - self.act = SiluAndMul() - self.down_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - prefix=f"{prefix}.down_proj", - quant_config=quant_config, - return_bias=False, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - h = self.gate_up_proj(hidden_states) - h = self.act(h) - return self.down_proj(h) - - -class Plamo2AttentionMixer(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - self.hidden_size = config.hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = config.num_attention_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = config.hidden_size_per_head - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - config.hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - config.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - max_position = config.max_position_embeddings - if hasattr(vllm_config.model_config, "max_model_len") and isinstance( - vllm_config.model_config.max_model_len, int - ): - max_position = min(max_position, vllm_config.model_config.max_model_len) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=config.rope_parameters, - ) - self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.q_norm.weight = torch.nn.Parameter( - torch.ones((self.num_heads, config.hidden_size_per_head)) - ) - set_weight_attrs( - self.q_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - self.k_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.k_norm.weight = torch.nn.Parameter( - torch.ones((self.num_kv_heads, config.hidden_size_per_head)) - ) - # Tensor-parallelism shards the K norm weights to the tp ranks - # in a head-wise manner. This approach does not work if there is only - # a single KV head, as is the case for PLaMo 2-1B. - if self.total_num_kv_heads != 1: - set_weight_attrs( - self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - - q_shape = q.shape - q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape) - q = self.q_norm.forward_native(q).reshape(q_shape) - k_shape = k.shape - k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape) - k = self.k_norm.forward_native(k).reshape(k_shape) - - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class Plamo2DecoderLayer(nn.Module): - def __init__( - self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs - ) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.is_mamba = is_mamba(config, layer_idx) - if self.is_mamba: - self.mixer = Plamo2MambaMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - else: - self.mixer = Plamo2AttentionMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - - self.mlp = DenseMLP( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" - ) - self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - **kwargs, - ): - if residual is None: - residual = hidden_states - hidden_states = self.pre_mixer_norm(hidden_states) - else: - hidden_states, residual = self.pre_mixer_norm(hidden_states, residual) - - if self.is_mamba: - # Plamo2MambaMixer writes output to this tensor - output = torch.empty_like(hidden_states) - mixer_kwargs = { - "output": output, - } - else: - mixer_kwargs = { - "positions": positions, - } - hidden_states = self.mixer( - hidden_states=hidden_states, - **mixer_kwargs, - ) - if self.is_mamba: - hidden_states = output - hidden_states = self.post_mixer_norm(hidden_states) - # Fully Connected - hidden_states, residual = self.pre_mlp_norm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - hidden_states = self.post_mlp_norm(hidden_states) - return hidden_states, residual - - -class Plamo2Decoder(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)} - - def get_layer(prefix: str): - layer_idx = int(prefix.rsplit(".", 1)[1]) - return Plamo2DecoderLayer( - vllm_config=vllm_config, - layer_idx=layer_idx, - prefix=prefix, - **extra_kwargs, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> torch.Tensor: - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - return hidden_states, residual - - -@support_torch_compile -class Plamo2Model(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - - self.config = config - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - prefix=f"{prefix}.embed_tokens", - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - self.layers = Plamo2Decoder(vllm_config=vllm_config, prefix=f"{prefix}.layers") - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - hidden_states, residual = self.layers( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - # Update the weight names to be compatible with the vllm version - # of the model. - # Do not change the order of the replacements. - replacements = { - # Rename incompatible weight names. - ".A_log": ".A", - ".B_norm_weight": ".B_norm.weight", - ".C_norm_weight": ".C_norm.weight", - ".dt_norm_weight": ".dt_norm.weight", - ".q_weight": ".q_norm.weight", - ".k_weight": ".k_norm.weight", - } - # Apply replacements based on the defined mappings - for old, new in replacements.items(): - if old in name: - name = name.replace(old, new) - - # Reshape the in_proj weights to match the shape expected - # by MergedColumnParallelLinear. - # This works both for unquantized weights and - # for quantized weights. - # In the quantized case, the weights are already transposed. - # Also, in addition to the quantized weights, - # the zero points and scales have to be reshaped as well. - # Packing should not be affected by this. - if ( - ".mixer.in_proj.weight" in name - or "mixer.in_proj.qweight" in name - or "mixer.in_proj.scales" in name - or "mixer.in_proj.qzeros" in name - ): - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - # for weight: - # loaded_weight.shape[0] == self.config.hidden_size - # for qweight: - # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa - # for scales and qzeros: - # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa - loaded_weight = loaded_weight.reshape( - loaded_weight.shape[0], self.config.mamba_num_heads, -1 - ) - gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1) - gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1) - hidden_states_weight = hidden_states_weight.reshape( - loaded_weight.shape[0], -1 - ) - loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1) - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - - # Offset parameter with vllm's RMSNorm haven't been supported yet. - if ".pre_mixer_norm" in name: - loaded_weight += 1.0 - elif ".post_mixer_norm" in name: - loaded_weight += 1.0 / 5 - elif ".pre_mlp_norm" in name: - loaded_weight += 1.0 - elif ".post_mlp_norm" in name: - loaded_weight += 1.0 / (5**1.5) - elif name == "norm.weight": - loaded_weight += 1.0 - - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class Plamo2ForCausalLM( - torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid -): - packed_modules_mapping = { - "qkv_proj": ["qkv_proj"], - "gate_up_proj": ["gate_up_proj"], - "in_proj": ["in_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - scheduler_config = vllm_config.scheduler_config - - self.config = config - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = scheduler_config - - # ModelConfig.get_head_size assumes head_dim is set or calculated as - # hidden_size // num_attention_heads. However, this is not always - # the case for PLaMo2, as indicated by the FIXME comment. - self.config.head_dim = self.config.hidden_size_per_head - - self.model = Plamo2Model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - self.vocab_size = self.config.vocab_size - self.lm_head = ParallelLMHead( - self.vocab_size, - self.config.hidden_size, - prefix=f"{prefix}.lm_head", - ) - if self.config.tie_word_embeddings: - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) - - self.logits_processor = LogitsProcessor( - config.vocab_size, self.config.vocab_size - ) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ): - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - @classmethod - def get_mamba_state_dtype_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( - vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, - ) - - @classmethod - def get_mamba_state_shape_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: - """Calculate shapes for Mamba's convolutional and state caches. - Args: - vllm_config: vLLM config - Returns: - Tuple containing: - - conv_state_shape: Shape for convolutional state cache - - temporal_state_shape: Shape for state space model cache - """ - parallel_config = vllm_config.parallel_config - hf_config = vllm_config.model_config.hf_config - intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head - - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=intermediate_size, - tp_world_size=parallel_config.tensor_parallel_size, - n_groups=0, - num_heads=hf_config.mamba_num_heads, - head_dim=hf_config.hidden_size_per_head, - state_size=hf_config.mamba_d_state, - conv_kernel=hf_config.mamba_d_conv, - ) - - @classmethod - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: - return MambaStateCopyFuncCalculator.mamba2_state_copy_func() - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 182b9758308d..f1d5f23a264c 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -55,7 +55,7 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta +from vllm.transformers_utils.config import set_default_rope_theta from vllm.v1.attention.backend import AttentionType from .interfaces import ( @@ -345,7 +345,7 @@ def __init__( quant_config = vllm_config.quant_config # TODO (@robertgshaw2): see if this can be moved out - if is_interleaved(vllm_config.model_config.hf_text_config): + if len(set(getattr(config, "layer_types", []))) > 1: assert config.max_window_layers == config.num_hidden_layers, ( "Sliding window for some but all layers is not supported. " "This model uses sliding window but `max_window_layers` = {} " diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index c987e07b43de..7957ca805ff6 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -67,12 +67,6 @@ ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -80,6 +74,12 @@ ) from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import PromptReplacement, PromptUpdate +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index a946159b8fcc..43274e4ba7e0 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -40,9 +40,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -165,7 +163,7 @@ def __init__( else: self.shared_expert = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_expert, num_experts=config.num_experts, top_k=config.num_experts_per_tok, diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 539f141cbaa7..e2e9f2452489 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -898,8 +898,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index bcd2576b74f1..50b47f16fecf 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -53,6 +53,7 @@ ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors +from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig from vllm.transformers_utils.configs.qwen3_5_moe import ( Qwen3_5MoeConfig, @@ -107,7 +108,11 @@ def get_hf_config(self): class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo): def get_hf_config(self): - return self.ctx.get_hf_config(Qwen3_5MoeConfig) + # transformers 5.x renames the top-level Qwen3.5-MoE config class to + # Qwen3_5MoeTextConfig for text-only models, while transformers ≤4.x + # returns Qwen3_5MoeConfig (the multimodal wrapper). Accept both so + # that vLLM works regardless of which transformers version is installed. + return self.ctx.get_hf_config((Qwen3_5MoeConfig, Qwen3_5MoeTextConfig)) class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): @@ -281,6 +286,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Qwen3_5ForCausalLMBase( nn.Module, HasInnerState, + IsHybrid, SupportsEagle3, SupportsLoRA, SupportsPP, @@ -360,6 +366,45 @@ def forward( return hidden_states + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + def compute_logits( self, hidden_states: torch.Tensor, @@ -397,8 +442,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dummy_inputs=Qwen3VLDummyInputsBuilder, ) class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid): - # Qwen3.5 does not support multimodal pruning (EVS). - supports_multimodal_pruning = False + supports_multimodal_pruning = True packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | { "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], @@ -416,8 +460,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - # Qwen3.5 does not support multimodal pruning (EVS). - self.is_multimodal_pruning_enabled = False + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( @@ -462,12 +519,6 @@ def embed_input_ids( return inputs_embeds - def recompute_mrope_positions(self, *args, **kwargs): - raise NotImplementedError( - "Qwen3.5 does not support multimodal pruning (EVS). " - "recompute_mrope_positions should never be called." - ) - def forward( self, input_ids: torch.Tensor, @@ -628,8 +679,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - # Qwen3.5 does not support multimodal pruning (EVS). - self.is_multimodal_pruning_enabled = False + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 0620509f2db5..6421c8b60d2d 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -103,6 +103,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.fc", ) + # GPTQ: quantized checkpoints may exclude MTP from quantization via + # quantization_config.dynamic with "-:pattern" entries. When detected, + # disable quantization for MTP layers so they use unquantized params. + original_quant = vllm_config.quant_config + if quant_config and quant_config.get_name() not in ("modelopt_fp4",): + hf_qc = getattr(model_config.hf_config, "quantization_config", None) + if isinstance(hf_qc, dict): + dynamic = hf_qc.get("dynamic", {}) + if any(k.startswith("-:") and "mtp" in k for k in dynamic): + vllm_config.quant_config = None self.layers = torch.nn.ModuleList( Qwen3_5DecoderLayer( vllm_config, @@ -111,11 +121,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) for idx in range(self.num_mtp_layers) ) - + vllm_config.quant_config = original_quant self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) - self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.pre_fc_norm_hidden = Qwen3_5RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -170,6 +179,7 @@ def forward( positions.shape[-1], self.config.hidden_size, ) + hidden_states, _ = self.norm(hidden_states, residual) return hidden_states diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index bf5ea501cc49..90b2f48faa82 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -35,6 +35,9 @@ from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) from .qwen2 import Qwen2MLP as Qwen3MLP from .qwen3 import Qwen3ForCausalLM @@ -69,6 +72,17 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool: ) +def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: + spec_config = vllm_config.speculative_config + config = spec_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(spec_config) + num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers + target_hidden_size = ( + getattr(config, "target_hidden_size", None) or config.hidden_size + ) + return target_hidden_size * num_features_to_use + + def _resolve_layer_attention( config: Qwen3Config, layer_idx: int ) -> tuple[int | None, bool]: @@ -395,17 +409,10 @@ def __init__( ] ) if self.use_aux_hidden_state: - num_features_to_use = self.config.num_hidden_layers - if "target_layer_ids" in drafter_config: - num_features_to_use = len(drafter_config["target_layer_ids"]) - elif "layer_ids" in drafter_config: - num_features_to_use = len(drafter_config["layer_ids"]) - if hasattr(self.config, "target_hidden_size"): - fc_input_size = self.config.target_hidden_size * num_features_to_use - else: - fc_input_size = self.config.hidden_size * num_features_to_use self.fc = ReplicatedLinear( - input_size=fc_input_size, + input_size=_get_dflash_fc_input_size( + vllm_config, + ), output_size=self.config.hidden_size, bias=False, params_dtype=vllm_config.model_config.dtype, diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 219819759ac0..04af45669498 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -22,9 +22,9 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, - VocabParallelEmbedding, ) from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model @@ -40,6 +40,10 @@ class DSparkMarkovHead(nn.Module): ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias (``draft_vocab_size``) added to the base draft logits. The two sizes coincide for full-vocab drafts. + + Both weights are replicated because the head runs sequentially for every + draft position. Sharding them would add an all-reduce and a full-vocab + gather to each position. """ def __init__( @@ -48,21 +52,28 @@ def __init__( draft_vocab_size: int, markov_rank: int, prefix: str, + quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() - # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard - self.markov_w1 = VocabParallelEmbedding( - vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1") - ) + self.markov_w1 = nn.Embedding(vocab_size, markov_rank) self.markov_w2 = ParallelLMHead( - draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2") + draft_vocab_size, + markov_rank, + bias=False, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "markov_w2"), + disable_tp=True, ) def embed(self, token_ids: torch.Tensor) -> torch.Tensor: """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r]).""" return self.markov_w1(token_ids) - def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor: + def bias( + self, + markov_embed: torch.Tensor, + logits_processor: LogitsProcessor, + ) -> torch.Tensor: """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).""" return logits_processor(self.markov_w2, markov_embed) @@ -89,6 +100,7 @@ def __init__( draft_vocab_size, config.markov_rank, prefix=maybe_prefix(prefix, "markov_head"), + quant_config=self.quant_config, ) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 16a2275cecd7..2c29f9cc616e 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -42,7 +42,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -201,7 +201,7 @@ def __init__( self.shared_expert_gate = None self.shared_expert = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_expert, gate=self.gate, num_experts=self.n_routed_experts, @@ -228,13 +228,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class + # In this case, the gate/router runs inside the MoERunner class final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: # Actually this will be dead code, since we always pass gate into - # FusedMoE in the current implementation. But we keep this code + # MoERunner in the current implementation. But we keep this code # here for clarity and future flexibility. router_logits, _ = self.gate(hidden_states) final_hidden_states = self.experts( diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 50bc4cce187c..b43aea8795f7 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -20,9 +20,7 @@ ) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate from vllm.model_executor.layers.layernorm import ( GemmaRMSNorm as Qwen3NextRMSNorm, @@ -175,7 +173,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.shared_expert", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_expert, gate=self.gate, num_experts=self.n_routed_experts, @@ -208,7 +206,7 @@ def forward( hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class + # In this case, the gate/router runs inside the MoERunner class final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 32a622567ceb..0afc6734afa8 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1530,52 +1530,6 @@ def _derive_audio_from_video_placeholders( result_placeholders["audio"] = audio_placeholders return result_placeholders - def _get_raw_input_ids( - self, - token_ids: list[int], - use_audio_in_video: bool = False, - ) -> list[int]: - tokenizer = self.info.get_tokenizer() - vision_bos_token = tokenizer.encode(tokenizer.vision_bos_token)[0] - vision_eos_token = tokenizer.encode(tokenizer.vision_eos_token)[0] - audio_bos_token = tokenizer.encode(tokenizer.audio_bos_token)[0] - audio_eos_token = tokenizer.encode(tokenizer.audio_eos_token)[0] - audio_token = tokenizer.encode("<|audio_pad|>")[0] - image_token = tokenizer.encode("<|image_pad|>")[0] - video_token = tokenizer.encode("<|video_pad|>")[0] - - result = token_ids[:] - if use_audio_in_video: - while True: - start = None - for i in range(len(result) - 1): - if result[i : i + 2] == [vision_bos_token, audio_bos_token]: - start = i - break - if start is not None: - end = None - for i in range(start + 2, len(result) - 1): - if result[i : i + 2] == [audio_eos_token, vision_eos_token]: - end = i - break - if end is not None: - result = ( - result[:start] - + [vision_bos_token, video_token, vision_eos_token] - + result[end + 2 :] - ) - else: - break - - for mm_token in [audio_token, image_token, video_token]: - compressed = [] - for x in result: - if x != mm_token or (not compressed or compressed[-1] != mm_token): - compressed.append(x) - result = compressed - - return result - class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin): def _process_audio_input( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index f86560e5f4e7..baf75fa2dd52 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -50,7 +50,12 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + MultiModalConfig, + VideoDummyOptions, + VideoPruningMethod, +) from vllm.distributed import get_pp_group, parallel_state from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger @@ -69,12 +74,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -92,6 +91,18 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count as vidcom2_compute_retained_tokens_count, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retention_mask as vidcom2_compute_retention_mask, +) from vllm.sequence import IntermediateTensors from vllm.tokenizers.protocol import TokenizerLike from vllm.tokenizers.registry import cached_tokenizer_from_config @@ -1256,7 +1267,7 @@ def _call_hf_processor( hf_config = self.info.get_hf_config() tokenizer = self.info.get_tokenizer() merge_size = hf_config.vision_config.spatial_merge_size - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() vision_start_token_id = hf_config.vision_start_token_id vision_end_token_id = hf_config.vision_end_token_id video_token_id = hf_config.video_token_id @@ -1339,11 +1350,18 @@ def _call_hf_processor( merge_size**2 ) - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + # Apply video pruning (EVS or VidCom2) if enabled. + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1459,16 +1477,22 @@ def get_video_replacement_qwen3vl(item_idx: int): f"video length ({grid_thw[0]})." ) - # Compute tokens per frame, with EVS support + # Compute tokens per frame, with EVS / VidCom2 support num_frames = int(grid_thw[0]) tokens_per_frame_base = int(grid_thw[1:].prod()) // merge_length - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1701,6 +1725,8 @@ class Qwen3VLForConditionalGeneration( supports_encoder_tp_data = True + supported_video_pruning_methods = ("evs", "vidcom2") + # To ensure correct weight loading and mapping. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -1719,6 +1745,17 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None: raise ValueError("Only image or video modality is supported") + def _init_video_pruning(self, multimodal_config: MultiModalConfig) -> None: + pruning_spec = multimodal_config.get_video_pruning_spec() + if pruning_spec is None: + self.video_pruning_method: VideoPruningMethod | None = None + self.video_pruning_rate = multimodal_config.video_pruning_rate + else: + self.video_pruning_method, self.video_pruning_rate = pruning_spec + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): super().__init__() config: Qwen3VLConfig = vllm_config.model_config.hf_config @@ -1730,10 +1767,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( @@ -1848,7 +1882,7 @@ def get_encoder_cudagraph_config(self): EncoderCudaGraphConfig, ) - # When EVS pruning is enabled, embed_multimodal post-processes both + # When video pruning is enabled, embed_multimodal post-processes both # image and video embeddings (mrope positions are appended for image, # prune+append for video). The encoder CUDA graph path bypasses that # post-process, producing inconsistent embedding formats vs eager. So @@ -2291,9 +2325,12 @@ def _postprocess_video_embeds_evs( t, h, w = size if self.is_multimodal_pruning_enabled: - # For each video, compute retention mask using EVS. - # retention_mask: [11424]. - retention_mask = compute_retention_mask( + # Compute the retention mask for each video (EVS or VidCom2). + if self.video_pruning_method == "vidcom2": + mask_fn = vidcom2_compute_retention_mask + else: + mask_fn = compute_retention_mask + retention_mask = mask_fn( emb, size, spatial_merge_size=self.visual.spatial_merge_size, diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 4413e1213bb6..f1409d23399e 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -220,10 +220,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 7ec320c5348d..213083392df4 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -331,54 +331,6 @@ def num_registers(self): def num_skip(self): return self.num_cls_tokens + self.num_registers - def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter): - if src_embed.shape != targ_embed.shape: - src_size = int(math.sqrt(src_embed.shape[1])) - - assert src_size**2 == src_embed.shape[1], ( - "Unable to interpolate non-square embedding" - ) - - src_embed = rearrange( - src_embed, "b (h w) c -> b c h w", h=src_size, w=src_size - ) - src_embed = F.interpolate( - src_embed, - size=(self.num_rows, self.num_cols), - mode="bicubic", - align_corners=True, - antialias=False, - ) - src_embed = rearrange(src_embed, "b c h w -> b (h w) c") - targ_embed.data.copy_(src_embed) - - def _load_projection( - self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor - ): - if src_proj_weight.shape != targ_proj_weight.shape: - src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3)) - - assert (src_patch_size**2) * 3 == src_proj_weight.shape[1], ( - "Unable to interpolate non-square patch size" - ) - - src_proj_weight = rearrange( - src_proj_weight, - "b (c h w) -> b c h w", - c=3, - h=src_patch_size, - w=src_patch_size, - ) - src_proj_weight = F.interpolate( - src_proj_weight, - size=(self.patch_size, self.patch_size), - mode="bicubic", - align_corners=True, - antialias=False, - ) - src_proj_weight = rearrange(src_proj_weight, "b c h w -> b (c h w)") - targ_proj_weight.data.copy_(src_proj_weight) - def embed_patches(self, x: torch.Tensor) -> torch.Tensor: patches = self.im_to_patches(x) patches = self.embedder(patches) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d28aac974b54..545b257794d0 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -55,6 +55,7 @@ supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, supports_pp, + supports_replayssm, supports_transcription, ) from .interfaces_base import ( @@ -96,7 +97,7 @@ "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), - "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "ExaoneMoeForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"), "FalconForCausalLM": ("falcon", "FalconForCausalLM"), "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), @@ -136,7 +137,10 @@ "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), "JambaForCausalLM": ("jamba", "JambaForCausalLM"), - "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), + "KimiLinearForCausalLM": ( + "vllm.models.kimi_k3", + "KimiLinearForCausalLM", + ), "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"), "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"), "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"), @@ -182,7 +186,6 @@ "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), "OPTForCausalLM": ("opt", "OPTForCausalLM"), "OrionForCausalLM": ("orion", "OrionForCausalLM"), - "OuroForCausalLM": ("ouro", "OuroForCausalLM"), "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"), "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), @@ -190,12 +193,13 @@ "PhiForCausalLM": ("phi", "PhiForCausalLM"), "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), - "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"), "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), + "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"), + "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"), "RWForCausalLM": ("falcon", "FalconForCausalLM"), "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), @@ -459,6 +463,10 @@ ), "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), + "KimiK3ForConditionalGeneration": ( + "vllm.models.kimi_k3", + "KimiK3ForConditionalGeneration", + ), "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), "MossTranscribeDiarizeForConditionalGeneration": ( "moss_transcribe_diarize", @@ -608,6 +616,10 @@ "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), + "K3DSparkModel": ( + "vllm.models.kimi_k3.nvidia.dspark_mla", + "K3DSparkForCausalLM", + ), "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"), "Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"), "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"), @@ -648,6 +660,7 @@ "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"), + "KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"), # Temporarily disabled. # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1. # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"), @@ -660,11 +673,16 @@ "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"), + "VaultGemmaForCausalLM": ("transformers", "TransformersForCausalLM"), # Multimodal models "Emu3ForConditionalGeneration": ( "transformers", "TransformersMultiModalForCausalLM", ), + "VibeVoiceAsrForConditionalGeneration": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), } _TRANSFORMERS_BACKEND_MODELS = { @@ -759,6 +777,8 @@ "TeleChatForCausalLM": "0.25.0", "PersimmonForCausalLM": "0.25.0", "FuyuForCausalLM": "0.25.0", + "Plamo2ForCausalLM": "0.26.0", + "OuroForCausalLM": "0.26.0", } _OOT_SUPPORTED_MODELS = { @@ -788,8 +808,10 @@ class _ModelInfo: is_hybrid: bool has_noops: bool supports_mamba_prefix_caching: bool + supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool + supported_video_pruning_methods: tuple[str, ...] @staticmethod def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": @@ -814,11 +836,15 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), + supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( supports_transcription(model) and model.supports_transcription_only ), has_noops=has_noops(model), + supported_video_pruning_methods=getattr( + model, "supported_video_pruning_methods", () + ), ) @@ -1068,7 +1094,7 @@ def register_model( else: msg = ( "`model_cls` should be a string or PyTorch model class, " - f"not a {type(model_arch)}" + f"not a {type(model_cls)}" ) raise TypeError(msg) @@ -1326,88 +1352,88 @@ def is_text_generation_model( architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.is_text_generation_model + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.is_text_generation_model def is_pooling_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.is_pooling_model + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.is_pooling_model def is_multimodal_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.supports_multimodal + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_multimodal def is_multimodal_raw_input_only_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.supports_multimodal_raw_input_only + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_multimodal_raw_input_only def is_pp_supported_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.supports_pp + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_pp def model_has_inner_state( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.has_inner_state + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.has_inner_state def is_attention_free_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.is_attention_free + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.is_attention_free def is_hybrid_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.is_hybrid + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.is_hybrid def is_noops_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.has_noops + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.has_noops def is_transcription_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.supports_transcription + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_transcription def is_transcription_only_model( self, architectures: str | list[str], model_config: ModelConfig, ) -> bool: - model_cls, _ = self.inspect_model_cls(architectures, model_config) - return model_cls.supports_transcription_only + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_transcription_only def _resolve_module_name(mod_relname: str) -> str: diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index b9f9532ae1b1..2b2fbfe70897 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -35,7 +35,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -334,7 +334,7 @@ def __init__( else: self.shared_experts = None - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, num_experts=self.num_experts, top_k=self.top_k, diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 7fb5a9170599..446e66ca8bf4 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -19,7 +19,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -68,7 +68,7 @@ def __init__( f"the number of experts {config.moe_num_experts}." ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.moe_num_experts, top_k=config.moe_top_k, hidden_size=config.hidden_size, diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 7b3bb93ad116..943549c46e32 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -725,6 +725,7 @@ def embed_input_ids( def get_encoder_cudagraph_config(self): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, + EncoderCudaGraphPathConfig, ) return EncoderCudaGraphConfig( @@ -734,9 +735,15 @@ def get_encoder_cudagraph_config(self): "patch_pixel_values", ], out_hidden_size=self.config.hidden_size, - enable_dual_path_graph=True, - global_token_per_image=self.img_output_tokens, - local_token_per_patch=self.patch_output_tokens, + paths={ + "global": EncoderCudaGraphPathConfig( + min_token_budget=self.img_output_tokens + ), + "local": EncoderCudaGraphPathConfig( + min_token_budget=self.patch_output_tokens, + allow_zero_tokens=True, + ), + }, ) def get_encoder_cudagraph_budget_range( @@ -771,8 +778,10 @@ def get_encoder_cudagraph_item_specs( output_tokens=( self.img_output_tokens + num_patch * self.patch_output_tokens ), - global_output_tokens=self.img_output_tokens, - local_output_tokens=num_patch * self.patch_output_tokens, + path_output_tokens={ + "global": self.img_output_tokens, + "local": num_patch * self.patch_output_tokens, + }, ) for num_patch in num_patches ] @@ -876,19 +885,21 @@ def encoder_eager_forward( def postprocess_encoder_output( self, - output: torch.Tensor, + outputs: dict[str, torch.Tensor], indices: list[int], per_item_out_tokens: list[int], dest: dict[int, torch.Tensor] | list[torch.Tensor | None], clone: bool = False, batch_mm_kwargs: dict[str, Any] | None = None, - local_output: torch.Tensor | None = None, ): """CPU-side per-item merge after dual-path graph replay. - ``output`` contains global-image features and ``local_output`` + ``outputs['global']`` contains global-image features and ``outputs['local']`` contains local-patch features (or ``None`` when there are no patches). """ + output = outputs["global"] + local_output = outputs.get("local") + assert batch_mm_kwargs is not None num_patches = batch_mm_kwargs["num_patches"] hidden = output.shape[-1] bsz = len(indices) @@ -899,7 +910,7 @@ def postprocess_encoder_output( patch_tokens = total_patches * self.patch_output_tokens global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden) - if total_patches > 0: + if total_patches > 0 and local_output is not None: patch_part = local_output[:patch_tokens].reshape( -1, self.patch_output_tokens, hidden ) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 07a25d23c8c4..0a699ad24179 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, MoERunner, fused_moe_make_expert_params_mapping, ) @@ -271,11 +271,11 @@ def forward( q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Add qk-norm inline similar to Qwen3 MOE attention q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) - q_by_head = self.q_norm(q_by_head.contiguous()) + q_by_head = self.q_norm(q_by_head) q = q_by_head.view(q.shape) k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) - k_by_head = self.k_norm(k_by_head.contiguous()) + k_by_head = self.k_norm(k_by_head) k = k_by_head.view(k.shape) if self.use_rope: q, k = self.rotary_emb(positions, q, k) @@ -376,7 +376,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.share_expert", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.share_expert, gate=self.gate, num_experts=config.moe_num_experts, @@ -404,7 +404,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states=hidden_states, router_logits=hidden_states ) else: - # TODO(bnell): this gate could be moved into the FusedMoE? + # TODO(bnell): this gate could be moved into the MoERunner? router_logits, _ = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits diff --git a/vllm/model_executor/models/step3p5_mtp.py b/vllm/model_executor/models/step3p5_mtp.py index b533a9111dcf..87b7aab49106 100644 --- a/vllm/model_executor/models/step3p5_mtp.py +++ b/vllm/model_executor/models/step3p5_mtp.py @@ -15,6 +15,9 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors @@ -29,11 +32,18 @@ def __init__( self, config: PretrainedConfig, quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: super().__init__() self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) + # Give the head its prefix so the quant config's exclude_modules matcher + # can skip it; without one it defaults to "" and never matches, so a + # checkpoint-excluded (BF16) MTP head gets quantized -> load crash. self.head = ParallelLMHead( - config.vocab_size, config.hidden_size, quant_config=quant_config + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.head", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -52,7 +62,9 @@ def __init__( self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) - self.shared_head = SharedHead(config=config, quant_config=quant_config) + self.shared_head = SharedHead( + config=config, quant_config=quant_config, prefix=f"{prefix}.shared_head" + ) self.mtp_block = Step3p5DecoderLayer( vllm_config, prefix=f"{prefix}.mtp_block", @@ -283,7 +295,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: and getattr(param, "requires_grad", False) is False } params_need_to_load -= optional_params - if params_need_to_load != loaded_params: + if params_need_to_load != loaded_params and is_mtp_completeness_check_enabled(): missing_params = list(params_need_to_load - loaded_params) param_name_example = missing_params[0] raise RuntimeError( diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index e863b0bb51c7..b49f7827dafe 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -217,7 +217,10 @@ def apply( ) with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) mm_placeholders = {"image": [PlaceholderRange(offset=0, length=0)]} diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 78a12876e662..ff4bb6cd433f 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING +import torch.nn.functional as F from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from vllm.model_executor.models.transformers.base import Base @@ -59,9 +60,22 @@ def vllm_attention_forward( if scaling is not None: self_attn.impl.scale = float(scaling) hidden = query.shape[-2] + head_dim_qk = query.shape[-1] + head_dim_v = value.shape[-1] 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 + # Pad `value` up to the query/key head size when it is smaller (expanded + # MLA). A larger last dim just means `value` isn't split per head, e.g. + # packed grouped/multi-query projections, and needs no padding. + pad_value = head_dim_v < head_dim_qk + if pad_value: + value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v)) + value = value.reshape(hidden, -1) + attn_output = self_attn.forward(query, key, value) + if pad_value: + attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v] + attn_output = attn_output.reshape(hidden, -1) + return attn_output, None ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 24cad1da20d2..35e5c9d1a0f8 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -104,6 +104,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): super().__init__() logger.info("Using Transformers modeling backend.") + self.vllm_config = vllm_config self.config = vllm_config.model_config.hf_config self.text_config = self.config.get_text_config() self.cache_config = vllm_config.cache_config @@ -357,7 +358,7 @@ def pipeline_parallel(self): tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - logger.warning( + logger.warning_once( "%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 " @@ -452,7 +453,7 @@ def recursive_replace(self): # 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) + fusers = Fusers(self.model, self.vllm_config) def register_fusion(fuser: BaseFuser, prefix: str): """Register a fused layer's mappings just before it is built.""" @@ -503,9 +504,7 @@ def _recursive_replace(module: nn.Module, prefix: str): new_module = replace_conv_class(child_module) 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 - ) + new_module = fuser.fuse(child_module, qual_name, self.vllm_config) logger.info_once(fuser.info(child_name)) _recursive_replace(new_module, prefix=qual_name) elif not isinstance(child_module, MoERunner): diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index 0aa7c419ceb8..cbcb24e2a95d 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/vllm/model_executor/models/transformers/fuser.py @@ -18,21 +18,27 @@ from vllm.model_executor.models.transformers.fusers import ( BaseFuser, GLUFuser, + PackedQKVFuser, QKVFuser, + RewriteFuser, RMSNormFuser, - StackedFuser, ) from vllm.model_executor.models.transformers.fx_utils import trace if TYPE_CHECKING: - from vllm.config.model import ModelConfig + from vllm.config import VllmConfig logger = init_logger(__name__) -@cached(cache={}, key=type) +def key(module: nn.Module) -> tuple: + """Cache key for `get_fuser`. Considers module type and its immediate children.""" + return (type(module), tuple(name for name, _ in module.named_children())) + + +@cached(cache={}, key=key) def get_fuser(module: nn.Module) -> BaseFuser | None: - """The fuser for `type(module)` (cached per class), or `None` if no match.""" + """The fuser for `module`'s class and shape (cached), 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. @@ -42,17 +48,20 @@ def get_fuser(module: nn.Module) -> BaseFuser | None: return None if (graph := trace(module)) is None: return None - for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser): + for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser): if (fuser := fuser_cls.match(graph, module)) is not None: - if isinstance(fuser, StackedFuser): + if isinstance(fuser, RewriteFuser): 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 + "Attempted to fuse %s using %s but failed " + "to update its forward method: %s", + type(module), + fuser_cls.__name__, + exc, ) - return None + continue return fuser # A norm we could not match structurally is left unfused; flag likely misses. if module.__class__.__name__.endswith("RMSNorm"): @@ -67,12 +76,12 @@ def get_fuser(module: nn.Module) -> BaseFuser | 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 + def __init__(self, model: nn.Module, vllm_config: "VllmConfig"): + self.vllm_config = vllm_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): + if fuser is not None and fuser.validate(m, self.vllm_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 index 58910b0ecc3c..ed4da5155cf1 100644 --- a/vllm/model_executor/models/transformers/fusers/__init__.py +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -2,17 +2,24 @@ # 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.base import ( + BaseFuser, + RewriteFuser, + 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.packed_qkv import PackedQKVFuser from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser __all__ = [ "BaseFuser", + "RewriteFuser", "StackedFuser", "GLUFuser", "MoEBlockFuser", + "PackedQKVFuser", "QKVFuser", "RMSNormFuser", ] diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index 54fb2d09165c..920e63101ac6 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -13,8 +13,7 @@ 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 + from vllm.config import VllmConfig @dataclass @@ -36,16 +35,12 @@ 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: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> 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", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Apply the fusion to an already-validated `module`, returning the module to install in its place (mutated in place, or freshly built).""" @@ -62,27 +57,61 @@ def packed_modules_mapping(self) -> dict[str, list[str]]: return {} +def local_output_sizes(merged_name: str) -> str: + """Source for the per-rank widths of the merged linear `self.`.""" + merged = f"self.{merged_name}" + return f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + + @dataclass -class StackedFuser(BaseFuser): - """A fuser that merges sibling projections into one stacked linear and - rewrites the forward to call it. +class RewriteFuser(BaseFuser): + """A fuser that rewrites the module's forward and rebinds 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. + `match` and `update_forward` analyse the class once; `fuse` swaps the + submodules 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`.""" + @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, vllm_config: "VllmConfig" + ) -> None: + """Replace `module`'s submodules with their vLLM equivalents.""" + + def fuse( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> 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, vllm_config) + module.forward = types.MethodType(self.fused_forward, module) + return module + + +@dataclass +class StackedFuser(RewriteFuser): + """A fuser that merges sibling projections into one stacked linear and + rewrites the forward to call it.""" + + 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).""" + def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( @@ -113,34 +142,3 @@ 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 index 951eb35777ae..2da84addea2f 100644 --- a/vllm/model_executor/models/transformers/fusers/glu.py +++ b/vllm/model_executor/models/transformers/fusers/glu.py @@ -33,8 +33,7 @@ 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 + from vllm.config import VllmConfig logger = init_logger(__name__) @@ -168,7 +167,7 @@ def update_forward(self, module: nn.Module) -> None: replace_expr(funcdef, muls[0], act_call) self.fused_forward = compile_forward(funcdef, fn) - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> 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)) @@ -176,12 +175,9 @@ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: return True def update_attrs( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: + quant_config = vllm_config.quant_config 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) diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index 6a3c7e85d6e5..af3a101a0ce1 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -6,14 +6,14 @@ import inspect import textwrap import types -from collections.abc import Iterator +from collections.abc import Iterable, 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.fused_moe import GateLinear from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, @@ -21,14 +21,10 @@ peel, trace, ) +from vllm.model_executor.models.transformers.utils import named_state 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)) @@ -79,6 +75,25 @@ def _is_scalar_gate(module: nn.Module) -> bool: ) +def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None: + """The floating dtype `nodes` cast to, if any. + + Computations that must run in higher precision say so in their forward (e.g. + `hidden_states.type(torch.float32)`), so the dtype is readable from the graph + even when the config does not name it.""" + for node in nodes: + if node.op not in ("call_method", "call_function"): + continue + name = str(node.target).rsplit(".", 1)[-1] + if name == "float": + return torch.float32 + if name in ("to", "type"): + for arg in (*node.args[1:], *node.kwargs.values()): + if isinstance(arg, torch.dtype) and arg.is_floating_point: + return arg + return None + + 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() @@ -131,18 +146,24 @@ class MoEBlockFuser: scoring_func: str shared_name: str | None shared_gate_name: str | None + router_dtype: torch.dtype | None = 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"]: + def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: + """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. + + Returns the scoring function and the dtype the router computes in.""" + state = {name for name, _ in named_state(gate)} + if "weight" not in state or state - {"weight", "e_score_correction_bias"}: 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: + # The routing top-k is the last one; any earlier one scores expert groups. + topks = [node for node in graph.nodes if is_op(node, "topk")] + if not topks: return None + topk = topks[-1] # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. scorers = [ n @@ -152,9 +173,11 @@ def _match_router(gate: nn.Module) -> str | None: if len(scorers) != 1: return None scorer = scorers[0] - if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")): + logits_cone = _reaches(scorer, "all_input_nodes") + if not any(is_op(n, "linear") for n in logits_cone): return None - return "softmax" if is_op(scorer, "softmax") else "sigmoid" + scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" + return scoring_func, _forced_dtype(logits_cone) @staticmethod def _match_shared_experts( @@ -198,10 +221,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None if _returns_tuple(type(moe_block)): return None # Router: the child that scores + top-k selects. - gate_name = scoring_func = None + gate_name = scoring_func = router_dtype = 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 + if ( + name != experts_name + and (router := cls._match_router(child)) is not None + ): + gate_name = name + scoring_func, router_dtype = router break if gate_name is None or scoring_func is None: return None @@ -229,17 +256,23 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None 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) + return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) - 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( + def gate( + self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None + ) -> GateLinear: + """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE.""" + hf_gate = getattr(moe_block, self.gate_name) + num_experts, hidden_size = hf_gate.weight.shape + gate = GateLinear( hidden_size, num_experts, bias=False, + out_dtype=out_dtype or self.router_dtype, prefix=maybe_prefix(prefix, self.gate_name), ) + if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None: + gate.register_buffer("e_score_correction_bias", bias) setattr(moe_block, self.gate_name, gate) return gate diff --git a/vllm/model_executor/models/transformers/fusers/packed_qkv.py b/vllm/model_executor/models/transformers/fusers/packed_qkv.py new file mode 100644 index 000000000000..ae5612fa816b --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/packed_qkv.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed-QKV fuser: `c_attn(x).split((q, kv, kv))` -> a `QKVParallelLinear`.""" + +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING + +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 ( + RewriteFuser, + local_output_sizes, +) +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + is_method, + recover_forward, + returned_linear, + upstream_linear, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import maybe_prefix + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@dataclass +class PackedQKVFuser(RewriteFuser): + """Fuser for attention with q, k and v packed into one projection.""" + + qkv_name: str + o_name: str | None + q_size: int + kv_size: int + + def info(self, name: str) -> str: + return ( + f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear" + ) + + @staticmethod + def _packed_sizes(node: fx.Node) -> tuple[int, int] | None: + """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one.""" + if not is_method(node, "split") or len(node.args) < 2: + return None + sizes = node.args[1] + if not isinstance(sizes, (tuple, list)) or len(sizes) != 3: + return None + if not all(isinstance(size, int) for size in sizes): + return None + q_size, k_size, v_size = sizes + if k_size != v_size or q_size < k_size: + return None + return q_size, k_size + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None": + for node in graph.nodes: + if (sizes := cls._packed_sizes(node)) is None: + continue + q_size, kv_size = sizes + qkv_node = upstream_linear(node.args[0], module) + if qkv_node is None: + continue + qkv_name = str(qkv_node.target) + # The split must consume the whole projection. + if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size: + continue + # o_proj produces the module's output and consumes the query width. + o_name = returned_linear(graph, module) + if o_name == qkv_name or ( + o_name is not None + and module.get_submodule(o_name).in_features != q_size + ): + o_name = None + return cls( + source_cls=type(module).__name__, + qkv_name=qkv_name, + o_name=o_name, + q_size=q_size, + kv_size=kv_size, + ) + return None + + def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call: + """The unique `self.(...)....split((a, b, c), ...)` call.""" + calls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "split" + and node.args + and isinstance(node.args[0], (ast.Tuple, ast.List)) + and len(node.args[0].elts) == 3 + and any( + isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name + for inner in ast.walk(node.func.value) + ) + ] + if len(calls) != 1: + raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits") + return calls[0] + + def update_forward(self, module: nn.Module) -> None: + """Rewrite the split sizes to the sharded projection's per-rank widths.""" + funcdef, fn = recover_forward(type(module)) + split = self._split_call(funcdef) + # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes] + sections = local_output_sizes(self.qkv_name) + split.args[0] = ast.parse(sections, mode="eval").body + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: + """Shapes must be compatible with a head-sharded packed GEMM.""" + head_size = vllm_config.model_config.get_head_size() + qkv = module.get_submodule(self.qkv_name) + compatible = ( + self.q_size % head_size == 0 + and self.kv_size % head_size == 0 + and qkv.out_features == self.q_size + 2 * self.kv_size + ) + if not compatible: + logger.debug("%s is not compatible with packed QKV fusion", type(module)) + return compatible + + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + quant_config = vllm_config.quant_config + head_size = vllm_config.model_config.get_head_size() + qkv_prefix = maybe_prefix(prefix, self.qkv_name) + qkv = module.get_submodule(self.qkv_name) + merged = QKVParallelLinear( + hidden_size=qkv.in_features, + head_size=head_size, + total_num_heads=self.q_size // head_size, + total_num_kv_heads=self.kv_size // head_size, + bias=qkv.bias is not None, + quant_config=quant_config, + prefix=qkv_prefix, + return_bias=False, + ) + setattr(module, self.qkv_name, merged) + log_replacement(qkv_prefix, qkv, merged) + # If there is an output projection, we know it must be rowwise. + if self.o_name is not None: + o_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_prefix + ) + setattr(module, self.o_name, new_o) + log_replacement(o_prefix, o_proj, new_o) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index 010a0018acc5..370c0cdf4fa5 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -10,13 +10,17 @@ 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.fusers.base import ( + StackedFuser, + local_output_sizes, +) from vllm.model_executor.models.transformers.fx_utils import ( compile_forward, innermost_block, is_linear, recover_forward, replace_expr, + returned_linear, single_self_call, ) from vllm.model_executor.models.transformers.utils import ( @@ -26,8 +30,7 @@ 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 + from vllm.config import VllmConfig logger = init_logger(__name__) @@ -84,15 +87,16 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | 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 + # o_proj produces the module's output. + o_name = returned_linear(graph, module) + # o_proj must be compatible with the q/k/v projections. + if o_name in names.values() or ( + o_name is not None + and module.get_submodule(o_name).in_features + != module.get_submodule(q.target).out_features + ): + o_name = None + names["o_name"] = o_name return cls(source_cls=type(module).__name__, **names) def update_forward(self, module: nn.Module) -> None: @@ -133,7 +137,7 @@ def update_forward(self, module: nn.Module) -> None: if names & set(temps): raise ValueError("fused temporaries would shadow existing names") merged = f"self.{self.merged_name}" - sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + sections = local_output_sizes(self.merged_name) template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)" assign = ast.parse(template).body[0] arg = next( @@ -149,12 +153,12 @@ def update_forward(self, module: nn.Module) -> None: 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: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> 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() + head_size = vllm_config.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 @@ -167,13 +171,10 @@ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: return compatible def update_attrs( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: - head_size = model_config.get_head_size() + quant_config = vllm_config.quant_config + head_size = vllm_config.model_config.get_head_size() q = module.get_submodule(self.q_name) k = module.get_submodule(self.k_name) merged = QKVParallelLinear( diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index 829fd9a541c8..781b417177aa 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -21,13 +21,13 @@ find_node, forward_input_count, is_op, + output_value, peel, trace, ) if TYPE_CHECKING: - from vllm.config.model import ModelConfig - from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.config import VllmConfig def _is_squared(node: object, x: fx.Node) -> bool: @@ -71,10 +71,8 @@ def _is_one_plus(node: object) -> bool: 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 + value = output_value(graph) + return value is not None and peel(value) is not node class TPAwareNormMixin(nn.Module): @@ -185,17 +183,14 @@ def _eps_from_graph(graph: fx.Graph) -> float | None: return eps return None - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: return True def fuse( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" + model_config = vllm_config.model_config weight = getattr(module, "weight", None) hidden_size = ( weight.size(0) if weight is not None else model_config.get_hidden_size() diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 0e043941d8ad..d665bcfe8cf5 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -9,10 +9,13 @@ """ import ast +import contextlib import inspect import operator import textwrap from collections.abc import Callable +from itertools import chain +from unittest import mock import torch from torch import fx, nn @@ -22,46 +25,64 @@ logger = init_logger(__name__) +_UNKNOWN = object() +"""Sentinel meta value for proxies whose concrete value could not be inferred. +Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`).""" -def _infer_len(node: fx.Node) -> int | None: - """Concrete length of a proxy's value, inferred from its node chain. +_MODULE_CALL = nn.Module.__call__ +"""The unpatched `nn.Module.__call__`. During tracing fx patches it to record +`call_module` nodes; meta execution must call modules for real.""" + + +def is_leaf_call(node: object) -> bool: + """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" + return isinstance(node, fx.Node) and node.meta.get("leaf_call", False) - 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 _reference_weight(module: nn.Module) -> torch.Tensor | None: + """A weight whose trailing dim is the module's hidden size. -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 + Linears and 2-D gate weights are `[out, hidden]`; norm weights are + `[hidden]`. Used to fabricate a placeholder input of matching size/dtype.""" + for child in module.modules(): + if isinstance(child, nn.Linear): + return child.weight + for param in module.parameters(): + if param.ndim in (1, 2): + return param return None -class _SizedProxy(fx.Proxy): - """Proxy whose `len` is inferred from the graph (see `_infer_len`).""" +class _MetaProxy(fx.Proxy): + """Proxy carrying the meta-tensor value of the traced expression. + + Shape questions (`len`, iteration, `.shape` unpacks) are answered by + executing each op on the meta values, so PyTorch's meta kernels are the + single source of shape inference — no per-op rules.""" + + meta: object = _UNKNOWN def __len__(self) -> int: - length = _infer_len(self.node) - if length is None: - return super().__len__() - return length + if self.meta is not _UNKNOWN: + return len(self.meta) + return super().__len__() # type: ignore[misc] + + def __getattr__(self, k: str) -> "_MetaAttribute": + return _MetaAttribute(self, k) + + +class _MetaAttribute(_MetaProxy, fx.proxy.Attribute): + """Attribute proxy (e.g. `x.shape`) carrying its meta value. + + `Proxy.__getattr__` constructs `Attribute` directly, bypassing + `Tracer.proxy`, so the meta value must be grafted on here too.""" + + def __init__(self, root: fx.Proxy, attr: str): + super().__init__(root, attr) + root_meta = getattr(root, "meta", _UNKNOWN) + if root_meta is not _UNKNOWN: + with contextlib.suppress(Exception): + self.meta = getattr(root_meta, attr) class _AllLeafTracer(fx.Tracer): @@ -69,31 +90,159 @@ class _AllLeafTracer(fx.Tracer): 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. + Every traced op is also executed on meta tensors (see `_MetaProxy`) so + shape unpacks and `*`-splats trace through; anything else untraceable ends + the trace early and the partial graph is matched. """ + varkw: str | None = None + """Name of the traced forward's `**kwargs` parameter, if any.""" + 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) + return _MetaProxy(node, self) + + def create_proxy(self, kind, target, args, kwargs, *extra, **extra_kwargs): + proxy = super().create_proxy(kind, target, args, kwargs, *extra, **extra_kwargs) + if isinstance(proxy, _MetaProxy) and proxy.meta is _UNKNOWN: + # A failure stays _UNKNOWN; only fatal if a shape question is asked. + with contextlib.suppress(Exception): + proxy.meta = self._infer_meta(kind, target, args, kwargs) + return proxy + + def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): + """Execute the op on meta tensors; PyTorch infers the output value.""" + if kind == "placeholder": + # vLLM always feeds the model [1, seq_len, hidden_size] hidden states. + weight = _reference_weight(self.root) + if str(target) == "hidden_states" and weight is not None: + return torch.empty( + 1, 8, weight.shape[-1], dtype=weight.dtype, device="meta" + ) + return _UNKNOWN + if kind == "get_attr": + value = operator.attrgetter(str(target))(self.root) + if isinstance(value, torch.Tensor): + value = torch.empty_like(value, device="meta") + return value + unknown = False + + def meta_of(arg: object) -> object: + nonlocal unknown + if isinstance(arg, fx.Proxy): + meta = getattr(arg, "meta", _UNKNOWN) + unknown = unknown or meta is _UNKNOWN + return meta + return arg + + meta_args = fx.node.map_aggregate(args, meta_of) + meta_kwargs = fx.node.map_aggregate(kwargs, meta_of) + if unknown: + return _UNKNOWN + if kind == "call_function": + return target(*meta_args, **meta_kwargs) + if kind == "call_method": + receiver, *rest = meta_args + return getattr(receiver, str(target))(*rest, **meta_kwargs) + if kind == "call_module": + # Run the child's forward with all its state on "meta", without + # mutating it (at match time params may be meta but buffers real). + # fx patches `nn.Module.__call__` while tracing; restore the real + # one so this execution is not itself recorded. + child = self.root.get_submodule(str(target)) + state = { + name: torch.empty_like(tensor, device="meta") + for name, tensor in chain( + child.named_parameters(), child.named_buffers() + ) + } + with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): + return torch.func.functional_call(child, state, meta_args, meta_kwargs) + return _UNKNOWN + + def _is_varkw(self, node: object) -> bool: + return ( + isinstance(node, fx.Node) + and node.op == "placeholder" + and str(node.target).lstrip("*") == self.varkw + ) def iter(self, obj: fx.Proxy): - length = _infer_len(obj.node) - if length is None: + # Assume kwargs is always empty to simplify tracing. + node = obj.node + if self._is_varkw(node) or ( + node.op == "call_method" + and node.target == "keys" + and self._is_varkw(node.args[0]) + ): + return iter(()) + meta = getattr(obj, "meta", _UNKNOWN) + if meta is _UNKNOWN: return super().iter(obj) - return iter([obj[i] for i in range(length)]) + return iter([obj[i] for i in range(len(meta))]) -def trace(module: nn.Module) -> fx.Graph | None: - """Trace `module.forward`, returning the partial graph on failure. +def _as_leaf_call(fn: Callable, length: int | None = None) -> Callable: + """Wrap any callable so tracing records it as one opaque `call_function` node. + + Lets the trace continue past untraceable bodies. Only the proxy arguments carry into + the node's dataflow; the rest are dropped rather than lifted into the graph. + `length` declares how many values the callable returns, so unpacking its result also + traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough. + """ + + def leaf(*args, **kwargs): + proxies = tuple(arg for arg in args if isinstance(arg, fx.Proxy)) + if not proxies: + return fn(*args, **kwargs) + proxy = proxies[0].tracer.create_proxy("call_function", fn, proxies, {}) + proxy.node.meta["leaf_call"] = True + if length is not None: + # The body never executes, so fabricate a value of the declared length. + proxy.meta = (_UNKNOWN,) * length + return proxy + + return leaf + + +def _leaf_attention_interfaces(): + """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node. - 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.""" + `vllm_attention_function` needs runtime context so it is untraceable. + Every interface returns `(attn_output, attn_weights)`.""" + from transformers.modeling_utils import AttentionInterface + + original = AttentionInterface.get_interface + + def get_interface(self, *args, **kwargs): + return _as_leaf_call(original(self, *args, **kwargs), length=2) + + return mock.patch.object(AttentionInterface, "get_interface", get_interface) + + +def trace(module: nn.Module) -> fx.Graph | None: + """Trace `module.forward`, returning the partial graph on failure.""" + parameters = forward_parameters(type(module)) + # vLLM never passes `past_key_values` so it is always the default value of `None`. + # Make this concrete to simplify tracing. + concrete_args = None + if "past_key_values" in parameters: + concrete_args = {"past_key_values": None} + # Get the name of the kwargs parameter passed to module.forward (usually "kwargs") tracer = _AllLeafTracer() + tracer.varkw = next( + ( + p.name + for p in parameters.values() + if p.kind is inspect.Parameter.VAR_KEYWORD + ), + None, + ) try: - return tracer.trace(module) + with _leaf_attention_interfaces(): + return tracer.trace(module, concrete_args=concrete_args) except Exception as exc: logger.debug("Could not fully trace %s: %s", type(module), exc) return getattr(tracer, "graph", None) @@ -129,20 +278,27 @@ def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]: return funcdef, fn +def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]: + """`cls.forward`'s signature parameters, or empty if uninspectable.""" + try: + return dict(inspect.signature(cls.forward).parameters) + except (TypeError, ValueError): + return {} + + 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): + params = list(forward_parameters(cls).values()) + if not params: 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) + return sum(1 for p in params[1:] if p.kind in fixed) def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable: @@ -227,6 +383,48 @@ def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node return next((n for n in graph.nodes if predicate(n)), None) +def output_value(graph: fx.Graph) -> object | None: + """The value the graph's `output` node returns, if the trace reached one.""" + output = find_node(graph, lambda n: n.op == "output") + if output is None or not output.args: + return None + return output.args[0] + + +def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: + """Nearest linear producing `node`, walking back through splits/reshapes. + + Non-linear submodules are transparent too (e.g. the dropout GPT-style + attentions apply after their output projection). Never walks through a leaf + call (e.g. an attention interface): its inputs are what attention consumes, + not what produced the value.""" + stack = [node] + seen: set[fx.Node] = set() + while stack: + current = stack.pop() + if not isinstance(current, fx.Node) or current in seen: + continue + seen.add(current) + if is_linear(current, module): + return current + if current.op in ( + "call_function", + "call_method", + "call_module", + ) and not is_leaf_call(current): + stack.extend(current.args) + return None + + +def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None: + """Name of the Linear producing the graph's (first) output value.""" + value = output_value(graph) + if isinstance(value, (tuple, list)) and value: + value = value[0] + linear = upstream_linear(value, module) + return None if linear is None else str(linear.target) + + def is_linear(node: fx.Node, module: nn.Module) -> bool: """Is node `nn.Linear.__call__()`.""" return node.op == "call_module" and isinstance( diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index d1f5dba03733..e1c3733a0f19 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -23,16 +23,22 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops 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.layers.fused_moe import ( + FusedMoEFactory, + MoERunner, + RoutedExperts, +) from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.model_executor.models.transformers.fuser import get_fuser 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 +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op from .utils import log_replacement @@ -51,7 +57,7 @@ class TransformersMoEState: # --8<-- [start:transformers_fused_moe] @PluggableLayer.register("transformers_fused_moe") class TransformersMoERunner(MoERunner): - """Custom FusedMoE for the Transformers modeling backend.""" + """Custom MoERunner for the Transformers modeling backend.""" # --8<-- [end:transformers_fused_moe] def __init__(self, *args, moe_state: TransformersMoEState, **kwargs): @@ -176,17 +182,31 @@ def recursive_replace(self): 0, ) - # Unused kwargs since we use custom_routing_function: - # - `scoring_func` and `e_score_correction_bias` only used for grouped - # topk routing inside vLLM and are non-trivial to infer - # and hard code `use_grouped_topk=False` - # - `renormalize` passed anyway because it's easy to infer - # - `num_expert_group` and `topk_group` used for inferring expert - # placement strategy in FusedMoE - # - `apply_router_weight_on_input` is already applied in Transformers + # Common kwargs renormalize = getattr(text_config, "norm_topk_prob", top_k > 1) + + # Routed scaling factor kwargs + routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0) + # aiter applies routed_scaling_factor internally + apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled() + routed_scaling_factor_kwargs = dict( + routed_scaling_factor=routed_scaling_factor, + apply_routed_scale_to_output=apply_routed_scale_to_output, + ) + + # Dtype the router computes in, if it is not the activation dtype. + config_router_dtype = getattr(text_config, "moe_router_dtype", None) + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) + + # Grouped topk routing kwargs num_expert_group = getattr(text_config, "n_group", None) topk_group = getattr(text_config, "topk_group", None) + use_grouped_topk = num_expert_group is not None and topk_group is not None + grouped_topk_routing_kwargs = dict( + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) # MoE activation function activation = "silu" @@ -211,6 +231,9 @@ def recursive_replace(self): self.num_shared_experts = num_shared_experts self.num_redundant_experts = num_redundant_experts + # Down projections of shared experts consumed by FusedMoE + shared_down_projs: list[tuple[nn.Module, str]] = [] + # Recursively fuse MoE layers def _recursive_replace(module: nn.Module, prefix: str): for child_name, child_module in module.named_children(): @@ -249,7 +272,6 @@ def _recursive_replace(module: nn.Module, prefix: str): hidden_size=hidden_size, intermediate_size=intermediate_size, renormalize=renormalize, - use_grouped_topk=False, quant_config=self.quant_config, prefix=qual_name, activation=activation, @@ -259,23 +281,61 @@ def _recursive_replace(module: nn.Module, prefix: str): routed_experts_cls=TransformersRoutedExperts, ) fuser = MoEBlockFuser.match(moe_block, experts_name) - if self.num_expert_groups <= 1 and fuser is not None: + # _maybe_apply_routed_scale_to_output edge case. Transformers + # decoder layers do not compensate for dividing by scaling factor. + reaches_fp16_trick = ( + routed_scaling_factor != 1.0 + and apply_routed_scale_to_output + and self.model_config.dtype == torch.float16 + and fuser is not None + and fuser.shared_name is not None + ) + if reaches_fp16_trick: + logger.warning_once( + "%s could be fused but routing it in vLLM would apply " + "`routed_scaling_factor` by dividing the shared expert " + "output, which only fp16 overflow protection expects the " + "decoder layer to compensate for. Falling back to routing " + "in Transformers; run in bfloat16 to fuse it.", + moe_block_cls, + ) + if fuser is not None and not reaches_fp16_trick: # MoE block forward is fully replaced. - # gate/router and shared expert (if any) runs in FusedMoE. + # gate/router and shared expert (if any) runs in MoERunner. + shared_experts = fuser.shared_experts(moe_block, prefix) + # Store shared experts for later down projection adjustment + if shared_experts is not None: + hf_shared = shared_experts.shared_experts + glu_fuser = get_fuser(hf_shared) + down_name = getattr(glu_fuser, "down_name", None) + if down_name is not None: + shared_down_projs.append((hf_shared, down_name)) + # Prefer config, otherwise read it from fuser. + router_dtype = config_router_dtype or fuser.router_dtype + gate = fuser.gate(moe_block, prefix, router_dtype) 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), + gate=gate, + shared_experts=shared_experts, ) + if router_dtype is not None: + kwargs["router_logits_dtype"] = router_dtype + if use_grouped_topk: + kwargs |= grouped_topk_routing_kwargs + if routed_scaling_factor != 1.0: + kwargs |= routed_scaling_factor_kwargs + bias = getattr(gate, "e_score_correction_bias", None) + if bias is not None: + kwargs["e_score_correction_bias"] = bias fuser.rewrite_forward(moe_block) routed = "gate + experts" if fuser.shared_name: routed += " + shared experts" logger.info_once( - "Fused: %s (%s) -> FusedMoE (internal routing)", + "Fused: %s (%s) -> MoERunner (internal routing)", routed, moe_block_cls, ) @@ -317,10 +377,10 @@ def custom_routing_function( runner_args={"moe_state": moe_state}, ) logger.info_once( - "Fused: experts (%s) -> FusedMoE (external routing)", + "Fused: experts (%s) -> MoERunner (external routing)", experts_cls, ) - fused_experts = FusedMoE(**kwargs) + fused_experts = FusedMoEFactory(**kwargs) moe_block.experts = fused_experts log_replacement(qual_name, experts, fused_experts) # Update MixtureOfExperts mixin state @@ -333,3 +393,7 @@ def custom_routing_function( self.num_moe_layers = len(self.moe_layers) # Continue with the replacement of layers in Base super().recursive_replace() + # GLUFuser likely fused shared_experts. The down projection after the GLU + # normally immediately reduces but we want FusedMoE to handle the reduction. + for hf_shared, down_name in shared_down_projs: + hf_shared.get_submodule(down_name).reduce_results = False diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 0f80d569754f..a87e28241fd4 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -17,7 +17,8 @@ """Transformers modeling backend mixin for multi-modal models.""" from collections.abc import Mapping -from typing import TYPE_CHECKING +from contextlib import nullcontext +from typing import TYPE_CHECKING, Any import torch from transformers import AutoModel @@ -27,6 +28,7 @@ from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal +from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MultiModalKwargsItems from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -36,6 +38,7 @@ from vllm.multimodal.parse import ( ImageProcessorItems, MultiModalDataItems, + MultiModalDataParser, ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, @@ -59,11 +62,90 @@ class MultiModalProcessingInfo(BaseProcessingInfo): + def _get_audio_processor(self) -> Any: + # TODO: drop feature_extractor branch once huggingface/transformers#44394 lands. + return getattr_iter( + self.get_hf_processor(), ("audio_processor", "feature_extractor") + ) + + def _is_audio_model(self) -> bool: + return self._get_audio_processor() is not None + + def _is_image_model(self) -> bool: + return hasattr(self.get_hf_processor(), "image_processor") + + def _is_video_model(self) -> bool: + return hasattr(self.get_hf_processor(), "video_processor") + + def _get_audio_token_id(self) -> int: + processor = self.get_hf_processor() + if hasattr(processor, "audio_token_id"): + return processor.audio_token_id + config = self.get_hf_config() + val = getattr_iter(config, ("audio_token_id", "audio_token_index")) + if val is not None: + return val + if hasattr(processor, "audio_token"): + tokenizer = self.get_tokenizer() + vocab = tokenizer.get_vocab() + if processor.audio_token in vocab: + return vocab[processor.audio_token] + raise ValueError("Cannot find audio_token_id on processor or model config") + + def _get_audio_sampling_rate(self) -> float: + sub = self._get_audio_processor() + if sub is not None and hasattr(sub, "sampling_rate"): + return sub.sampling_rate + return 16000.0 + + def get_data_parser(self) -> MultiModalDataParser: + if self._is_audio_model(): + return MultiModalDataParser( + target_sr=self._get_audio_sampling_rate(), + expected_hidden_size=self._get_expected_hidden_size(), + ) + return super().get_data_parser() + def get_supported_mm_limits(self): - return {"image": None} + limits = {} + if self._is_audio_model(): + limits["audio"] = None + if self._is_image_model(): + limits["image"] = None + if not limits: + raise ValueError( + f"Unable to detect a supported modality on " + f"{type(self.get_hf_processor()).__name__}. " + "Checked `_is_audio_model` and `_is_image_model`." + ) + return limits def get_mm_max_tokens_per_item(self, seq_len, mm_counts): - return {"image": self.get_max_image_tokens()} + result = {} + if self._is_audio_model(): + result["audio"] = self.get_max_audio_tokens() + if self._is_image_model(): + result["image"] = self.get_max_image_tokens() + if not result: + raise ValueError( + f"Unable to detect a supported modality on " + f"{type(self.get_hf_processor()).__name__}. " + "Checked `_is_audio_model` and `_is_image_model`." + ) + return result + + def get_max_audio_tokens(self) -> int: + config = self.get_hf_config() + audio_config_names = ("audio_config", "encoder_config") + names = ("max_source_positions", "max_position_embeddings", "max_pos_emb") + audio_config = getattr_iter(config, audio_config_names, default=config) + val = getattr_iter(audio_config, names) + if val is not None: + return int(val) + raise ValueError( + f"Unable to get max input length from {type(audio_config).__name__}. " + f"The following attribute names were checked: {names}." + ) def get_max_image_tokens(self) -> int: width, height = self.get_max_image_size() @@ -82,14 +164,19 @@ def get_max_image_size(self): class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - num_images = mm_counts.get("image", 0) - - processor = self.info.get_hf_processor() - if "gemma3" in processor.__class__.__name__.lower(): - image_token = processor.boi_token - else: - image_token = getattr(processor, "image_token", "") - return image_token * num_images + text = "" + if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): + processor = self.info.get_hf_processor() + audio_token = getattr(processor, "audio_token", "") + text += audio_token * num_audios + if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): + processor = self.info.get_hf_processor() + if "gemma3" in processor.__class__.__name__.lower(): + image_token = processor.boi_token + else: + image_token = getattr(processor, "image_token", "") + text += image_token * num_images + return text def get_dummy_mm_data( self, @@ -97,20 +184,28 @@ def get_dummy_mm_data( mm_counts: Mapping[str, int], mm_options: Mapping[str, "BaseDummyOptions"], ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - - target_width, target_height = self.info.get_max_image_size() - - image_overrides = mm_options.get("image") - - return { - "image": self._get_dummy_images( + data: MultiModalDataDict = {} + if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): + sampling_rate = self.info._get_audio_sampling_rate() + sub = self.info._get_audio_processor() + chunk_length = getattr(sub, "chunk_length", None) if sub else None + if chunk_length is None: + chunk_length = 30 + audio_len = int(chunk_length * sampling_rate) + data["audio"] = self._get_dummy_audios( + length=audio_len, + num_audios=num_audios, + overrides=mm_options.get("audio"), + ) + if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): + target_width, target_height = self.info.get_max_image_size() + data["image"] = self._get_dummy_images( width=target_width, height=target_height, num_images=num_images, - overrides=image_overrides, - ), - } + overrides=mm_options.get("image"), + ) + return data class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): @@ -142,21 +237,39 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: # HF Processors always return a mask but vLLM doesn't need it hf_inputs.pop("attention_mask", None) - num_image_patches = hf_inputs.get("num_image_patches") - mm_fields = { - key: MultiModalFieldConfig.flat_from_sizes("image", num_image_patches) - for key in hf_inputs - } - mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes( - "image", num_image_patches - ) - # Keep these as batched, as they always have batch size as first dim - mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") - mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") - mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( - "image", keep_on_cpu=True - ) + mm_fields: dict[str, MultiModalFieldConfig] = {} + if self.info._is_audio_model(): + num_audio_tokens = hf_inputs.get("num_audio_tokens") + mm_fields.update( + { + key: MultiModalFieldConfig.flat_from_sizes( + "audio", num_audio_tokens + ) + for key in hf_inputs + } + ) + mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched("audio") + if self.info._is_image_model(): + num_image_patches = hf_inputs.get("num_image_patches") + mm_fields.update( + { + key: MultiModalFieldConfig.flat_from_sizes( + "image", num_image_patches + ) + for key in hf_inputs + } + ) + mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes( + "image", num_image_patches + ) + + # Keep these as batched, as they always have batch size as first dim + mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) return mm_fields def _get_hf_mm_data( @@ -164,13 +277,93 @@ def _get_hf_mm_data( mm_items: MultiModalDataItems, ) -> tuple[Mapping[str, object], Mapping[str, object]]: """ - In contrast to the base class, this method always adds - `return_mm_token_type_ids` to the processor data + In contrast to the base class, this method requests + `return_mm_token_type_ids` and remaps the `audios` key to `audio` for + audio models. """ processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) + if self.info._is_audio_model() and "audios" in processor_data: + processor_data["audio"] = processor_data.pop("audios") processor_data["return_mm_token_type_ids"] = True return processor_data, passthrough_data + def _apply_audio( + self, + prompt_ids: list[int], + processed_data: "BatchFeature", + ) -> dict[str, list[PlaceholderRange]]: + audio_token_id = self.info._get_audio_token_id() + prompt_tensor = torch.tensor(prompt_ids) + is_audio = prompt_tensor == audio_token_id + + if not is_audio.any(): + return {} + + padded = torch.cat([torch.tensor([False]), is_audio, torch.tensor([False])]) + transitions = padded.int().diff() + starts = torch.where(transitions == 1)[0] + ends = torch.where(transitions == -1)[0] + lengths = ends - starts + + ranges = [ + PlaceholderRange( + offset=s.item(), + length=ln.item(), + is_embed=torch.ones(ln.item(), dtype=torch.bool), + ) + for s, ln in zip(starts, lengths) + ] + processed_data["num_audio_tokens"] = lengths + return {"audio": ranges} + + def _apply_vision( + self, + prompt_ids: list[int], + processed_data: "BatchFeature", + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + mm_token_type_ids: torch.Tensor | None, + ) -> dict[str, list[PlaceholderRange]]: + if mm_token_type_ids is None: + return {} + + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + # We can infer vLLM style placeholder from token type ids, if we split + # it for each input `mm_data`. + mm_positions = torch.where(mm_token_type_ids == 1)[1] + images = mm_items.get_items("image", ImageProcessorItems) + image_sizes = [] + for item_idx in range(len(images)): + image_size = images.get_image_size(item_idx) + image_sizes.append((image_size.height, image_size.width)) + + mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens( + image_sizes=image_sizes, + **self.info.ctx.get_merged_mm_kwargs({}), + ) + + mm_placeholders: dict[str, list[PlaceholderRange]] = {} + split_sizes = mm_tokens_per_modality["num_image_tokens"] + if split_sizes: + chunked_mm_positions = torch.split(mm_positions, split_sizes) + mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()] + chunked_mm_tokens = torch.split(mm_tokens, split_sizes) + ranges = [ + PlaceholderRange( + offset=positions[0].item(), + length=positions.shape[0], + is_embed=(mm_tokens == hf_processor.image_token_id).bool(), + ) + for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens) + ] + mm_placeholders = {"image": ranges} + + processed_data["num_image_patches"] = torch.tensor( + mm_tokens_per_modality["num_image_patches"] + ) + return mm_placeholders + def apply( self, inputs: ProcessorInputs, @@ -198,8 +391,8 @@ def apply( # Bypass cached processor and always apply to the full set of mm inputs # NOTE: we can't just set caching=False because base class method # transforms outputs to `MultiModalKwargs` which is not going to - # work for Transformers. We have a lot of logic tied to - # `mm_tokens_per_modality` below + # work for Transformers. The vision path has logic tied to + # `mm_tokens_per_modality` in _apply_vision() prompt_ids, processed_data, _ = self._apply_hf_processor_text_mm( prompt_text=prompt, mm_items=mm_items, @@ -207,52 +400,36 @@ def apply( tokenization_kwargs=tokenization_kwargs, ) + # Use overrides if provided; fallback to data-dependent hashing. + with timing_ctx.record("get_mm_hashes"): + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) + # For gemma3 we check `token_type_ids` as the key mm_token_type_ids = processed_data.pop("token_type_ids", None) mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids) - # We can infer vLLM style placeholder from token type ids, if we split - # it for each input `mm_data`. - mm_positions = torch.where(mm_token_type_ids == 1)[1] - images = mm_items.get_items("image", ImageProcessorItems) - image_sizes = [] - for item_idx in range(len(images)): - image_size = images.get_image_size(item_idx) - image_sizes.append((image_size.height, image_size.width)) - - mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens( - image_sizes=image_sizes, - **self.info.ctx.get_merged_mm_kwargs({}), - ) - - mm_placeholders = {} - split_sizes = mm_tokens_per_modality["num_image_tokens"] - if split_sizes: - chunked_mm_positions = torch.split(mm_positions, split_sizes) - mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()] - chunked_mm_tokens = torch.split(mm_tokens, split_sizes) - ranges = [ - PlaceholderRange( - offset=positions[0].item(), - length=positions.shape[0], - is_embed=(mm_tokens == hf_processor.image_token_id).bool(), + mm_placeholders: dict[str, list[PlaceholderRange]] = {} + if self.info._is_audio_model(): + mm_placeholders.update(self._apply_audio(prompt_ids, processed_data)) + if self.info._is_image_model(): + mm_placeholders.update( + self._apply_vision( + prompt_ids, + processed_data, + mm_items, + hf_processor_mm_kwargs, + mm_token_type_ids, ) - for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens) - ] - mm_placeholders = {"image": ranges} + ) - processed_data["num_image_patches"] = torch.tensor( - mm_tokens_per_modality["num_image_patches"] - ) mm_kwargs = MultiModalKwargsItems.from_hf_inputs( processed_data, self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs), ) - # Use overrides if provided; fallback to data-dependent hashing. - with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) - return mm_input( prompt_token_ids=prompt_ids, mm_kwargs=mm_kwargs, @@ -337,7 +514,7 @@ def forward( # Positions shape handling for MRoPE models if self.model_config.uses_mrope: # [3, seq_len] -> [3, 1, seq_len] - positions = positions[:, None] + positions = positions[:, None].contiguous() model_output = super().forward( input_ids, positions, intermediate_tensors, inputs_embeds ) @@ -364,7 +541,79 @@ def __init__(self, multimodal_model): return LanguageModel(self) - def embed_multimodal(self, **kwargs): + def get_mm_mapping(self) -> MultiModelKeys: + """ + Get the module prefix in multimodal models + """ + for name in ("language_model", "text_model"): + if getattr(self.model, name, None) is not None: + return MultiModelKeys.from_string_field(language_model=f"model.{name}") + raise ValueError( + "Could not locate the language model submodule for LoRA support" + ) + + def _split_embeddings( + self, embeddings: torch.Tensor, split_sizes: list[int] + ) -> list[torch.Tensor]: + total_expected = sum(split_sizes) + + # Flatten to 2D: [total_tokens, hidden_dim] + if embeddings.ndim == 3: + embeddings = embeddings.view(-1, embeddings.shape[-1]) + + total_tokens = embeddings.shape[0] + if total_tokens == total_expected: + # Direct match: split_sizes are actual token counts + token_split_sizes = split_sizes + elif total_expected > 0 and total_tokens % total_expected == 0: + # Uniform expansion: each item expands to N tokens + tokens_per_item = total_tokens // total_expected + token_split_sizes = [s * tokens_per_item for s in split_sizes] + elif total_expected > 0: + # Mismatch (profiling with dummy data) - pad/truncate + if total_tokens == 0: + raise ValueError( + "Encoder returned empty embeddings. " + f"Expected {total_expected} tokens from " + f"split_sizes={split_sizes}" + ) + if total_tokens < total_expected: + repeat_factor = (total_expected + total_tokens - 1) // total_tokens + embeddings = embeddings.repeat(repeat_factor, 1) + embeddings = embeddings[:total_expected] + token_split_sizes = split_sizes + else: + return [] + + return list(torch.split(embeddings, token_split_sizes, dim=0)) + + def _embed_audio(self, **kwargs) -> list[torch.Tensor] | None: + self.check_version("5.13.0", "audio models support") + input_features: torch.Tensor | None = kwargs.pop("input_features", None) + if input_features is None: + input_features = kwargs.pop("input_values", None) + if input_features is None: + return None + + num_audio_tokens = kwargs.pop("num_audio_tokens") + kwargs.pop("token_type_ids", None) + kwargs.pop("mm_token_type_ids", None) + + context = nullcontext() + if current_platform.is_rocm(): + context = torch.nn.attention.sdpa_kernel( + backends=[torch.nn.attention.SDPBackend.MATH] + ) + with context: + audio_output = self.model.get_audio_features( + input_features, return_dict=True, **kwargs + ) + audio_embeddings = audio_output.pooler_output + + split_sizes = num_audio_tokens.flatten().tolist() + return self._split_embeddings(audio_embeddings, split_sizes) + + def _embed_vision(self, **kwargs) -> list[torch.Tensor] | torch.Tensor | None: pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None) image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None) # Model might use `image_patches` instead of `pixel_values` @@ -379,84 +628,57 @@ def embed_multimodal(self, **kwargs): num_image_patches = kwargs.pop("num_image_patches") - if pixel_values is not None: + context = nullcontext() + if current_platform.is_rocm(): # ROCm: Force math SDP backend for vision encoder to avoid accuracy issues # with flash_sdp and mem_efficient_sdp - if current_platform.is_rocm(): - # TODO: [ROCm] Fix accuracy issues with flash backend - logger.debug( - "ROCm platform detected. Forcing math SDP backend " - "for vision encoder. Currently ROCm platform has " - "accuracy issues with `flash_sdp` and" - "`mem_efficient_sdp` backends. See issue: " - "https://github.com/vllm-project/vllm/issues/30167" - ) - with torch.nn.attention.sdpa_kernel( - backends=[torch.nn.attention.SDPBackend.MATH] - ): - vision_embeddings = self.model.get_image_features( - pixel_values, **kwargs - ) - else: - vision_embeddings = self.model.get_image_features( - pixel_values, **kwargs - ) + # TODO: [ROCm] Fix accuracy issues with flash backend + logger.debug( + "ROCm platform detected. Forcing math SDP backend " + "for vision encoder. Currently ROCm platform has " + "accuracy issues with `flash_sdp` and" + "`mem_efficient_sdp` backends. See issue: " + "https://github.com/vllm-project/vllm/issues/30167" + ) + context = torch.nn.attention.sdpa_kernel( + backends=[torch.nn.attention.SDPBackend.MATH] + ) + with context: + vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) - # Transformers `v5`, `self.get_image_features` returns a tuple - # containing the features and optionally attentions/hidden_states - # After v5 is settled, we can enable qwen3-vl with several outputs - # from `self.get_image_features` - if isinstance(vision_embeddings, tuple): - vision_embeddings = vision_embeddings[0] - elif isinstance(vision_embeddings, dict): - vision_embeddings = vision_embeddings.pooler_output - - if isinstance(vision_embeddings, torch.Tensor): - split_sizes = num_image_patches.flatten().tolist() - total_patches = sum(split_sizes) - - # Flatten to 2D: [total_tokens, hidden_dim] - if vision_embeddings.ndim == 3: - vision_embeddings = vision_embeddings.view( - -1, vision_embeddings.shape[-1] - ) + # Transformers `v5`, `self.get_image_features` returns a tuple + # containing the features and optionally attentions/hidden_states + # After v5 is settled, we can enable qwen3-vl with several outputs + # from `self.get_image_features` + if isinstance(vision_embeddings, tuple): + vision_embeddings = vision_embeddings[0] + elif isinstance(vision_embeddings, dict): + vision_embeddings = vision_embeddings.pooler_output - total_tokens = vision_embeddings.shape[0] - if total_tokens == total_patches: - # Direct match: num_image_patches are actual token counts - # (e.g., Qwen2.5-VL style) - token_split_sizes = split_sizes - elif total_patches > 0 and total_tokens % total_patches == 0: - # Uniform expansion: each patch expands to N tokens - # (e.g., Idefics3 style) - tokens_per_patch = total_tokens // total_patches - token_split_sizes = [s * tokens_per_patch for s in split_sizes] - elif total_patches > 0: - # Mismatch (profiling with dummy data) - pad/truncate - if total_tokens == 0: - raise ValueError( - "Vision encoder returned empty embeddings. " - f"Expected {total_patches} patches from " - f"num_image_patches={split_sizes}" - ) - if total_tokens < total_patches: - repeat_factor = ( - total_patches + total_tokens - 1 - ) // total_tokens - vision_embeddings = vision_embeddings.repeat(repeat_factor, 1) - vision_embeddings = vision_embeddings[:total_patches] - token_split_sizes = split_sizes - else: - return [] + if isinstance(vision_embeddings, torch.Tensor): + split_sizes = num_image_patches.flatten().tolist() + return self._split_embeddings(vision_embeddings, split_sizes) - return list(torch.split(vision_embeddings, token_split_sizes, dim=0)) + return vision_embeddings - return vision_embeddings - else: - logger.debug( - "No pixel values or image embeddings provided for multimodal embedding." - ) - return None + def embed_multimodal(self, **kwargs): + embeddings: tuple[torch.Tensor, ...] = () + if "input_features" in kwargs or "input_values" in kwargs: + audio_embeddings = self._embed_audio(**kwargs) + if audio_embeddings is not None: + embeddings += tuple(audio_embeddings) + if ( + "pixel_values" in kwargs + or "image_embeds" in kwargs + or "image_patches" in kwargs + ): + vision_embeddings = self._embed_vision(**kwargs) + if vision_embeddings is not None: + if isinstance(vision_embeddings, torch.Tensor): + embeddings += (vision_embeddings,) + else: + embeddings += tuple(vision_embeddings) + return embeddings def get_mrope_input_positions( self, @@ -481,12 +703,8 @@ def get_mrope_input_positions( image_grid_thw = kwargs.get("image_grid_thw", []) video_grid_thw = kwargs.get("video_grid_thw", []) - image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)( - image_grid_thw - ) - video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)( - video_grid_thw - ) + image_grid_thw = torch.stack(image_grid_thw) if image_grid_thw else None + video_grid_thw = torch.stack(video_grid_thw) if video_grid_thw else None # `get_rope_index` doesn't always accept arbitrary `kwargs` kwargs = {} diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 4d9b01ce3930..9000ace90516 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -16,7 +16,9 @@ # limitations under the License. """Transformers modeling backend utilities.""" +from collections.abc import Iterator from contextlib import contextmanager +from itertools import chain from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -209,6 +211,11 @@ def _recursive_replace(module: nn.Module, prefix: str): _recursive_replace(model, prefix=prefix) +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 log_replacement(name: str, old_module: nn.Module, new_module: nn.Module): logger.debug("%s: %s -> %s", name, old_module, new_module) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index a6f39f6ebf8b..f92aa4f08662 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -1074,3 +1074,25 @@ def scatter_output_slices( sliced = output[offset : offset + n_tok] dest[idx] = sliced.clone() if clone else sliced offset += n_tok + + +def parse_diarized_timestamp(marker: str) -> float | None: + if ( + not marker + or not marker.isascii() + or marker.count(".") > 1 + or not marker.replace(".", "").isdigit() + ): + return None + return float(marker) + + +def parse_diarized_speaker(speaker: str) -> str | None: + if ( + len(speaker) < 2 + or speaker[0] != "S" + or not speaker[1:].isascii() + or not speaker[1:].isdigit() + ): + return None + return speaker diff --git a/vllm/model_executor/warmup/cutedsl_warmup.py b/vllm/model_executor/warmup/cutedsl_warmup.py index 991eab22b8a8..d2b8142f3b29 100644 --- a/vllm/model_executor/warmup/cutedsl_warmup.py +++ b/vllm/model_executor/warmup/cutedsl_warmup.py @@ -2,9 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Run registered CuTeDSL warmup compile units.""" +# TODO(roberto): Remove this compatibility registry after registered CuTeDSL +# warmups are migrated to the shared JIT warmup infrastructure. +# https://github.com/vllm-project/vllm/pull/47451 + from __future__ import annotations -import time import weakref from collections.abc import Callable, Hashable, Iterable from dataclasses import dataclass @@ -93,25 +96,17 @@ def _compile_cutedsl_warmup_units( def cutedsl_warmup() -> None: """Run CuTeDSL compile providers before serving.""" if not current_platform.is_cuda(): - logger.info("Skipping CuTeDSL warmup on non-CUDA platform.") + logger.debug("Skipping CuTeDSL warmup on non-CUDA platform.") return compile_units = _collect_unique_compile_units(_iter_cutedsl_warmup_compile_units()) if not compile_units: - logger.info("Skipping CuTeDSL warmup because no compile units were requested.") + logger.debug("Skipping CuTeDSL warmup because no compile units were requested.") return - unit_names = list(dict.fromkeys(unit.name for unit in compile_units)) - logger.info( + logger.info_once( "Warming up CuTeDSL compile_units=%d names=%s.", len(compile_units), - unit_names, - ) - - start_time = time.perf_counter() - compiled_count = _compile_cutedsl_warmup_units(compile_units) - logger.info( - "CuTeDSL warmup compiled %d units in %.2f s.", - compiled_count, - time.perf_counter() - start_time, + tuple(dict.fromkeys(unit.name for unit in compile_units)), ) + _compile_cutedsl_warmup_units(compile_units) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 4b05cad1b1ec..44f2121659ea 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -108,7 +108,7 @@ def _extract_data_from_fused_moe_module( m_: torch.nn.Module, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: """ - Extract weights, weight scales and num_topk from FusedMoE module. + Extract weights, weight scales and num_topk from MoERunner module. """ assert isinstance(m_, MoERunner) m = m_.routed_experts @@ -234,7 +234,7 @@ def _deepgemm_fp8_gemm_nt_warmup( device = w.device a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.empty( + a1q_scales = torch.zeros( (max_tokens, k // block_m), device=device, dtype=torch.float32 ) out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16) @@ -336,7 +336,7 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( def _warmup(w: torch.Tensor, w_scale: torch.Tensor): _, n, k = w.size() a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.empty( + a1q_scales = torch.zeros( (MAX_M, k // block_m), device=device, dtype=torch.float32 ) out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16) diff --git a/vllm/model_executor/warmup/fa4_cutedsl_config.py b/vllm/model_executor/warmup/fa4_cutedsl_config.py deleted file mode 100644 index 916ffd8a67f7..000000000000 --- a/vllm/model_executor/warmup/fa4_cutedsl_config.py +++ /dev/null @@ -1,204 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FA4 MLA prefill CuTeDSL compile warmup config.""" - -from __future__ import annotations - -from collections.abc import Hashable, Iterator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal - -import torch - -if TYPE_CHECKING: - from vllm.v1.attention.backends.fa_utils import ( - FlashAttentionCuTeDSLCompileSpec, - ) - -FA4ArchitectureFamily = Literal["sm90", "sm100f", "sm120"] - -FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16) - -# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan -# covers qhead_per_kvhead=1. -# Batch is not a current FA4 MLA-prefill key field. Use b1 for compile-only -# specs because it is the conservative case for Split-KV shape heuristics. -# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but vLLM -# does not use them in this backend yet; they need a separate -# num_kv_heads/qv/top-k-aware warmup plan if wired in later. -FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1 -FA4_MLA_PREFILL_Q_TILE = 128 -FA4_MLA_PREFILL_K_TILE = 128 -FA4_MLA_PREFILL_LONG_K_BLOCKS = 32 -FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64 -FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True) -FA4_MLA_PREFILL_LSE_OPTIONS = (False, True) - - -@dataclass(frozen=True) -class FA4MLAPrefillCompileContext: - dtype: torch.dtype - num_heads: int - qk_head_dim: int - v_head_dim: int - kv_nope_head_dim: int - requires_v_padding: bool - scale: float - num_splits: int - fa_version: int - - # Return the V head dim FA4 sees. - @property - def effective_v_head_dim(self) -> int: - if self.requires_v_padding: - return self.qk_head_dim - return self.v_head_dim - - -@dataclass(frozen=True) -class FA4MLAPrefillCompileRequest: - """One compile-only FA4 MLA prefill request.""" - - key: Hashable - compile_spec: FlashAttentionCuTeDSLCompileSpec - - # Compile this request. - def compile(self) -> None: - self.compile_spec.compile() - - -# Yield deduped compile requests. -def iter_fa4_mla_prefill_compile_requests( - ctx: FA4MLAPrefillCompileContext, -) -> Iterator[FA4MLAPrefillCompileRequest]: - """Yield compile requests for this fixed MLA backend. - - FA4 dedupes duplicate atomic kernel selections in its own JIT cache. - """ - seen: set[Hashable] = set() - for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx): - key = compile_spec.request_key() - if key in seen: - continue - seen.add(key) - yield FA4MLAPrefillCompileRequest( - key=key, - compile_spec=compile_spec, - ) - - -# Build compile specs for this setup. -def iter_fa4_mla_prefill_compile_specs( - ctx: FA4MLAPrefillCompileContext, -) -> Iterator[FlashAttentionCuTeDSLCompileSpec]: - """Yield compile-only FA4 MLA prefill requests for this fixed setup.""" - - arch_family = _fa4_architecture_family_from_compute_capability( - *torch.cuda.get_device_capability() - ) - if not _supports_fa4_mla_prefill(ctx, arch_family): - return - - from vllm.v1.attention.backends.fa_utils import ( - FlashAttentionCuTeDSLCompileSpec, - ) - - batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE - v_stride = None - if not ctx.requires_v_padding: - v_stride = ( - ctx.num_heads * ctx.kv_nope_head_dim, - ctx.kv_nope_head_dim, - 1, - ) - - for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family): - total_q_tokens = batch_size * max_seqlen_q - total_kv_tokens = batch_size * max_seqlen_k - for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS: - for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS: - yield FlashAttentionCuTeDSLCompileSpec( - q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim), - k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim), - v_shape=( - total_kv_tokens, - ctx.num_heads, - ctx.effective_v_head_dim, - ), - v_stride=v_stride, - q_dtype=ctx.dtype, - cu_seqlens_q_shape=(batch_size + 1,), - cu_seqlens_k_shape=(batch_size + 1,), - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - softmax_scale=ctx.scale, - causal=causal, - return_softmax_lse=return_lse, - num_splits=ctx.num_splits, - fa_version=ctx.fa_version, - ) - - -# Pick one q/k point per current FA4 MLA-prefill shape regime. -def _shape_probes_for_context( - ctx: FA4MLAPrefillCompileContext, - arch_family: FA4ArchitectureFamily, -) -> tuple[tuple[str, int, int], ...]: - q_stage1_q = 1 - q_stage2_q = FA4_MLA_PREFILL_Q_TILE + 1 - # FA4 never auto-splits when ceil(max_seqlen_k / tile_n) <= 4. - no_split_k = 4 * FA4_MLA_PREFILL_K_TILE - long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE - # Diff-head-dim Blackwell Split-KV switches tile_n at 64 K blocks. - very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE - - base_probes = ( - ("q_stage1", q_stage1_q, FA4_MLA_PREFILL_K_TILE), - ("q_stage2", q_stage2_q, no_split_k), - ) - # SM120 currently rejects Split-KV in FA4; num_splits=1 also has no split - # shape regimes on any architecture. - if ctx.num_splits == 1 or arch_family == "sm120": - return base_probes - - long_k_probes = ( - ("q_stage1_long_k", q_stage1_q, long_k), - ("q_stage2_long_k", q_stage2_q, long_k), - ) - - # SM90 does not have the SM100 q_stage or diff-head-dim tile_n=64 branch. - # Same-dim SM100-family MLA also does not need the very-long-K probe. - if arch_family == "sm90" or ctx.qk_head_dim == ctx.effective_v_head_dim: - return (*base_probes, *long_k_probes) - - very_long_k_probes = ( - ("q_stage1_very_long_k", q_stage1_q, very_long_k), - ("q_stage2_very_long_k", q_stage2_q, very_long_k), - ) - return (*base_probes, *long_k_probes, *very_long_k_probes) - - -# Check whether this setup can use FA4 MLA prefill. -def _supports_fa4_mla_prefill( - ctx: FA4MLAPrefillCompileContext, - arch_family: FA4ArchitectureFamily, -) -> bool: - return ( - ctx.dtype in FA4_STANDARD_DTYPES - and ctx.num_heads > 0 - and (arch_family != "sm120" or ctx.num_splits == 1) - ) - - -# Map CUDA capability to the FA4 arch family used by warmup checks. -def _fa4_architecture_family_from_compute_capability( - major: int, - minor: int, -) -> FA4ArchitectureFamily: - if (major, minor) == (9, 0): - return "sm90" - if major == 10: - return "sm100f" - if (major, minor) == (12, 0): - return "sm120" - raise ValueError(f"FA4 warmup does not know CUDA capability {major}.{minor}") diff --git a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py new file mode 100644 index 000000000000..7f1042338e4d --- /dev/null +++ b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up FA4 CuTeDSL MLA prefill compile keys.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + + +def fa4_cutedsl_warmup(worker: Worker) -> None: + runner = worker.model_runner + if runner.is_pooling_model: + return + + vllm_config = runner.vllm_config + if not vllm_config.model_config.use_mla: + return + + try: + backend_cls = get_mla_prefill_backend(vllm_config) + except ValueError: + # fall back to top-k MQA prefill path. + return + if backend_cls.get_name() != "FLASH_ATTN": + return + + from vllm.v1.attention.backends.mla.prefill import flash_attn + + flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py new file mode 100644 index 000000000000..c5411b5efb48 --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared interfaces and tracing helpers for explicit JIT warmup keys.""" + +from __future__ import annotations + +import ast +import inspect +import itertools +import operator +import textwrap +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +__all__ = [ + "VllmJitKernel", + "WarmupIntRange", + "get_ast_full_name", + "get_function_source_node", + "zip_inputs", +] + + +CompileKeyT = TypeVar("CompileKeyT") + + +@dataclass(frozen=True) +class WarmupIntRange: + start: int + stop: int + step: int = 1 + + +WarmupValues = Any +CompileKeyDispatchFn = Callable[..., CompileKeyT] + + +@dataclass(frozen=True) +class _WarmupInputRows: + """Warmup dispatch inputs expanded in lockstep.""" + + rows: tuple[Mapping[str, WarmupValues], ...] + + +def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]: + if isinstance(values, WarmupIntRange): + return tuple(range(values.start, values.stop, values.step)) + if isinstance(values, (list, tuple)): + return tuple(values) + return (values,) + + +def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows: + """Group row-wise dispatch inputs that should be expanded in lockstep.""" + if not rows: + raise ValueError("zip_inputs requires at least one dispatch input row") + if not all(isinstance(row, Mapping) for row in rows): + raise ValueError("zip_inputs rows must be mappings") + + first_names = frozenset(rows[0]) + if not first_names: + raise ValueError("zip_inputs rows require at least one dispatch input name") + if not all(isinstance(name, str) for name in first_names): + raise ValueError("zip_inputs dispatch input names must be strings") + + input_rows: list[Mapping[str, WarmupValues]] = [] + for row in rows: + names = frozenset(row) + if names != first_names: + raise ValueError("zip_inputs rows must use the same dispatch input names") + input_rows.append(dict(row)) + + return _WarmupInputRows(rows=tuple(input_rows)) + + +def _expand_warmup_value_grid( + values: Mapping[str, WarmupValues], + input_names: frozenset[str], +) -> tuple[dict[str, Any], ...]: + names = tuple(name for name in values if name in input_names) + if not names: + return ({},) + + expanded_values = tuple(_expand_warmup_values(values[name]) for name in names) + return tuple( + dict(zip(names, value_set)) for value_set in itertools.product(*expanded_values) + ) + + +def _expand_warmup_input_rows( + rows: tuple[Mapping[str, WarmupValues], ...], + input_names: frozenset[str], +) -> tuple[dict[str, Any], ...]: + active_names = frozenset(name for name in rows[0] if name in input_names) + if not active_names: + return ({},) + + return tuple( + {name: value for name, value in row.items() if name in active_names} + for row in rows + ) + + +def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]: + merged: dict[str, Any] = {} + for part in parts: + for name, value in part.items(): + if name in merged: + raise ValueError( + f"Warmup dispatch input '{name}' is specified more than once" + ) + merged[name] = value + return merged + + +@dataclass(frozen=True) +class _CompileKeyDispatchTrace: + local_exprs: tuple[tuple[str, ast.AST], ...] + field_exprs: tuple[tuple[str, ast.AST], ...] + globals: Mapping[str, Any] + input_names: frozenset[str] + defaults: Mapping[str, Any] + + def compile_key( + self, + compile_key_type: type[CompileKeyT], + kwargs: Mapping[str, Any], + ) -> CompileKeyT: + dispatch_values = {**self.defaults, **kwargs} + for name, expr in self.local_exprs: + dispatch_values[name] = _eval_dispatch_expr( + expr, dispatch_values, self.globals + ) + return compile_key_type( + **{ + field: _eval_dispatch_expr(expr, dispatch_values, self.globals) + for field, expr in self.field_exprs + } + ) + + +_BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, +} +_CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = { + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, +} + + +def _dispatch_expr_source(node: ast.AST) -> str: + try: + return ast.unparse(node) + except Exception: + return ast.dump(node) + + +def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError: + return ValueError( + f"{reason}: {_dispatch_expr_source(node)}. " + "Supported dispatch expressions are names, constants, attributes, " + "tuple/list literals, conditional expressions, comparisons, boolean " + "operators, unary not/minus, arithmetic, and calls without **kwargs." + ) + + +class _DispatchExprEvaluator(ast.NodeVisitor): + def __init__( + self, + values: Mapping[str, Any], + globals_: Mapping[str, Any], + ) -> None: + self.values = values + self.globals = globals_ + + def eval(self, node: ast.AST) -> Any: + return self.visit(node) + + def generic_visit(self, node: ast.AST) -> Any: + raise _dispatch_expr_error(node, "Unsupported dispatch expression") + + def visit_Name(self, node: ast.Name) -> Any: + if node.id in self.values: + return self.values[node.id] + if node.id in self.globals: + return self.globals[node.id] + raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'") + + def visit_Constant(self, node: ast.Constant) -> Any: + return node.value + + def visit_IfExp(self, node: ast.IfExp) -> Any: + return self.visit(node.body if self.visit(node.test) else node.orelse) + + def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]: + return tuple(self.visit(elt) for elt in node.elts) + + def visit_List(self, node: ast.List) -> list[Any]: + return [self.visit(elt) for elt in node.elts] + + def visit_BoolOp(self, node: ast.BoolOp) -> Any: + if isinstance(node.op, ast.And): + result = None + for value in node.values: + result = self.visit(value) + if not result: + return result + return result + if isinstance(node.op, ast.Or): + result = None + for value in node.values: + result = self.visit(value) + if result: + return result + return result + raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator") + + def visit_Compare(self, node: ast.Compare) -> bool: + left = self.visit(node.left) + for op_node, comparator in zip(node.ops, node.comparators): + right = self.visit(comparator) + op = _CMP_OPS.get(type(op_node)) + if op is None: + raise _dispatch_expr_error( + node, "Unsupported dispatch comparison operator" + ) + if not op(left, right): + return False + left = right + return True + + def visit_UnaryOp(self, node: ast.UnaryOp) -> Any: + operand = self.visit(node.operand) + if isinstance(node.op, ast.Not): + return not operand + if isinstance(node.op, ast.USub): + return -operand + raise _dispatch_expr_error(node, "Unsupported dispatch unary operator") + + def visit_BinOp(self, node: ast.BinOp) -> Any: + op = _BIN_OPS.get(type(node.op)) + if op is None: + raise _dispatch_expr_error(node, "Unsupported dispatch binary operator") + return op(self.visit(node.left), self.visit(node.right)) + + def visit_Call(self, node: ast.Call) -> Any: + args = [self.visit(arg) for arg in node.args] + fn = self.visit(node.func) + call_kwargs: dict[str, Any] = {} + for keyword in node.keywords: + if keyword.arg is None: + raise _dispatch_expr_error( + node, "Dispatch helper calls cannot use **kwargs" + ) + call_kwargs[keyword.arg] = self.visit(keyword.value) + return fn(*args, **call_kwargs) + + def visit_Attribute(self, node: ast.Attribute) -> Any: + return getattr(self.visit(node.value), node.attr) + + +def get_ast_full_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + parent = get_ast_full_name(node.value) + if parent is not None: + return f"{parent}.{node.attr}" + return None + + +def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef: + source_fn = getattr(fn, "fn", fn) + source = textwrap.dedent(inspect.getsource(source_fn)) + tree = ast.parse(source) + function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)] + if len(function_defs) != 1: + name = getattr(source_fn, "__name__", type(source_fn).__name__) + raise ValueError(f"Expected one function in {name}, found {len(function_defs)}") + return function_defs[0] + + +def _eval_dispatch_expr( + node: ast.AST, + kwargs: Mapping[str, Any], + globals_: Mapping[str, Any], +) -> Any: + return _DispatchExprEvaluator(kwargs, globals_).eval(node) + + +def _collect_input_names( + node: ast.AST, + candidate_names: set[str], + local_names: set[str] | None = None, +) -> set[str]: + if local_names is None: + local_names = set() + return { + child.id + for child in ast.walk(node) + if ( + isinstance(child, ast.Name) + and child.id in candidate_names + and child.id not in local_names + ) + } + + +def _collect_dispatch_body( + fn: CompileKeyDispatchFn[Any], + function_def: ast.FunctionDef, +) -> tuple[list[tuple[str, ast.AST]], ast.Call]: + local_exprs: list[tuple[str, ast.AST]] = [] + for statement in function_def.body: + if ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str) + ): + continue + + if isinstance(statement, ast.Assign): + if len(statement.targets) != 1 or not isinstance( + statement.targets[0], ast.Name + ): + raise _dispatch_expr_error( + statement, "Dispatch assignments must target one local name" + ) + local_exprs.append((statement.targets[0].id, statement.value)) + continue + + if isinstance(statement, ast.AnnAssign): + if statement.value is None: + raise _dispatch_expr_error( + statement, "Dispatch annotations must assign a value" + ) + if not isinstance(statement.target, ast.Name): + raise _dispatch_expr_error( + statement, "Dispatch assignments must target one local name" + ) + local_exprs.append((statement.target.id, statement.value)) + continue + + if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call): + return local_exprs, statement.value + + if isinstance(statement, ast.Return): + raise _dispatch_expr_error( + statement, "Dispatch must return one CompileKey(...) call" + ) + + raise _dispatch_expr_error( + statement, + "Dispatch may only contain local assignments before CompileKey return", + ) + + raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call") + + +def _trace_compile_key_dispatch( + fn: CompileKeyDispatchFn[Any], +) -> _CompileKeyDispatchTrace: + source_fn = getattr(fn, "__func__", fn) + globals_ = source_fn.__globals__ + function_def = get_function_source_node(fn) + + local_exprs, return_call = _collect_dispatch_body(fn, function_def) + + field_exprs: list[tuple[str, ast.AST]] = [] + signature = inspect.signature(fn) + defaults = { + name: parameter.default + for name, parameter in signature.parameters.items() + if parameter.default is not inspect.Parameter.empty + } + candidate_names = set(signature.parameters) + input_names: set[str] = set() + local_names = {name for name, _ in local_exprs} + for _, expr in local_exprs: + input_names.update(_collect_input_names(expr, candidate_names)) + for keyword in return_call.keywords: + if keyword.arg is None: + raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey") + field_exprs.append((keyword.arg, keyword.value)) + input_names.update( + _collect_input_names(keyword.value, candidate_names, local_names) + ) + + return _CompileKeyDispatchTrace( + tuple(local_exprs), + tuple(field_exprs), + globals_, + frozenset(input_names), + defaults, + ) + + +class VllmJitKernel(Generic[CompileKeyT], ABC): + """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" + + CompileKey: type[CompileKeyT] + + def __init__(self) -> None: + self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + + def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: + return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs) + + def _trace_dispatch( + self, dispatch: CompileKeyDispatchFn[CompileKeyT] + ) -> Callable[..., list[CompileKeyT]]: + compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch) + + def traced( + *input_groups: _WarmupInputRows, + **kwargs: WarmupValues, + ) -> list[CompileKeyT]: + for group in input_groups: + if not isinstance(group, _WarmupInputRows): + raise TypeError( + "_trace_dispatch positional arguments must be " + "zip_inputs(...) groups" + ) + expanded_input_groups = tuple( + _expand_warmup_input_rows( + group.rows, compile_key_dispatch_trace.input_names + ) + for group in input_groups + ) + expanded_kwargs = _expand_warmup_value_grid( + kwargs, compile_key_dispatch_trace.input_names + ) + dispatch_value_groups = (*expanded_input_groups, expanded_kwargs) + return list( + dict.fromkeys( + compile_key_dispatch_trace.compile_key( + self.CompileKey, _merge_warmup_kwargs(dispatch_values) + ) + for dispatch_values in itertools.product(*dispatch_value_groups) + ) + ) + + return traced + + @abstractmethod + def dispatch(self, **kwargs: Any) -> CompileKeyT: + """Build one compile key from one concrete dispatch point.""" + raise NotImplementedError + + @abstractmethod + def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]: + """Return compile keys that should be warmed for this kernel.""" + raise NotImplementedError + + @abstractmethod + def compile(self, compile_key: CompileKeyT) -> None: + """Compile one warmup key.""" + raise NotImplementedError + + def warmup(self, *args: Any, **kwargs: Any) -> None: + """Compile this kernel's warmup keys.""" + for compile_key in self.get_warmup_keys(*args, **kwargs): + self.compile(compile_key) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py new file mode 100644 index 000000000000..b90975762175 --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import ast +import inspect +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from vllm.model_executor.warmup.jit_warmup import ( + get_ast_full_name, + get_function_source_node, +) + + +@dataclass(frozen=True) +class TritonWarmupTensor: + # Compile-only tensor descriptor for Triton pointer specialization. + dtype: Any + aligned: bool = True + shape: tuple[int, ...] = (1,) + + def data_ptr(self) -> int: + return 0 if self.aligned else 1 + + def ptr_range(self) -> int: + return 0 + + def stride(self) -> tuple[int, ...]: + strides: list[int] = [] + stride = 1 + for size in reversed(self.shape): + strides.append(stride) + stride *= size + return tuple(reversed(strides)) + + +@dataclass(frozen=True) +class TritonPointerInputVariant: + # Named pointer-alignment variant for compile-only Triton warmup. + alignments: tuple[tuple[str, bool], ...] + + @classmethod + def from_alignment(cls, **aligned: bool) -> "TritonPointerInputVariant": + return cls(tuple(aligned.items())) + + def is_aligned(self, name: str) -> bool: + for alignment_name, aligned in self.alignments: + if alignment_name == name: + return aligned + raise KeyError(f"Unknown Triton pointer input variant: {name}") + + def pointer( + self, + name: str, + dtype: Any, + shape: tuple[int, ...] = (1,), + ) -> TritonWarmupTensor: + return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape) + + +def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]: + if isinstance(node, ast.Constant) and isinstance(node.value, str | int): + return (node.value,) + if isinstance(node, ast.List | ast.Tuple): + refs: list[str | int] = [] + for elt in node.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int): + refs.append(elt.value) + else: + raise ValueError( + f"Unsupported Triton specialization ref: {ast.dump(elt)}" + ) + return tuple(refs) + raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}") + + +def _normalize_arg_refs( + refs: tuple[str | int, ...], + arg_names: tuple[str, ...], +) -> frozenset[str]: + names: set[str] = set() + for ref in refs: + if isinstance(ref, int): + names.add(arg_names[ref]) + else: + names.add(ref) + return frozenset(names) + + +def _decorator_keyword_refs( + function_def: ast.FunctionDef, + keyword_name: str, +) -> tuple[str | int, ...]: + for decorator in function_def.decorator_list: + if not isinstance(decorator, ast.Call): + continue + decorator_name = get_ast_full_name(decorator.func) + if decorator_name not in ("triton.jit", "jit"): + continue + for keyword in decorator.keywords: + if keyword.arg == keyword_name: + return _literal_str_refs(keyword.value) + return () + + +def _triton_do_not_specialize_args( + kernel: Callable[..., Any], + function_def: ast.FunctionDef, + arg_names: tuple[str, ...], +) -> frozenset[str]: + refs = getattr(kernel, "do_not_specialize", None) + if refs is not None: + return _normalize_arg_refs(tuple(refs), arg_names) + return _normalize_arg_refs( + _decorator_keyword_refs(function_def, "do_not_specialize"), + arg_names, + ) + + +def _triton_constexpr_arg_names( + kernel: Callable[..., Any], + function_def: ast.FunctionDef, + arg_names: tuple[str, ...], +) -> frozenset[str]: + constexprs = getattr(kernel, "constexprs", None) + if constexprs is not None: + return frozenset(arg_names[index] for index in constexprs) + + names: set[str] = set() + for arg in function_def.args.args + function_def.args.kwonlyargs: + if arg.annotation is None: + continue + annotation = get_ast_full_name(arg.annotation) + if annotation in ("tl.constexpr", "triton.language.constexpr", "constexpr"): + names.add(arg.arg) + return frozenset(names) + + +def _leftmost_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.BinOp): + return _leftmost_name(node.left) + return None + + +def _pointer_arg_names( + function_def: ast.FunctionDef, + arg_names: tuple[str, ...], +) -> frozenset[str]: + candidate_names = set(arg_names) + pointer_names = {name for name in arg_names if name.endswith("_ptr")} + for node in ast.walk(function_def): + if not isinstance(node, ast.Call): + continue + if get_ast_full_name(node.func) not in ("tl.load", "tl.store"): + continue + if not node.args: + continue + name = _leftmost_name(node.args[0]) + if name in candidate_names: + pointer_names.add(name) + return frozenset(pointer_names) + + +def trace_triton_kernel_specialization_args( + kernel: Callable[..., Any], +) -> tuple[str, ...]: + function_def = get_function_source_node(kernel) + source_fn = getattr(kernel, "fn", kernel) + arg_names = tuple(inspect.signature(source_fn).parameters) + constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names) + do_not_specialize_args = _triton_do_not_specialize_args( + kernel, function_def, arg_names + ) + pointer_args = _pointer_arg_names(function_def, arg_names) + + return tuple( + name + for name in arg_names + if name in constexpr_args + or (name not in pointer_args and name not in do_not_specialize_args) + ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d4b9a735d2cd..c35e1ac30c9a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -17,6 +17,9 @@ from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( deepseek_v4_mhc_warmup, ) +from vllm.model_executor.warmup.fa4_cutedsl_warmup import ( + fa4_cutedsl_warmup, +) from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, write_flashinfer_autotune_cache, @@ -25,9 +28,12 @@ deepseek_v4_sparse_mla_attention_warmup, flashinfer_sparse_mla_decode_autotune_warmup, ) +from vllm.model_executor.warmup.kimi_k3_triton_warmup import ( + kimi_k3_triton_warmup, +) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( - sparse_mla_triton_warmup_if_needed, + sparse_mla_triton_warmup, ) from vllm.model_executor.warmup.v1_block_table_warmup import ( warm_v1_block_table_kernels, @@ -42,16 +48,30 @@ logger = init_logger(__name__) -_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = ( - (6144, 264), # Inkling - (7168, 256), # DSV3 - (7168, 384), # DSV4-Pro - (14400, 256), # DSV4-Flash -) _LL_BF16_WARMUP_M_RANGE = range(1, 17) -def _warmup_ll_bf16_router_gemm() -> None: +def _ll_bf16_router_shapes_from_model( + model: torch.nn.Module, +) -> tuple[tuple[int, int], ...]: + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + shapes: set[tuple[int, int]] = set() + for module in model.modules(): + if not isinstance(module, GateLinear): + continue + weight = getattr(module, "weight", None) + if not isinstance(weight, torch.Tensor): + continue + if weight.dim() != 2 or weight.dtype != torch.bfloat16: + continue + n, k = weight.shape + if k % 8 == 0: + shapes.add((int(k), int(n))) + return tuple(sorted(shapes)) + + +def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( is_available as is_ll_bf16_gemm_available, ) @@ -62,24 +82,35 @@ def _warmup_ll_bf16_router_gemm() -> None: if not is_ll_bf16_gemm_available(): return - logger.info("Warming up ll_bf16 router GEMM kernels.") + shapes = _ll_bf16_router_shapes_from_model(model) + if not shapes: + logger.debug_once( + "Skipping ll_bf16 router GEMM warmup: no bf16 GateLinear shapes found." + ) + return + + logger.info_once("Warming up ll_bf16 router GEMM kernels for shapes: %s.", shapes) ll_bf16_gemm_kernel.warmup( - shapes=_LL_BF16_WARMUP_MODEL_SHAPES, + shapes=shapes, m_values=_LL_BF16_WARMUP_M_RANGE, ) -def kernel_warmup(worker: "Worker"): +def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( minimax_m3_msa_warmup, ) - # Pooling models do not use the generation slot-mapping path. - if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels( - getattr(worker.model_runner, "device", torch.device("cuda")), - worker.scheduler_config.max_num_batched_tokens, - ) + if not worker.use_v2_model_runner: + # Pooling models do not use the generation slot-mapping path. + if not worker.model_runner.is_pooling_model: + warm_v1_block_table_kernels(worker.model_runner) + # The KV-block zeroing kernel is driven by the scheduler's + # `new_block_ids_to_zero`, so no dummy run ever reaches it. + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + if zeroer is not None: + zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder @@ -94,7 +125,23 @@ def kernel_warmup(worker: "Worker"): ) # Run next so input-prep kernels JIT against pristine runner state. - sparse_mla_triton_warmup_if_needed(worker) + if worker.vllm_config.kernel_config.enable_jit_warmup: + kimi_k3_triton_warmup(worker) + fa4_cutedsl_warmup(worker) + sparse_mla_triton_warmup(worker) + + if current_platform.has_device_capability(90): + _warmup_ll_bf16_router_gemm(worker.get_model()) + + if worker.vllm_config.kernel_config.enable_cutedsl_warmup: + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 + cutedsl_warmup() + + if process_local_only: + return + flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) @@ -116,13 +163,10 @@ def kernel_warmup(worker: "Worker"): ) # FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs if enable_flashinfer_autotune is False: - logger.info("Skipping FlashInfer autotune because it is disabled.") + logger.info_once("Skipping FlashInfer autotune because it is disabled.") elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) - if current_platform.has_device_capability(90): - _warmup_ll_bf16_router_gemm() - # FlashInfer attention warmup # Only warmup if the model has FlashInfer attention groups # and is not a pooling model @@ -144,7 +188,7 @@ def _is_flashinfer_backend(backend): for group in groups ) ): - logger.info("Warming up FlashInfer attention.") + logger.info_once("Warming up FlashInfer attention.") # Warmup with mixed batch containing both prefill and decode tokens # This is to warm up both prefill and decode attention kernels worker.model_runner._dummy_run( @@ -155,9 +199,6 @@ def _is_flashinfer_backend(backend): create_mixed_batch=True, ) - if worker.vllm_config.kernel_config.enable_cutedsl_warmup: - cutedsl_warmup() - def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None: @@ -188,85 +229,70 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower. - Tuning is performed only on rank 0. The resulting cache is broadcast - to every rank so all ranks dispatch the same kernel tactic. + Every rank profiles the same tactics. When distributed, per-tactic + timings are averaged over the world CPU group so all ranks select the + same tactic. """ + from flashinfer.autotuner import AutoTuner, set_autotune_process_group + import vllm.utils.flashinfer as fi_utils from vllm.distributed.parallel_state import get_world_group + world = get_world_group() + is_leader = world.rank_in_group == 0 + tuner = AutoTuner.get() + autotune_kwargs: dict = {} skip_ops = _flashinfer_autotune_skip_ops(runner) if skip_ops: - logger.info( + logger.info_once( "Skipping FlashInfer autotuning for ops %s", - sorted(skip_ops), + tuple(sorted(skip_ops)), ) autotune_kwargs["skip_ops"] = skip_ops - use_persistent_cache = True - - # 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: - with torch.inference_mode(), fi_utils.autotune(**autotune_kwargs): - runner._dummy_run( - num_tokens=runner.scheduler_config.max_num_batched_tokens, - skip_eplb=True, - is_profile=True, - ) - get_world_group().barrier() - return - - world = get_world_group() - is_leader = world.rank_in_group == 0 - cache_path = resolve_flashinfer_autotune_file(runner) if is_leader: - logger.info("Using FlashInfer autotune cache file: %s", cache_path) + logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) # We skip EPLB here since we don't want to record dummy metrics. # When autotuning with number of tokens m, flashinfer will autotune # operations for all number of tokens up to m, so we only need to # run with the max number of tokens. + # Randomize inputs to avoid every token pick the same experts, + # which lead to some EP ranks receiving no tokens and skipping their + # MoE kernel entirely, and cause hang due to all-reduce collective + # during synchronized autotuning. dummy_run_kwargs = dict( num_tokens=runner.scheduler_config.max_num_batched_tokens, skip_eplb=True, is_profile=True, + randomize_inputs=True, ) - with torch.inference_mode(): - if is_leader: - with fi_utils.autotune( - tune_mode=True, cache=str(cache_path), **autotune_kwargs - ): - runner._dummy_run(**dummy_run_kwargs) - else: - runner._dummy_run(**dummy_run_kwargs) - - # Broadcast autotune cache from rank 0 to all other ranks so every - # rank loads the same set of chosen tactics. - tune_results: bytes | None = None + # Read cached autotune results and broadcast to all ranks. + cached_results: bytes | None = None if is_leader and cache_path.exists(): with open(cache_path, "rb") as f: - tune_results = f.read() - - tune_results = world.broadcast_object(tune_results, src=0) - - if tune_results is None: - logger.warning( - "No FlashInfer autotune cache entries found." - "Falling back to default tactics." - ) - else: - write_flashinfer_autotune_cache(cache_path, tune_results) + cached_results = f.read() + cached_results = world.broadcast_object(cached_results, src=0) + if cached_results is not None: + write_flashinfer_autotune_cache(cache_path, cached_results) world.barrier() - from flashinfer.autotuner import AutoTuner + tuner.load_configs(str(cache_path)) + + group = world.cpu_group if world.world_size > 1 else None + set_autotune_process_group(group) + try: + with ( + torch.inference_mode(), + fi_utils.autotune(tune_mode=True, **autotune_kwargs), + ): + runner._dummy_run(**dummy_run_kwargs) + finally: + set_autotune_process_group(None) - AutoTuner.get().load_configs(str(cache_path)) - logger.info( - "FlashInfer autotune cache loaded on rank %d from %s.", - world.rank_in_group, - cache_path, - ) + if world.world_size > 1: + world.barrier() + if is_leader: + tuner.save_configs(str(cache_path)) diff --git a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py new file mode 100644 index 000000000000..4c4a0327c7f5 --- /dev/null +++ b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up Kimi-K3 Triton kernels.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + compilation_config = getattr( + worker.model_runner, + "compilation_config", + None, + ) + static_context = getattr(compilation_config, "static_forward_context", None) + if not isinstance(static_context, dict): + return None + return next( + ( + layer + for layer in static_context.values() + if isinstance(layer, KimiK3DeltaAttention) + ), + None, + ) + + +def _warm_attn_res(worker: Worker) -> None: + from vllm.models.kimi_k3.nvidia.ops.attn_res import ( + attn_res, + get_attn_res_triton_warmup_profiles, + ) + + config = worker.model_config.hf_text_config + block_size = getattr(config, "attn_res_block_size", None) + if block_size is None: + return + + hidden_size = int(config.hidden_size) + max_blocks = (int(config.num_hidden_layers) + block_size - 1) // block_size + if max_blocks < 2: + return + + dtype = worker.model_config.dtype + device = torch.device("cuda") + eps = float(config.rms_norm_eps) + prefix = torch.zeros((1, hidden_size), dtype=dtype, device=device) + delta = torch.zeros_like(prefix) + blocks = torch.zeros( + (1, max_blocks, hidden_size), + dtype=dtype, + device=device, + ) + norm_weight = torch.zeros(hidden_size, dtype=dtype, device=device) + qk_weight = torch.zeros_like(norm_weight) + output_norm_weight = torch.zeros_like(norm_weight) + + for ( + num_blocks, + has_delta, + block_write_idx, + apply_output_norm, + ) in get_attn_res_triton_warmup_profiles(max_blocks): + attn_res( + prefix, + delta if has_delta else None, + blocks, + norm_weight, + qk_weight, + output_norm_weight if apply_output_norm else None, + num_blocks=num_blocks, + block_write_idx=block_write_idx, + eps=eps, + output_norm_eps=eps if apply_output_norm else 0.0, + ) + + +def _warm_recurrent_kda( + layer: KimiK3DeltaAttention, + input_dtype: torch.dtype, +) -> None: + from vllm.models.kimi_k3.nvidia.ops.third_party.kda.fused_recurrent import ( + fused_recurrent_kda, + get_fused_recurrent_kda_fwd_warmup_profiles, + ) + + num_speculative_tokens = int(layer.num_spec) + # fused_recurrent_kda_fwd_kernel is only used by speculative decode. + if num_speculative_tokens <= 0: + return + + kv_cache = layer.kv_cache + if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2: + return + state = kv_cache[1] + if not isinstance(state, torch.Tensor) or not state.numel(): + return + + logger.info("Warming up Kimi-K3 speculative KDA kernels.") + h = int(layer.local_num_heads) + d = int(layer.head_dim) + tokens_per_sequence = num_speculative_tokens + 1 + for num_sequences in get_fused_recurrent_kda_fwd_warmup_profiles(h): + num_tokens = num_sequences * tokens_per_sequence + packed_qkv = torch.empty( + (num_tokens, 3 * h * d), + dtype=input_dtype, + device=state.device, + ) + q, k, v = ( + tensor.view(1, num_tokens, h, d) + for tensor in packed_qkv.split(h * d, dim=-1) + ) + fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=torch.empty( + (1, num_tokens, h, d), + dtype=input_dtype, + device=state.device, + ), + raw_beta=torch.empty( + (1, num_tokens, h), + dtype=input_dtype, + device=state.device, + ), + A_log=layer.A_log, + dt_bias=layer.dt_bias, + lower_bound=layer.gate_lower_bound, + initial_state=state[:1], + cu_seqlens=torch.arange( + 0, + num_tokens + 1, + tokens_per_sequence, + dtype=torch.int32, + device=state.device, + ), + ssm_state_indices=torch.zeros( + (num_sequences, tokens_per_sequence), + dtype=torch.int32, + device=state.device, + ), + num_accepted_tokens=torch.ones( + num_sequences, + dtype=torch.int32, + device=state.device, + ), + out=torch.empty( + (1, num_tokens, h, d), + dtype=input_dtype, + device=state.device, + ), + ) + + +@torch.inference_mode() +def kimi_k3_triton_warmup(worker: Worker) -> None: + """Warm Kimi-K3 Triton kernels reachable by this server.""" + if not current_platform.is_cuda(): + return + + layer = _get_kda_layer(worker) + if layer is None: + return + + _warm_attn_res(worker) + _warm_recurrent_kda(layer, worker.model_config.dtype) diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 9b7d76207b41..06a016a5a06c 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -24,23 +24,10 @@ } ) -_ZERO_KV_N_BLOCKS = (1, 2) - -_SLOT_MAPPING_KV_BLOCK_SIZE = 16 -_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3) - # Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L. _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) -@dataclass(frozen=True) -class _ZeroKvWarmupConfig: - page_size_el: int - block_size: int - n_segs: int - - @dataclass(frozen=True) class _QwenGDNWarmupConfig: h: int @@ -147,93 +134,6 @@ def _qwen_gdn_warmup_config( return None -def _get_kv_block_zeroer(runner: object) -> object | None: - zeroer = getattr(runner, "kv_block_zeroer", None) - if zeroer is None: - zeroer = getattr(runner, "_kv_block_zeroer", None) - return zeroer - - -def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: - zeroer = _get_kv_block_zeroer(runner) - meta = getattr(zeroer, "_meta", None) - if meta is None: - return None - - _, page_size_el, block_size, n_segs = meta - return _ZeroKvWarmupConfig( - page_size_el=int(page_size_el), - block_size=int(block_size), - n_segs=int(n_segs), - ) - - -def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool: - zeroer = _get_kv_block_zeroer(runner) - zero_block_ids = getattr(zeroer, "zero_block_ids", None) - if not callable(zero_block_ids): - return False - - for n_blocks in _ZERO_KV_N_BLOCKS: - zero_block_ids(list(range(n_blocks))) - return True - - -def _warm_zero_kv_blocks_kernel( - device: torch.device, config: _ZeroKvWarmupConfig -) -> None: - from vllm.v1.worker.utils import _zero_kv_blocks_kernel - - max_n_blocks = max(_ZERO_KV_N_BLOCKS) - scratch = torch.empty( - max_n_blocks * config.page_size_el, - dtype=torch.int32, - device=device, - ) - seg_addrs = torch.tensor( - [scratch.data_ptr()] * config.n_segs, - dtype=torch.uint64, - device=device, - ) - - for n_blocks in _ZERO_KV_N_BLOCKS: - block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) - grid = (n_blocks * config.n_segs * (config.page_size_el // config.block_size),) - _zero_kv_blocks_kernel[grid]( - seg_addrs, - block_ids, - n_blocks, - N_SEGS=config.n_segs, - PAGE_SIZE_EL=config.page_size_el, - BLOCK_SIZE=config.block_size, - ) - - -def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: - from vllm.v1.worker.block_table import BlockTable - - # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny. - num_tokens = 1 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - - for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES: - # Use BlockTable so the JIT key matches the production slot-mapping call. - block_table = BlockTable( - block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - max_num_reqs=1, - max_num_blocks_per_req=block_table_stride, - max_num_batched_tokens=num_tokens, - pin_memory=False, - device=device, - kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE, - ) - block_table.add_row(list(range(block_table_stride)), 0) - block_table.commit_block_table(num_reqs=1) - block_table.compute_slot_mapping(1, query_start_loc, positions) - - def _warm_causal_conv1d_fwd_kernel( device: torch.device, config: _QwenGDNWarmupConfig ) -> None: @@ -369,16 +269,6 @@ def qwen_triton_warmup( device = getattr(runner, "device", torch.device("cuda")) logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) - zero_config = _zero_kv_warmup_config(runner) - warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner) - if zero_config is not None: - _warm_zero_kv_blocks_kernel(device, zero_config) - elif not warmed_zeroer: - logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") - - _warm_compute_slot_mapping_kernel(device) - _synchronize_device(device) - compilation_config = getattr(runner, "compilation_config", None) static_forward_context = getattr(compilation_config, "static_forward_context", None) gdn_config = _qwen_gdn_warmup_config(static_forward_context) diff --git a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py index 185449cb8ee6..cf35bf40e840 100644 --- a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py +++ b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py @@ -4,11 +4,10 @@ from typing import TYPE_CHECKING -import torch - from vllm.logger import init_logger if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker @@ -31,47 +30,7 @@ ) _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"}) -_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8) -_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2) -_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128) -_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3) -_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = ( - # query_slice_start offset, query_slice_stop offset - (0, 0), - (0, -1), - (1, 0), - (1, -1), -) -_COMBINE_TOPK_SWA_INPUT_VARIANTS = ( - # offset_topk, offset_query_and_seq, offset_gather - (False, False, False), - (False, True, False), - (True, True, True), -) -_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = ( - # compress_ratio, topk, topk_width, N - (1, 0, 512, 512), - (4, 512, 512, 512 * 4), - # DSv4-Pro C4A traffic uses top-k 1024 with N=1024. - (4, 1024, 1024, 1024), - (128, 8192, 8192, 8192 * 128), - # Real C128A traffic also specializes N=1 in one call path. - (128, 8192, 8192, 1), -) - - -def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int: - return max(0, min(num_tokens, max_tokens)) - - -def _next_power_of_2(x: int) -> int: - return 1 << (x - 1).bit_length() - - -def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int: - model_config = getattr(runner.vllm_config, "model_config", None) - hf_config = getattr(model_config, "hf_config", None) - return int(getattr(hf_config, name, default) or default) +_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"}) def _attention_backend_name(backend: object) -> str | None: @@ -96,242 +55,57 @@ def _has_attention_backend( return False -def _warm_sparse_swa_prefill_metadata_kernel( - device: torch.device, - window_size: int, - prefill_tokens: int, +def _compile_sparse_swa_prefill_metadata_kernel( + vllm_config: "VllmConfig", ) -> None: from vllm.v1.attention.backends.mla.sparse_swa import ( - _compute_prefill_metadata_kernel, + _COMPUTE_PREFILL_METADATA_KERNEL, ) - for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS: - for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES: - query_lens = [1] * num_decodes - query_lens += [prefill_tokens] * num_prefills - query_start_locs = [0] - for query_len in query_lens: - query_start_locs.append(query_start_locs[-1] + query_len) - query_start_loc = torch.tensor( - query_start_locs, - dtype=torch.int32, - device=device, - ) - seq_lens = torch.tensor( - [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]], - dtype=torch.int32, - device=device, - ) - prefill_gather_lens = torch.empty( - num_prefills, dtype=torch.int32, device=device - ) - _compute_prefill_metadata_kernel[(1,)]( - prefill_gather_lens, - seq_lens, - query_start_loc, - num_prefills, - num_decodes, - window_size, - BLOCK_SIZE=_next_power_of_2(num_prefills), - ) + _COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config) -def _warm_prefill_chunk_metadata_kernel( - device: torch.device, - compress_ratio: int, - query_len: int, +def _compile_prefill_chunk_metadata_kernel( + vllm_config: "VllmConfig", ) -> None: - from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata - - num_reqs = 2 - query_start_loc_cpu = torch.arange( - 0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32 + from vllm.v1.attention.backends.mla.indexer import ( + _BUILD_PREFILL_CHUNK_METADATA_KERNEL, ) - query_start_loc = query_start_loc_cpu.to(device=device) - uncompressed_seq_lens_cpu = torch.tensor( - [ - compress_ratio * multiplier + query_len - for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS - ], - dtype=torch.int32, - ) - compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio - uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device) - compressed_seq_lens = compressed_seq_lens_cpu.to(device=device) - block_table = torch.zeros( - (num_reqs, int(compressed_seq_lens_cpu.max().item())), - dtype=torch.int32, - device=device, - ) - - offset_uncompressed_seq_lens = torch.empty( - num_reqs + 1, dtype=torch.int32, device=device - )[1:] - offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens) - query_slices = tuple( - slice(start, num_reqs * query_len + stop) - for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS - ) - for warmup_uncompressed_seq_lens in ( - uncompressed_seq_lens, - offset_uncompressed_seq_lens, - ): - for query_slice in query_slices: - build_prefill_chunk_metadata( - 0, - num_reqs, - query_start_loc, - query_start_loc_cpu, - warmup_uncompressed_seq_lens, - compressed_seq_lens, - compressed_seq_lens_cpu, - block_table, - compress_ratio, - query_slice=query_slice, - ) + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config) -def _warm_combine_topk_swa_indices_kernel( - device: torch.device, - num_tokens: int, - window_size: int, - compress_ratio: int, - topk: int, - topk_width: int, - n: int, +def _compile_combine_topk_swa_indices_kernel( + vllm_config: "VllmConfig", ) -> None: - from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices - - if num_tokens <= 0: - return - - def _make_topk_indices(*, offset: bool) -> torch.Tensor: - if offset: - topk_storage = torch.full( - (num_tokens * topk_width + 1,), - -1, - dtype=torch.int32, - device=device, - ) - topk_indices = topk_storage[1:].reshape(num_tokens, topk_width) - else: - topk_indices = torch.full( - (num_tokens, topk_width), -1, dtype=torch.int32, device=device - ) - if topk > 0: - topk_indices.copy_( - torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device) - .reshape(num_tokens, topk_width) - .remainder(topk_width) - ) - return topk_indices - - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - seq_lens = torch.tensor( - [window_size + num_tokens], dtype=torch.int32, device=device + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + _COMBINE_TOPK_SWA_INDICES_KERNEL, ) - gather_lens = torch.tensor( - [min(window_size + num_tokens, window_size + num_tokens - 1)], - dtype=torch.int32, - device=device, - ) - offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:] - offset_query_start_loc.copy_(query_start_loc) - offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] - offset_seq_lens.copy_(seq_lens) - offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] - offset_gather_lens.copy_(gather_lens) - - for ( - offset_topk, - offset_query_and_seq, - offset_gather, - ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS: - warmup_topk_indices = _make_topk_indices(offset=offset_topk) - warmup_query_start_loc = ( - offset_query_start_loc if offset_query_and_seq else query_start_loc - ) - warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens - warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens - n_values = (n,) if n == 1 else (n, n + 1) - for m in (window_size + num_tokens, topk_width): - for n_value in n_values: - combine_topk_swa_indices( - warmup_topk_indices, - warmup_query_start_loc, - warmup_seq_lens, - warmup_gather_lens, - window_size, - compress_ratio, - topk, - M=m, - N=n_value, - ) - - -@torch.inference_mode() -def sparse_mla_triton_warmup( - runner: "GPUModelRunner", - num_tokens: int, - *, - compress_ratios: tuple[int, ...], - combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (), -) -> None: - device = getattr(runner, "device", torch.device("cuda")) - window_size = _hf_config_int(runner, "sliding_window", 128) - - _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens) - for compress_ratio in compress_ratios: - _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens) - for compress_ratio, topk, topk_width, n in combine_topk_swa_cases: - _warm_combine_topk_swa_indices_kernel( - device, - num_tokens, - window_size, - compress_ratio, - topk, - topk_width, - n, - ) - -def deepseek_v4_sparse_triton_warmup( - runner: "GPUModelRunner", - num_tokens: int, -) -> None: - sparse_mla_triton_warmup( - runner, - num_tokens, - compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS, - combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES, - ) + _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config) -def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None: +def sparse_mla_triton_warmup(worker: "Worker") -> None: runner = worker.model_runner if runner.is_pooling_model: return max_tokens = worker.scheduler_config.max_num_batched_tokens - num_tokens = _clamp_warmup_tokens(8, max_tokens) - if num_tokens <= 0: + max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens) + if max_tokens <= 0 or max_num_prefills <= 0: return + vllm_config = runner.vllm_config try: if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS): - deepseek_v4_sparse_triton_warmup(runner, num_tokens) + _compile_sparse_swa_prefill_metadata_kernel(vllm_config) + _compile_prefill_chunk_metadata_kernel(vllm_config) + _compile_combine_topk_swa_indices_kernel(vllm_config) elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS): - sparse_mla_triton_warmup( - runner, - num_tokens, - compress_ratios=(1,), - ) + _compile_sparse_swa_prefill_metadata_kernel(vllm_config) + _compile_prefill_chunk_metadata_kernel(vllm_config) elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS): - _warm_prefill_chunk_metadata_kernel( - getattr(runner, "device", torch.device("cuda")), - compress_ratio=1, - query_len=num_tokens, - ) + _compile_prefill_chunk_metadata_kernel(vllm_config) + except Exception: logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True) diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py index 8d2328432ebc..d49e1ba7cc89 100644 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -2,42 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm up v1 block-table Triton kernels.""" +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + _SLOT_MAPPING_WARMUP_TOKENS = 8 -_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16) -_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -def warm_v1_block_table_kernels( - device: torch.device, - max_tokens: int, -) -> None: - from vllm.v1.worker.block_table import BlockTable +def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: + """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens)) + device = runner.device + block_table = runner.input_batch.block_table + num_tokens = min( + _SLOT_MAPPING_WARMUP_TOKENS, + runner.scheduler_config.max_num_batched_tokens, + ) if num_tokens <= 0: return query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES: - max_num_blocks_per_req = max( - 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size - ) - max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16 - block_table = BlockTable( - block_size=block_size, - max_num_reqs=1, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max(num_tokens, max_tokens), - pin_memory=False, - device=device, - kernel_block_size=block_size, - cp_kv_cache_interleave_size=( - _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE - ), - ) - block_table.add_row(list(range(max_num_blocks_per_req)), 0) - block_table.commit_block_table(1) - block_table.compute_slot_mapping(1, query_start_loc, positions) + block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/models/common/__init__.py b/vllm/models/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/common/ops/__init__.py b/vllm/models/common/ops/__init__.py new file mode 100644 index 000000000000..0d0fe8de26ec --- /dev/null +++ b/vllm/models/common/ops/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ops shared across model implementations.""" + +from .fused_qk_rmsnorm import fused_q_kv_rmsnorm + +__all__ = [ + "fused_q_kv_rmsnorm", +] diff --git a/vllm/models/deepseek_v32/nvidia/fused_ops.py b/vllm/models/common/ops/fused_allreduce_rms_norm.py similarity index 91% rename from vllm/models/deepseek_v32/nvidia/fused_ops.py rename to vllm/models/common/ops/fused_allreduce_rms_norm.py index 6a795e2a1530..443c8f3eec3b 100644 --- a/vllm/models/deepseek_v32/nvidia/fused_ops.py +++ b/vllm/models/common/ops/fused_allreduce_rms_norm.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path). +"""Fused all-reduce + residual-add + RMSNorm for eager model paths. -These recover fusions that vLLM's torch.compile passes would normally do but -that don't fire when running eager under the breakable CUDA graph. +This recovers a fusion that vLLM's torch.compile passes would normally do but +that doesn't fire for models running eager (or under a breakable CUDA graph). """ import torch diff --git a/vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py similarity index 92% rename from vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py rename to vllm/models/common/ops/fused_qk_rmsnorm.py index 0dd348a46e26..ac7d0dfd4a14 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -21,6 +22,7 @@ def _fused_q_kv_rmsnorm_kernel( Q_SIZE: tl.constexpr, KV_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, ): # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y. # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes @@ -41,6 +43,10 @@ def _fused_q_kv_rmsnorm_kernel( weight_ptr = kv_weight_ptr row_out = kv_out_ptr + token_idx * kv_out_stride + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which # keep x, rrms, and w all in fp32 and perform a single cast at store. @@ -92,5 +98,6 @@ def fused_q_kv_rmsnorm( Q_SIZE=q_size, KV_SIZE=kv_size, BLOCK_SIZE=block_size, + launch_pdl=current_platform.is_arch_support_pdl(), ) return qr_out, kv_out diff --git a/vllm/models/common/ops/sequence_parallel.py b/vllm/models/common/ops/sequence_parallel.py new file mode 100644 index 000000000000..8631405eaec7 --- /dev/null +++ b/vllm/models/common/ops/sequence_parallel.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) + + +def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None: + device_communicator = get_tp_group().device_communicator + if device_communicator is None: + return None + collective = getattr(device_communicator, name, None) + return None if collective is None else collective(x) + + +def sp_all_gather(x: torch.Tensor) -> torch.Tensor: + output = _custom_collective("custom_all_gather", x) + if output is not None: + return output + return tensor_model_parallel_all_gather(x, 0) + + +def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + output = _custom_collective("custom_reduce_scatter", x) + if output is not None: + return output + return tensor_model_parallel_reduce_scatter(x, 0) + + +def sp_shard(x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + pad = (0, 0) * (x.ndim - 1) + (0, sp_pad) + x = torch.nn.functional.pad(x, pad) + chunk = x.shape[0] // tp_size + return x[tp_rank * chunk : (tp_rank + 1) * chunk] + + +def sp_padding_mask( + is_padding: torch.Tensor | None, + hidden_states: torch.Tensor, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + if is_padding is None: + is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool) + assert is_padding.shape[0] == num_tokens + + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-num_tokens) % tp_size + if sp_pad > 0: + is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True) + chunk = is_padding.shape[0] // tp_size + tp_rank = get_tensor_model_parallel_rank() + return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk] diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py index 1b0aa64262fe..f3eba73142b4 100644 --- a/vllm/models/deepseek_v32/__init__.py +++ b/vllm/models/deepseek_v32/__init__.py @@ -10,11 +10,15 @@ from vllm.platforms import current_platform -if current_platform.is_rocm() or current_platform.is_xpu(): - raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.") - -from .nvidia.model import DeepseekV32ForCausalLM -from .nvidia.mtp import DeepseekV32MTP +if current_platform.is_rocm(): + from .amd.model import DeepseekV32ForCausalLM + from .amd.mtp import DeepseekV32MTP +elif current_platform.is_xpu(): + raise NotImplementedError("deepseek_v32 does not yet support XPU.") +else: + # Covers Blackwell (sm100) and all other CUDA devices. + from .nvidia.model import DeepseekV32ForCausalLM + from .nvidia.mtp import DeepseekV32MTP __all__ = [ "DeepseekV32ForCausalLM", diff --git a/vllm/models/deepseek_v32/amd/__init__.py b/vllm/models/deepseek_v32/amd/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v32/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py new file mode 100644 index 000000000000..f5067bd0b99c --- /dev/null +++ b/vllm/models/deepseek_v32/amd/model.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import torch + +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.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV2ForCausalLM, + DeepseekV2MLP, + DeepseekV2MoE, + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, +) +from vllm.model_executor.models.utils import ( + PPMissingLayer, + get_pp_missing_layer_names, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, +) +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.sequence import IntermediateTensors + +from .rocm import DeepseekV32MLAAttention + + +class DeepseekV32DecoderLayer(torch.nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + ) -> None: + super().__init__() + + if config is None: + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + + self.hidden_size = config.hidden_size + moe_layer_freq = getattr(config, "moe_layer_freq", 1) + layer_idx = int(prefix.split(sep=".")[-1]) + self.layer_idx = layer_idx + self.use_mha = False + + self.self_attn = DeepseekV32MLAAttention( + vllm_config=vllm_config, + config=config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + ) + + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + ): + self.mlp = DeepseekV2MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.mlp.experts.moe_config.skip_final_all_reduce = True + else: + self.mlp = DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + reduce_results=False, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class DeepseekV32Model(torch.nn.Module): + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + from vllm.platforms import current_platform + + self.device = current_platform.device_type + + self.vocab_size = config.vocab_size + self.is_v32 = True + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV32DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + self.aux_hidden_state_layers = tuple[int, ...]() + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + assert input_ids is not None + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + aux_hidden_states = [] + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + if idx in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states + residual) + hidden_states, residual = layer(positions, hidden_states, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + expert_params_mapping = 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.n_routed_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if get_spec_layer_idx_from_weight_name(self.config, name) is not None: + continue + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, + ): + continue + + 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 + name_mapped = name.replace(weight_name, param_name) + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + name = name_mapped + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment] + if name is None: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): + model_cls = DeepseekV32Model + + def set_moe_parameters(self): + self.num_expert_groups = getattr(self.config, "n_group", 1) + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, DeepseekV2MoE): + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + self.extract_moe_parameters(example_moe) diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py new file mode 100644 index 000000000000..109bd2662265 --- /dev/null +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import typing +from collections.abc import Callable, Iterable + +import torch +import torch.nn as nn + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_reduce +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 ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV2MixtureOfExperts, + DeepseekV2MoE, + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, +) +from vllm.model_executor.models.utils import ( + get_pp_missing_layer_names, + maybe_prefix, +) +from vllm.models.deepseek_v32.common.kernels import fused_eh_norm +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from .model import DeepseekV32DecoderLayer + + +class DeepseekV32MultiTokenPredictorLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + + 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) + + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=current_platform.device_type, + ) + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + self.mtp_block = DeepseekV32DecoderLayer( + vllm_config, + prefix, + config=config, + topk_indices_buffer=topk_indices_buffer, + ) + + 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 + eh_input = fused_eh_norm( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + hidden_states = self.eh_proj(eh_input) + hidden_states, residual = self.mtp_block( + positions=positions, hidden_states=hidden_states, residual=None + ) + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states = residual + hidden_states + hidden_states = self.shared_head.norm(hidden_states) + return hidden_states, hidden_states + + +class DeepseekV32MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.layers = torch.nn.ModuleDict( + { + str(idx): DeepseekV32MultiTokenPredictorLayer( + vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def set_skip_topk(self, skip: bool): + for layer in self.layers.values(): + self_attn = getattr(layer.mtp_block, "self_attn", None) + if self_attn is not None and hasattr(self_attn, "skip_topk"): + self_attn.skip_topk = skip + + 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, + ) -> 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, hidden_states) + + +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = DeepseekV32MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.set_moe_parameters() + + def set_moe_parameters(self): + self.num_moe_layers = self.config.num_nextn_predict_layers + self.num_expert_groups = self.config.n_group + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers.values(): + mlp = layer.mtp_block.mlp + if isinstance(mlp, DeepseekV2MoE): + example_moe = mlp + self.moe_mlp_layers.append(mlp) + self.moe_layers.append(mlp.experts) + self.extract_moe_parameters(example_moe) + + 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 | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + rocm_aiter_moe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + expert_params_mapping = 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.n_routed_experts + + ( + self.config.n_shared_experts + if rocm_aiter_moe_shared_expert_enabled + else 0 + ), + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + is_fusion_moe_shared_experts_layer = ( + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + ) + name = self._rewrite_spec_layer_name(spec_layer, name) + + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, + ): + continue + + 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 + if is_fusion_moe_shared_experts_layer: + continue + name_mapped = name.replace(weight_name, param_name) + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + else: + name = name_mapped + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + num_chunks = 1 + if is_fusion_moe_shared_experts_layer: + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 + split_dim = ( + 1 + if ("down_proj.weight" in name and loaded_weight.ndim > 1) + else 0 + ) + total = loaded_weight.shape[split_dim] + assert total % num_chunks == 0 + chunk_size = total // num_chunks + + for j in range(num_chunks): + chunk_name = name + weight_to_load = loaded_weight + if is_fusion_moe_shared_experts_layer: + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) + if loaded_weight.ndim == 1: + weight_to_load = loaded_weight[chunk_slice] + elif split_dim == 0: + weight_to_load = loaded_weight[chunk_slice, :] + else: + weight_to_load = loaded_weight[:, chunk_slice] + chunk_name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts + j}", + ) + + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] + if weight_name not in chunk_name: + continue + is_expert_weight = True + name_mapped = chunk_name.replace(weight_name, param_name) + param = params_dict[name_mapped] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + weight_to_load, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fusion_moe_shared_experts_layer: + name = name_mapped + else: + loaded_params.add(name_mapped) + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment] + if name is None: + continue + if ( + spec_layer != self.model.mtp_start_layer_idx + and ".layers" not in name + ): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) + + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_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_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint." + ) + return loaded_params diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py new file mode 100644 index 000000000000..b71fc9ddfc30 --- /dev/null +++ b/vllm/models/deepseek_v32/amd/rocm.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache +from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer +from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, +) + + +class DeepseekV32MLASparseBackend(ROCMAiterMLASparseBackend): + @staticmethod + def get_supported_kernel_block_sizes() -> list: + return [16, 32] + + +class DeepseekV32ROCmIndexerBackend(DeepseekV32IndexerBackend): + @staticmethod + def get_supported_kernel_block_sizes() -> list: + return [16, 32] + + +class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache): + def get_attn_backend(self): + return DeepseekV32ROCmIndexerBackend + + +class DeepseekV32ROCmIndexer(DeepseekV32Indexer): + indexer_cache_cls = DeepseekV32ROCmIndexerCache + + +class DeepseekV32MLAAttention(DeepseekV32Attention): + require_fp8_kv_cache: bool = False + indexer_cls = DeepseekV32ROCmIndexer + + def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None): + super().__init__( + vllm_config, + config, + prefix, + topk_indices_buffer, + attn_backend=DeepseekV32MLASparseBackend, + ) + + self.indexer_op: SparseAttnIndexer | None = None + if self.indexer is not None: + self.indexer_op = SparseAttnIndexer( + self.indexer.k_cache, + self.indexer.quant_block_size, + self.indexer.scale_fmt, + self.indexer.topk_tokens, + self.indexer.head_dim, + self.indexer.max_model_len, + self.indexer.max_total_seq_len, + topk_indices_buffer, + skip_k_cache_insert=True, + ) + self._fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype) + self._fp8_kv_needs_view = self._fp8_kv and self.kv_cache_dtype != "fp8_ds_mla" + + def _compute_ql_nope(self, q_c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope = q_nope.transpose(0, 1) # (N, tokens, P) + + if self.is_aiter_triton_fp4_bmm_enabled: + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4 + + ql_nope = batched_gemm_a16wfp4( + q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True + ) + elif self.is_aiter_triton_fp8_bmm_enabled: + from vllm._aiter_ops import rocm_aiter_ops + + ql_nope = rocm_aiter_ops.triton_fp8_bmm( + q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True + ) + else: + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) + + return ql_nope, q_pe + + def _run_indexer( + self, + q_c: torch.Tensor, + index_q_fp8: torch.Tensor | None, + index_weights_out: torch.Tensor | None, + ) -> None: + """Run the ROCm sparse indexer (forward_hip) if this layer has an indexer.""" + if self.indexer_op is not None: + self.indexer_op.forward_hip(q_c, index_q_fp8, None, index_weights_out) + + def _build_q_for_attn( + self, + ql_nope: torch.Tensor, + mqa_q: torch.Tensor, + num_actual: int, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if self._fp8_kv: + # fp8 KV: mqa_q is the full [ql_nope; q_pe] packed as fp8. + return mqa_q[:num_actual] + + return (ql_nope[:num_actual], mqa_q[:num_actual]) + + def _compute_uv_out( + self, + attn_out: torch.Tensor, + output: torch.Tensor, + num_actual: int, + ) -> None: + x = attn_out.view( + num_actual, self.num_local_heads, self.kv_lora_rank + ).transpose(0, 1) # (N, tokens, L) + out_view = output[:num_actual].view( + num_actual, self.num_local_heads, self.v_head_dim + ) + + if self.is_aiter_triton_fp4_bmm_enabled: + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4 + + batched_gemm_a16wfp4( + x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True + ) + elif self.is_aiter_triton_fp8_bmm_enabled: + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.triton_fp8_bmm( + x, + self.W_V, + self.W_V_scale, + group_size=128, + transpose_bm=True, + YQ=out_view, + ) + else: + torch.bmm(x, self.W_UV, out=out_view.transpose(0, 1)) + + @eager_break_during_capture + def _fused_attention( + self, + positions: torch.Tensor, + q_c: torch.Tensor, + kv_c: torch.Tensor, + k_pe: torch.Tensor, + index_k: torch.Tensor | None, + index_weights: torch.Tensor | None, + output: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if isinstance(attn_metadata_raw, dict): + attn_metadata = attn_metadata_raw.get(self.layer_name) + elif isinstance(attn_metadata_raw, list): + attn_metadata = attn_metadata_raw[0].get(self.layer_name) + else: + attn_metadata = attn_metadata_raw + + slot_mapping = forward_context.slot_mapping + assert isinstance(slot_mapping, dict) + mla_slot = slot_mapping.get(self.layer_name) + + if self.indexer is not None: + has_indexer = True + indexer_k_norm_w = self.indexer.k_norm.weight + indexer_k_norm_bias = self.indexer.k_norm.bias + indexer_k_norm_eps = self.indexer.k_norm.eps + indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache + indexer_k_cache = self.indexer.k_cache.kv_cache + indexer_softmax_scale = self.indexer.softmax_scale + indexer_n_head_scale = self.indexer.n_head**-0.5 + else: + has_indexer = False + indexer_k_norm_w = None + indexer_k_norm_bias = None + indexer_k_norm_eps = 1e-6 + indexer_k_rope_cos_sin_cache = None + indexer_k_cache = None + indexer_softmax_scale = 0.0 + indexer_n_head_scale = 0.0 + + if attn_metadata is None: + mla_kv_cache = None + mla_k_scale = None + indexer_k_cache = None + mla_slot = None + else: + mla_kv_cache = self.kv_cache + mla_k_scale = self._k_scale + + q_c = fused_norm_rope( + positions, + q_c, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + kv_c, + self.kv_a_layernorm.weight, + self.kv_a_layernorm.variance_epsilon, + k_pe, + self.rotary_emb.cos_sin_cache, + index_k, + indexer_k_norm_w, + indexer_k_norm_bias, + indexer_k_norm_eps, + indexer_k_rope_cos_sin_cache, + self.topk_indices_buffer, + slot_mapping=mla_slot, + indexer_k_cache=indexer_k_cache, + mla_kv_cache=mla_kv_cache, + mla_kv_cache_dtype=self.kv_cache_dtype, + mla_k_scale=mla_k_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, + ) + + ql_nope, q_pe = self._compute_ql_nope(q_c) + + if self.indexer is not None: + index_q = self.indexer.wq_b(q_c)[0] + index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) + else: + index_q = None + + index_q_fp8, index_weights_out, mqa_q = fused_q( + positions, + q_pe, + self.rotary_emb.cos_sin_cache, + index_q, + self.indexer_rope_emb.cos_sin_cache if has_indexer else None, + ql_nope, + self._q_scale, + index_weights, + indexer_softmax_scale, + indexer_n_head_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, + quantize_mqa=self._fp8_kv, + ) + + self._run_indexer(q_c, index_q_fp8, index_weights_out) + + if attn_metadata is None: + output.zero_() + return + + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] + kv_cache = self.kv_cache + if self._fp8_kv_needs_view: + kv_cache = kv_cache.view(torch.float8_e4m3fn) + + q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual) + attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] + q_for_attn, kv_cache, attn_metadata, self + ) + + self._compute_uv_out(attn_out, output, num_actual) diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/attention.py similarity index 80% rename from vllm/models/deepseek_v32/nvidia/attention.py rename to vllm/models/deepseek_v32/attention.py index dcf955ad59ba..79d350b1bcc4 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -31,12 +31,13 @@ yarn_get_mscale, ) from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q from vllm.utils.torch_utils import is_quantized_kv_cache -from .kernels import fused_norm_rope, fused_q - class DeepseekV32Indexer(nn.Module): + indexer_cache_cls = DeepseekV32IndexerCache + def __init__( self, vllm_config: VllmConfig, @@ -55,7 +56,6 @@ def __init__( self.rope_dim = config.qk_rope_head_dim self.q_lora_rank = q_lora_rank - # No tensor parallel, just replicated. self.wq_b = ReplicatedLinear( self.q_lora_rank, self.head_dim * self.n_head, @@ -63,8 +63,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. - # FP8 wk weights are upcasted to BF16 during loading to keep this fused. + self.wk_weights_proj = MergedColumnParallelLinear( hidden_size, [self.head_dim, self.n_head], @@ -80,9 +79,8 @@ def __init__( self.quant_block_size = 128 self.topk_indices_buffer = topk_indices_buffer - # fp8 naive cache: value in fp8 + fp32 scale per quant_block_size element. assert cache_config is not None, "DeepSeek V3.2 indexer requires cache_config" - self.k_cache = DeepseekV32IndexerCache( + self.k_cache = type(self).indexer_cache_cls( head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, dtype=torch.uint8, prefix=f"{prefix}.k_cache", @@ -120,7 +118,7 @@ def forward( q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - # Fused wk + weights_proj: one GEMM, then split. + kw, _ = self.wk_weights_proj(hidden_states) k = kw[:, : self.head_dim] weights = kw[:, self.head_dim :] @@ -131,14 +129,13 @@ def forward( ) q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) - # RoPE (NeoX) can introduce extra leading dims; reshape back to flat. + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) k_pe = k_pe.reshape(-1, 1, self.rope_dim) q = torch.cat([q_pe, q_nope], dim=-1) k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) - # Only quant q here; k quant is fused with cache insertion. q = q.view(-1, self.head_dim) q_fp8, q_scale = per_token_group_quant_fp8( q, @@ -158,9 +155,10 @@ def forward( class DeepseekV32Attention(MLAAttention): - # Narrow the base's broadly-typed `indexer` to the concrete type so the - # `if self.indexer is not None` guards below type-check its attributes. indexer: "DeepseekV32Indexer | None" + indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer + + require_fp8_kv_cache: bool = True def __init__( self, @@ -168,6 +166,7 @@ def __init__( config: DeepseekV2Config | DeepseekV3Config, prefix: str, topk_indices_buffer: torch.Tensor | None = None, + attn_backend: "type | None" = None, ) -> None: quant_config = vllm_config.quant_config cache_config = vllm_config.cache_config @@ -200,11 +199,6 @@ def __init__( mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) scaling = scaling * mscale * mscale - # DSA "shared indexer" pattern: only some layers carry an indexer; the - # rest reuse the top-k written by the previous indexer layer into the - # shared topk_indices_buffer. DeepSeek-V3.2 builds it on every layer - # (index_topk_freq defaults to 1); GLM-5.2 uses index_topk_freq=4 so - # only layers [0,1,2,6,10,...] (+ MTP) carry one. layer_id = extract_layer_index(prefix) index_topk_freq = getattr(config, "index_topk_freq", 1) index_topk_pattern = getattr(config, "index_topk_pattern", None) @@ -217,12 +211,10 @@ def __init__( skip_topk = index_topk_pattern[layer_id] == "S" else: skip_topk = False - # MTP/nextn layers always build a full indexer (they toggle at runtime). + num_hidden_layers = getattr(config, "num_hidden_layers", None) is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers - # Build kv_b_proj + indexer first; they are passed to MLAAttention.__init__ - # (which runs nn.Module.__init__ and registers them). kv_b_proj = ColumnParallelLinear( kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), @@ -232,7 +224,7 @@ def __init__( ) indexer = None if not skip_topk or is_mtp_layer: - indexer = DeepseekV32Indexer( + indexer = type(self).indexer_cls( vllm_config, config, hidden_size, @@ -243,8 +235,6 @@ def __init__( prefix=f"{prefix}.indexer", ) - # Set up the MLA engine (impl, KV cache, scales, backend, registration, - # and process_weights_after_loading) via the MLAAttention base. super().__init__( num_heads=num_local_heads, scale=scaling, @@ -260,41 +250,32 @@ def __init__( use_sparse=True, indexer=indexer, topk_indices_buffer=topk_indices_buffer, + attn_backend=attn_backend, ) self.num_local_heads = num_local_heads self.qk_head_dim = qk_head_dim self.indexer = indexer self.topk_indices_buffer = topk_indices_buffer - # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0 - # computes the top-k, steps 1+ set this True to reuse it. + self.skip_topk = False - # Fused fp8 paths: Triton fused norm/rope/cache + fused-q. Two layouts, - # picked by the sparse MLA backend's query support: - # * supports_quant_query_input (FlashInfer sparse, SM100): per-tensor - # fp8 cache + a single packed fp8 MQA query. - # * not supported (FlashMLA sparse, SM90/SM100): fp8_ds_mla cache - # (per-128 block-scaled fp8 NoPE + unquantized bf16 RoPE) + a bf16 - # (ql_nope, q_pe) query tuple. FA3 cannot mix a bf16 query with an - # fp8 KV cache, so FlashMLA (which dequantizes internally) is used. - # FlashMLA sparse runs on both Hopper and Blackwell, so this is the - # only DSA path on SM90 and an opt-in alternative on SM100. - assert is_quantized_kv_cache(self.kv_cache_dtype), ( - "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse " - "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) " - "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)." - ) - self._fp8_query = self.impl.supports_quant_query_input - if not self._fp8_query: - assert self.kv_cache_dtype == "fp8_ds_mla", ( - "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend " - "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. " - "Launch with --kv-cache-dtype fp8_ds_mla." + + if self.require_fp8_kv_cache: + assert is_quantized_kv_cache(self.kv_cache_dtype), ( + "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse " + "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) " + "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)." ) - # The paged KV cache is stored as uint8 and viewed as fp8 for the decode - # (per-tensor fp8). The fp8_ds_mla layout is consumed as raw bytes. - self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla" - # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX. + self._fp8_query = self.impl.supports_quant_query_input + if not self._fp8_query: + assert self.kv_cache_dtype == "fp8_ds_mla", ( + "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend " + "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. " + "Launch with --kv-cache-dtype fp8_ds_mla." + ) + + self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla" + self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) # Remaining MLA projections (registered on this module). @@ -313,9 +294,7 @@ def __init__( prefix=f"{prefix}.q_b_proj", ) self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps) - # reduce_results=False: the attention all-reduce is fused with the - # following post_attention_layernorm in the decoder layer via - # fused_allreduce_rms_norm. + self.o_proj = RowParallelLinear( num_heads * v_head_dim, hidden_size, @@ -380,11 +359,6 @@ def _fused_attention( index_weights: torch.Tensor | None, output: torch.Tensor, ) -> None: - # One eager break for the whole attention. In FULL cudagraph mode (pure - # decode) this decorator is a no-op, so everything here is captured; in - # PIECEWISE (prefill) it runs eagerly. The cache writes, sparse indexer, - # and forward_mqa all depend on per-step metadata and must not be split - # out (PIECEWISE capture would otherwise miss them). forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata if isinstance(attn_metadata_raw, dict): @@ -494,8 +468,10 @@ def _fused_attention( self.indexer.max_model_len, self.indexer.max_total_seq_len, self.topk_indices_buffer, - True, # skip_k_cache_insert - False, # use_fp4_cache + skip_k_cache_insert=True, + use_pcp=False, + dense_mha_metadata_layer_name="", + use_fp4_cache=False, # fused_norm_rope already cleared the topk buffer this forward. skip_topk_buffer_clear=True, ) @@ -514,8 +490,6 @@ def _fused_attention( :num_actual ] else: - # FlashMLA sparse: bf16 (ql_nope, q_pe) tuple. mqa_q is the RoPE'd - # q_pe; ql_nope is consumed directly. mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual]) attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] mqa_q_arg, kv_cache, attn_metadata, self diff --git a/vllm/models/deepseek_v32/common/__init__.py b/vllm/models/deepseek_v32/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v32/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v32/nvidia/kernels.py b/vllm/models/deepseek_v32/common/kernels.py similarity index 100% rename from vllm/models/deepseek_v32/nvidia/kernels.py rename to vllm/models/deepseek_v32/common/kernels.py diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 353aedc8ceed..240e97329ea7 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -38,11 +38,10 @@ make_layers, sequence_parallel_chunk, ) +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.models.deepseek_v32.attention import DeepseekV32Attention from vllm.sequence import IntermediateTensors -from .attention import DeepseekV32Attention -from .fused_ops import fused_allreduce_rms_norm - def _all_gather_sp_states( hidden_states: torch.Tensor, diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index d3d8e5aae9c5..828cdd02080d 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -8,7 +8,10 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig -from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed import ( + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -17,14 +20,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) -from vllm.model_executor.models.deepseek_mtp import ( - SharedHead, - _restore_full_token_layout_if_needed, -) +from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import ( DeepseekV2MixtureOfExperts, DeepseekV2MoE, @@ -35,10 +38,10 @@ get_pp_missing_layer_names, maybe_prefix, ) +from vllm.models.deepseek_v32.common.kernels import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from .kernels import fused_eh_norm from .model import DeepseekV32DecoderLayer @@ -92,13 +95,8 @@ def forward( hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) - hidden_states, residual = _restore_full_token_layout_if_needed( - hidden_states, - residual, - positions.shape[0], - is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, - ) - if not self.mtp_block.use_sequence_parallel_moe: + is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe + if not is_sequence_parallel: # Without sequence parallelism, the MoE output is left un-reduced. hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The @@ -112,6 +110,9 @@ def forward( # the legacy proposer (model_returns_tuple is True for the # DeepSeekMTPModel architecture). hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[0]] return hidden_states, hidden_states @@ -416,7 +417,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled(): raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint." diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 9093bd412899..becfd1d4f708 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -8,7 +8,9 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -16,7 +18,7 @@ ) from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -103,13 +105,100 @@ def __init__( else: self.act_fn = SiluAndMul() + # gate_up_proj B-preshuffle (ColumnParallel -> no all-reduce); set at load. + self._gateup = rocm_aiter_ops.is_enabled() + # Block scale for the preshuffled gate_up weight; None = not preshuffled. + self._gateup_scale: torch.Tensor | None = None + + def prepare_gateup_preshuffle(self) -> None: + # B-preshuffle the gate_up_proj weight in place (single weight). + if not self._gateup: + return + from vllm.model_executor.utils import replace_parameter + + w = getattr(self.gate_up_proj, "weight", None) + ws = getattr(self.gate_up_proj, "weight_scale_inv", None) # per-block scale + if w is None or ws is None or w.dim() != 2: + return + # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold. + if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0: + return + if ws.dtype == torch.float8_e8m0fnu: + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, + ) + + ws = _upcast_e8m0_to_fp32(ws).contiguous() + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) + self._gateup_scale = ws + def forward(self, x): - gate_up, _ = self.gate_up_proj(x) + if self._gateup_scale is not None and x.dim() == 2: + # gate_up via fp8 group-quant (col-major) + B-preshuffle GEMM. + x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True) + gate_up = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + x_fp8, + self.gate_up_proj.weight, + x_scale, + self._gateup_scale, + output_dtype=x.dtype, + ) + else: + gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x +def _shared_experts_are_fp4(config, layer_idx: int | None = None) -> bool: + """Whether the shared experts are MXFP4 and thus fusable. + + ``layer_idx=None`` resolves the model-wide default (global scheme), used by + the main-model weight loader / mapper callers that operate per-model. + """ + quant_cfg = getattr(config, "quantization_config", None) + if quant_cfg is None: + return False + if layer_idx is None: + base = None + elif layer_idx >= config.num_hidden_layers: + base = f"mtp.{layer_idx - config.num_hidden_layers}.ffn.shared_experts" + else: + base = f"layers.{layer_idx}.ffn.shared_experts" + if base and any(e.startswith(base) for e in (quant_cfg.get("exclude") or [])): + return False + entry = ( + (quant_cfg.get("layer_quant_config") or {}).get(f"{base}.w1") if base else None + ) + if entry is None: + entry = quant_cfg.get("global_quant_config") + return ((entry or {}).get("weight") or {}).get("dtype") == "fp4" + + +def _fuse_shared_experts_enabled(config, prefix: str = "") -> bool: + """Whether to fuse the shared expert into the routed MXFP4 grouped GEMM. + + Fusion fuses the shared expert into the routed experts' MXFP4 grouped GEMM, + so it only applies where the shared expert is the same precision as the + routed experts. Some layers may carry a shared expert in a different quantization + than the routed experts; when so, it runs as its own linear and must not be fused. + """ + if not ( + current_platform.is_rocm() + and getattr(config, "n_shared_experts", None) + and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS + and not get_current_vllm_config().parallel_config.enable_expert_parallel + ): + return False + return _shared_experts_are_fp4( + config, extract_layer_index(prefix) if prefix else None + ) + + class DeepseekV4MoE(nn.Module): def __init__( self, @@ -164,7 +253,11 @@ def __init__( requires_grad=False, ) - if config.n_shared_experts is None: + self.n_shared_experts = config.n_shared_experts + + self.fuse_shared_experts = _fuse_shared_experts_enabled(config, prefix) + + if config.n_shared_experts is None or self.fuse_shared_experts: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -186,8 +279,11 @@ def __init__( self.experts_start_idx = self.tp_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, + n_shared_experts=( + config.n_shared_experts if self.fuse_shared_experts else None + ), gate=self.gate, num_experts=config.n_routed_experts, top_k=config.num_experts_per_tok, @@ -212,7 +308,7 @@ def forward( org_shape = hidden_states.shape if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class + # In this case, the gate/router runs inside the MoERunner class final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states, @@ -520,13 +616,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): requires_grad=False, ) self.hc_head_op = HCHeadOp() - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address (outside the cudagraph pool) so the copy_ in forward() - # refreshes it correctly across captured shapes. - # refreshes it correctly across captured shapes. Only allocated on - # the last PP rank — that's where MTP target hidden states are - # produced. - if get_pp_group().is_last_rank: + spec_config = vllm_config.speculative_config + needs_mtp_hidden_states = spec_config is not None and ( + spec_config.use_eagle() or spec_config.uses_draft_model() + ) + if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, @@ -626,9 +720,9 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + if self._mtp_hidden_buffer is not None: + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) hidden_states = self.hc_head_op( hidden_states, @@ -667,7 +761,38 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() + # Use each MoE's own per-layer fusion decision (computed with its prefix + # at init) as the single source of truth, so the redirect below cannot + # diverge from how the module was built if per-layer quantization ever + # mixes fused and non-fused layers. + fuse_by_layer = { + extract_layer_index(mod_name): mod.fuse_shared_experts + for mod_name, mod in self.named_modules() + if isinstance(mod, DeepseekV4MoE) + } + n_routed = self.config.n_routed_experts + # The redirect below maps the single shared-expert tensor group to one + # appended slot; multiple shared experts would need per-expert slicing + # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1. + if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1: + raise NotImplementedError( + "deepseek-v4 fused shared-expert loading supports only " + f"n_shared_experts == 1, got {self.config.n_shared_experts}" + ) + for name, loaded_weight in weights: + # Shared-expert fusion: redirect ``.ffn.shared_experts.w{1,2,3}`` + # into appended routed-expert slot ``.ffn.experts.{n_routed}`` + # so the MXFP4-quantized shared expert loads through the routed + # expert loader (grouped GEMM). Single shared expert only. + if ".ffn.shared_experts.w" in name and fuse_by_layer.get( + extract_layer_index(name), False + ): + name = name.replace( + ".ffn.shared_experts.w", + f".ffn.experts.{n_routed}.w", + ) + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -745,24 +870,41 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + # When fusing shared experts, include the appended slots + # (ids n_routed_experts .. n_routed_experts + n_shared - 1) so the + # redirected shared-expert weights route through the expert loader. + n_shared = getattr(self.config, "n_shared_experts", 0) or 0 + num_experts = self.config.n_routed_experts + ( + n_shared if _fuse_shared_experts_enabled(self.config) else 0 + ) return fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", - num_experts=self.config.n_routed_experts, + num_experts=num_experts, ) -def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: +def _make_deepseek_v4_weights_mapper( + expert_dtype: str, fuse_shared_experts: bool = False +) -> WeightsMapper: if expert_dtype == "fp4": # MXFP4 experts use Mxfp4MoEMethod, which registers scales as # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and - # shared experts use Fp8LinearMethod's block scales, which - # register as ``weight_scale_inv``. + # (non-fused) shared experts use Fp8LinearMethod's block scales, + # which register as ``weight_scale_inv``. + # + # - DeepSeek native ``.scale``: expert scales -> ``.weight_scale``, + # everything else -> ``.weight_scale_inv``. + # - AMD-Quark ``.weight_scale``: linear/attn scales -> + # ``.weight_scale_inv``. Expert and shared-expert + # ``w{1,2,3}.weight_scale`` are left untouched (consumed as-is by + # the MXFP4 expert loader, which produces ``w{13,2}_weight_scale``); scale_regex = { re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", re.compile(r"\.scale$"): ".weight_scale_inv", + re.compile(r"(? WeightsMapper: scale_regex = { re.compile(r"\.scale$"): ".weight_scale_inv", } + # When shared experts are fused into the routed MXFP4 grouped GEMM, the + # shared_experts tensors are redirected to routed expert slots ; leave + # their names untouched here. + substr_map = ( + {} + if fuse_shared_experts + else {".shared_experts.w2": ".shared_experts.down_proj"} + ) return WeightsMapper( orig_to_new_prefix={ "layers.": "model.layers.", @@ -785,9 +935,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: "embed.weight": "embed_tokens.weight", ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", }, - orig_to_new_substr={ - ".shared_experts.w2": ".shared_experts.down_proj", - }, + orig_to_new_substr=substr_map, ) @@ -804,8 +952,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config expert_dtype = getattr(config, "expert_dtype", "fp4") - if expert_dtype != "fp4": - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) + fuse_shared_experts = _fuse_shared_experts_enabled(config) + if expert_dtype != "fp4" or fuse_shared_experts: + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( + expert_dtype, fuse_shared_experts=fuse_shared_experts + ) self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") @@ -853,8 +1004,15 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) - loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - return loaded_params + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + # After per-layer quant finalize, so we preshuffle the final fp8 weights. + for module in self.modules(): + if isinstance(module, DeepseekV4ROCMAiterMLAAttention): + module.prepare_attn_preshuffle() + elif isinstance(module, DeepseekV4MLP): + module.prepare_gateup_preshuffle() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index f5ef4bb06746..51479af238b5 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -32,6 +32,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name @@ -48,7 +51,7 @@ logger = init_logger(__name__) # MoE expert scales are fused into per-layer w13/w2 tensors. The exact -# parameter suffix depends on which FusedMoE method handles the experts: +# parameter suffix depends on which MoERunner method handles the experts: # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; # - fp8 experts (Fp8MoEMethod with block_quant=True) register # ``w{1,2,3}_weight_scale_inv``. @@ -331,6 +334,21 @@ def _find_mtp_layer_idx(name: str) -> int: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + def _resolve_scale_name(name: str) -> str: + # Quark checkpoints name FP8 block scales ``.weight_scale``, + # but block-FP8 layers register them as ``.weight_scale_inv`` + # while MXFP4 experts register ``.weight_scale``. Auto-detect: + # rename to ``_inv`` only when that variant exists and the plain + # one does not. + if name.endswith(".weight_scale") and name not in params_dict: + inv = name.removesuffix(".weight_scale") + ".weight_scale_inv" + if inv in params_dict: + return inv + # Otherwise leave the name unchanged: either it already matches a + # param, or it is genuinely unknown and should surface the normal + # KeyError downstream rather than be silently rewritten. + return name + # TP for attention tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() @@ -390,6 +408,7 @@ def _find_mtp_layer_idx(name: str) -> int: if weight_name not in name: continue name = name.replace(weight_name, param_name) + name = _resolve_scale_name(name) param = params_dict[name] weight_loader = param.weight_loader @@ -444,6 +463,7 @@ def _find_mtp_layer_idx(name: str) -> int: ) if name.endswith(".ffn.gate.bias"): name = name.replace(".bias", ".e_score_correction_bias") + name = _resolve_scale_name(name) param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader @@ -461,7 +481,7 @@ def _find_mtp_layer_idx(name: str) -> int: self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled(): raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint. The checkpoint may have " diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index d4aef6ec9ea6..7c232fe9e528 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -6,6 +6,10 @@ import torch +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache @@ -442,10 +446,73 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Block scale for the preshuffled weight; None = not preshuffled. + self._wqa_wkv_scale: torch.Tensor | None = None + self._wo_b_scale: torch.Tensor | None = None + @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return num_heads + def prepare_attn_preshuffle(self) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_enabled(): + return + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, + ) + from vllm.model_executor.utils import replace_parameter + + def _prep(linear) -> torch.Tensor | None: + w = getattr(linear, "weight", None) + if w is None or w.dim() != 2: + return None + # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold. + if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0: + return None + ws = getattr(linear, "weight_scale_inv", None) # per-block scale + if ws is None: + return None + if ws.dtype == torch.float8_e8m0fnu: + ws = _upcast_e8m0_to_fp32(ws).contiguous() + # Shuffle the weight in place (single weight, no unshuffled copy). + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) + return ws + + self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) + self._wo_b_scale = _prep(self.wo_b) + + def _bpre_attn_gemm( + self, + weight: torch.Tensor, + scale: torch.Tensor, + x: torch.Tensor, + reduce_tp: bool, + ) -> torch.Tensor: + from vllm._aiter_ops import rocm_aiter_ops + + x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True) + out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + x_fp8, weight, x_scale, scale, output_dtype=x.dtype + ) + if reduce_tp and get_tensor_model_parallel_world_size() > 1: + out = tensor_model_parallel_all_reduce(out) + return out + + def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self._wqa_wkv_scale is not None and hidden_states.dim() == 2: + return self._bpre_attn_gemm( + self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False + ) + return super()._fused_wqa_wkv_gemm(hidden_states) + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. z = rocm_inv_rope_einsum( @@ -457,7 +524,10 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: self.o_lora_rank, self.wo_a, ) - return self.wo_b(z.flatten(1)) + zf = z.flatten(1) + if self._wo_b_scale is not None and zf.dim() == 2: + return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) + return self.wo_b(zf) def forward_mqa( self, diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 5628a6d0d728..c56a36bcd2ff 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -22,12 +22,14 @@ RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v4.common.ops import ( fused_indexer_q_rope_quant, - fused_q_kv_rmsnorm, ) +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE if TYPE_CHECKING: + from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) @@ -46,6 +48,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, maybe_execute_in_parallel, @@ -65,6 +68,25 @@ logger = init_logger(__name__) +@triton.jit +def _fill_short_context_topk_indices( + output, + positions, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + PADDED_TOP_K: tl.constexpr, +): + # small triton kernel that selects every candidate, -1 otherwise + row = tl.program_id(0) + offsets = tl.arange(0, PADDED_TOP_K) + num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO + tl.store( + output + row * TOP_K + offsets, + tl.where(offsets < num_compressed, offsets, -1), + mask=offsets < TOP_K, + ) + + def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, @@ -160,6 +182,7 @@ def __init__( prefix: str, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ) -> None: super().__init__() config = vllm_config.model_config.hf_config @@ -248,6 +271,7 @@ def __init__( ) self.indexer_rotary_emb = self.rotary_emb self.topk_indices_buffer = topk_indices_buffer + self.eager_scratch_pool = eager_scratch_pool self.indexer = None if self.compress_ratio == 4: @@ -269,6 +293,7 @@ def __init__( compress_ratio=self.compress_ratio, prefix=f"{prefix}.indexer", aux_stream=indexer_aux_stream, + eager_scratch_pool=eager_scratch_pool, ) # Will be None on ROCm for now. @@ -319,6 +344,7 @@ def __init__( rotate=True, prefix=f"{prefix}.compressor", k_cache_prefix=self.prefix, + eager_scratch_pool=eager_scratch_pool, ) def forward( @@ -369,6 +395,11 @@ def forward( # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). return self._o_proj(o, positions) + def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: aux_streams = self.aux_stream_list if aux_streams is not None: @@ -413,9 +444,7 @@ def indexer_compressor_kv_score() -> torch.Tensor: aux_fns[2] = indexer_compressor_kv_score def fused_wqa_wkv() -> torch.Tensor: - # MergedColumnParallelLinear returns (output, bias); bias is None. - qr_kv, _ = self.fused_wqa_wkv(hidden_states) - return qr_kv + return self._fused_wqa_wkv_gemm(hidden_states) qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( fused_wqa_wkv, @@ -547,10 +576,24 @@ def _fused_qnorm_rope_kv_insert( if cache_dtype == torch.uint8: # fp8_ds_mla UE8M0 paged path. Horizontally fused: # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling - # the padding head slots; the kernel allocates and returns - # the padded q tensor. + # the padding head slots. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + if self.eager_scratch_pool is not None: + q_out = self.eager_scratch_pool.q_out(q.shape[0]) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( + q, + kv, + q_out, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + return q_out return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, @@ -599,6 +642,13 @@ def _fused_qnorm_rope_kv_insert( ) return q_fp8 + def _global_topk_output_buffers( + self, topk_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if self.compress_ratio != 4 or self.eager_scratch_pool is None: + return None + return self.eager_scratch_pool.global_topk_outputs(topk_indices) + def get_attn_backend(self) -> type[AttentionBackend]: return self.backend_cls @@ -678,6 +728,7 @@ def __init__( compress_ratio: int = 1, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.vllm_config = vllm_config @@ -690,6 +741,7 @@ def __init__( self.rope_dim = config.qk_rope_head_dim # 64 self.q_lora_rank = q_lora_rank # 1536 self.compress_ratio = compress_ratio + self.eager_scratch_pool = eager_scratch_pool self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache logger.info_once( "Using %s indexer cache for Lightning Indexer.", @@ -727,11 +779,16 @@ def __init__( ) assert cache_config is not None, "Deepseek V4 indexer requires cache_config" - # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: - # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. - # For FP4 indexer cache, we still allocate the same amount of memory as FP8, - # but only use the first half of the memory. - k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 + if self.use_fp4_kv: + # MXFP4 stores two values per byte plus one UE8M0 byte per 32 values. + # head_dim bytes = 64 packed values + 4 UE8M0 scales = 68. + k_cache_head_dim = self.head_dim // 2 + self.head_dim // MXFP4_BLOCK_SIZE + else: + # NOTE(yifan): FP8 indexer cache uses the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + k_cache_head_dim = ( + self.head_dim + self.head_dim // self.quant_block_size * 4 + ) self.k_cache = DeepseekV4IndexerCache( head_dim=k_cache_head_dim, dtype=torch.uint8, @@ -748,6 +805,7 @@ def __init__( prefix=f"{prefix}.compressor", k_cache_prefix=self.k_cache.prefix, use_fp4_cache=self.use_fp4_kv, + eager_scratch_pool=eager_scratch_pool, ) self.indexer_op = SparseAttnIndexer( @@ -781,10 +839,36 @@ def forward( ) -> torch.Tensor: compressor = self.compressor + attn_metadata = get_forward_context().attn_metadata + if isinstance(attn_metadata, dict): + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: + # candidates num smaller than topk, every candidate is selected + # but we still need to build k cache + compressor(compressed_kv_score, positions, rotary_emb) + assert self.topk_indices_buffer is not None + num_tokens = ( + indexer_metadata.num_decode_tokens + + indexer_metadata.num_prefill_tokens + ) + if num_tokens > 0: + _fill_short_context_topk_indices[(num_tokens,)]( + self.topk_indices_buffer, + positions, + TOP_K=self.topk_tokens, + COMPRESS_RATIO=self.compress_ratio, + PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens), + num_warps=8, + ) + return self.topk_indices_buffer + def wq_b_and_q_quant(): # ReplicatedLinear returns (output, bias); bias is None. q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) + outputs = None + if self.eager_scratch_pool is not None and self.use_fp4_kv: + outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0]) return fused_indexer_q_rope_quant( positions, q, @@ -793,6 +877,7 @@ def wq_b_and_q_quant(): self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + output_buffers=outputs, ) # compressor returns None and writes K to the indexer KV cache; the diff --git a/vllm/models/deepseek_v4/common/ops/__init__.py b/vllm/models/deepseek_v4/common/ops/__init__.py index ff6ee22996d6..b7899f01de7e 100644 --- a/vllm/models/deepseek_v4/common/ops/__init__.py +++ b/vllm/models/deepseek_v4/common/ops/__init__.py @@ -11,7 +11,6 @@ from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm -from .fused_qk_rmsnorm import fused_q_kv_rmsnorm from .save_partial_states import save_partial_states __all__ = [ @@ -23,7 +22,6 @@ "fused_indexer_q_rope_quant", "fused_inv_rope_fp8_quant", "fused_mtp_input_rmsnorm", - "fused_q_kv_rmsnorm", "mtp_shared_head_rmsnorm", "quantize_and_insert_k_cache", "save_partial_states", diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 55106d06af90..dc06dc91b22f 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -14,14 +14,23 @@ window indices for sparse prefill. """ +from dataclasses import dataclass +from typing import Any + import torch from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonPointerInputVariant, + TritonWarmupTensor, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl +from vllm.utils.math_utils import next_power_of_2 @triton.jit @@ -429,6 +438,7 @@ def compute_global_topk_indices_and_lens( block_table: torch.Tensor, block_size: int, is_valid_token: torch.Tensor, + output_buffers: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Map local topk indices to global KV cache slots and count valid entries. @@ -438,8 +448,15 @@ def compute_global_topk_indices_and_lens( 3. Masking padding tokens to length 0 """ num_tokens = topk_indices.shape[0] - global_topk_indices = torch.empty_like(topk_indices) - topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + if output_buffers is None: + global_topk_indices = torch.empty_like(topk_indices) + topk_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + else: + global_topk_indices, topk_lens = output_buffers + assert global_topk_indices.shape == topk_indices.shape + assert topk_lens.shape == (num_tokens,) _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( global_topk_indices, global_topk_indices.stride(0), @@ -526,31 +543,31 @@ def combine_topk_swa_indices( topk: int, M: int, N: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] combined_topk = ( (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) - combined_indices = torch.full( - (num_tokens, combined_topk), - fill_value=-1, - dtype=torch.int32, - device=topk_indices.device, - ) - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) + if out is None: + combined_indices = torch.full( + (num_tokens, combined_topk), + fill_value=-1, + dtype=torch.int32, + device=topk_indices.device, + ) + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + else: + combined_indices, combined_lens = out - NUM_WORKERS = 128 - _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)]( + _COMBINE_TOPK_SWA_INDICES_KERNEL( combined_indices, - combined_indices.stride(0), combined_lens, topk_indices, - topk_indices.stride(0), query_start_loc, seq_lens, gather_lens, @@ -559,82 +576,245 @@ def combine_topk_swa_indices( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, - PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), ) return combined_indices, combined_lens -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - PADDED_TOP_K: tl.constexpr, +_COMBINE_TOPK_SWA_NUM_WORKERS = 128 + + +# Representative pointer alignment variants for Triton pointer specialization. +_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs( + dict( + topk_indices=True, + query_start_loc=True, + seq_lens=True, + gather_lens=True, + ), + dict( + topk_indices=True, + query_start_loc=False, + seq_lens=False, + gather_lens=True, + ), + dict( + topk_indices=False, + query_start_loc=False, + seq_lens=False, + gather_lens=False, + ), +) + + +_DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS = zip_inputs( + # DSv4-Flash / SWA-only and C4A. + dict(compress_ratio=1, topk=0, topk_width=512), + dict(compress_ratio=4, topk=512, topk_width=512), + # DSv4-Pro C4A. + dict(compress_ratio=4, topk=1024, topk_width=1024), + # DSv4-Pro C128A. + dict(compress_ratio=128, topk=8192, topk_width=8192), +) + + +def _hf_config_int(vllm_config: Any, name: str, default: int) -> int: + model_config = getattr(vllm_config, "model_config", None) + hf_config = getattr(model_config, "hf_config", None) + return int(getattr(hf_config, name, default) or default) + + +def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int: + scheduler_config = getattr(vllm_config, "scheduler_config", None) + return int(getattr(scheduler_config, name, default) or default) + + +class CombineTopkSwaIndicesKernel( + VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"] ): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) + @dataclass(frozen=True) + class CompileKey: + TOP_K: int + COMPRESS_RATIO: int + WINDOW_SIZE: int + PADDED_TOP_K: int + input_variant: TritonPointerInputVariant + + @staticmethod + @triton.jit( + do_not_specialize=[ + "combined_indices_stride", + "topk_indices_stride", + "M", + "N", + ] + ) + def kernel( + combined_indices_ptr, + combined_indices_stride, + combined_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + query_start_loc_ptr, + seq_lens_ptr, + gather_lens_ptr, + M, + N, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + PADDED_TOP_K: tl.constexpr, + ): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + # query_start_loc is a global tensor; rebase to chunk-local offsets + # by subtracting the chunk's starting value. + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + gather_len = tl.load(gather_lens_ptr + batch_idx) + start_pos = seq_len - query_len + # The SWA portion of the gathered buffer starts from position + # (seq_len - gather_len), not position 0. We need this offset + # to correctly index into the gathered buffer. + gather_start = seq_len - gather_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + # topk_len is fully determined by the query token's absolute position: + # both the C4A indexer and the C128A metadata builder emit + # min((pos + 1) // compress_ratio, topk_tokens) valid entries. + # Caller passes TOP_K=0 for SWA-only layers to zero this out. + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + + offset = tl.arange(0, PADDED_TOP_K) + mask = offset < topk_len + topk_indices = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + ) + tl.store( + combined_indices_ptr + token_idx * combined_indices_stride + offset, + topk_indices + M * batch_idx, + mask=mask, + ) + offset = tl.arange(0, WINDOW_SIZE) + # Index into gathered buffer: N + (position - gather_start) + # For positions [pos - swa_len + 1, pos], the buffer indices are: + # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + offset, + M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, + mask=offset < swa_len, + ) - # query_start_loc is a global tensor; rebase to chunk-local offsets - # by subtracting the chunk's starting value. - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - # The SWA portion of the gathered buffer starts from position - # (seq_len - gather_len), not position 0. We need this offset - # to correctly index into the gathered buffer. - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - # topk_len is fully determined by the query token's absolute position: - # both the C4A indexer and the C128A metadata builder emit - # min((pos + 1) // compress_ratio, topk_tokens) valid entries. - # Caller passes TOP_K=0 for SWA-only layers to zero this out. - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - - offset = tl.arange(0, PADDED_TOP_K) - mask = offset < topk_len - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, + combined_len = topk_len + swa_len + tl.store(combined_lens_ptr + token_idx, combined_len) + + def dispatch( # type: ignore[override] + self, + *, + topk_width: int, + topk_indices: bool, + query_start_loc: bool, + seq_lens: bool, + gather_lens: bool, + topk: int, + compress_ratio: int, + WINDOW_SIZE: int, + ) -> CompileKey: + padded_topk = next_power_of_2(topk_width) + input_variant = TritonPointerInputVariant.from_alignment( + topk_indices=topk_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + gather_lens=gather_lens, ) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + offset, - topk_indices + M * batch_idx, - mask=mask, + return self.CompileKey( + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=WINDOW_SIZE, + PADDED_TOP_K=padded_topk, + input_variant=input_variant, ) - offset = tl.arange(0, WINDOW_SIZE) - # Index into gathered buffer: N + (position - gather_start) - # For positions [pos - swa_len + 1, pos], the buffer indices are: - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + offset, - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, - mask=offset < swa_len, + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0: + return [] + + window_size = _hf_config_int(vllm_config, "sliding_window", 128) + return self._trace_dispatch(self.dispatch)( + _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS, + _COMBINE_TOPK_SWA_POINTER_INPUTS, + WINDOW_SIZE=window_size, + ) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + input_variant = compile_key.input_variant + warmup( + int32_ptr, + 1, # do not specialize combined_indices_stride + int32_ptr, + input_variant.pointer("topk_indices", torch.int32), + 1, # do not specialize topk_indices_stride + input_variant.pointer("query_start_loc", torch.int32), + input_variant.pointer("seq_lens", torch.int32), + input_variant.pointer("gather_lens", torch.int32), + 1, # do not specialize M + 1, # do not specialize N + TOP_K=compile_key.TOP_K, + COMPRESS_RATIO=compile_key.COMPRESS_RATIO, + WINDOW_SIZE=compile_key.WINDOW_SIZE, + PADDED_TOP_K=compile_key.PADDED_TOP_K, + grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), + ) + + def __call__( + self, + combined_indices: torch.Tensor, + combined_lens: torch.Tensor, + topk_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor, + M: int, + N: int, + *, + TOP_K: int, + COMPRESS_RATIO: int, + WINDOW_SIZE: int, + ) -> None: + num_reqs = seq_lens.shape[0] + self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( + combined_indices, + combined_indices.stride(0), + combined_lens, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + TOP_K=TOP_K, + COMPRESS_RATIO=COMPRESS_RATIO, + WINDOW_SIZE=WINDOW_SIZE, + PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), ) - combined_len = topk_len + swa_len - tl.store(combined_lens_ptr + token_idx, combined_len) + +_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel() def build_flashinfer_mixed_sparse_indices( diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 128746dda8c6..5aa00174079a 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -295,6 +295,7 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + output_buffers: tuple[torch.Tensor, ...] | None = None, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, @@ -332,7 +333,13 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + if output_buffers is None: + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + else: + expected_num_buffers = 3 if use_fp4 else 2 + assert len(output_buffers) == expected_num_buffers + index_weights_out = output_buffers[-1] + assert index_weights_out.shape == index_weights.shape if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -340,16 +347,23 @@ def fused_indexer_q_rope_quant( f"size {MXFP4_BLOCK_SIZE}" ) num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE - index_q_packed = torch.empty( - (num_tokens, num_index_q_heads, index_q_head_dim // 2), - dtype=torch.uint8, - device=index_q.device, - ) - index_q_scale = torch.empty( - (num_tokens, num_index_q_heads, num_scale_blocks), - dtype=torch.uint8, - device=index_q.device, - ) + packed_shape = (num_tokens, num_index_q_heads, index_q_head_dim // 2) + scale_shape = (num_tokens, num_index_q_heads, num_scale_blocks) + if output_buffers is None: + index_q_packed = torch.empty( + packed_shape, + dtype=torch.uint8, + device=index_q.device, + ) + index_q_scale = torch.empty( + scale_shape, + dtype=torch.uint8, + device=index_q.device, + ) + else: + index_q_packed, index_q_scale, _ = output_buffers + assert index_q_packed.shape == packed_shape + assert index_q_scale.shape == scale_shape if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( @@ -367,6 +381,18 @@ def fused_indexer_q_rope_quant( index_q_scale, index_weights_out, ) + elif current_platform.is_xpu(): + torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( + index_q, + positions, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_packed, + index_q_scale, + index_weights_out, + ) else: _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)]( positions, @@ -406,7 +432,11 @@ def fused_indexer_q_rope_quant( fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 - index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) + if output_buffers is None: + index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) + else: + index_q_fp8, _ = output_buffers + assert index_q_fp8.shape == index_q.shape if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( @@ -423,6 +453,17 @@ def fused_indexer_q_rope_quant( index_q_fp8, index_weights_out, ) + elif current_platform.is_xpu(): + torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8( + index_q, + positions, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_fp8, + index_weights_out, + ) else: _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)]( positions, diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py index b667b87679cd..aa8750580d04 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -257,7 +257,6 @@ def _fused_inv_rope_fp8_quant_kernel_impl( ) grid = (tma_aligned_T, n_groups * heads_per_group) use_gdc = current_platform.is_arch_support_pdl() - pdl_kwargs = {"launch_pdl": True} if use_gdc else {} _fused_inv_rope_fp8_quant_per_head[grid]( o, positions, @@ -281,8 +280,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl( HALF_ROPE=half_rope, TMA_ALIGNED_SCALES=tma_aligned_scales, USE_GDC=use_gdc, + launch_pdl=use_gdc, num_stages=1, - **pdl_kwargs, num_warps=1, ) return fp8_buf, scale_buf diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 13f327f6bc19..93ae1e89c919 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -2,12 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import torch from torch import nn -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -35,6 +35,9 @@ SlidingWindowMLASpec, ) +if TYPE_CHECKING: + from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool + def _prefer_two_stage_compressor() -> bool: # Platforms that favor the triton variant of two-stage compressor split. @@ -42,6 +45,19 @@ def _prefer_two_stage_compressor() -> bool: return current_platform.is_rocm() +def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None: + starts = metadata._num_computed_tokens_cpu + if starts is None: + return None + + starts_list = starts.tolist() + query_start_loc = metadata.query_start_loc_cpu.tolist() + return any( + start % 128 + query_start_loc[i + 1] - query_start_loc[i] >= 128 + for i, start in enumerate(starts_list) + ) + + class CompressorBackend(AttentionBackend): def __init__(self): super().__init__() @@ -90,6 +106,7 @@ class CompressorMetadata: token_to_req_indices: torch.Tensor | None = None # [num_tokens] num_decode_tokens: int | None = None + c128_boundary: bool | None = None class CompressorMetadataBuilder(AttentionMetadataBuilder): @@ -127,6 +144,11 @@ def build( block_size=self.block_size, token_to_req_indices=token_to_req_indices, num_decode_tokens=num_decode_tokens, + c128_boundary=( + _get_c128_boundary(common_attn_metadata) + if self.block_size == 8 + else None + ), ) @@ -207,6 +229,7 @@ def __init__( prefix: str = "", k_cache_prefix="", use_fp4_cache: bool = False, + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.compress_ratio = compress_ratio @@ -216,6 +239,7 @@ def __init__( self.prefix = prefix self.k_cache_prefix = k_cache_prefix self.use_fp4_cache = use_fp4_cache + self.eager_scratch_pool = eager_scratch_pool config = vllm_config.model_config.hf_config self.rope_head_dim = config.qk_rope_head_dim @@ -317,7 +341,8 @@ def forward( ) # Get the metadata and handle dummy profiling run. - attn_metadata = get_forward_context().attn_metadata + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata if not isinstance(attn_metadata, dict): return @@ -359,6 +384,16 @@ def forward( pdl_kwargs=pdl_kwargs, ) + # full graph cannot branch on per-step CPU metadata after capture + if ( + current_platform.is_cuda() + and self.head_dim == 512 + and self.compress_ratio == 128 + and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL + and state_metadata.c128_boundary is False + ): + return + # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write. # RoPE requirements (kernel applies forward GPT-J style rotation): # - is_neox_style=False (interleaved pairs, NOT split-half) @@ -398,6 +433,10 @@ def forward( store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) + if not self.overlap and self.eager_scratch_pool is not None: + extra_kwargs["compress_scratch"] = ( + self.eager_scratch_pool.compressor_scratch(num_actual) + ) elif self._use_two_stage_fused_compressor: # head=512 cr>=128 (no overlap): two-pass split compressor on the # prefill suffix, single-pass on the decode prefix. diff --git a/vllm/models/deepseek_v4/eager_scratch.py b/vllm/models/deepseek_v4/eager_scratch.py new file mode 100644 index 000000000000..bc46239bfe34 --- /dev/null +++ b/vllm/models/deepseek_v4/eager_scratch.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from math import prod + +import torch + +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE +from vllm.utils.math_utils import round_up + + +class DeepseekV4EagerScratchPool: + """Model-wide outputs and scratch used inside the attention eager break.""" + + _ALIGNMENT = 256 + + def __init__( + self, + max_num_tokens: int, + padded_q_heads: int, + q_head_dim: int, + index_q_heads: int, + index_q_head_dim: int, + index_topk: int, + device: torch.device | str, + ) -> None: + self.max_num_tokens = max_num_tokens + self.index_topk = index_topk + self._q = torch.empty( + (max_num_tokens, padded_q_heads, q_head_dim), + dtype=torch.bfloat16, + device=device, + ) + + fp4_specs = ( + ((max_num_tokens, index_q_heads, index_q_head_dim // 2), torch.uint8), + ( + ( + max_num_tokens, + index_q_heads, + index_q_head_dim // MXFP4_BLOCK_SIZE, + ), + torch.uint8, + ), + ((max_num_tokens, index_q_heads), torch.float32), + ) + global_specs = ( + ((max_num_tokens, index_topk), torch.int32), + ((max_num_tokens,), torch.int32), + ) + compressor_specs = (((max_num_tokens, q_head_dim), torch.float32),) + # FP4 indexer is C4 only, global mapping after FP4 indexer + # compressor scratch is C128 only + # so here we use max instead of sum + aux_bytes = max( + self._packed_size(specs) + for specs in (fp4_specs, global_specs, compressor_specs) + ) + storage = torch.empty(aux_bytes, dtype=torch.uint8, device=device) + + self._q_outputs: dict[int, torch.Tensor] = {} + fp4_values, fp4_scales, fp4_weights = self._views(storage, fp4_specs) + self._fp4_template = (fp4_values, fp4_scales, fp4_weights) + self._fp4_outputs: dict[ + int, tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ] = {} + global_indices, global_lens = self._views(storage, global_specs) + self._global_template = (global_indices, global_lens) + self._global_outputs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + self._compressor_template = self._views(storage, compressor_specs)[0] + self._compressor_outputs: dict[int, torch.Tensor] = {} + self._storage = storage + + @classmethod + def _packed_size( + cls, specs: tuple[tuple[tuple[int, ...], torch.dtype], ...] + ) -> int: + offset = 0 + for shape, dtype in specs: + offset = round_up(offset, cls._ALIGNMENT) + prod(shape) * dtype.itemsize + return round_up(offset, cls._ALIGNMENT) + + @classmethod + def _views( + cls, + storage: torch.Tensor, + specs: tuple[tuple[tuple[int, ...], torch.dtype], ...], + ) -> list[torch.Tensor]: + offset = 0 + views = [] + for shape, dtype in specs: + offset = round_up(offset, cls._ALIGNMENT) + num_bytes = prod(shape) * dtype.itemsize + views.append(storage[offset : offset + num_bytes].view(dtype).view(shape)) + offset += num_bytes + return views + + def q_out(self, num_tokens: int) -> torch.Tensor: + output = self._q_outputs.get(num_tokens) + if output is None: + output = self._q[:num_tokens] + self._q_outputs[num_tokens] = output + return output + + def compressor_scratch(self, num_tokens: int) -> torch.Tensor: + output = self._compressor_outputs.get(num_tokens) + if output is None: + output = self._compressor_template[:num_tokens] + self._compressor_outputs[num_tokens] = output + return output + + def indexer_q_outputs( + self, + num_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + output = self._fp4_outputs.get(num_tokens) + if output is None: + values, scales, weights = self._fp4_template + output = ( + values[:num_tokens], + scales[:num_tokens], + weights[:num_tokens], + ) + self._fp4_outputs[num_tokens] = output + return output + + def global_topk_outputs( + self, topk_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens, topk = topk_indices.shape + assert topk == self.index_topk + output = self._global_outputs.get(num_tokens) + if output is None: + indices, lens = self._global_template + output = (indices[:num_tokens], lens[:num_tokens]) + self._global_outputs[num_tokens] = output + return output diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index e4e258372a29..5478c423a223 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -15,11 +15,13 @@ import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, @@ -40,9 +42,16 @@ DSparkMarkovHead, ) from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from .model import ( DeepseekV4DecoderLayer, + DeepseekV4Model, + _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, ) @@ -65,6 +74,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) + self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 @@ -85,12 +95,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) + current_vllm_config = get_current_vllm_config() self.layers = nn.ModuleList( [ DeepseekV4DecoderLayer( current_vllm_config, prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + topk_indices_buffer=self.topk_indices_buffer, ) for i in range(self.num_dspark_layers) ] @@ -171,6 +188,15 @@ def forward( ) -> torch.Tensor: if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) + full_num_tokens = positions.shape[0] + if self.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, inputs_embeds + ) + inputs_embeds = sp_shard(inputs_embeds) + input_ids = sp_shard(input_ids) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) @@ -185,6 +211,8 @@ def forward( residual, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] # hc_head reduces the hc copies; return the PRE-norm head hidden hidden_states = hc_head_fused_kernel_tilelang( hidden_states, @@ -277,6 +305,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: assert vllm_config.speculative_config is not None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config + self.quant_config = vllm_config.quant_config + self.pad_shared_expert = getattr( + self.quant_config, "weight_block_size", None + ) is not None and not _use_sequence_parallel(vllm_config) self.model = DSparkDeepseekV4Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -396,6 +428,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if self.pad_shared_expert and ".shared_experts." in name: + loaded_weight = DeepseekV4Model._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) # E8M0 expert scales: keep raw exponent bytes. if ".experts." in name: @@ -440,10 +478,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict[name][: narrow.shape[0]].copy_(narrow) loaded_params.add(name) continue - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) if name.endswith(".ffn.gate.bias"): name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 1848c1930db0..d8742e610342 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -60,6 +60,20 @@ def _packed_block_span(pool: torch.Tensor) -> int: return block_stride // token_stride +# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545). +_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128) + + +def _pad_to_supported_q_heads(num_heads: int) -> int: + for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS: + if num_heads <= supported: + return supported + raise ValueError( + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " + "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})." + ) + + class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend): """FlashInfer backend using the DSv4 sparse metadata/cache layout. @@ -164,13 +178,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: - # FP8 decode kernel only supports h_q = 64 or 128. - if num_heads > 128: - raise ValueError( - f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " - "(FP8 decode kernel requires h_q in {64, 128})." - ) - return 64 if num_heads <= 64 else 128 + return _pad_to_supported_q_heads(num_heads) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: return deep_gemm_fp8_o_proj( @@ -545,18 +553,7 @@ def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: - if num_heads <= 16: - return 16 - if num_heads <= 32: - return 32 - if num_heads <= 64: - return 64 - if num_heads <= 128: - return 128 - raise ValueError( - f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " - "(SM120 kernel requires h_q in {16, 32, 64, 128})." - ) + return _pad_to_supported_q_heads(num_heads) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: return deep_gemm_fp8_o_proj( @@ -748,6 +745,9 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, + output_buffers=self._global_topk_output_buffers( + self.topk_indices_buffer[:num_decode_tokens] + ), ) ) extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) @@ -837,6 +837,7 @@ def _forward_prefill( attn_metadata.block_table, block_size, swa_metadata.is_valid_token[prefill_token_slice], + output_buffers=self._global_topk_output_buffers(local_topk_indices), ) ) diff --git a/vllm/models/deepseek_v4/nvidia/flashmla.py b/vllm/models/deepseek_v4/nvidia/flashmla.py index 9fa4e1c11b94..29ec2d6f2abf 100644 --- a/vllm/models/deepseek_v4/nvidia/flashmla.py +++ b/vllm/models/deepseek_v4/nvidia/flashmla.py @@ -20,6 +20,7 @@ DeepseekV4FlashMLABackend, DeepseekV4FlashMLAMetadata, ) +from vllm.utils.math_utils import round_up from vllm.v1.attention.ops.flashmla import ( flash_mla_sparse_fwd, flash_mla_with_kvcache, @@ -95,8 +96,13 @@ def forward_mqa( // self.compress_ratio ) M = N + self.window_size + self.max_num_batched_tokens + assert self.topk_indices_buffer is not None + top_k = 0 if swa_only else self.topk_indices_buffer.shape[-1] + combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), ) output.zero_() return @@ -170,6 +176,9 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, + output_buffers=self._global_topk_output_buffers( + self.topk_indices_buffer[:num_decode_tokens] + ), ) topk_indices = global_indices.view(num_decode_tokens, 1, -1) else: @@ -284,11 +293,15 @@ def _forward_prefill( ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() + combined_topk = round_up(top_k + self.window_size, 128) for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: chunk_size = chunk_end - chunk_start - kv = workspace_manager.get_simultaneous( + workspace = workspace_manager.get_simultaneous( ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16), - )[0] + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), + ) + kv, combined_indices_out, combined_lens_out = workspace if not swa_only: # Gather compressed KV assert attn_metadata is not None @@ -322,6 +335,8 @@ def _forward_prefill( query_end = ( query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base ) + combined_indices_out = combined_indices_out[: query_end - query_start] + combined_lens_out = combined_lens_out[: query_end - query_start] combined_indices, combined_lens = combine_topk_swa_indices( topk_indices[query_start:query_end], @@ -335,6 +350,7 @@ def _forward_prefill( top_k, chunk_M, chunk_N, + out=(combined_indices_out, combined_lens_out), ) flash_mla_sparse_fwd( q=q[query_start:query_end], diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ddc2fe0f4bc2..8925890a3cd4 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -27,7 +27,7 @@ ) from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.router.base_router import ( @@ -65,7 +65,14 @@ maybe_prefix, ) from vllm.model_executor.utils import set_weight_attrs +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -514,6 +521,7 @@ def __init__( self, vllm_config: VllmConfig, prefix: str = "", + use_sequence_parallel: bool = False, ): super().__init__() @@ -521,6 +529,7 @@ def __init__( config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.prefix = prefix + self.use_sequence_parallel = use_sequence_parallel self.use_mega_moe = ( vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" ) @@ -594,6 +603,7 @@ def __init__( swiglu_limit=self.swiglu_limit, quant_config=quant_config, reduce_results=self.use_mega_moe, + is_sequence_parallel=use_sequence_parallel, prefix=f"{prefix}.shared_experts", ) @@ -670,7 +680,7 @@ def _init_fused_moe_experts( self.physical_expert_start = self.experts_start_idx self.physical_expert_end = self.experts_end_idx - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=config.n_routed_experts, @@ -688,6 +698,7 @@ def _init_fused_moe_experts( router_logits_dtype=torch.float32, enable_eplb=parallel_config.enable_eplb, num_redundant_experts=eplb_config.num_redundant_experts, + is_sequence_parallel=self.use_sequence_parallel, ) def forward( @@ -736,7 +747,7 @@ def _forward_fused_moe( ) -> torch.Tensor: org_shape = hidden_states.shape if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class + # In this case, the gate/router runs inside the MoERunner class final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states, @@ -791,6 +802,17 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: return DeepseekV4FlashMLAAttention +def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: + parallel_config = vllm_config.parallel_config + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + return ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -798,11 +820,13 @@ def __init__( prefix, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + eager_scratch_pool: DeepseekV4EagerScratchPool | None = None, ): super().__init__() config = vllm_config.model_config.hf_config self.hidden_size = config.hidden_size + self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.rms_norm_eps = config.rms_norm_eps self.attn = _select_dsv4_attn_cls(vllm_config)( @@ -810,8 +834,15 @@ def __init__( prefix=f"{prefix}.attn", topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, + eager_scratch_pool=eager_scratch_pool, + ) + if self.use_sequence_parallel: + self.attn.wo_b.reduce_results = False + self.ffn = DeepseekV4MoE( + vllm_config, + prefix=f"{prefix}.ffn", + use_sequence_parallel=self.use_sequence_parallel, ) - self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) @@ -929,8 +960,12 @@ def forward( norm_eps=attn_norm_eps, ) - # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above. + if self.use_sequence_parallel: + x = sp_all_gather(x)[: positions.shape[0]] + x = self.attn(positions, x, None) + if self.use_sequence_parallel: + x = sp_reduce_scatter(x) ffn_norm_weight = self.ffn_norm.weight.data ffn_norm_eps = self.ffn_norm.variance_epsilon @@ -969,6 +1004,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.use_mega_moe = ( vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" ) + self.use_sequence_parallel = _use_sequence_parallel(vllm_config) if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " @@ -986,6 +1022,22 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + padded_heads = _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads( + config.num_attention_heads // get_tensor_model_parallel_world_size() + ) + self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None + if not vllm_config.parallel_config.use_ubatching: + # TODO: support dbo if needed + # this requires the buffer to have ubatch dim + self.eager_scratch_pool = DeepseekV4EagerScratchPool( + vllm_config.scheduler_config.max_num_batched_tokens, + padded_heads, + config.head_dim, + config.index_n_heads, + config.index_head_dim, + config.index_topk, + current_platform.device_type, + ) # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -1011,6 +1063,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=prefix, topk_indices_buffer=self.topk_indices_buffer, aux_stream_list=aux_stream_list, + eager_scratch_pool=self.eager_scratch_pool, ), prefix=f"{prefix}.layers", ) @@ -1039,13 +1092,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): torch.empty(1, dtype=torch.float32), requires_grad=False, ) - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address (outside the cudagraph pool) so the copy_ in forward() - # refreshes it correctly across captured shapes. - # refreshes it correctly across captured shapes. Only allocated on - # the last PP rank — that's where MTP target hidden states are - # produced. - if get_pp_group().is_last_rank: + spec_config = vllm_config.speculative_config + needs_mtp_hidden_states = spec_config is not None and ( + spec_config.use_eagle() or spec_config.uses_draft_model() + ) + if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, @@ -1096,6 +1147,16 @@ def forward( if self.use_mega_moe: input_ids = input_ids.to(torch.int64) + full_num_tokens = positions.shape[0] + if self.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + input_ids = sp_shard(input_ids) + residual, post_mix, res_mix = None, None, None aux_hidden_states: list[torch.Tensor] = [] final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call @@ -1116,7 +1177,10 @@ def forward( aux_recon = mhc_post_tilelang( hidden_states, residual, post_mix, res_mix ) - aux_hidden_states.append(aux_recon.mean(dim=1)) + aux_hidden_state = aux_recon.mean(dim=1) + if self.use_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) final_aux_recon = aux_recon if layer is not None: # Reuse if the last layer was captured as an aux hidden state @@ -1130,9 +1194,12 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] + + if self._mtp_hidden_buffer is not None: + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) hidden_states = hc_head_fused_kernel_tilelang( hidden_states, @@ -1176,12 +1243,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # ranks land on the zero pad). SP / unquantized ones need no padding. pad_shared_expert = ( getattr(self.quant_config, "weight_block_size", None) is not None - and not self.parallel_config.use_sequence_parallel_moe + and not self.use_sequence_parallel ) for name, loaded_weight in weights: if pad_shared_expert and ".shared_experts." in name: - loaded_weight = self._pad_shared_expert_weight(name, loaded_weight) + loaded_weight = self._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -1256,15 +1325,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params + @staticmethod def _pad_shared_expert_weight( - self, name: str, loaded_weight: torch.Tensor + quant_config: QuantizationConfig | None, + name: str, + loaded_weight: torch.Tensor, ) -> torch.Tensor: """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate axis so the standard TP loaders split it into even, block-aligned shards (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0; down (w2 -> down_proj) [H, I] pads dim 1. """ - block_size = getattr(self.quant_config, "weight_block_size", None) + block_size = getattr(quant_config, "weight_block_size", None) assert block_size is not None # Round the intermediate axis up to a whole number of TP shards. The axis # is in elements for weights (step = block) and in blocks for scales. diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 64715deae99b..eaf105205c1c 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -18,11 +18,13 @@ import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, @@ -37,10 +39,18 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from vllm.models.deepseek_v4.common.ops import ( fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm, @@ -49,13 +59,15 @@ from .model import ( DeepseekV4DecoderLayer, + DeepseekV4Model, + _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, ) logger = init_logger(__name__) # MoE expert scales are fused into per-layer w13/w2 tensors. The exact -# parameter suffix depends on which FusedMoE method handles the experts: +# parameter suffix depends on which FusedMoEFactory method handles the experts: # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; # - fp8 experts (Fp8MoEMethod with block_quant=True) register # ``w{1,2,3}_weight_scale_inv``. @@ -154,6 +166,14 @@ def forward( self.enorm.variance_epsilon, self.hc_mult, ) + if self.mtp_block.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, inputs_embeds + ) + inputs_embeds = sp_shard(inputs_embeds) + previous_hidden_states = sp_shard(previous_hidden_states) hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) @@ -161,6 +181,8 @@ def forward( positions=positions, x=hidden_states, input_ids=None ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) + if self.mtp_block.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] # Return the flat pre-hc_head residual so it can be re-fed as the # next spec step's `previous_hidden_states` when # num_speculative_tokens > 1. hc_head is deferred to compute_logits. @@ -262,6 +284,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config + self.pad_shared_expert = getattr( + self.quant_config, "weight_block_size", None + ) is not None and not _use_sequence_parallel(vllm_config) self.model = DeepSeekV4MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -384,6 +409,12 @@ def _find_mtp_layer_idx(name: str) -> int: else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if self.pad_shared_expert and ".shared_experts." in name: + loaded_weight = DeepseekV4Model._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -439,10 +470,6 @@ def _find_mtp_layer_idx(name: str) -> int: loaded_params.add(name) continue else: - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) if name.endswith(".ffn.gate.bias"): # ``e_score_correction_bias`` lives on the gate # under a different attribute name. @@ -467,7 +494,7 @@ def _find_mtp_layer_idx(name: str) -> int: self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled(): raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint. The checkpoint may have " diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py index 18e3b10562bd..b064fbe9cde6 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py @@ -60,10 +60,13 @@ def deep_gemm_fp8_o_proj( device=o.device, dtype=torch.bfloat16, ) + weight_scale = ( + wo_a.weight_scale if hasattr(wo_a, "weight_scale") else wo_a.weight_scale_inv + ) fp8_einsum( "bhr,hdr->bhd", (o_fp8, o_scale), - (wo_a.weight, wo_a.weight_scale_inv), + (wo_a.weight, weight_scale), z, recipe=einsum_recipe, ) diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index 4ff4b232d10f..960c516c4522 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -2097,6 +2097,7 @@ def compress_norm_rope_store_cutedsl( store_full_kv: bool = False, store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, + compress_scratch: torch.Tensor | None = None, ) -> None: if compress_ratio == 4: # For C4A, the single fused kernel is faster than the two-kernel version. @@ -2129,11 +2130,15 @@ def compress_norm_rope_store_cutedsl( ) else: # For C128, the two-kernel version is faster than the single fused kernel. - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, - ) + if compress_scratch is None: + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) + else: + assert compress_scratch.shape == (num_actual, head_dim) + compressed_kv = compress_scratch split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( state_cache, token_to_req_indices, diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index 721a9138914c..293d71f2f414 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from vllm.config import get_current_vllm_config from vllm.model_executor.layers.fused_moe import ( @@ -117,13 +117,33 @@ def _get_nvfp4_config(self) -> ModelOptNvFp4Config: def get_name(cls) -> QuantizationMethods: return "deepseek_v4_fp8" + @staticmethod + def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool: + """True for AMD-Quark exports whose global scheme is MXFP4.""" + weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") + # A non-dict weight (e.g. a list of multiple specs) means not an OCP + # MXFP4 scheme (e.g. NVFP4 with 2-level scale). + if not isinstance(weight, dict): + return False + return ( + weight.get("dtype") == "fp4" + and weight.get("qscheme") == "per_group" + and weight.get("group_size") == 32 + ) + @classmethod def override_quantization_method( cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: if not ( isinstance(hf_quant_cfg, dict) - and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") + and ( + hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") + or ( + hf_quant_cfg.get("quant_method") == "quark" + and cls._is_quark_mxfp4_ocp(hf_quant_cfg) + ) + ) ): return None model_type = getattr(hf_config, "model_type", None) @@ -131,6 +151,25 @@ def override_quantization_method( return "deepseek_v4_fp8" return None + @classmethod + def from_config(cls, config: dict) -> DeepseekV4FP8Config: + # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8 + # path: the runtime layout matches the DeepSeek-native fp8 checkpoint, + # so translate the schema into format Fp8Config.from_config expects. + if config.get("quant_method") == "quark": + quark_exclude = config.get("exclude") or [] + config = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "fmt": "e4m3", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + "ignored_layers": [ + name for name in quark_exclude if isinstance(name, str) + ], + } + return cast("DeepseekV4FP8Config", super().from_config(config)) + def get_quant_method(self, layer, prefix): if isinstance(layer, RoutedExperts): if is_layer_skipped( diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 4523d1875ebf..4ac27bd59a8d 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -261,6 +261,13 @@ def _build_c128a_metadata( assert cm.positions is not None, ( "positions is required for C128A metadata build" ) + active_topk_width = min( + max( + triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), + _C128A_TOPK_ALIGNMENT, + ), + self.c128a_max_compressed, + ) block_size = self.kv_cache_spec.block_size // self.compress_ratio global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( cm.positions[:num_total], @@ -273,7 +280,7 @@ def _build_c128a_metadata( self.c128a_global_decode_buffer, self.c128a_decode_lens_buffer, self.c128a_prefill_buffer, - max_compressed_tokens=self.c128a_max_compressed, + max_compressed_tokens=active_topk_width, ) result: dict[str, torch.Tensor | None] = {} @@ -305,25 +312,30 @@ def build_c128a_topk_metadata( Decode tokens: position → block_table lookup → global slot ids + topk_lens. Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. - Writes into pre-allocated buffers for CUDA graph address stability. - Returns slices of the buffers. + Writes into packed views of pre-allocated buffers for CUDA graph stability. """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens - global_decode = global_decode_buffer[:num_decode_tokens] + # view(-1) as 1-d array and then expanded to + # [num_decode_tokens, max_compressed_tokens] + global_decode = global_decode_buffer.view(-1)[ + : num_decode_tokens * max_compressed_tokens + ].view(num_decode_tokens, max_compressed_tokens) decode_lens = decode_lens_buffer[:num_decode_tokens] - prefill_local = prefill_buffer[:num_prefill_tokens] + prefill_local = prefill_buffer.view(-1)[ + : num_prefill_tokens * max_compressed_tokens + ].view(num_prefill_tokens, max_compressed_tokens) if num_tokens == 0: return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( global_decode_buffer, - global_decode_buffer.stride(0), + max_compressed_tokens, decode_lens_buffer, prefill_buffer, - prefill_buffer.stride(0), + max_compressed_tokens, positions, compress_ratio, max_compressed_tokens, diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index e8449b9c058b..0c10404d1069 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -8,7 +8,6 @@ import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_ep_group, @@ -19,7 +18,7 @@ from vllm.forward_context import get_forward_context from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -730,7 +729,7 @@ def _init_fused_moe_experts( self.n_local_experts = config.n_routed_experts // self.tp_size self.experts_start_idx = self.tp_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts - self.experts = FusedMoE( + self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=config.n_routed_experts, @@ -978,7 +977,6 @@ def forward( return x, residual, post_mix, res_mix -@support_torch_compile class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1059,11 +1057,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): requires_grad=False, ) self.hc_head_op = HCHeadOp() - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address so the copy_ in forward() refreshes it correctly across - # captured shapes. Only allocated on the last PP rank — that's - # where MTP target hidden states are produced. - if get_pp_group().is_last_rank: + spec_config = vllm_config.speculative_config + needs_mtp_hidden_states = spec_config is not None and ( + spec_config.use_eagle() or spec_config.uses_draft_model() + ) + if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, @@ -1141,9 +1139,9 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + if self._mtp_hidden_buffer is not None: + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) hidden_states = self.hc_head_op( hidden_states, diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py index 8baca78b8ba1..5c905ac22e6d 100644 --- a/vllm/models/deepseek_v4/xpu/mtp.py +++ b/vllm/models/deepseek_v4/xpu/mtp.py @@ -32,6 +32,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name @@ -51,7 +54,7 @@ logger = init_logger(__name__) # MoE expert scales are fused into per-layer w13/w2 tensors. The exact -# parameter suffix depends on which FusedMoE method handles the experts: +# parameter suffix depends on which FusedMoEFactory method handles the experts: # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; # - fp8 experts (Fp8MoEMethod with block_quant=True) register # ``w{1,2,3}_weight_scale_inv``. @@ -469,7 +472,7 @@ def _find_mtp_layer_idx(name: str) -> int: self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled(): raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint. The checkpoint may have " diff --git a/vllm/models/deepseek_v4/xpu/xpu_sparse.py b/vllm/models/deepseek_v4/xpu/xpu_sparse.py index 77cc35cf492a..73439d86461d 100644 --- a/vllm/models/deepseek_v4/xpu/xpu_sparse.py +++ b/vllm/models/deepseek_v4/xpu/xpu_sparse.py @@ -89,19 +89,37 @@ def get_padded_num_q_heads(cls, num_heads: int) -> int: return num_heads def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - # XPU uses BF16 reference wo_a path (same as ROCm). - from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum + from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant, + ) - z = rocm_inv_rope_einsum( - self.rotary_emb, + o_fp8, o_scale = fused_inv_rope_fp8_quant( o, positions, - self.rope_head_dim, - self.n_local_groups, - self.o_lora_rank, - self.wo_a, + self.rotary_emb.cos_sin_cache, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + tma_aligned_scales=False, ) - return self.wo_b(z.flatten(1)) + + # Precomputed contiguous [G, K, N] weight and [G, K/bs, N/bs] scale. + wo_a_weight = self.wo_a.bmm_weight + wo_a_scale = self.wo_a.bmm_scale + + # TODO: optimize fused_inv_rope_fp8_quant for xpu bmm to + # eliminate o_scale transpose + contiguous + z = torch.ops.vllm.xpu_fp8_bmm( + o_fp8.transpose(0, 1), + wo_a_weight, + torch.bfloat16, + o_scale.transpose(0, 1).contiguous(), + wo_a_scale, + None, + ) + + return self.wo_b(z.transpose(0, 1).flatten(1)) def forward_mqa( self, diff --git a/vllm/models/inkling/__init__.py b/vllm/models/inkling/__init__.py index 32e58905e255..4726d68e4806 100644 --- a/vllm/models/inkling/__init__.py +++ b/vllm/models/inkling/__init__.py @@ -2,12 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING +from vllm.platforms import current_platform + if TYPE_CHECKING: - from .nvidia.model import ( - InklingForCausalLM, - InklingForConditionalGeneration, - ) - from .nvidia.mtp import InklingMTP + if current_platform.is_rocm(): + from .amd.model import ( + InklingForCausalLM, + InklingForConditionalGeneration, + ) + from .amd.mtp import InklingMTP as InklingMTP + else: + from .nvidia.model import ( + InklingForCausalLM, + InklingForConditionalGeneration, + ) + from .nvidia.mtp import InklingMTP as InklingMTP __all__ = [ "InklingForConditionalGeneration", @@ -18,11 +27,23 @@ def __getattr__(name: str): if name == "InklingMTP": - from .nvidia import mtp + if current_platform.is_rocm(): + from .amd import mtp as amd_mtp + + return amd_mtp.InklingMTP + + from .nvidia import mtp as nvidia_mtp + + return nvidia_mtp.InklingMTP - return mtp.InklingMTP if name in __all__: - from .nvidia import model + if current_platform.is_rocm(): + from .amd import model as amd_model + + return getattr(amd_model, name) + + from .nvidia import model as nvidia_model + + return getattr(nvidia_model, name) - return getattr(model, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/amd/__init__.py b/vllm/models/inkling/amd/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/inkling/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/inkling/amd/attention.py b/vllm/models/inkling/amd/attention.py new file mode 100644 index 000000000000..edbc64bd817d --- /dev/null +++ b/vllm/models/inkling/amd/attention.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionMetadata, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + SlidingWindowSpec, +) + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_num_splits, + inkling_fa4_rel_attention, +) +from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup +from .ops.qkvr_prep import fused_qkvr_prep +from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata +from .short_conv import InklingShortConv + + +def compute_log_scaling_tau( + positions: torch.Tensor, n_floor: int, alpha: float +) -> torch.Tensor: + effective_n = (positions + 1).to(torch.float32) + return 1.0 + alpha * torch.log(torch.clamp(effective_n / float(n_floor), min=1.0)) + + +class RelLogitsProj(nn.Module): + """Project the per-head relative branch ``r`` to per-distance logits.""" + + def __init__(self, d_rel: int, rel_extent: int) -> None: + super().__init__() + self.d_rel = d_rel + self.rel_extent = rel_extent + self.proj = nn.Parameter(torch.empty(d_rel, rel_extent), requires_grad=False) + + def forward(self, r_out: torch.Tensor) -> torch.Tensor: + # r_out: (T, num_heads, d_rel) -> (T, num_heads, rel_extent) + return torch.einsum("thd,de->the", r_out, self.proj) + + +class InklingAttention(nn.Module, AttentionLayerBase): + def __init__( + self, + config: InklingModelConfig, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rel_extent: int, + local_extent: int, + is_local: bool, + prefix: str, + quant_config: QuantizationConfig | None = None, + conv_owner: InklingConvState, + ) -> None: + super().__init__() + self.prefix = prefix + self.is_local = is_local + self.hidden_size = config.hidden_size + self.head_dim = head_dim + self.d_rel = config.d_rel + self.log_scaling_n_floor = config.log_scaling_n_floor + self.log_scaling_alpha = config.log_scaling_alpha + # q/k are per-head RMS-normed (unit norm), so Inkling scales by 1/head_dim. + self.scaling = 1.0 / head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_total_heads = num_heads + self.num_total_kv_heads = num_kv_heads + assert self.num_total_heads % tp_size == 0 + self.num_heads = self.num_total_heads // tp_size + if self.num_total_kv_heads >= tp_size: + assert self.num_total_kv_heads % tp_size == 0 + else: + assert tp_size % self.num_total_kv_heads == 0 + self.num_kv_heads = max(1, self.num_total_kv_heads // tp_size) + # When tp_size > num_kv_heads the K/V projections are padded up to + # tp_size heads so each rank gets at least one (GQA replication). + kv_total_for_sizing = max(self.num_total_kv_heads, tp_size) + + self.qkvr = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[ + head_dim * self.num_total_heads, + head_dim * kv_total_for_sizing, + head_dim * kv_total_for_sizing, + self.d_rel * self.num_total_heads, + ], + bias=config.q_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkvr", + ) + self.wo_ud = RowParallelLinear( + input_size=head_dim * self.num_total_heads, + output_size=config.hidden_size, + bias=config.o_bias, + quant_config=quant_config, + # reduce_results=False: the partial output is all-reduced below + # (one-shot custom AR) so the attention-output sconv can run on the + # full hidden width fused with the residual add + rmsnorm. + reduce_results=False, + prefix=f"{prefix}.wo_ud", + ) + self.rel_extent = local_extent if is_local else rel_extent + self.local_extent = local_extent if is_local else None + self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent) + self.q_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + self.k_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + + # Short convolution on the K/V streams (per-head-width, TP sharded), + # applied after the qkvr projection and before q/k norm. + kv_conv_dim = self.num_kv_heads * head_dim + self.conv_owner = conv_owner + self.k_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_K + ) + self.v_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_V + ) + + # FA4 left/right window; right=0 keeps it causal. local_extent-1 mirrors + # the source (sliding_window_size - 1). + self.window_size: tuple[int, int] = ( + (local_extent - 1, 0) if is_local else (-1, -1) + ) + # Static per-layer-type KV length bound for the split heuristic: local + # layers never see more than the sliding window. + vllm_config = get_current_vllm_config() + self._max_kv_len = ( + local_extent if is_local else vllm_config.model_config.max_model_len + ) + + # ---- KV-cache wiring (reuse FlashAttentionBackend for metadata) ---- + cache_config = vllm_config.cache_config + self.kv_cache_dtype = ( + cache_config.cache_dtype if cache_config is not None else "auto" + ) + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + self.register_buffer("k_scale", torch.ones((), dtype=torch.float32)) + self.register_buffer("v_scale", torch.ones((), dtype=torch.float32)) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + + register_fa4_warmup( + InklingFA4WarmupConfig( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + rel_extent=self.rel_extent, + window_size=self.window_size, + is_local=self.is_local, + max_kv_len=self._max_kv_len, + dtype=vllm_config.model_config.dtype, + kv_dtype=self.kv_cache_torch_dtype, + block_size=vllm_config.cache_config.block_size, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=( + vllm_config.scheduler_config.max_num_batched_tokens + ), + ) + ) + + def get_attn_backend(self) -> type[AttentionBackend]: + return FlashAttentionBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + block_size = vllm_config.cache_config.block_size + if self.is_local: + assert self.local_extent is not None + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + sliding_window=self.local_extent, + ) + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + ) + + def _split_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]: + key_cache, value_cache = self.kv_cache.transpose(1, 2).split( + self.head_dim, dim=-1 + ) + return ( + canonicalize_singleton_dim_strides(key_cache), + canonicalize_singleton_dim_strides(value_cache), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + log_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + qkvr, _ = self.qkvr(hidden_states) + + attn_metadata = get_forward_context().attn_metadata + attn_output = torch.empty( + (num_tokens, self.num_heads, self.head_dim), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if not isinstance(attn_metadata, dict): + attn_output.zero_() + else: + conv_meta = attn_metadata[self.conv_owner.prefix] + md = attn_metadata[self.prefix] + assert isinstance(conv_meta, InklingSconvMetadata) + fa_md = cast(FlashAttentionMetadata, md) + assert self.kv_cache.numel() > 0 + assert self.conv_owner.kv_cache.numel() > 0 + # One launch: K/V sconv (conv-cache insert + conv + residual), + # Q/K per-head rmsnorm, and the attention KV-cache write. K/V are + # consumed via the KV cache; only normed q is materialized. + key_cache, value_cache = self._split_kv_cache() + off_k, _ = self.conv_owner.stream_ranges[_K] + off_v, _ = self.conv_owner.stream_ranges[_V] + q, rel_logits = fused_qkvr_prep( + qkvr, + self.k_sconv.weight.squeeze(1), + self.v_sconv.weight.squeeze(1), + self.q_norm.weight, + self.k_norm.weight, + self.rel_logits_proj.proj, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.d_rel, + self.conv_owner.kv_cache, + key_cache, + value_cache, + positions, + conv_meta.block_table, + conv_meta.seq_idx, + conv_meta.slot_mapping, + conv_meta.query_start, + fa_md.slot_mapping, + off_k, + off_v, + self.conv_owner.cache_block_size, + log_scaling if not self.is_local else None, + ) + q = q.view(num_tokens, self.num_heads, self.head_dim) + self._attention(q, rel_logits, attn_output) + + flat = attn_output.view(num_tokens, -1) + output, _ = self.wo_ud(flat) + return output + + @eager_break_during_capture + def _attention( + self, + q: torch.Tensor, + rel_logits: torch.Tensor, + output: torch.Tensor, + ) -> None: + attn_metadata = get_forward_context().attn_metadata + assert isinstance(attn_metadata, dict) + md = cast(FlashAttentionMetadata, attn_metadata[self.prefix]) + + nt = md.num_actual_tokens + key_cache, value_cache = self._split_kv_cache() + max_seqlen_q = bucket_max_seqlen_q(md.max_query_len) + num_splits = inkling_fa4_num_splits( + is_local=self.is_local, + batch_size=md.seq_lens.shape[0], + max_query_len=max_seqlen_q, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + max_kv_len=md.max_seq_len, + ) + inkling_fa4_rel_attention( + q[:nt], + key_cache, + value_cache, + block_table=md.block_table, + cache_seqlens=md.seq_lens, + cu_seqlens_q=md.query_start_loc, + max_seqlen_q=max_seqlen_q, + softmax_scale=self.scaling, + causal=True, + window_size=self.window_size, + rel_extent=self.rel_extent, + rel_logits=rel_logits[:nt], + num_splits=num_splits, + max_kv_len=md.max_seq_len, + out=output[:nt], + ) diff --git a/vllm/models/inkling/amd/layernorm.py b/vllm/models/inkling/amd/layernorm.py new file mode 100644 index 000000000000..18f9951f65d4 --- /dev/null +++ b/vllm/models/inkling/amd/layernorm.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling RMSNorm (no bias, weight-scaled), backed by the vendored Triton kernel.""" + +from __future__ import annotations + +import torch +from torch import nn + +from .ops import rmsnorm + + +class InklingRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return x + original_shape = x.shape + x_2d = x.contiguous().view(-1, self.hidden_size) + y = rmsnorm(x_2d, self.weight, self.variance_epsilon) + return y.view(original_shape) diff --git a/vllm/models/inkling/amd/logits_processor.py b/vllm/models/inkling/amd/logits_processor.py new file mode 100644 index 000000000000..2e142c5961ef --- /dev/null +++ b/vllm/models/inkling/amd/logits_processor.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling logits processor (muP + LoRA aware). + +Inkling divides the final logits by a muP width multiplier +(``logits_mup_width_multiplier``). This applies it two ways, depending on +whether an lm_head LoRA is attached: + +* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no + separate elementwise kernel, no extra rounding, no weight mutation. +* LoRA attached: the LoRA manager wraps this layer in + ``LogitsProcessorWithLoRA``, whose ``forward`` calls + ``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with + ``self`` bound to the wrapper. We detect that via ``base_layer`` and take the + LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA + delta), then divide the full logits by the multiplier so the delta is scaled + too. muP thus composes with the LoRA delta, with the dispatch as the only + model-side branching. +""" + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + +class InklingLogitsProcessor(LogitsProcessor): + """``LogitsProcessor`` that applies Inkling's muP logits width multiplier. + + Args: + vocab_size: Padded vocabulary size. + org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``). + scale: Base logits scale (kept ``1.0`` for the served checkpoint). + logits_as_input: Whether the input is already logits. + soft_cap: Optional logit soft cap (``None`` for the served checkpoint). + logits_mup_width_multiplier: muP width divisor for the final logits; + ``None`` or ``0`` disables it. + """ + + def __init__( + self, + vocab_size: int, + org_vocab_size: int | None = None, + scale: float = 1.0, + logits_as_input: bool = False, + soft_cap: float | None = None, + logits_mup_width_multiplier: float | None = None, + ) -> None: + super().__init__( + vocab_size=vocab_size, + org_vocab_size=org_vocab_size, + scale=scale, + logits_as_input=logits_as_input, + soft_cap=soft_cap, + ) + self.logits_mup_width_multiplier = logits_mup_width_multiplier + self._logits_zero: torch.Tensor | None = None + + def forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper, + # which calls this forward with ``self`` bound to the wrapper. The + # wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch + # ``_lora_forward`` explicitly through the base_layer's class (it + # provides ``_get_logits``/``logits_as_input``; only ``_lora_forward`` + # lives on this class). + if hasattr(self, "base_layer"): + return type(self.base_layer)._lora_forward( + self, lm_head, hidden_states, embedding_bias + ) + return self._base_forward(lm_head, hidden_states, embedding_bias) + + def _lora_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits`` + # returns the base logits plus the lm_head LoRA delta. Apply the muP + # divisor on the full logits so the LoRA delta is scaled too. + mup_multiplier = self.base_layer.logits_mup_width_multiplier + mup = 1.0 / mup_multiplier if mup_multiplier else None + if self.logits_as_input: + logits = hidden_states + else: + logits = self._get_logits(hidden_states, lm_head, embedding_bias) + # TODO: fuse this multiplication + if logits is not None and mup: + assert self.base_layer.soft_cap is None + assert self.base_layer.scale == 1.0 + logits = logits * mup + return logits + + def _base_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + mup = self.logits_mup_width_multiplier + if not mup: + return super().forward(lm_head, hidden_states, embedding_bias) + # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue): + # no separate elementwise kernel, no bf16 rounding of scaled logits, and + # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no + # soft cap, unit logits scale. + assert self.soft_cap is None + assert self.scale == 1.0 + w = lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.org_vocab_size] + return logits diff --git a/vllm/models/inkling/amd/mlp.py b/vllm/models/inkling/amd/mlp.py new file mode 100644 index 000000000000..51deb38e64ef --- /dev/null +++ b/vllm/models/inkling/amd/mlp.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling dense SwiGLU MLP (also used as the MoE shared expert). + +The checkpoint stores the gate/up projection as a single fused, *interleaved* +weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain +``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair +together, and an interleaved SwiGLU activation. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig + + +class InklingDenseMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + use_global_scale: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = ColumnParallelLinear( + hidden_size, + 2 * intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.down_proj", + ) + if use_global_scale: + self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False) + else: + self.global_scale = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from .ops import silu_and_mul_triton + + gate_up, _ = self.gate_up_proj(x) + x = silu_and_mul_triton(gate_up) + x, _ = self.down_proj(x) + if self.global_scale is not None: + x.mul_(self.global_scale) + # TP-partial output: the layer's reduce-scatter fallback consumes it. + return x diff --git a/vllm/models/inkling/amd/model.py b/vllm/models/inkling/amd/model.py new file mode 100644 index 000000000000..e37a54938d8c --- /dev/null +++ b/vllm/models/inkling/amd/model.py @@ -0,0 +1,669 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling model implementation for AMD GPUs.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.models.inkling.common.mm_preprocess import ( + InklingDummyInputsBuilder, + InklingMultiModalProcessor, + InklingProcessingInfo, + inkling_audio_enabled, + inkling_vision_enabled, +) +from vllm.models.inkling.common.towers import InklingAudio, InklingVision +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from ..configs import InklingMMConfig, InklingModelConfig +from .attention import InklingAttention, compute_log_scaling_tau +from .layernorm import InklingRMSNorm +from .logits_processor import InklingLogitsProcessor +from .mlp import InklingDenseMLP +from .moe import InklingMoE +from .ops.norm import add_rmsnorm, embed_rmsnorm +from .sconv_swa_attn import _ATTN, _MLP, InklingConvState +from .short_conv import InklingShortConv + + +def _layer_id(name: str) -> int | None: + m = re.search(r"\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +def _sconv_add_norm( + delta: torch.Tensor, + hidden: torch.Tensor, + sconv: InklingShortConv, + norm: InklingRMSNorm | None, + positions: torch.Tensor, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``. + + ROCm uses the portable collective path. The Lamport P2P implementation in + the NVIDIA model relies on CUDA GDC and CUDA-specific symmetric-memory + publication semantics which cannot be linked into a gfx950 HSACO.""" + norm_w = norm.weight if norm is not None else None + eps = norm.variance_epsilon if norm is not None else 0.0 + + # RCCL RS -> shard sconv -> AG -> fused add(+rmsnorm). + shard = tensor_model_parallel_reduce_scatter(delta, dim=-1) + shard = sconv(shard.contiguous(), positions) + full = tensor_model_parallel_all_gather(shard, dim=-1) + if norm is None: + return None, hidden + full + return add_rmsnorm(hidden, full, norm_w, eps) + + +class InklingDecoderLayer(nn.Module): + def __init__( + self, + config: InklingModelConfig, + layer_id: int, + is_local: bool, + quant_config: QuantizationConfig | None, + prefix: str, + force_dense_mlp: bool = False, + ) -> None: + super().__init__() + # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv + # streams (K/V/attn/mlp) are packed head-major into one block and share + # it. Built first so the attention layer can wire its K/V sconv to it. + self.conv_state = InklingConvState( + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + hidden_size=config.hidden_size, + kernel_size=config.sconv_kernel_size, + prefix=f"{prefix}.conv_state", + ) + self.attn_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attn = InklingAttention( + config, + num_heads=( + config.swa_num_attention_heads + if is_local + else config.num_attention_heads + ), + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + rel_extent=config.rel_extent, + local_extent=config.sliding_window_size, + is_local=is_local, + prefix=f"{prefix}.attn", + quant_config=quant_config, + conv_owner=self.conv_state, + ) + self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if force_dense_mlp or layer_id < config.dense_mlp_idx: + self.mlp: nn.Module = InklingDenseMLP( + hidden_size=config.hidden_size, + intermediate_size=config.dense_intermediate_size, + use_global_scale=config.use_global_scale, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + # InklingMoE decides per layer (from the checkpoint exclude list) + # whether the routed experts are NVFP4 or bf16; the shared sink + # experts are always bf16. + self.mlp = InklingMoE( + config, + prefix=f"{prefix}.mlp", + quant_config=quant_config, + ) + + # Short convolution on the attention-output and MLP-output residual + # streams, hidden-sharded: the sublayer outputs are reduce-scattered + # to [T, H/tp], the sconv runs on the shard, and an all-gather + # restores the full residual — all fused with the residual add + next + # rmsnorm via the Lamport P2P kernels for decode-sized batches. + tp_size = get_tensor_model_parallel_world_size() + sconv_dim = config.hidden_size // tp_size + self.attn_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_ATTN + ) + self.mlp_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_MLP + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None, + defer_mlp_add: bool = False, + attn_in: torch.Tensor | None = None, + log_scaling: torch.Tensor | None = None, + ) -> ( + torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]] + ): + # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in + # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm. + # A None delta means the partials sit in the NVLS symm buffer. + if pending is None: + if attn_in is None: + # First layer; on the text path attn_norm comes fused with + # the embedding gather (chain_weight in embed_rmsnorm). + attn_in = self.attn_norm(hidden_states) + else: + attn_in, hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.attn_norm, positions + ) + attn_output = self.attn(positions, attn_in, log_scaling) + mlp_in, hidden_states = _sconv_add_norm( + attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions + ) + mlp_output = self.mlp(mlp_in) + if defer_mlp_add: + # Caller folds mlp_output (pre-reduce, pre-sconv) into the next + # fused sconv+add+rmsnorm. + return hidden_states, (mlp_output, self.mlp_sconv) + # Standalone (MTP) tail: finish the sublayer without a norm. + return _sconv_add_norm( + mlp_output, hidden_states, self.mlp_sconv, None, positions + )[1] + + +class InklingReplicatedEmbedding(nn.Module): + """Full-vocab embedding table replicated on every TP rank. + + Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a + 1/tp shard) for no masked lookup and no per-lookup TP all-reduce — one + all-reduce per MTP draft step plus one per verify pass — and keeps the + full table on-rank for the fused gather+norm kernels (``embed_rmsnorm``, + ``embed_dual_rmsnorm_cat``). Bit-exact vs vocab-parallel: the all-reduce + there only ever summed one real row against exact zeros. The LM head + stays vocab-sharded. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()), + requires_grad=False, + ) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return embed_rmsnorm(input_ids, self.weight, None, 0.0) + + +class InklingModel(nn.Module): + def __init__( + self, + *, + config: InklingModelConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.config = config + self.embed_tokens = InklingReplicatedEmbedding( + config.padded_vocab_size, config.hidden_size + ) + self.embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + local_ids = set(config.local_layer_ids) + + def get_layer(prefix: str) -> InklingDecoderLayer: + idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1]) + return InklingDecoderLayer( + config, idx, idx in local_ids, quant_config, prefix + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + self.norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + # Row gather + embed_norm in one launch. + norm = self.embed_norm + return embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + attn_in0: torch.Tensor | None = None + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + # embed_norm was already applied when producing inputs_embeds. + hidden_states = inputs_embeds + else: + # Gather + embed_norm + the first layer's attn_norm, one launch. + norm = self.embed_norm + hidden_states, attn_in0 = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + self.config.rms_norm_eps, + chain_weight=self.layers[self.start_layer].attn_norm.weight, + ) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + log_scaling = None + if self.config.log_scaling_n_floor is not None: + log_scaling = compute_log_scaling_tau( + positions, + self.config.log_scaling_n_floor, + self.config.log_scaling_alpha, + ) + + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, pending = layer( + positions, + hidden_states, + pending=pending, + defer_mlp_add=True, + attn_in=attn_in0, + log_scaling=log_scaling, + ) + attn_in0 = None + + if not get_pp_group().is_last_rank: + if pending is not None: + hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], None, positions + )[1] + return IntermediateTensors({"hidden_states": hidden_states}) + if pending is not None: + # Final RS/sconv/AG + residual add fused with the final rmsnorm. + norm_out = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.norm, positions + )[0] + assert norm_out is not None + return norm_out + return self.norm(hidden_states) + + +class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA): + """Shared text-backbone causal-LM scaffolding for both entry classes.""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + ".w13_dn": ".gate_up_proj", + ".w2_md": ".down_proj", + }, + orig_to_new_stacked={ + ".attn.wq_du.": (".attn.qkvr.", 0), + ".attn.wk_dv.": (".attn.qkvr.", 1), + ".attn.wv_dv.": (".attn.qkvr.", 2), + ".attn.wr_du.": (".attn.qkvr.", 3), + }, + orig_to_new_prefix={ + "model.llm.layers.": "model.layers.", + "model.llm.embed_norm": "model.embed_norm", + "model.llm.embed": "model.embed_tokens", + "model.llm.norm": "model.norm", + "model.llm.unembed": "lm_head", + "language_model.layers.": "model.layers.", + "language_model.lm_head.": "lm_head.", + }, + orig_to_new_suffix={ + # NVFP4 scale + ".w13_weight.scale": ".w13_weight_scale", + ".w13_weight.scale2": ".w13_weight_scale_2", + ".w2_weight.scale": ".w2_weight_scale", + ".w2_weight.scale2": ".w2_weight_scale_2", + }, + ) + # Quark uses this mapping when resolving quantization exclusions for the + # four checkpoint attention projections fused into qkvr. + packed_modules_mapping = { + "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"], + "w13": ["w1", "w3"], + } + embedding_modules = { + "lm_head": "output_embeddings", + } + + def _build( + self, + vllm_config: VllmConfig, + text_config: InklingModelConfig, + prefix: str, + ) -> None: + quant_config = vllm_config.quant_config + self.config = text_config + # ROCm checkpoints use Quark OCP MXFP4. The global Quark config is + # passed into each routed MoE and its exclusion list keeps the rest of + # Inkling in bf16. + # Read by the MRV2 runner to publish per-request short-conv metadata. + # Short convolution is intrinsic to Inkling, so this is always set. + self.uses_sconv = True + self.model = InklingModel( + config=text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.lm_head = ParallelLMHead( + text_config.padded_vocab_size, + text_config.hidden_size, + org_num_embeddings=text_config.padded_vocab_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = InklingLogitsProcessor( + text_config.padded_vocab_size, + org_vocab_size=text_config.vocab_size, + soft_cap=text_config.final_logit_softcapping, + logits_mup_width_multiplier=text_config.logits_mup_width_multiplier, + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_weights(self, weights, self.config) + + +class InklingForCausalLM(_TmlForCausalLMBase): + """Text-only entry point (``inkling_model`` checkpoints).""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self._build(vllm_config, vllm_config.model_config.hf_config, prefix) + + +@MULTIMODAL_REGISTRY.register_processor( + InklingMultiModalProcessor, + info=InklingProcessingInfo, + dummy_inputs=InklingDummyInputsBuilder, +) +class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal): + """Top-level (multimodal) entry point. + + Builds the vision + audio towers on top of the shared text backbone. Inkling has + NO cross-modal fusion (the vision tower emits one token per patch, the audio + tower one token per frame), so generation reuses the inherited backbone + ``forward`` / ``compute_logits`` (the latter already applies muP) and this + class only adds multimodal embedding + merge. + """ + + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={ + "model.audio.": "audio.", + "model.visual.": "visual.vision_encoder.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|content_image|>" + if modality.startswith("audio"): + return "<|content_audio_input|>" + raise ValueError("Only image or audio modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: InklingMMConfig = vllm_config.model_config.hf_config + + self.visual = ( + InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual")) + if inkling_vision_enabled(config) + else None + ) + self.audio = ( + InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio")) + if inkling_audio_enabled(config) + else None + ) + + self._build(vllm_config, config.text_config, prefix) + + # -- multimodal embedding ------------------------------------------- + + def _process_image_input( + self, pixel_values: Any, num_patches: Any + ) -> tuple[torch.Tensor, ...]: + assert self.visual is not None + # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors, + # or a single concatenated tensor. Normalize to a flat batch, run the + # tower once, then split back per item. + if isinstance(pixel_values, (list, tuple)): + if not pixel_values: + return () + sizes = [int(p.shape[0]) for p in pixel_values] + patches = torch.cat(list(pixel_values), dim=0) + else: + patches = pixel_values + sizes = self._sizes_from(num_patches, patches.shape[0]) + + patches = patches.to(device=self.visual.device, dtype=self.visual.dtype) + embeds = self.visual(patches) # [total_patches, D] + return tuple(embeds.split(sizes)) + + def _process_audio_input( + self, input_audio_features: Any, num_audio_tokens: Any + ) -> tuple[torch.Tensor, ...]: + assert self.audio is not None + if isinstance(input_audio_features, (list, tuple)): + if not input_audio_features: + return () + sizes = [int(d.shape[0]) for d in input_audio_features] + dmel = torch.cat(list(input_audio_features), dim=0) + else: + dmel = input_audio_features + sizes = self._sizes_from(num_audio_tokens, dmel.shape[0]) + + dmel = dmel.to(device=self.audio.device) + embeds = self.audio(dmel) # [total_frames, D] + return tuple(embeds.split(sizes)) + + @staticmethod + def _sizes_from(counts: Any, total: int) -> list[int]: + if counts is None: + return [total] + if isinstance(counts, torch.Tensor): + return [int(c) for c in counts.flatten().tolist()] + if isinstance(counts, (list, tuple)): + flat: list[int] = [] + for c in counts: + flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c)) + return flat + return [int(counts)] + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + # Iterate modalities in a stable order so the returned per-item tensors + # line up with their appearance order; the positional merge in + # embed_input_ids handles actual placement. + pixel_values = kwargs.get("pixel_values") + num_patches = kwargs.get("num_patches") + input_audio_features = kwargs.get("input_audio_features") + num_audio_tokens = kwargs.get("num_audio_tokens") + + embeddings: tuple[torch.Tensor, ...] = () + if pixel_values is not None and self.visual is not None: + embeddings += self._process_image_input(pixel_values, num_patches) + if input_audio_features is not None and self.audio is not None: + embeddings += self._process_audio_input( + input_audio_features, num_audio_tokens + ) + return embeddings + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg + # signature for multimodal models. Text embeddings come from the shared + # backbone (which applies embed_norm); MM embeddings are scattered in. + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # Placeholder ids use unused vocabulary slots and these positions are + # overwritten by MM embeds below. + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def get_language_model(self) -> nn.Module: + # This class IS the causal LM (the towers are side branches), so the + # language model is self — callers expect a module exposing ``.model`` + # / ``.lm_head`` (e.g. the MTP/eagle loader shares embeddings via + # ``get_language_model().model.embed_tokens``). + return self + + +# =========================================================================== +# Weight loading +# =========================================================================== + + +_MOE_EXPERT_WEIGHT_RE = re.compile( + r"^(?P.*\.mlp)\.(?P(?:shared_)?experts\..+)$" +) + + +def _is_peft_adapter_weight(name: str) -> bool: + return ".lora_A." in name or ".lora_B." in name + + +def _load_inkling_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + config: InklingModelConfig, +) -> set[str]: + moe_modules = { + name: mod for name, mod in module.named_modules() if isinstance(mod, InklingMoE) + } + loaded: set[str] = set() + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + local_ids = set(config.local_layer_ids) + + def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in module.hf_to_vllm_mapper.apply(weights): + # LightSeek's MXFP4 conversion bundles a quantized copy of the + # standalone PEFT adapter in the base safetensors index. These are + # not base-model parameters; adapters remain opt-in at serving. + if _is_peft_adapter_weight(name): + continue + shard_id = getattr(weight, "shard_id", None) + # Replicate K/V conv-free GQA heads when tp_size > num_kv_heads. + if ( + shard_id in (1, 2) + and name.endswith(".attn.qkvr.weight") + and weight.shape[0] > 0 + ): + lid = _layer_id(name) + if lid is not None: + is_local = lid in local_ids + n_kv = ( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ) + head_dim = config.swa_head_dim if is_local else config.head_dim + if tp_size > n_kv and weight.shape[0] == n_kv * head_dim: + kv_idx = (tp_rank * n_kv) // tp_size + weight = weight.narrow(0, kv_idx * head_dim, head_dim) + weight.shard_id = shard_id + + # MoE expert tensors (fused stacked, routed + shared sink): translate + # the checkpoint layout to per-expert MoERunner loads. + moe_match = _MOE_EXPERT_WEIGHT_RE.match(name) + if moe_match is not None and moe_match.group("mlp") in moe_modules: + moe = moe_modules[moe_match.group("mlp")] + for rel in moe.load_expert_weight(moe_match.group("rest"), weight): + loaded.add(f"{moe_match.group('mlp')}.{rel}") + continue + + yield name, weight + + loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) + loaded |= loader.load_weights(_iter_loadable_weights()) + + # Post-load MoE fixups (default input scales, zeroed EP-padding experts). + for moe_name, moe in moe_modules.items(): + for rel in moe.finalize_load(): + loaded.add(f"{moe_name}.{rel}") + return loaded + + +EntryClass = [InklingForCausalLM, InklingForConditionalGeneration] diff --git a/vllm/models/inkling/amd/moe.py b/vllm/models/inkling/amd/moe.py new file mode 100644 index 000000000000..d8141eb4cf08 --- /dev/null +++ b/vllm/models/inkling/amd/moe.py @@ -0,0 +1,693 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling mixture-of-experts on vLLM's MoERunner abstraction. + +Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over +the routed experts, log-sigmoid renormalization over the k routed + S shared +"sink" logits, scaled by route_scale * global_scale. The routed top-k goes +through vLLM's MoERunner (which handles TP/EP); the sink experts run in +:class:`InklingSinkExperts` -- replicated across EP ranks (every token +activates every sink) and always bf16 (the checkpoint excludes every +``shared_experts`` from quantization). + +MXFP4 routed experts reuse vLLM's Quark OCP MXFP4 fused-MoE method; excluded +(bf16) layers fall back to the unquantized method. The checkpoint's fused +stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` / +``.input_amax`` aux tensors) are translated to the standard per-expert loads +in :meth:`InklingMoE.load_expert_weight`. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.triton_utils import tl, tldevice, triton +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream + +from ..configs import InklingModelConfig + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import ( + RoutedExperts, + ) + +# --------------------------------------------------------------------------- +# Gate / expert selection +# --------------------------------------------------------------------------- + +_INKLING_LL_BF16_MAX_TOKENS = 64 +_MXFP4_INPUT_SCALE_DENOMINATOR = torch.finfo(torch.float8_e4m3fn).max * 6.0 + + +def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + leading = list(x.shape[:-1]) + flat = x.flatten(0, -2) + if ( + flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS + and flat.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and flat.is_cuda + and flat.is_contiguous() + and weight.is_contiguous() + and flat.shape[1] % 8 == 0 + and current_platform.has_device_capability(90) + and ll_bf16.is_available() + ): + out = ll_bf16.ll_bf16_gemm(flat, weight) + else: + out = torch.mm(flat, weight.T, out_dtype=torch.float32) + return out.view(*leading, weight.shape[0]) + + +@triton.jit(do_not_specialize=["T", "route_scale"]) +def _inkling_gate_select_kernel( + logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad) + bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False) + global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False) + ids_ptr, # [T, K + S] int32 out: selected expert ids + weights_ptr, # [T, K + S] fp32 out: renormalized weights + route_scale, + T, + G: tl.constexpr, # total gate experts (routed + shared) + stride_logits_0, + R: tl.constexpr, # routed experts + K: tl.constexpr, # top-k routed + S: tl.constexpr, # shared (sink) experts + HAS_BIAS: tl.constexpr, + HAS_GSCALE: tl.constexpr, + BLOCK_G: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + if pid >= T: + return + offs = tl.arange(0, BLOCK_G) + mask_r = offs < R + logits = tl.load( + logits_ptr + pid * stride_logits_0 + offs, + mask=offs < G, + other=float("-inf"), + ).to(tl.float32) + + # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf. + sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf")) + if HAS_BIAS: + bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32) + sel = tl.where(mask_r, sel + bias, float("-inf")) + + scale = route_scale + if HAS_GSCALE: + scale = scale * tl.load(global_scale_ptr).to(tl.float32) + + # Iterative top-K (K is small); argmax tie-breaks to the lowest index + # (stable ordering). + A: tl.constexpr = K + S + offs_a = tl.arange(0, A) + top_ids = tl.zeros([A], dtype=tl.int32) + active = tl.zeros([A], dtype=tl.float32) + for kk in tl.static_range(K): + idx = tl.argmax(sel, axis=0).to(tl.int32) + raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == kk, idx, top_ids) + active = tl.where(offs_a == kk, raw, active) + sel = tl.where(offs == idx, float("-inf"), sel) + if S > 0: + # Shared sink logits sit at the tail of the gate output; their expert + # ids continue after the routed range (R + j). + for jj in tl.static_range(S): + raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids) + active = tl.where(offs_a == K + jj, raw, active) + + # Log-sigmoid renormalization over the K + S active logits. + abs_l = tl.abs(active) + min_l = tl.minimum(active, 0.0) + log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l)) + max_lp = tl.max(log_probs, axis=0) + exp_shifted = tldevice.exp(log_probs - max_lp) + sum_exp = tl.sum(exp_shifted, axis=0) + weights = exp_shifted / sum_exp * scale + + tl.store(ids_ptr + pid * A + offs_a, top_ids) + tl.store(weights_ptr + pid * A + offs_a, weights) + + +def inkling_gate_select( + logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding) + n_gate_experts: int, + n_routed_experts: int, + topk: int, + n_shared_experts: int, + bias: torch.Tensor | None, + route_scale: float, + global_scale: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids).""" + assert logits.dtype == torch.float32 + tokens = logits.shape[0] + active = topk + n_shared_experts + topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device) + topk_weights = torch.empty( + (tokens, active), dtype=torch.float32, device=logits.device + ) + if tokens == 0: + return topk_weights, topk_ids + _inkling_gate_select_kernel[(tokens,)]( + logits, + bias if bias is not None else logits, + global_scale if global_scale is not None else logits, + topk_ids, + topk_weights, + route_scale, + tokens, + n_gate_experts, + logits.stride(0), + n_routed_experts, + topk, + n_shared_experts, + HAS_BIAS=bias is not None, + HAS_GSCALE=global_scale is not None, + BLOCK_G=triton.next_power_of_2(n_gate_experts), + ) + return topk_weights, topk_ids + + +class InklingGate(nn.Module): + """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and + global scale (the served checkpoint's only configuration).""" + + def __init__( + self, + d_model: int, + n_routed_experts: int, + n_shared_experts: int, + experts_per_token: int, + route_scale: float, + *, + use_global_scale: bool = False, + use_gate_bias: bool = False, + ) -> None: + super().__init__() + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.n_total_experts = n_routed_experts + n_shared_experts + self.topk = experts_per_token + self.route_scale = route_scale + + padded_experts = self.n_total_experts + (-self.n_total_experts) % 8 + self.weight = Parameter( + torch.empty(padded_experts, d_model), requires_grad=False + ) + set_weight_attrs(self.weight, {"weight_loader": self._load_weight}) + self.global_scale: Parameter | None + if use_global_scale: + self.global_scale = Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + else: + self.global_scale = None + self.bias: Parameter | None + if use_gate_bias: + self.bias = Parameter( + torch.empty(n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + else: + self.bias = None + + @staticmethod + def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None: + param.data.zero_() + param.data[: loaded_weight.shape[0]].copy_(loaded_weight) + + def compute_logits(self, x: torch.Tensor) -> torch.Tensor: + """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk).""" + return _linear_with_fp32_out(x, self.weight) + + def select_experts( + self, gating_output: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Full selection: (weights, ids) of [T, K + S]. The first K entries + are the routed top-k; the S trailing entries are the sink gammas.""" + return inkling_gate_select( + gating_output, + self.n_total_experts, + self.n_routed_experts, + self.topk, + self.n_shared_experts, + self.bias, + self.route_scale, + self.global_scale, + ) + + +# --------------------------------------------------------------------------- +# MoE layer +# --------------------------------------------------------------------------- + + +def _inkling_moe_ep_size() -> int: + """EP size the MoERunner layer will run with (mirrors + FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when + expert parallelism is enabled).""" + parallel_config = get_current_vllm_config().parallel_config + if not parallel_config.enable_expert_parallel: + return 1 + world = ( + get_tensor_model_parallel_world_size() + * get_dp_group().world_size + * get_pcp_group().world_size + ) + return world if world > 1 else 1 + + +class InklingSinkExperts(nn.Module): + """Shared "sink" experts with per-token gammas, in bf16. + + Replicated across EP ranks (every token activates every sink, so + EP-sharding them would hotspot the owning rank) and TP-sharded on the + intermediate dim so the output remains a TP-partial sum like the routed + output. The sinks are always bf16 (the checkpoint excludes every + ``shared_experts`` from quantization): the experts concatenate into two + plain dense GEMMs with the fused sink epilogue between them. + """ + + def __init__( + self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = "" + ) -> None: + super().__init__() + self.n_experts = n_experts + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + intermediate_pp = d_mlp // tp_size + self.w13_weight = Parameter( + torch.empty(n_experts, 2 * intermediate_pp, d_model), + requires_grad=False, + ) + self.w2_weight = Parameter( + torch.empty(d_model, n_experts * intermediate_pp), + requires_grad=False, + ) + self._unit: torch.Tensor | None = None + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint sink tensor (stacked over the S experts).""" + if key == "w13_weight": + if weight.shape != self.w13_weight.shape: + shard = self.w13_weight.shape[1] + weight = weight.narrow(1, self.tp_rank * shard, shard) + self.w13_weight.data.copy_(weight) + return [key] + + assert key == "w2_weight" + shard = self.w2_weight.shape[1] // self.n_experts + shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard + for expert_idx, expert_weight in enumerate(weight): + local_weight = expert_weight.narrow(1, shard_start, shard) + start = expert_idx * shard + self.w2_weight.data[:, start : start + shard].copy_(local_weight) + return [key] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp).""" + from .ops import sink_silu_mul_epilogue + + # One GEMM over the experts' stacked w13 (a view), fused epilogue, + # then one GEMM whose K-reduction over the K-concatenated w2 performs + # the expert sum. + if self._unit is None or self._unit.device != x.device: + self._unit = torch.ones( + self.n_experts, dtype=torch.float32, device=x.device + ) + raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F) + h = sink_silu_mul_epilogue( + raw, self._unit, gammas, self._unit, self.n_experts, x.dtype + ) + return h @ self.w2_weight.T # (T, D) + + +class InklingSinkExpertsLinear(nn.Module): + """LoRA-capable implementation of the Inkling sink experts.""" + + def __init__( + self, + n_experts: int, + d_model: int, + d_mlp: int, + *, + prefix: str = "", + ) -> None: + super().__init__() + from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, + ) + + self.n_experts = n_experts + self.d_mlp = d_mlp + total = n_experts * d_mlp + self.w13 = MergedColumnParallelLinear( + input_size=d_model, + output_sizes=[total, total], + bias=False, + prefix=f"{prefix}.w13", + ) + self.w2 = RowParallelLinear( + input_size=total, + output_size=d_model, + bias=False, + reduce_results=False, + prefix=f"{prefix}.w2", + ) + self._w2_input_pp = self.w2.input_size_per_partition + self._col_expert: torch.Tensor | None = None + + def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor: + if self._col_expert is None or self._col_expert.device != gammas.device: + local = self._w2_input_pp + start = get_tensor_model_parallel_rank() * local + cols = torch.arange(start, start + local, device=gammas.device) + self._col_expert = (cols // self.d_mlp).long() + return gammas[:, self._col_expert] + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + if key == "w13_weight": + d_model = weight.shape[-1] + gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous() + up = weight[:, 1::2, :].reshape(-1, d_model).contiguous() + self.w13.weight_loader(self.w13.weight, gate, 0) + self.w13.weight_loader(self.w13.weight, up, 1) + return ["w13.weight"] + w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous() + self.w2.weight_loader(self.w2.weight, w) + return ["w2.weight"] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.w13(x) + gate, up = gate_up.chunk(2, dim=-1) + hidden_states = torch.nn.functional.silu(gate) * up + hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype) + output, _ = self.w2(hidden_states) + return output + + +class InklingMoE(nn.Module): + def __init__( + self, + config: InklingModelConfig, + *, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + # Overfit to the served checkpoint: sigmoid gate renormalized after + # top-k, shared sink experts, interleaved gate/up checkpoint rows. + assert config.gate_activation == "sigmoid" and config.norm_after_topk + assert config.n_shared_experts > 0 and config.shared_expert_sink + assert config.inference_moe_w13_interleaved + n_routed = config.n_routed_experts + n_shared = config.n_shared_experts + self.n_routed_experts = n_routed + self.gate = InklingGate( + d_model=config.hidden_size, + n_routed_experts=n_routed, + n_shared_experts=n_shared, + experts_per_token=config.num_experts_per_tok, + route_scale=config.route_scale, + use_global_scale=config.use_global_scale, + use_gate_bias=config.use_gate_bias, + ) + + # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs + # (local_expert_offset = ep_rank * local_num_experts), so pad the + # expert count to a multiple of the EP size. A no-op for the usual + # power-of-two EP sizes (n_routed is a power of two). + num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size() + + # The released MXFP4 checkpoint keeps the first routed-expert layer + # in bf16 and lists its two expert weights explicitly in Quark's + # exclusion list. FusedMoEFactory asks for a quant method at the module + # prefix, so an exact weight exclusion would otherwise be missed and + # the bf16 tensors would be loaded into MXFP4 parameters. + routed_quant_config = quant_config + if quant_config is not None: + excluded = set(getattr(quant_config, "quant_config", {}).get("exclude", ())) + expert_prefix = f"{prefix}.experts" + routed_weights = { + f"{expert_prefix}.w13_weight", + f"{expert_prefix}.w2_weight", + } + if routed_weights <= excluded: + routed_quant_config = None + + self.experts = FusedMoEFactory( + num_experts=num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + renormalize=False, + quant_config=routed_quant_config, + prefix=f"{prefix}.experts", + custom_routing_function=self._select_routed, + router_logits_dtype=torch.float32, + activation="silu", + ) + # The decoder layer reduce-scatters the MoE delta into the sconv + # stream itself (RS -> shard sconv -> AG); the runner must return the + # per-rank partial sum instead of all-reducing. + self.experts.moe_config.skip_final_all_reduce = True + + self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + sink_experts_cls = ( + InklingSinkExpertsLinear + if get_current_vllm_config().lora_config is not None + else InklingSinkExperts + ) + self.sink_experts = sink_experts_cls( + n_experts=n_shared, + d_model=config.hidden_size, + d_mlp=config.intermediate_size, + prefix=f"{prefix}.shared_experts", + ) + + # Sink chain overlaps the routed MoE call on the aux stream for + # decode-sized batches (same pattern as the runner's SharedExperts + # multi-stream overlap). The routed GEMM runs on the default stream and + # the sink chain on the aux stream, joined via these two events by + # ``maybe_execute_in_parallel``. + self._sink_stream: torch.cuda.Stream | None = aux_stream() + self._sink_events = (torch.cuda.Event(), torch.cuda.Event()) + + def _select_routed( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """MoERunner ``custom_routing_function``: the routed top-k slice of the + full (routed + sink) selection. + + forward() stashes its selection (keyed by logits identity) so the + gate select runs once per layer; the fallback covers paths where the + runner re-derives the logits (e.g. naive DP dispatch). + """ + del hidden_states, renormalize + assert topk == self.gate.topk + cached = self._routed_sel + self._routed_sel = None + if cached is not None and cached[0] is gating_output: + return cached[1], cached[2] + weights, ids = self.gate.select_experts(gating_output) + return weights[:, :topk].contiguous(), ids[:, :topk].contiguous() + + def forward(self, x: torch.Tensor) -> torch.Tensor | None: + router_logits = self.gate.compute_logits(x) + num_tokens = x.shape[0] + # One gate select per layer: the routed slice is stashed for the + # routing function inside the MoERunner op; the sink gammas are the + # trailing columns. + k = self.gate.topk + weights, ids = self.gate.select_experts(router_logits) + self._routed_sel = ( + router_logits, + weights[:, :k].contiguous(), + ids[:, :k].contiguous(), + ) + gammas = weights[:, k:] + + out, sink_out = maybe_execute_in_parallel( + lambda: self.experts(hidden_states=x, router_logits=router_logits), + lambda: self.sink_experts(x, gammas), + self._sink_events[0], + self._sink_events[1], + self._sink_stream + if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + else None, + ) + self._routed_sel = None + + return out.add_(sink_out) + + # -- weight loading ---------------------------------------------------- + + def _local_expert_slots(self) -> dict[int, int]: + """Global expert id -> local slot for this rank's expert partition.""" + manager = self.experts.routed_experts.expert_map_manager + if manager.expert_map is None: + return {g: g for g in range(manager.global_num_experts)} + emap = manager.expert_map.tolist() + return {g: slot for g, slot in enumerate(emap) if slot >= 0} + + def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint expert tensor. + + ``name`` is relative to the mlp module: ``experts.`` (routed + stack) or ``shared_experts.shared_`` (sink experts). Returns the + loaded param names (relative to this module). + """ + if name.startswith("shared_experts."): + key = name.split(".", 1)[1].replace("shared_", "", 1) + return [ + f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight) + ] + + experts: RoutedExperts = self.experts.routed_experts + key = name.split(".", 1)[1] + + # original_shape is unused by the vLLM serving layout. + if key.endswith(".original_shape"): + return [] + if key.endswith(".input_amax"): + projection = "w13" if key.startswith("w13") else "w2" + amax = float(weight.max()) + assert math.isfinite(amax) and amax > 0, ( + f"bad {projection} input_amax: {amax}" + ) + input_scale = getattr(experts, f"{projection}_input_scale") + input_scale.data.fill_(amax / _MXFP4_INPUT_SCALE_DENOMINATOR) + return [f"experts.routed_experts.{projection}_input_scale"] + + # Quark's OCP MXFP4 converter stores block scales as a two-dimensional + # tensor with the expert and projection-row dimensions flattened: + # w13: [E * 2I, H / 32], w2: [E * H, I / 32]. + # Restore the expert dimension before applying the EP/TP slicing used + # for both packed weights and scales. + if key.endswith("_weight_scale") and weight.ndim == 2: + if weight.shape[0] % self.n_routed_experts != 0: + raise ValueError( + f"cannot unflatten {name} with shape {tuple(weight.shape)} " + f"over {self.n_routed_experts} experts" + ) + weight = weight.view( + self.n_routed_experts, + weight.shape[0] // self.n_routed_experts, + weight.shape[1], + ) + + param = getattr(experts, key) + slots = self._local_expert_slots() + gids = sorted(slots) + lids = [slots[g] for g in gids] + tp_rank = experts.moe_config.moe_parallel_config.tp_rank + tp_size = experts.moe_config.moe_parallel_config.tp_size + + if key.endswith("_scale_2"): + # Per-expert scalars, vectorized over the local experts. The + # fused w13 param carries one slot per gate/up half. + vals = weight[gids].float().to(param.device) + param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals + elif key.startswith("w13"): + # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the + # fused param layout is [w1(gate); w3(up)]. The TP-local rows form + # one contiguous slab of the interleaved tensor, so upload just + # that slab (a single bounded synchronous H2D; pre-uploading whole + # untrimmed tensors pins the mmap pages of the entire checkpoint + # and OOMs the host) and de-interleave on device. + dst_half = param.shape[1] // 2 + if weight.shape[1] % (2 * tp_size) != 0: + raise ValueError( + f"cannot TP-shard {name} with shape {tuple(weight.shape)} " + f"over {tp_size} ranks" + ) + logical_half = weight.shape[1] // (2 * tp_size) + if logical_half > dst_half: + raise ValueError( + f"checkpoint shard for {name} has {logical_half} rows per " + f"gate/up half, but destination only has {dst_half}" + ) + for gid, lid in slots.items(): + slab = weight[gid].narrow( + 0, + tp_rank * 2 * logical_half, + 2 * logical_half, + ) + slab = slab.to(param.device) + param.data[lid, :logical_half].copy_(slab[0::2]) + param.data[ + lid, + dst_half : dst_half + logical_half, + ].copy_(slab[1::2]) + else: + # w2: shard the packed intermediate (last) dim. AITER rounds the + # destination intermediate width to 256, so derive the logical TP + # slice from the checkpoint and leave the destination tail at its + # initialized padding value (zero for weights, one for scales). + if weight.shape[2] % tp_size != 0: + raise ValueError( + f"cannot TP-shard {name} with shape {tuple(weight.shape)} " + f"over {tp_size} ranks" + ) + shard = weight.shape[2] // tp_size + if shard > param.shape[2]: + raise ValueError( + f"checkpoint shard for {name} has width {shard}, but " + f"destination only has {param.shape[2]}" + ) + for gid, lid in slots.items(): + param.data[lid, :, :shard].copy_( + weight[gid].narrow(1, tp_rank * shard, shard) + ) + return [f"experts.routed_experts.{key}"] + + def finalize_load(self) -> list[str]: + """Post-load fixups for zeroed padding experts.""" + experts = self.experts.routed_experts + out: list[str] = [] + # Zero the EP-alignment padding experts (if any) so their + # (never-routed) slots hold defined values. + slots = self._local_expert_slots() + for gid in range(self.n_routed_experts, experts.global_num_experts): + lid = slots.get(gid) + if lid is None: + continue + for pname in ( + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_weight_scale_2", + "w2_weight_scale_2", + ): + p = getattr(experts, pname, None) + if p is not None: + p.data[lid].zero_() + return out diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py new file mode 100644 index 000000000000..34cde020ed38 --- /dev/null +++ b/vllm/models/inkling/amd/mtp.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). + +Implements the first MTP depth from the reference ``mtp_model.py`` shipped with +the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` +input projection, and a full Inkling transformer block with a dense bf16 MLP. + +The draft shares the target's token embedding table and LM head +(``load_eagle_model`` wires those references) and applies the backbone +``embed_norm`` on top: the depth layers were trained on the same normed +embeddings the backbone consumes (their own ``embed_norm`` weights are +near-identity trims, unlike the backbone's whitening ``embed_norm``). +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import maybe_prefix +from vllm.sequence import IntermediateTensors + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm + +# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. +# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a +# local (pname, wname, shard) list since the MTP loader remaps by hand. +_ATTENTION_PARAMS_MAPPING = [ + ("qkvr", "wq_du", 0), + ("qkvr", "wk_dv", 1), + ("qkvr", "wv_dv", 2), + ("qkvr", "wr_du", 3), +] + + +def _mtp_depth_from_name(name: str) -> int | None: + m = re.search(r"\.mtp\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +class InklingMTPDepthLayer(nn.Module): + """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" + + def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + super().__init__() + self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + return_bias=False, + prefix=f"{prefix}.input_proj", + ) + # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # full attention (the swa_* head config and sliding_window window) to + # match this depth's checkpoint transformer_block weights. + self.transformer_block = InklingDecoderLayer( + config, + layer_id=0, + is_local=is_local, + quant_config=None, + prefix=f"{prefix}.transformer_block", + force_dense_mlp=True, + ) + + def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)] + # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch. + hidden = self.input_proj(combined) + # The short conv self-fetches its paged SWA-cache metadata from the + # forward context (via its conv_owner prefix); no conv_meta to thread. + return self.transformer_block(positions, hidden) + + +class InklingMultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + if vllm_config.speculative_config.num_speculative_tokens != 1: + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + self.chain_hidden_post_norm = config.chain_hidden_post_norm + local_ids = set(config.local_layer_ids) + self.layers = nn.ModuleDict( + {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + ) + self.chain_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.chain_hidden_post_norm + else None + ) + # The target's raw token embedding (pre embed_norm), attached by + # load_eagle_model. Never materialized here: building our own + # replicated copy would transiently double the 2.3 GiB table. + self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment] + # The depth layers consume the *backbone-normed* embedding + # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights + # are near-identity (trained on already-normalized inputs), and + # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70. + # Weight loaded from the target's embed_norm.weight; gated like the + # target's InklingModel.embed_norm. + self.backbone_embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + """Draft-prefill embedding: fused gather + backbone embed_norm, then + the target's tower embeddings scattered in unnormed (the backbone + convention — MM embeds are merged after embed_norm).""" + norm = self.backbone_embed_norm + embeds = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type] + return embeds + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def fused_input_cat( + self, + layer: InklingMTPDepthLayer, + previous_hidden: torch.Tensor, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one + launch: embedding row gather + the backbone embed_norm + the depth + embed_norm chain on one side, hidden_norm on the other, written + straight into the cat buffer.""" + hidden_w = layer.hidden_norm.weight + embed_w = layer.embed_norm.weight + eps = layer.hidden_norm.variance_epsilon + if inputs_embeds is not None: + # Draft prefill with target-merged MM embeddings (already + # backbone-normed via embed_input_ids); only the depth embed_norm + # remains. + return embed_dual_rmsnorm_cat( + previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds + ) + return embed_dual_rmsnorm_cat( + previous_hidden, + hidden_w, + embed_w, + eps, + input_ids=input_ids, + embed_table=self.embed_tokens.weight, + pre_norm_weight=( + self.backbone_embed_norm.weight + if self.backbone_embed_norm is not None + else None + ), + ) + + 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: + # The draft's short conv is a paged SWA-cache layer (its conv_owner is + # auto-enumerated as a draft attention layer); its per-token metadata is + # built by the speculator's build_attn_metadata and read from the + # forward context, so nothing extra is threaded here. + if spec_step_idx != 0: + raise ValueError("Inkling MTP only supports spec_step_idx=0") + layer = self.layers["0"] + combined = self.fused_input_cat( + layer, previous_hidden_states, input_ids, inputs_embeds + ) + hidden = layer(combined, positions) + if self.chain_norm is not None: + hidden = self.chain_norm(hidden) + return hidden + + +class InklingMTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + self.model = InklingMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # The target's (vocab-sharded) LM head, attached by load_eagle_model; + # never materialized here (same reasoning as model.embed_tokens). + self.lm_head: ParallelLMHead = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor( + config.padded_vocab_size, + org_vocab_size=config.vocab_size, + soft_cap=config.final_logit_softcapping, + ) + self._logits_zero: torch.Tensor | None = None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.embed_input_ids( + input_ids, multimodal_embeddings, is_multimodal=is_multimodal + ) + + def forward( + self, + input_ids: torch.Tensor, + 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 | None: + # The MTP shares the base model's LM head, which is trained on + # ``hidden / mup``-scaled inputs, so apply the same mup scaling here + # for a matching logit scale — folded into the lm_head GEMM alpha + # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant + # for greedy draft sampling, but it matters for the gumbel sampling + # distribution at temperature > 0.) + mup = self.config.logits_mup_width_multiplier + if not mup: + return self.logits_processor(self.lm_head, hidden_states) + assert self.logits_processor.soft_cap is None + assert self.logits_processor.scale == 1.0 + w = self.lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self.logits_processor._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + return logits + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Greedy draft tokens via rank-local argmax + tiny (value, index) + reduction — no full-vocab logits all-gather. The muP divisor is a + positive scalar, so the argmax is invariant and the scaling is + skipped entirely.""" + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_mtp_weights(self, weights) + + +def _load_inkling_mtp_weights( + module: InklingMTP, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Load ``model.mtp.*`` weights into the MTP module. + + Checkpoint keys look like ``model.mtp.chain_norm.weight`` and + ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone + layer's fused-projection layout, so we apply the same qkvr / gate_up / down + remapping as ``_load_inkling_weights``. Token embedding and LM head are shared + (provided by ``load_eagle_model``) and are not present in mtp.safetensors. + """ + # Per-depth attention is full or sliding-window (config.local_layer_ids); + # each depth's qkvr MergedColumnParallelLinear is built with the matching + # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding. + # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never + # exceeds a layer's kv-head count and no GQA K/V replication is needed here. + params = dict(module.named_parameters()) + loaded: set[str] = set() + + def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: + param = params.get(name) + if param is None: + return False + loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + if loader is default_weight_loader or param.shape == weight.shape: + default_weight_loader(param, weight) + else: + loader(param, weight) + else: + loader(param, weight, shard_id) # type: ignore[call-arg] + loaded.add(name) + return True + + for name, weight in weights: + depth = _mtp_depth_from_name(name) + # Token embedding and LM head are never materialized on the draft + # (no params to load into); load_eagle_model attaches the target's. + if name in ("model.llm.embed.weight", "model.llm.unembed.weight"): + continue + # The backbone embed_norm, applied to the shared embedding before the + # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The + # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded + # below. Only the shared backbone key routes here. + if name == "model.llm.embed_norm.weight": + _load("model.backbone_embed_norm.weight", weight) + continue + # Only consume the MTP weights; everything else belongs to the target. + if ".mtp." not in name: + continue + # Only the first checkpoint depth is used for MTP=1. + if depth is not None and depth != 0: + continue + # model.mtp.chain_norm.weight -> model.chain_norm.weight + # model.mtp.layers.{i}.X -> model.layers.{i}.X + original_name = name + name = name.replace(".mtp.layers.", ".layers.").replace( + ".mtp.chain_norm.", ".chain_norm." + ) + + if ".chain_norm." in name and module.model.chain_norm is None: + raise ValueError( + "Inkling checkpoint contains chain_norm weights but " + "chain_hidden_post_norm is disabled." + ) + + # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr). + matched = False + for pname, wname, shard in _ATTENTION_PARAMS_MAPPING: + if f".attn.{wname}." in name: + mapped_name = name.replace(f".{wname}.", f".{pname}.") + if not _load(mapped_name, weight, shard): + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + matched = True + break + if matched: + continue + + # Dense MLP fused gate/up + down. + if ".mlp.w13_dn.weight" in name: + loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) + elif ".mlp.w2_md.weight" in name: + loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) + else: + if name.endswith(".bias") and name not in params: + continue + loaded_weight = _load(name, weight) + if not loaded_weight: + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + required = { + name + for name in params + if name.startswith("model.layers.") or name.startswith("model.chain_norm.") + } + if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled(): + raise ValueError( + "Inkling MTP checkpoint is missing required parameters: " + + ", ".join(missing) + ) + return loaded + + +EntryClass = [InklingMTP] diff --git a/vllm/models/inkling/amd/ops/__init__.py b/vllm/models/inkling/amd/ops/__init__.py new file mode 100644 index 000000000000..3c6ea5fe8f61 --- /dev/null +++ b/vllm/models/inkling/amd/ops/__init__.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling kernels (NVIDIA). + +``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4 +relative-attention wrapper are exposed lazily to keep this package's +import path lightweight. +""" + +from typing import TYPE_CHECKING + +from .norm import add_rmsnorm, rmsnorm +from .sconv import fused_sconv + +_LAZY_EXPORTS = { + "silu_and_mul_triton": "silu_and_mul", + "sink_silu_mul_epilogue": "silu_and_mul", + "inkling_fa4_rel_attention": "fa4_rel_attention", + "inkling_rel_attention_split_kv_decode": "rel_attention_decode", +} + +if TYPE_CHECKING: + from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401 + from .silu_and_mul import ( # noqa: F401 + silu_and_mul_triton, + sink_silu_mul_epilogue, + ) + +__all__ = [ + "add_rmsnorm", + "rmsnorm", + "fused_sconv", + *sorted(_LAZY_EXPORTS), +] + + +def __getattr__(name: str): + module = _LAZY_EXPORTS.get(name) + if module is not None: + import importlib + + mod = importlib.import_module(f".{module}", __name__) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/amd/ops/fa4_rel_attention.py b/vllm/models/inkling/amd/ops/fa4_rel_attention.py new file mode 100644 index 000000000000..d448ac8e2c5f --- /dev/null +++ b/vllm/models/inkling/amd/ops/fa4_rel_attention.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm paged attention with Inkling's query-dependent relative bias. + +The NVIDIA implementation uses the score-mod hook in tml-fa4. ROCm Flash +Attention and AITER do not expose an equivalent hook, so this module implements +the same operation directly in Triton. Query heads belonging to one KV head +are processed together and KV pages are gathered through vLLM's block table. +""" + +from __future__ import annotations + +import os +from typing import cast + +import torch + +from vllm.models.inkling.amd.ops.rel_attention_decode import ( + inkling_rel_attention_split_kv_decode, + use_split_kv_decode, +) +from vllm.platforms.rocm import on_gfx950 +from vllm.triton_utils import tl, triton + + +def bucket_max_seqlen_q(max_seqlen_q: int) -> int: + """Round the scheduling bound up to a power of two.""" + return 1 << max(0, max_seqlen_q - 1).bit_length() + + +def use_gfx950_gluon_decode( + *, max_query_len: int, page_size: int, head_dim: int +) -> bool: + """Use the vendored TokenSpeed CDNA4 decode kernel where it is supported.""" + return ( + os.getenv("INKLING_GFX950_GLUON", "1") == "1" + and on_gfx950() + and max_query_len == 1 + and page_size in (64, 128, 256) + and head_dim in (64, 128) + ) + + +def use_gfx950_gluon_extend( + *, + max_query_len: int, + max_kv_len: int, + page_size: int, + head_dim: int, + window_left: int, +) -> bool: + """Use Gluon only for the long full-attention extend regime it wins.""" + return ( + os.getenv("INKLING_GFX950_GLUON", "1") == "1" + and on_gfx950() + and max_query_len > 1 + and max_kv_len >= 8192 + and window_left < 0 + and page_size in (64, 128, 256) + and head_dim in (64, 128) + ) + + +def inkling_fa4_num_splits( + *, + is_local: bool, + batch_size: int, + max_query_len: int, + num_heads: int, + num_kv_heads: int, + max_kv_len: int, +) -> int: + """Keep the NVIDIA-facing split heuristic as API-compatible metadata. + + The ROCm Triton implementation performs online softmax in one program and + does not consume the result. Keeping this function unchanged avoids + platform-specific scheduling branches in :mod:`inkling.amd.attention`. + """ + if is_local: + return 1 + + q_rows = max_query_len * (num_heads // num_kv_heads) + q_tiles = (q_rows + 255) // 256 + base_ctas = batch_size * num_kv_heads * q_tiles + target_ctas = ( + 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64) + ) + max_splits = 128 + if q_tiles == 1 and batch_size == 1: + if num_kv_heads == 8: + max_splits = 16 + elif num_kv_heads == 4 or max_kv_len <= 8192: + max_splits = 32 + elif max_kv_len <= 65536: + max_splits = 64 + return max( + 1, + min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128), + ) + + +@triton.heuristics( + { + "BLOCK_D": lambda args: triton.next_power_of_2(args["head_dim"]), + "BLOCK_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), + "BLOCK_QH": lambda args: ( + args["BLOCK_Q"] * triton.next_power_of_2(args["gqa_group_size"]) + ), + } +) +@triton.jit(do_not_specialize_on_alignment=["cache_seqlens", "cu_seqlens_q"]) +def _inkling_rel_attention_kernel( + q_ptr, # [total_q, Hq, D] + k_ptr, # [blocks, page, Hkv, D] + v_ptr, # [blocks, page, Hkv, D] + rel_ptr, # [total_q, Hq, rel_extent] + out_ptr, # [total_q, Hq, D] + block_table_ptr, # [batch, max_pages] + cache_seqlens, + cu_seqlens_q, + gqa_group_size, + head_dim, + softmax_scale, + stride_q_t, + stride_q_h, + stride_q_d, + stride_k_b, + stride_k_p, + stride_k_h, + stride_k_d, + stride_v_b, + stride_v_p, + stride_v_h, + stride_v_d, + stride_r_t, + stride_r_h, + stride_r_e, + stride_o_t, + stride_o_h, + stride_o_d, + stride_bt_b, + page_size: tl.constexpr, + rel_extent: tl.constexpr, + window_left: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_QH: tl.constexpr, +): + # A program owns BLOCK_Q query positions and every query head sharing one + # KV head. This makes the QK/PV operations MFMA-friendly on CDNA. + pid_q = tl.program_id(0) + pid_kh = tl.program_id(1) + pid_b = tl.program_id(2) + + q_start = tl.load(cu_seqlens_q + pid_b) + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start + q_block = pid_q * BLOCK_Q + if q_block >= q_len: + return + + kv_len = tl.load(cache_seqlens + pid_b) + prefix_len = kv_len - q_len + q_head_start = pid_kh * gqa_group_size + bt_row = block_table_ptr + pid_b * stride_bt_b + + q_block_ptr = tl.make_block_ptr( + base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_q_t, stride_q_h, stride_q_d), + offsets=(q_block, 0, 0), + block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), + order=(2, 1, 0), + ) + q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") + q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) + + q_rows = q_block + tl.arange(0, BLOCK_Q) + q_abs = prefix_len + q_rows + q_valid = q_rows < q_len + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + + m_i = tl.full((BLOCK_QH,), float("-inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_QH,), dtype=tl.float32) + acc = tl.zeros((BLOCK_QH, BLOCK_D), dtype=tl.float32) + log2e: tl.constexpr = 1.4426950408889634 + + for k_start in range(0, kv_len, BLOCK_K): + k_pos = k_start + tl.arange(0, BLOCK_K) + k_valid = k_pos < kv_len + # Masked lanes still need an in-range page-table load. + safe_pos = tl.minimum(k_pos, tl.maximum(kv_len - 1, 0)) + page = tl.load(bt_row + safe_pos // page_size).to(tl.int64) + page_off = safe_pos % page_size + + k = tl.load( + k_ptr + + page[None, :] * stride_k_b + + page_off[None, :] * stride_k_p + + pid_kh * stride_k_h + + off_d[:, None] * stride_k_d, + mask=d_valid[:, None] & k_valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k) * (softmax_scale * log2e) + + dist = q_abs[:, None] - k_pos[None, :] + score_valid = q_valid[:, None] & k_valid[None, :] & (dist >= 0) + if window_left >= 0: + score_valid &= dist <= window_left + + # rel_logits is query- and head-dependent. Expand [Q, K] distance + # indices across the GQA heads, then flatten to match QK's row order. + off_h = tl.arange(0, BLOCK_H) + rel_dist = dist[:, None, :] + rel_valid = ( + q_valid[:, None, None] + & (off_h[None, :, None] < gqa_group_size) + & (rel_dist >= 0) + & (rel_dist < rel_extent) + ) + safe_dist = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) + bias = tl.load( + rel_ptr + + (q_start + q_rows[:, None, None]) * stride_r_t + + (q_head_start + off_h[None, :, None]) * stride_r_h + + safe_dist * stride_r_e, + mask=rel_valid, + other=0.0, + ) + bias = tl.reshape(bias, (BLOCK_QH, BLOCK_K)) + scores += bias.to(tl.float32) * log2e + score_valid = tl.reshape( + score_valid[:, None, :] & (off_h[None, :, None] < gqa_group_size), + (BLOCK_QH, BLOCK_K), + ) + scores = tl.where(score_valid, scores, float("-inf")) + + m_ij = tl.maximum(m_i, tl.max(scores, axis=1)) + has_scores = tl.sum(score_valid.to(tl.int32), axis=1) > 0 + # A sliding-window query can have entire early KV tiles masked. Avoid + # the -inf - -inf NaN in online softmax until its first live tile. + alpha = tl.where(has_scores, tl.exp2(m_i - m_ij), 1.0) + p = tl.where( + score_valid, + tl.exp2(scores - m_ij[:, None]), + 0.0, + ) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc *= alpha[:, None] + + v = tl.load( + v_ptr + + page[:, None] * stride_v_b + + page_off[:, None] * stride_v_p + + pid_kh * stride_v_h + + off_d[None, :] * stride_v_d, + mask=k_valid[:, None] & d_valid[None, :], + other=0.0, + ) + acc += tl.dot(p.to(v.dtype), v) + m_i = m_ij + + acc /= l_i[:, None] + acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) + out_block_ptr = tl.make_block_ptr( + base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_o_t, stride_o_h, stride_o_d), + offsets=(q_block, 0, 0), + block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), + order=(2, 1, 0), + ) + tl.store( + out_block_ptr, + acc.to(out_ptr.dtype.element_ty), + boundary_check=(0, 1, 2), + ) + + +@torch.no_grad() +def inkling_fa4_rel_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + max_kv_len: int | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Paged varlen attention with Inkling's relative score modification.""" + del num_splits + if not causal or window_size[1] not in (0, -1): + raise NotImplementedError("Inkling ROCm attention requires causal masking") + if q.ndim != 3 or key_cache.ndim != 4 or value_cache.ndim != 4: + raise ValueError("expected q [T,H,D] and paged K/V [B,P,Hkv,D]") + if rel_logits.shape != (q.shape[0], q.shape[1], rel_extent): + raise ValueError( + f"relative logits have shape {tuple(rel_logits.shape)}, expected " + f"{(q.shape[0], q.shape[1], rel_extent)}" + ) + + num_kv_heads = key_cache.shape[2] + if q.shape[1] % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + if out is None: + out = torch.empty_like(q) + gqa_group_size = q.shape[1] // num_kv_heads + batch = cu_seqlens_q.shape[0] - 1 + if max_kv_len is None: + max_kv_len = key_cache.shape[0] * key_cache.shape[1] + if use_gfx950_gluon_decode( + max_query_len=max_seqlen_q, + page_size=key_cache.shape[1], + head_dim=q.shape[2], + ): + # Import only after the architecture guard. The implementation uses + # gfx950-only CDNA4 Gluon layouts and async-copy operations. + from vllm.models.inkling.amd.ops.gluon.rel_mha_decode_gfx950 import ( + gluon_rel_mha_decode_gfx950, + ) + + return gluon_rel_mha_decode_gfx950( + q, + key_cache, + value_cache, + block_table, + cache_seqlens, + max_kv_len, + rel_logits, + cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + window_left=window_size[0], + softmax_scale=softmax_scale, + out=out, + ) + if use_gfx950_gluon_extend( + max_query_len=max_seqlen_q, + max_kv_len=max_kv_len, + page_size=key_cache.shape[1], + head_dim=q.shape[2], + window_left=window_size[0], + ): + from vllm.models.inkling.amd.ops.gluon.rel_mha_extend_gfx950 import ( + gluon_rel_mha_extend_gfx950, + ) + + # The copied kernel retains TokenSpeed's cu_seqlens_kv parameter for + # API compatibility but does not consume it. + return cast( + torch.Tensor, + gluon_rel_mha_extend_gfx950( + q, + cu_seqlens_q, + cu_seqlens_q, + key_cache, + value_cache, + block_table, + cache_seqlens, + window_left=window_size[0], + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_kv_len, + rel_logits=rel_logits, + softmax_scale=softmax_scale, + out=out, + ), + ) + if use_split_kv_decode( + max_query_len=max_seqlen_q, + max_kv_len=max_kv_len, + page_size=key_cache.shape[1], + window_left=window_size[0], + ): + return inkling_rel_attention_split_kv_decode( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + softmax_scale=softmax_scale, + window_left=window_size[0], + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + out=out, + ) + block_q = 1 if max_seqlen_q == 1 else 4 + grid = (triton.cdiv(max_seqlen_q, block_q), num_kv_heads, batch) + _inkling_rel_attention_kernel[grid]( + q, + key_cache, + value_cache, + rel_logits, + out, + block_table, + cache_seqlens, + cu_seqlens_q, + gqa_group_size, + q.shape[2], + softmax_scale, + q.stride(0), + q.stride(1), + q.stride(2), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + key_cache.stride(3), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + value_cache.stride(3), + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + block_table.stride(0), + page_size=key_cache.shape[1], + rel_extent=rel_extent, + window_left=window_size[0], + BLOCK_Q=block_q, + BLOCK_K=64, + num_warps=4, + num_stages=1, + ) + return out diff --git a/vllm/models/inkling/amd/ops/fa4_warmup.py b/vllm/models/inkling/amd/ops/fa4_warmup.py new file mode 100644 index 000000000000..998acdc09737 --- /dev/null +++ b/vllm/models/inkling/amd/ops/fa4_warmup.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility shim for the ROCm relative-attention implementation. + +ROCm uses a Triton kernel which is compiled by vLLM's normal model warmup. The +NVIDIA path registers ahead-of-time CuTeDSL units; importing that provider on +ROCm would pull in CUDA-only tml-fa4 code. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class InklingFA4WarmupConfig: + num_heads: int + num_kv_heads: int + head_dim: int + rel_extent: int + window_size: tuple[int, int] + is_local: bool + max_kv_len: int + dtype: torch.dtype + kv_dtype: torch.dtype + block_size: int + max_num_reqs: int + max_num_batched_tokens: int + + +def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None: + """Triton compilation is triggered by the ordinary vLLM warmup forward.""" + del config diff --git a/vllm/models/inkling/amd/ops/gluon/__init__.py b/vllm/models/inkling/amd/ops/gluon/__init__.py new file mode 100644 index 000000000000..54494063da99 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""GFX950-only Gluon relative-attention kernels for Inkling.""" diff --git a/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py new file mode 100644 index 000000000000..987508f91f2a --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py @@ -0,0 +1,1037 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""rel_mha decode Gluon kernel for AMD GFX950.""" + +import math +from typing import NamedTuple + +import torch + +from vllm.models.inkling.amd.ops.gluon.utils import ( + _INV_LN2_VALUE, + InputStrides, + PagedKVStrides, + max, + maximum, +) +from vllm.triton_utils import aggregate, gl, gluon + +cdna4 = gl.amd.cdna4 +async_copy = cdna4.async_copy +cdiv = gl.cdiv +_GFX950_SM_COUNT = 256 + + +# ===-----------------------------------------------------------------------===# +# Kernel Config +# ===-----------------------------------------------------------------------===# + + +@aggregate +class AttentionConfig: + SM_SCALE: gl.constexpr + PAGE_TABLE_STRIDE: gl.constexpr + PAGE_SIZE: gl.constexpr + NUM_KV_SPLITS: gl.constexpr + MAX_SEQLEN_Q: gl.constexpr + NUM_Q_HEADS: gl.constexpr + NUM_KV_HEADS: gl.constexpr + HEAD_DIM: gl.constexpr + BLOCK_M: gl.constexpr + BLOCK_N: gl.constexpr + IS_SLIDING: gl.constexpr + WINDOW_LEFT: gl.constexpr + REL_EXTENT: gl.constexpr + REL_BIAS_QK_SCALE: gl.constexpr + IS_FP8: gl.constexpr + GROUP_SIZE: gl.constexpr + NUM_GROUPS: gl.constexpr + q_strides: InputStrides + rel_strides: InputStrides + k_strides: PagedKVStrides + v_strides: PagedKVStrides + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + load_layout: gl.constexpr + store_layout: gl.constexpr + reduce_layout: gl.constexpr + k_smem_layout: gl.constexpr + v_smem_layout: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + q_strides, + rel_strides, + k_strides, + v_strides, + ): + assert NUM_Q_HEADS % NUM_KV_HEADS == 0 + assert HEAD_DIM in (64, 128) + assert BLOCK_N == PAGE_SIZE + if IS_SLIDING: + assert WINDOW_LEFT >= 0 + else: + assert WINDOW_LEFT == -1 + + mfma_layout = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[1, 1], + ) + qk_layout = mfma_layout + pv_layout = mfma_layout + # qk_kw is derived from a 128-bit load / dtype bitwidth. + # pv_kw is empirically tuned. + qk_kw = 16 if IS_FP8 else 8 + pv_kw = 8 if IS_FP8 else 4 + q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw) + k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw) + p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw) + v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw) + # Elements loaded per lane depend on the input dtype, matching qk_kw. + # load_threads is how many lanes span HEAD_DIM. + load_vec = 16 if IS_FP8 else 8 + load_threads = HEAD_DIM // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [1, 1], [1, 0] + ) + # Output is always 16-bit, so a 128-bit store has 8 elements. + # store_threads is how many lanes span HEAD_DIM. + store_vec = 8 + store_threads = HEAD_DIM // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [1, 1], [1, 0] + ) + reduce_layout = gl.BlockedLayout([1, HEAD_DIM // 64], [1, 64], [1, 1], [1, 0]) + # Padding interval is 64 lanes * load_vec elems. + pad_interval = 64 * load_vec + # Empirically tuned. + pad_k = 16 if IS_FP8 else 8 + pad_v = 32 + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + + self.SM_SCALE = gl.constexpr(SM_SCALE) + self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) + self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) + self.NUM_KV_SPLITS = gl.constexpr(NUM_KV_SPLITS) + self.MAX_SEQLEN_Q = gl.constexpr(MAX_SEQLEN_Q) + self.NUM_Q_HEADS = gl.constexpr(NUM_Q_HEADS) + self.NUM_KV_HEADS = gl.constexpr(NUM_KV_HEADS) + self.HEAD_DIM = gl.constexpr(HEAD_DIM) + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_N = gl.constexpr(BLOCK_N) + self.IS_SLIDING = gl.constexpr(IS_SLIDING) + self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) + self.REL_EXTENT = gl.constexpr(REL_EXTENT) + self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) + self.IS_FP8 = gl.constexpr(IS_FP8) + self.GROUP_SIZE = gl.constexpr(NUM_Q_HEADS // NUM_KV_HEADS) + self.NUM_GROUPS = gl.constexpr((self.GROUP_SIZE + BLOCK_M - 1) // BLOCK_M) + self.q_strides = q_strides + self.rel_strides = rel_strides + self.k_strides = k_strides + self.v_strides = v_strides + self.qk_layout = gl.constexpr(qk_layout) + self.pv_layout = gl.constexpr(pv_layout) + self.q_layout = gl.constexpr(q_layout) + self.k_layout = gl.constexpr(k_layout) + self.p_layout = gl.constexpr(p_layout) + self.v_layout = gl.constexpr(v_layout) + self.load_layout = gl.constexpr(load_layout) + self.store_layout = gl.constexpr(store_layout) + self.reduce_layout = gl.constexpr(reduce_layout) + self.k_smem_layout = gl.constexpr(k_smem_layout) + self.v_smem_layout = gl.constexpr(v_smem_layout) + + +# ===-----------------------------------------------------------------------===# +# Kernel Program +# ===-----------------------------------------------------------------------===# + + +@aggregate +class AttentionProgram: + cfg: gl.constexpr + q_ptr: gl.tensor + rel_logits_ptr: gl.tensor + k_cache_ptr: gl.tensor + v_cache_ptr: gl.tensor + page_table_ptr: gl.tensor + cache_seqlens_ptr: gl.tensor + mid_o_ptr: gl.tensor + mid_lse_ptr: gl.tensor + q_index: gl.tensor + batch: gl.tensor + kv_head: gl.tensor + group_start: gl.tensor + split_id: gl.tensor + cache_len: gl.tensor + kv_start: gl.tensor + split_start: gl.tensor + split_end: gl.tensor + + @gluon.constexpr_function + def __init__( + self, + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + q_index, + batch, + kv_head, + group_start, + split_id, + cache_len, + kv_start, + split_start, + split_end, + ): + self.cfg = gl.constexpr(cfg) + self.q_ptr = q_ptr + self.rel_logits_ptr = rel_logits_ptr + self.k_cache_ptr = k_cache_ptr + self.v_cache_ptr = v_cache_ptr + self.page_table_ptr = page_table_ptr + self.cache_seqlens_ptr = cache_seqlens_ptr + self.mid_o_ptr = mid_o_ptr + self.mid_lse_ptr = mid_lse_ptr + self.q_index = q_index + self.batch = batch + self.kv_head = kv_head + self.group_start = group_start + self.split_id = split_id + self.cache_len = cache_len + self.kv_start = kv_start + self.split_start = split_start + self.split_end = split_end + + @gluon.jit + def create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + ): + # Gluon treats ``cfg`` as the constexpr factory argument, while mypy + # models the first parameter of this aggregate method as ``self``. + q_index = gl.program_id(0) + batch = q_index // cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] + q_pos = q_index - batch * cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] + head_block = gl.program_id(1) + kv_head = head_block // cfg.NUM_GROUPS # type: ignore[attr-defined] + group_block = head_block - kv_head * cfg.NUM_GROUPS # type: ignore[attr-defined] + group_start = group_block * cfg.BLOCK_M # type: ignore[attr-defined] + split_id = gl.program_id(2) + cache_len = gl.load(cache_seqlens_ptr + batch) + cache_len = cache_len - ( + cfg.MAX_SEQLEN_Q - 1 - q_pos # type: ignore[attr-defined] + ) + cache_len = maximum(cache_len, 0) + if cfg.IS_SLIDING: # type: ignore[attr-defined] + # WINDOW_LEFT is defined as exclusive (keys strictly to the left, not + # counting the current token), so the window is WINDOW_LEFT + 1 keys + # once the current token is included. E.g. with + # WINDOW_LEFT = 127 and cache_len = 500 (current token at index 499), + # kv_start = 500 - (127 + 1) = 372 keeps keys [372, 499] = 128 keys. + # Without the + 1, kv_start = 373 would drop the leftmost key. + kv_start = cache_len - min( + cache_len, + cfg.WINDOW_LEFT + 1, # type: ignore[attr-defined] + ) + else: + kv_start = cache_len - cache_len + first_page = kv_start // cfg.PAGE_SIZE # type: ignore[attr-defined] + end_page = cdiv(cache_len, cfg.PAGE_SIZE) # type: ignore[attr-defined] + num_pages = end_page - first_page + pages_per_split = cdiv( + num_pages, + cfg.NUM_KV_SPLITS, # type: ignore[attr-defined] + ) + split_start_page = first_page + split_id * pages_per_split + split_end_page = min(split_start_page + pages_per_split, end_page) + split_start = ( + split_start_page * cfg.PAGE_SIZE # type: ignore[attr-defined] + ) + split_end = min( + split_end_page * cfg.PAGE_SIZE, # type: ignore[attr-defined] + cache_len, + ) + return AttentionProgram( + gl.constexpr(cfg), + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + q_index, + batch, + kv_head, + group_start, + split_id, + cache_len, + kv_start, + split_start, + split_end, + ) + + @gluon.jit + def load_q(self): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + offsets = cfg.q_strides.offsets(self.q_index, q_heads[:, None], offs_d[None, :]) + return cdna4.buffer_load(self.q_ptr, offsets, mask=valid[:, None], other=0.0) + + @gluon.jit + def init_state(self): + cfg = self.cfg + m_i = gl.full( + [cfg.BLOCK_M], + value=-float("inf"), + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + l_i = gl.full( + [cfg.BLOCK_M], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + acc = gl.zeros( + [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout + ) + return m_i, l_i, acc + + @gluon.jit + def load_page(self, start_n): + cfg = self.cfg + page_index = start_n // cfg.PAGE_SIZE + valid = start_n < self.split_end + return gl.load( + self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, + mask=valid, + other=0, + ) + + @gluon.jit + def issue_load_k(self, physical_page, k_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.k_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + # can't use buffer_load: paged KV offsets may exceed its 32-bit range. + async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def issue_load_v(self, physical_page, v_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.v_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + # can't use buffer_load: paged KV offsets may exceed its 32-bit range. + async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def shared_load_k(self, k_smem): + return k_smem.permute([1, 0]).load(self.cfg.k_layout) + + @gluon.jit + def shared_load_v(self, v_smem): + return v_smem.load(self.cfg.v_layout) + + @gluon.jit + def compute_qk(self, q, k): + cfg = self.cfg + qk = gl.zeros( + [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout + ) + return cdna4.mfma(q, k, qk) + + @gluon.jit + def apply_rel_bias(self, qk, start_n): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout)) + offs_n = start_n + gl.arange( + 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) + ) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + rel_dist = (self.cache_len - 1) - offs_n + rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) + rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) + rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) + offsets = cfg.rel_strides.offsets( + self.q_index, q_heads[:, None], rel_idx[None, :] + ) + head_valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + rel_bias = cdna4.buffer_load( + self.rel_logits_ptr, + offsets, + mask=head_valid[:, None] & rel_valid[None, :], + other=0.0, + ).to(gl.float32) + return qk + rel_bias * cfg.REL_BIAS_QK_SCALE + + @gluon.jit + def apply_kv_mask(self, qk, start_n): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) + tokens = start_n + offs_n[None, :] + mask = (tokens >= self.kv_start) & (tokens < self.split_end) + return gl.where(mask, qk, -float("inf")) + + @gluon.jit + def softmax(self, qk, m_i, l_i, acc): + cfg = self.cfg + row_max = max(qk, axis=1) + row_max = gl.convert_layout(row_max, gl.SliceLayout(1, cfg.pv_layout)) + m_new = maximum(m_i, row_max) + m_new_scaled = m_new * cfg.SM_SCALE + qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] + p = gl.exp2(qk_shifted) + m_diff = m_i * cfg.SM_SCALE - m_new_scaled + alpha = gl.exp2(m_diff) + l_ij = gl.sum(p, axis=1) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + p = p.to(self.q_ptr.dtype.element_ty) + p = gl.convert_layout(p, cfg.p_layout) + return p, m_new, l_i, acc + + @gluon.jit + def compute_pv(self, p, v, acc): + return cdna4.mfma(p, v, acc) + + @gluon.jit + def store_split(self, acc, l_i, m_i): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = ((self.group_start + offs_m) < cfg.GROUP_SIZE) & ( + self.split_start < self.split_end + ) + acc = gl.convert_layout(acc, cfg.store_layout) + l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout)) + m_i = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.store_layout)) + recip_l_i = 1.0 / l_i + part_o = acc * recip_l_i[:, None] + part_lse = m_i * cfg.SM_SCALE + gl.log2(l_i) + mid_o_offsets = ( + (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.NUM_KV_SPLITS + + self.split_id + ) * cfg.HEAD_DIM + offs_d[None, :] + mid_lse_offsets = ( + self.q_index * cfg.NUM_Q_HEADS + q_heads + ) * cfg.NUM_KV_SPLITS + self.split_id + cdna4.buffer_store(part_o, self.mid_o_ptr, mid_o_offsets, mask=valid[:, None]) + cdna4.buffer_store(part_lse, self.mid_lse_ptr, mid_lse_offsets, mask=valid) + + @gluon.jit + def store_output(self, acc, l_i): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + acc = gl.convert_layout(acc, cfg.store_layout) + l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout)) + output = acc * (1.0 / l_i)[:, None] + output = output.to(self.mid_o_ptr.dtype.element_ty) + offsets = (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.HEAD_DIM + offsets += offs_d[None, :] + cdna4.buffer_store(output, self.mid_o_ptr, offsets, mask=valid[:, None]) + + +# ===-----------------------------------------------------------------------===# +# Entry Point +# ===-----------------------------------------------------------------------===# + + +@gluon.jit +def _rel_mha_decode_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + Q_STRIDE_B: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + NUM_KV_SPLITS: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = AttentionProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + ) + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_state() + + physical_page = program.load_page(program.split_start) + + for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N): + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + physical_page = program.load_page(start_n + cfg.BLOCK_N) + + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + qk = program.apply_kv_mask(qk, start_n) + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + program.store_split(acc, l_i, m_i) + + +@gluon.jit +def _rel_mha_decode_sliding_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + out_ptr, + Q_STRIDE_B: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + 1, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = AttentionProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + out_ptr, + out_ptr, + ) + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_state() + + for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N): + physical_page = program.load_page(start_n) + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + qk = program.apply_kv_mask(qk, start_n) + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + program.store_output(acc, l_i) + + +@gluon.jit +def _rel_mha_decode_reduce_fp16( + mid_o_ptr, + mid_lse_ptr, + out_ptr, + cache_seqlens_ptr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + NUM_KV_SPLITS: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + PAGE_SIZE: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + 1, # REL_EXTENT + 1.0, # REL_BIAS_QK_SCALE + IS_FP8, + InputStrides(1, 1, 1), + InputStrides(1, 1, 1), + PagedKVStrides(1, 1, 1, 1), + PagedKVStrides(1, 1, 1, 1), + ) + q_index = gl.program_id(0) + batch = q_index // MAX_SEQLEN_Q + q_pos = q_index - batch * MAX_SEQLEN_Q + q_head = gl.program_id(1) + cache_len = gl.load(cache_seqlens_ptr + batch) + cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) + cache_len = maximum(cache_len, 0) + if cfg.IS_SLIDING: + kv_start = cache_len - min(cache_len, cfg.WINDOW_LEFT + 1) + else: + kv_start = cache_len - cache_len + first_page = kv_start // cfg.PAGE_SIZE + end_page = cdiv(cache_len, cfg.PAGE_SIZE) + num_pages = end_page - first_page + pages_per_split = cdiv(num_pages, cfg.NUM_KV_SPLITS) + + # SPLIT_TILE pads NUM_KV_SPLITS up to a power of 2. + SPLIT_TILE: gl.constexpr = 1 << (NUM_KV_SPLITS - 1).bit_length() + offs_s = gl.arange(0, SPLIT_TILE, layout=gl.SliceLayout(1, cfg.reduce_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.reduce_layout)) + # split_valid masks out empty splits and the power-of-2 padding tail. + split_start_page = first_page + offs_s * pages_per_split + split_end_page_raw = split_start_page + pages_per_split + split_end_page = gl.where( + split_end_page_raw < end_page, split_end_page_raw, end_page + ) + split_start_tok = split_start_page * cfg.PAGE_SIZE + split_end_raw = split_end_page * cfg.PAGE_SIZE + split_end_tok = gl.where(split_end_raw < cache_len, split_end_raw, cache_len) + split_valid = (split_start_tok < split_end_tok) & (offs_s < cfg.NUM_KV_SPLITS) + # Load every split's partial output and lse. + base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.NUM_KV_SPLITS + offs_s + part_lse = gl.load(mid_lse_ptr + base, mask=split_valid, other=-float("inf")) + o_off = base[:, None] * cfg.HEAD_DIM + offs_d[None, :] + part_o = cdna4.buffer_load(mid_o_ptr, o_off, mask=split_valid[:, None], other=0.0) + + # Global softmax max over all splits. + m_i = max(part_lse, axis=0) + # Weighted sum of the split partials, normalized by the total softmax mass. + beta = gl.exp2(part_lse - m_i) + l_i = gl.sum(beta, axis=0) + acc = gl.sum(part_o * beta[:, None], axis=0) + + out_base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.HEAD_DIM + output = acc * (1.0 / l_i) + output = output.to(out_ptr.dtype.element_ty) + cdna4.buffer_store(output, out_ptr, out_base + offs_d) + + +def _select_num_kv_splits( + *, + batch: int, + num_kv_heads: int, + num_groups: int, + num_pages: int, + sm_count: int, +) -> int: + """Pick num_kv_splits to balance occupancy against reduce overhead. + + The launch grid is (batch * num_kv_heads * num_groups) * num_kv_splits + work-groups. Too few splits under-fill the machine at low batch; too many + leave each split with a handful of pages, so the reduce kernel dominates. + + Return the smaller of two candidate counts: splits_for_occupancy (enough to + fill ~wave_target waves of CUs) and splits_for_pages (~min_pages_per_split + pages per split), with the pages candidate clamped to [min_page_splits, + max_page_splits] so a short context still splits without launching empty work + and a long one does not over-split where reduce cost outgrows the decode win. + """ + wave_target = 2 + min_pages_per_split = 2 + min_page_splits = 8 + max_page_splits = 32 + + base_ctas = batch * num_kv_heads * num_groups + target_ctas = sm_count * wave_target + splits_for_occupancy = (target_ctas + base_ctas - 1) // base_ctas + + splits_for_pages = num_pages // min_pages_per_split + min_page_splits = min(min_page_splits, num_pages) + if splits_for_pages < min_page_splits: + splits_for_pages = min_page_splits + if splits_for_pages > max_page_splits: + splits_for_pages = max_page_splits + return min(splits_for_occupancy, splits_for_pages) + + +class LaunchConfig(NamedTuple): + num_q_heads: int + num_kv_heads: int + num_groups: int + head_dim: int + page_size: int + num_kv_splits: int + block_m: int + block_n: int + sm_scale: float + rel_bias_qk_scale: float + is_sliding: bool + window_left: int + + +def get_config( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + max_seqlen_k: int, + window_left: int, + softmax_scale: float | None, +) -> LaunchConfig: + head_dim = q.shape[2] + page_size = k_cache.shape[1] + block_m = 16 + block_n = page_size + group_size = q.shape[1] // k_cache.shape[2] + num_groups = math.ceil(group_size / block_m) + is_sliding = window_left >= 0 + window_left = window_left if is_sliding else -1 + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + sm_scale = softmax_scale + effective_seqlen_k = ( + min(max_seqlen_k, window_left + 1) if is_sliding else max_seqlen_k + ) + num_pages = (effective_seqlen_k + page_size - 1) // page_size + if is_sliding: + num_kv_splits = 4 + else: + num_kv_splits = _select_num_kv_splits( + batch=q.shape[0], + num_kv_heads=k_cache.shape[2], + num_groups=num_groups, + num_pages=num_pages, + sm_count=_GFX950_SM_COUNT, + ) + return LaunchConfig( + num_q_heads=q.shape[1], + num_kv_heads=k_cache.shape[2], + num_groups=num_groups, + head_dim=head_dim, + page_size=page_size, + num_kv_splits=num_kv_splits, + block_m=block_m, + block_n=block_n, + sm_scale=sm_scale * _INV_LN2_VALUE, + rel_bias_qk_scale=1.0 / softmax_scale, + is_sliding=is_sliding, + window_left=window_left, + ) + + +def gluon_rel_mha_decode_gfx950( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + rel_logits: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int = 1, + window_left: int = -1, + softmax_scale: float | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + total_q = q.shape[0] + + config = get_config( + q=q, + k_cache=k_cache, + max_seqlen_k=max_seqlen_k, + window_left=window_left, + softmax_scale=softmax_scale, + ) + + is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + out_dtype = torch.bfloat16 if is_fp8 else q.dtype + if k_cache.shape[1] not in (64, 128, 256): + raise ValueError( + "gfx950 Gluon relative-attention decode requires page size " + f"64, 128, or 256; got {k_cache.shape[1]}" + ) + output = ( + torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out + ) + + # Always use split-k for full attention and for the small-batch sliding + # decode path. Sliding uses a fixed 8 splits, one page per split for the + # TP-4 local-attention shape with 512-token window and 64-token pages. + mid_o = torch.empty( + (total_q, config.num_q_heads, config.num_kv_splits, config.head_dim), + device=q.device, + dtype=torch.float32, + ) + mid_lse = torch.empty( + (total_q, config.num_q_heads, config.num_kv_splits), + device=q.device, + dtype=torch.float32, + ) + + grid = ( + total_q, + config.num_kv_heads * config.num_groups, + config.num_kv_splits, + ) + _rel_mha_decode_fp16[grid]( + q, + rel_logits, + k_cache, + v_cache, + page_table, + cache_seqlens, + mid_o, + mid_lse, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), + config.sm_scale, + page_table.stride(0), + config.page_size, + config.num_kv_splits, + max_seqlen_q, + config.num_q_heads, + config.num_kv_heads, + config.head_dim, + config.block_m, + config.block_n, + config.is_sliding, + config.window_left, + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + rel_logits.shape[2], + config.rel_bias_qk_scale, + is_fp8, + num_warps=1, + ) + + reduce_grid = (total_q, config.num_q_heads) + _rel_mha_decode_reduce_fp16[reduce_grid]( + mid_o, + mid_lse, + output, + cache_seqlens, + config.sm_scale, + page_table.stride(0), + config.num_kv_splits, + max_seqlen_q, + config.page_size, + config.num_q_heads, + config.num_kv_heads, + config.head_dim, + config.block_m, + config.block_n, + config.is_sliding, + config.window_left, + is_fp8, + num_warps=1, + ) + return output diff --git a/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py b/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py new file mode 100644 index 000000000000..224524e76b95 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py @@ -0,0 +1,760 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""rel_mha extend Gluon kernel for AMD GFX950. + +This handles ragged, multi-token queries against a paged KV cache. The query +axis is tiled into the MFMA ``M`` dimension (prefill-style): ``BLOCK_M`` query +rows of a single q-head share each paged KV tile, so every KV tile is loaded +once and reused across all rows in the tile. The grid is +``(blocks_per_req, batch, n_heads)`` -- all host-known sizes -- and each program +self-locates its request from ``cu_seqlens_q`` / ``cache_seqlens`` in-kernel, so +the launch stays CUDA-graph static with no device->host sync. + +Visibility per query row depends on ``is_causal`` and the optional sliding +window: + +* ``is_causal=False``: every query token attends the full visible cache, i.e. + ``visible_kv = cache_seqlens[batch]``. +* ``is_causal=True``: query tokens are a causal suffix, so the ``i``-th query + token of a request (0-indexed) attends ``prefix + i + 1`` tokens, where + ``prefix = cache_seqlens[batch] - query_len[batch]``. + +Causal masking only touches the few KV tiles that reach the diagonal; the long +prefix is a mask-free fast path. Sinks and sliding windows (causal or not) are +applied per tile. This is the sole Gluon extend implementation. +""" + +import math + +import torch + +from vllm.models.inkling.amd.ops.gluon.utils import ( + _INV_LN2_VALUE, + _LN2, + InputStrides, + PagedKVStrides, + max, + maximum, +) +from vllm.triton_utils import aggregate, gl, gluon + +cdna4 = gl.amd.cdna4 +async_copy = cdna4.async_copy + + +# ===-----------------------------------------------------------------------===# +# Query-batched path +# +# This path tiles BLOCK_M query rows (of a single q-head) into the MFMA M +# dimension -- like the prefill kernel -- so each paged KV tile is loaded once +# and reused across all BLOCK_M rows. KV still comes from the paged cache +# (decode-style async page loads). Causal masking only touches the few KV tiles +# that reach the diagonal; the long prefix is a mask-free fast path. +# ===-----------------------------------------------------------------------===# + +_EXTEND_SHORT_Q_BLOCK_M = 64 +_EXTEND_SHORT_Q_NUM_WARPS = 2 +_EXTEND_LONG_Q_BLOCK_M = 128 +_EXTEND_LONG_Q_NUM_WARPS = 4 + + +def _select_extend_tile(max_seqlen_q: int, page_size: int) -> tuple[int, int, int]: + """Return (BLOCK_M, BLOCK_N, NUM_WARPS) for the given max query length. + + Queries that fit in a single short tile use it (least padding, most + occupancy); longer ones use the tall tile that covers more rows per + shared-KV pass. + """ + if max_seqlen_q <= _EXTEND_SHORT_Q_BLOCK_M: + return _EXTEND_SHORT_Q_BLOCK_M, page_size, _EXTEND_SHORT_Q_NUM_WARPS + return _EXTEND_LONG_Q_BLOCK_M, page_size, _EXTEND_LONG_Q_NUM_WARPS + + +@aggregate +class ExtendConfig: + N_HEADS: gl.constexpr + N_KV_HEADS: gl.constexpr + GROUP_SIZE: gl.constexpr + HEAD_DIM: gl.constexpr + SM_SCALE: gl.constexpr + BLOCK_M: gl.constexpr + BLOCK_N: gl.constexpr + NUM_WARPS: gl.constexpr + PAGE_SIZE: gl.constexpr + PAGE_TABLE_STRIDE: gl.constexpr + IS_CAUSAL: gl.constexpr + HAS_LSE: gl.constexpr + WINDOW_LEFT: gl.constexpr + REL_EXTENT: gl.constexpr + REL_BIAS_QK_SCALE: gl.constexpr + IS_FP8: gl.constexpr + q_strides: InputStrides + rel_strides: InputStrides + k_strides: PagedKVStrides + v_strides: PagedKVStrides + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + load_layout: gl.constexpr + store_layout: gl.constexpr + k_smem_layout: gl.constexpr + v_smem_layout: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + N_HEADS, + N_KV_HEADS, + HEAD_DIM, + SM_SCALE, + BLOCK_M, + BLOCK_N, + NUM_WARPS, + PAGE_SIZE, + PAGE_TABLE_STRIDE, + IS_CAUSAL, + HAS_LSE, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + q_strides, + rel_strides, + k_strides, + v_strides, + ): + assert HEAD_DIM in (64, 128) + assert BLOCK_N == PAGE_SIZE + + instr_shape = [32, 32, 16] + mfma_layout = gl.amd.AMDMFMALayout( + version=4, + instr_shape=instr_shape, + transposed=True, + warps_per_cta=[NUM_WARPS, 1], + ) + qk_layout = mfma_layout + pv_layout = mfma_layout + # Elements loaded per lane depend on the input dtype, matching qk_kw. + # load_threads is how many lanes span HEAD_DIM. + load_vec = 16 if IS_FP8 else 8 + load_threads = HEAD_DIM // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [NUM_WARPS, 1], [1, 0] + ) + # Output is always 16-bit, so a 128-bit store has 8 elements. + # store_threads is how many lanes span HEAD_DIM. + store_vec = 8 + store_threads = HEAD_DIM // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [NUM_WARPS, 1], [1, 0] + ) + # Padding interval is 64 lanes * load_vec elems. + pad_interval = 64 * load_vec + # Empirically tuned. + pad_k = 16 if IS_FP8 else 8 + pad_v = 16 if IS_FP8 else 32 + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + + self.N_HEADS = gl.constexpr(N_HEADS) + self.N_KV_HEADS = gl.constexpr(N_KV_HEADS) + self.GROUP_SIZE = gl.constexpr(N_HEADS // N_KV_HEADS) + self.HEAD_DIM = gl.constexpr(HEAD_DIM) + self.SM_SCALE = gl.constexpr(SM_SCALE) + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_N = gl.constexpr(BLOCK_N) + self.NUM_WARPS = gl.constexpr(NUM_WARPS) + self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) + self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) + self.IS_CAUSAL = gl.constexpr(IS_CAUSAL) + self.HAS_LSE = gl.constexpr(HAS_LSE) + self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) + self.REL_EXTENT = gl.constexpr(REL_EXTENT) + self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) + self.IS_FP8 = gl.constexpr(IS_FP8) + self.q_strides = q_strides + self.rel_strides = rel_strides + self.k_strides = k_strides + self.v_strides = v_strides + self.qk_layout = gl.constexpr(qk_layout) + self.pv_layout = gl.constexpr(pv_layout) + # qk_kw is derived from a 128-bit load / dtype bitwidth. + # pv_kw is empirically tuned. + qk_kw = 16 if IS_FP8 else 8 + pv_kw = 16 if IS_FP8 else 4 + self.q_layout = gl.constexpr(gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)) + self.k_layout = gl.constexpr(gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)) + self.p_layout = gl.constexpr(gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)) + self.v_layout = gl.constexpr(gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)) + self.load_layout = gl.constexpr(load_layout) + self.store_layout = gl.constexpr(store_layout) + self.k_smem_layout = gl.constexpr(k_smem_layout) + self.v_smem_layout = gl.constexpr(v_smem_layout) + + +@aggregate +class ExtendProgram: + cfg: gl.constexpr + q_ptr: gl.tensor + rel_logits_ptr: gl.tensor + k_cache_ptr: gl.tensor + v_cache_ptr: gl.tensor + page_table_ptr: gl.tensor + output_ptr: gl.tensor + lse_ptr: gl.tensor + batch: gl.tensor + q_head: gl.tensor + kv_head: gl.tensor + q_start: gl.tensor + seq_base: gl.tensor + seq_len: gl.tensor + prefix: gl.tensor + cache_len: gl.tensor + + @gluon.constexpr_function + def __init__( + self, + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + batch, + q_head, + kv_head, + q_start, + seq_base, + seq_len, + prefix, + cache_len, + ): + self.cfg = gl.constexpr(cfg) + self.q_ptr = q_ptr + self.rel_logits_ptr = rel_logits_ptr + self.k_cache_ptr = k_cache_ptr + self.v_cache_ptr = v_cache_ptr + self.page_table_ptr = page_table_ptr + self.output_ptr = output_ptr + self.lse_ptr = lse_ptr + self.batch = batch + self.q_head = q_head + self.kv_head = kv_head + self.q_start = q_start + self.seq_base = seq_base + self.seq_len = seq_len + self.prefix = prefix + self.cache_len = cache_len + + @gluon.jit + def create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + ): + # Gluon treats ``cfg`` as the constexpr factory argument, while mypy + # models the first parameter of this aggregate method as ``self``. + block_in_req = gl.program_id(0) + batch = gl.program_id(1) + q_head = gl.program_id(2) + q_start = block_in_req * cfg.BLOCK_M # type: ignore[attr-defined] + seq_base = gl.load(cu_seqlens_q_ptr + batch) + seq_end = gl.load(cu_seqlens_q_ptr + batch + 1) + seq_len = seq_end - seq_base + cache_len = gl.load(cache_seqlens_ptr + batch) + prefix = cache_len - seq_len + kv_head = q_head // cfg.GROUP_SIZE # type: ignore[attr-defined] + return ExtendProgram( + gl.constexpr(cfg), + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + batch, + q_head, + kv_head, + q_start, + seq_base, + seq_len, + prefix, + cache_len, + ) + + @gluon.jit + def load_q(self): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout) + ) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) + row = self.seq_base + offs_m + offsets = cfg.q_strides.offsets(row[:, None], self.q_head, offs_d[None, :]) + mask = offs_m[:, None] < self.seq_len + return cdna4.buffer_load(self.q_ptr, offsets, mask=mask, other=0.0) + + @gluon.jit + def load_page(self, start_n): + cfg = self.cfg + page_index = start_n // cfg.PAGE_SIZE + valid = start_n < self.cache_len + return gl.load( + self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, + mask=valid, + other=0, + ) + + @gluon.jit + def issue_load_k(self, physical_page, k_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.k_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def issue_load_v(self, physical_page, v_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.v_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def shared_load_k(self, k_smem): + cfg = self.cfg + k_buffer = k_smem.permute([1, 0]) + return k_buffer.load(cfg.k_layout) + + @gluon.jit + def shared_load_v(self, v_smem): + cfg = self.cfg + return v_smem.load(cfg.v_layout) + + @gluon.jit + def compute_qk(self, q, k): + cfg = self.cfg + qk = gl.zeros( + [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout + ) + return cdna4.mfma(q, k, qk) + + @gluon.jit + def apply_rel_bias(self, qk, start_n): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) + ) + offs_n_abs = start_n + gl.arange( + 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) + ) + q_pos = self.prefix + offs_m + rel_dist = q_pos[:, None] - offs_n_abs[None, :] + rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) + rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) + rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) + offsets = cfg.rel_strides.offsets( + self.seq_base + offs_m[:, None], self.q_head, rel_idx + ) + mask = (offs_m[:, None] < self.seq_len) & rel_valid + rel_bias = cdna4.buffer_load( + self.rel_logits_ptr, offsets, mask=mask, other=0.0 + ).to(gl.float32) + return qk + rel_bias * cfg.REL_BIAS_QK_SCALE + + @gluon.jit + def compute_pv(self, p, v, acc): + return cdna4.mfma(p, v, acc) + + @gluon.jit + def init_attention_state(self): + cfg = self.cfg + m_i = gl.full( + [cfg.BLOCK_M], + value=-float("inf"), + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + l_i = gl.full( + [cfg.BLOCK_M], + value=0, + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + acc = gl.zeros( + [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout + ) + return m_i, l_i, acc + + @gluon.jit + def softmax(self, qk, m_i, l_i, acc): + cfg = self.cfg + # In sliding window case, some rows can see fully masked tiles before + # any valid KV. Guard the online softmax state so `-inf - -inf` does not + # produce NaNs. + HAS_INVALID: gl.constexpr = cfg.WINDOW_LEFT >= 0 + + row_max = max(qk, 1) + m_new = maximum(m_i, row_max) + m_new_scaled = m_new * cfg.SM_SCALE + if HAS_INVALID: + invalid = m_new == -float("inf") + m_new_scaled = gl.where(invalid, 0.0, m_new_scaled) + + qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] + p = gl.exp2(qk_shifted) + m_diff = m_i * cfg.SM_SCALE - m_new_scaled + if HAS_INVALID: + m_diff = gl.where(invalid, 0.0, m_diff) + + alpha = gl.exp2(m_diff) + l_ij = gl.sum(p, axis=1) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + p = p.to(self.q_ptr.dtype.element_ty) + p = gl.convert_layout(p, cfg.p_layout) + return p, m_new, l_i, acc + + @gluon.jit + def store_output(self, output): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout) + ) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + offsets = ( + ((self.seq_base + offs_m[:, None]) * cfg.N_HEADS + self.q_head) + * cfg.HEAD_DIM + + offs_d[None, :] + ).to(gl.int32) + mask = offs_m[:, None] < self.seq_len + output = output.to(self.output_ptr.dtype.element_ty) + cdna4.buffer_store(output, self.output_ptr, offsets, mask=mask) + + @gluon.jit + def store_lse(self, l_i, m_i): + cfg = self.cfg + if cfg.HAS_LSE: + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout) + ) + offsets = ((self.seq_base + offs_m) * cfg.N_HEADS + self.q_head).to( + gl.int32 + ) + mask = offs_m < self.seq_len + lse_l_i = gl.where(l_i > 0.0, l_i, 1.0) + # Softmax runs in base-2 (exp2 hardware fast path), so m_i*SM_SCALE + + # log2(l_i) is the LSE in base-2 units. Convert to natural log (the + # public op contract / torch.logsumexp convention) by scaling by ln2. + lse = (m_i * cfg.SM_SCALE + gl.log2(lse_l_i)) * _LN2 + cdna4.buffer_store(lse, self.lse_ptr, offsets, mask=mask) + + +@gluon.jit +def _rel_mha_extend_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + Q_STRIDE_T: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + N_HEADS: gl.constexpr, + N_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + NUM_WARPS: gl.constexpr, + IS_CAUSAL: gl.constexpr, + HAS_LSE: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = ExtendConfig( + N_HEADS, + N_KV_HEADS, + HEAD_DIM, + SM_SCALE, + BLOCK_M, + BLOCK_N, + NUM_WARPS, + PAGE_SIZE, + PAGE_TABLE_STRIDE, + IS_CAUSAL, + HAS_LSE, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_T, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = ExtendProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + ) + # Over-provisioned tile past this request's real query rows: nothing to do. + if program.q_start >= program.seq_len: + return + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_attention_state() + + offs_m_q = program.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) + ) + offs_n_q = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) + diag_row = program.prefix + offs_m_q + + if IS_CAUSAL: + kv_end = min(program.cache_len, program.prefix + program.q_start + cfg.BLOCK_M) + else: + kv_end = program.cache_len + + # Sliding window (inclusive-left, matches flash-attn window_size=(W, 0)): + # skip KV tiles entirely below the window's lower edge. The tile's top query + # row sits at absolute position pos_top; its window opens at + # pos_top - WINDOW_LEFT (that key is visible -> W + 1 keys total). The min() + # form clamps to 0 without the shadowed builtin max(). + if cfg.WINDOW_LEFT >= 0: + pos_top = program.prefix + program.q_start + kv_start = pos_top - min(pos_top, cfg.WINDOW_LEFT) + kv_start = (kv_start // cfg.BLOCK_N) * cfg.BLOCK_N + else: + kv_start = 0 + + for start_n in range(kv_start, kv_end, cfg.BLOCK_N): + physical_page = program.load_page(start_n) + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + + if cfg.WINDOW_LEFT >= 0: + # Window lower edge + cache bound always apply; causal upper edge only + # when IS_CAUSAL (independent layering keeps non-causal + window correct). + offs_n_abs = start_n + offs_n_q + mask = (offs_n_abs[None, :] >= diag_row[:, None] - cfg.WINDOW_LEFT) & ( + offs_n_abs[None, :] < program.cache_len + ) + if IS_CAUSAL: + mask &= offs_n_abs[None, :] <= diag_row[:, None] + qk = gl.where(mask, qk, -float("inf")) + elif IS_CAUSAL: + if start_n + cfg.BLOCK_N > program.prefix + program.q_start: + offs_n_abs = start_n + offs_n_q + mask = (offs_n_abs[None, :] <= diag_row[:, None]) & ( + offs_n_abs[None, :] < program.cache_len + ) + qk = gl.where(mask, qk, -float("inf")) + else: + if start_n + cfg.BLOCK_N > program.cache_len: + offs_n_abs = start_n + offs_n_q + qk = gl.where( + offs_n_abs[None, :] < program.cache_len, qk, -float("inf") + ) + + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + denom = gl.where(l_i > 0.0, l_i, 1.0) + output = acc * (1.0 / denom)[:, None] + output = gl.convert_layout(output, cfg.store_layout) + program.store_output(output) + program.store_lse(l_i, m_i) + + +def gluon_rel_mha_extend_gfx950( + q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + window_left: int = -1, + return_lse: bool = False, + max_seqlen_q: int = 1, + max_seqlen_k: int = 1, + rel_logits: torch.Tensor | None = None, + softmax_scale: float | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert rel_logits is not None + head_dim = q.shape[2] + n_heads = q.shape[1] + n_kv_heads = k_cache.shape[2] + page_size = k_cache.shape[1] + block_m, block_n, num_warps = _select_extend_tile(max_seqlen_q, page_size) + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + sm_scale = softmax_scale * _INV_LN2_VALUE + rel_bias_qk_scale = 1.0 / softmax_scale + + # max_seqlen_q must be >= the true max query length; extra tiles early-exit. + batch = cu_seqlens_q.shape[0] - 1 + safe_max_q = max_seqlen_q if max_seqlen_q > 0 else 1 + blocks_per_req = (safe_max_q + block_m - 1) // block_m + cu_q_i32 = cu_seqlens_q.to(torch.int32).contiguous() + cache_i32 = cache_seqlens.to(torch.int32).contiguous() + + is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + out_dtype = torch.bfloat16 if is_fp8 else q.dtype + if page_size not in (64, 128, 256): + raise ValueError( + "gfx950 Gluon relative-attention extend requires page size " + f"64, 128, or 256; got {page_size}" + ) + output = ( + torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out + ) + if return_lse: + lse = torch.empty((q.shape[0], n_heads), device=q.device, dtype=torch.float32) + lse_arg = lse + else: + lse = None + lse_arg = q + grid = (blocks_per_req, batch, n_heads) + _rel_mha_extend_fp16[grid]( + q, + rel_logits, + k_cache, + v_cache, + page_table, + output, + lse_arg, + cu_q_i32, + cache_i32, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), + sm_scale, + page_table.stride(0), + page_size, + n_heads, + n_kv_heads, + head_dim, + block_m, + block_n, + num_warps, + True, + return_lse, + window_left, + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + rel_logits.shape[2], + rel_bias_qk_scale, + is_fp8, + num_warps=num_warps, + ) + if return_lse: + assert lse is not None + return output, lse + return output diff --git a/vllm/models/inkling/amd/ops/gluon/utils.py b/vllm/models/inkling/amd/ops/gluon/utils.py new file mode 100644 index 000000000000..90f2e586bc6e --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/utils.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from vllm.triton_utils import aggregate, gl, gluon, tl + +_INV_LN2_VALUE = 1.4426950408889634 +_INV_LN2 = tl.constexpr(_INV_LN2_VALUE) +_LN2_VALUE = 0.6931471805599453 +_LN2 = tl.constexpr(_LN2_VALUE) +_PROPAGATE_NAN_ALL = gl.constexpr(tl.PropagateNan.ALL) + + +@gluon.jit +def maximum(a, b, propagate_nan: gl.constexpr = _PROPAGATE_NAN_ALL): + return gl.maximum(a, b, propagate_nan=propagate_nan) + + +@gluon.jit +def max(input, axis=None, keep_dims=False): + return gl.reduce(input, axis, maximum, keep_dims=keep_dims) + + +@gluon.constexpr_function +def attention_layouts(head_dim, block_n, is_fp8, dtype, num_warps, instr_shape): + mfma = gl.amd.AMDMFMALayout( + version=4, + instr_shape=instr_shape, + transposed=True, + warps_per_cta=[num_warps, 1], + ) + qk_layout = mfma + pv_layout = mfma + # qk_kw is derived from a 128-bit load / dtype bitwidth; pv_kw is tuned. + qk_kw = 16 if is_fp8 else 8 + pv_kw = 8 if is_fp8 else 4 + q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw) + k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw) + p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw) + v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw) + # load_vec = elems/lane (dtype-dependent, == qk_kw); load_threads span HEAD_DIM. + load_vec = 16 if is_fp8 else 8 + load_threads = head_dim // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [num_warps, 1], [1, 0] + ) + # store_vec is always 16-bit (128 / 16 = 8) regardless of input dtype. + store_vec = 8 + store_threads = head_dim // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [num_warps, 1], [1, 0] + ) + # Take only the built-in's padding, not its swizzle: the swizzle scatters + # head_dim across banks, making the LDS write stride non-constant, so it can't + # lower through async_copy's affine [1, 0] load. Padding-only is affine and + # DMA-legal. + # TODO(perf): to also use the swizzle, co-design a matched load layout so the + # DMA stays legal. + k_api = gl.amd.cdna4.compute_efficient_padded_shared_layout( + k_layout, [block_n, head_dim], dtype, is_k_contig=True + ) + v_api = gl.amd.cdna4.compute_efficient_padded_shared_layout( + v_layout, [block_n, head_dim], dtype, is_k_contig=False + ) + assert k_api is not None and v_api is not None, ( + "no CDNA4 padded shared layout for this operand/dtype" + ) + k_pairs = list(k_api.interval_padding_pairs) + v_pairs = list(v_api.interval_padding_pairs) + assert len(k_pairs) == 1 and len(v_pairs) == 1, ( + "expected a single interval padding pair from the built-in" + ) + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[int(k_pairs[0][0]), int(k_pairs[0][1])]], [block_n, head_dim], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[int(v_pairs[0][0]), int(v_pairs[0][1])]], [block_n, head_dim], [1, 0] + ) + return ( + qk_layout, + pv_layout, + q_layout, + k_layout, + p_layout, + v_layout, + load_layout, + store_layout, + k_smem_layout, + v_smem_layout, + ) + + +@aggregate +class InputStrides: + stride_t: gl.constexpr + stride_h: gl.constexpr + stride_d: gl.constexpr + + @gluon.constexpr_function + def __init__(self, stride_t, stride_h, stride_d): + self.stride_t = gl.constexpr(stride_t) + self.stride_h = gl.constexpr(stride_h) + self.stride_d = gl.constexpr(stride_d) + + @gluon.jit + def offsets(self, token, head, dim): + return (token * self.stride_t + head * self.stride_h + dim * self.stride_d).to( + gl.int32 + ) + + +@aggregate +class PagedKVStrides: + stride_b: gl.constexpr + stride_p: gl.constexpr + stride_h: gl.constexpr + stride_d: gl.constexpr + + @gluon.constexpr_function + def __init__(self, stride_b, stride_p, stride_h, stride_d): + self.stride_b = gl.constexpr(stride_b) + self.stride_p = gl.constexpr(stride_p) + self.stride_h = gl.constexpr(stride_h) + self.stride_d = gl.constexpr(stride_d) + + @gluon.jit + def offsets(self, page, token, head, dim): + # KV pools can exceed the 32-bit buffer-offset range. Keep the page + # arithmetic in int64, matching the original kernel's large-cache path. + return ( + page.to(gl.int64) * self.stride_b + + token.to(gl.int64) * self.stride_p + + head * self.stride_h + + dim * self.stride_d + ) diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py new file mode 100644 index 000000000000..9705a9a3df37 --- /dev/null +++ b/vllm/models/inkling/amd/ops/lamport.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. + +The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. + +Liveness +-------- +Large grids are deliberately split at the two communication dependencies: + + 1. ``_publish_input_kernel`` only publishes rank partials. + 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. + 3. ``_sconv_publish_kernel`` only computes and publishes the local result. + 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. + +Without PDL, CUDA stream order completes a producer before its consumer. With +PDL, every producer CTA posts all peer stores before triggering its dependent, +and the consumer executes ``gdc_wait`` before polling. A consumer therefore +waits only for stores from a producer that is already running (or complete) on +another GPU. It never waits for another block in its own grid. Consequently +spinning consumers cannot occupy resources needed by any producer, and the +wait-for graph has no cycle. The proof is independent of grid size, block +dispatch order, and occupancy. + +For one token, the first three phases use eight independent channel slices to +expose enough CTA parallelism for decode latency. Each slice has exclusive +ownership of its cache and Lamport columns; the gather/RMSNorm phase retains +one CTA per token so no cross-CTA reduction or completion counter is needed. + +Immediate buffer reuse (including replay of a captured CUDA graph) is also +safe. Rank R cannot republish input for call n+1 until its gather for n has +finished; that gather waited for owner O's output, which O publishes only +after consuming R's input for n. Likewise, R cannot republish output for n+1 +until its reduction for n+1 has observed destination D's input; D publishes +that input only after its gather consumed R's output for n. Thus every prior +read happens-before a same-slot rewrite. Three generations reduce incidental +coupling for ordinary launches, but correctness does not depend on rotation. + +The payload itself is the Lamport flag. A 32-bit store publishes two bf16s +atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real +negative zeroes are changed to positive zero before publication. Consumers +use volatile 32-bit loads and restore the sentinel after consuming a slot. +""" + +from __future__ import annotations + +import os + +import torch + +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton + +logger = init_logger(__name__) + +_MAX_TOKENS = 16384 +_EMPTY_PAIR = tl.constexpr(0x80008000) + + +@triton.jit +def _pack_bf16_pairs(values): + """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" + lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) + lo = lo.to(tl.uint16, bitcast=True) + hi = hi.to(tl.uint16, bitcast=True) + lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) + hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) + return lo | (hi << 16) + + +@triton.jit +def _unpack_bf16_pairs(values): + lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + return tl.interleave(lo, hi) + + +@triton.jit +def _wait_pairs(ptr, offsets, mask): + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + return values + + +@triton.jit +def _publish_input_kernel( + stage_ptr, + peer_ptrs, + peer_offset_u32, + stride_stage_t, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Publish this rank's full partial row into every shard owner's slots.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + elem = tl.arange(0, CS_P2) + elem_mask = elem < CSS + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + # Address generation is independent of the preceding stream kernel. The + # acquire remains before stage/peer loads, which may consume its writes. + if USE_PDL: + tl.extra.cuda.gdc_wait() + + for owner in tl.static_range(WORLD): + values = tl.load( + stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, + mask=elem_mask, + other=0.0, + ) + packed = _pack_bf16_pairs(values) + base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) + dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair + tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _reduce_insert_kernel( + input_peer_ptrs, + input_peer_offset_u32, + cache_ptr, + slot_ptr, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume all partials for this rank and insert the reduced cache row.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + source = tl.arange(0, WORLD) + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + offsets = ( + (token * WORLD + source)[:, None] * (CS // 2) + + split * (CSS // 2) + + pair[None, :] + ) + mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] + input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + input_u32 += input_peer_offset_u32 + # Slot metadata and the cache destination do not depend on input publish. + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + safe_slot = tl.maximum(slot, 0).to(tl.int64) + dst = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(input_u32, offsets, mask) + + lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + reduced = tl.interleave( + tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), + tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), + ) + tl.store( + input_u32 + offsets, + tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=mask, + ) + + tl.store(dst, reduced, mask=valid & channel_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _sconv_publish_kernel( + peer_ptrs, + peer_offset_u32, + residual_ptr, + weight_ptr, + cache_ptr, + position_ptr, + sequence_ptr, + slot_ptr, + block_table_ptr, + stride_residual_t, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + stride_block_table_r, + max_blocks, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + WINDOW: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Compute this rank's shard, then publish it to every rank.""" + tl.static_assert(WINDOW == 4) + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + position = tl.load(position_ptr + token) + sequence = tl.load(sequence_ptr + token) + safe_slot = tl.maximum(slot, 0).to(tl.int64) + own_ptr = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + # These loads were made visible before reduce/insert could signal us and + # are independent of its cache store. Hoisting them gives PDL useful work + # to overlap while retaining the acquire before every cache read. + residual = tl.load( + residual_ptr + token * stride_residual_t + RANK * CS + global_channel, + mask=channel_mask, + other=0.0, + ) + weight0 = tl.load( + weight_ptr + global_channel * WINDOW, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight1 = tl.load( + weight_ptr + global_channel * WINDOW + 1, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight2 = tl.load( + weight_ptr + global_channel * WINDOW + 2, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight3 = tl.load( + weight_ptr + global_channel * WINDOW + 3, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + if USE_PDL: + tl.extra.cuda.gdc_wait() + current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) + + conv = tl.zeros([CS_P2], tl.float32) + for tap_idx in tl.static_range(WINDOW): + source_position = position - (WINDOW - 1) + tap_idx + take = valid & (source_position >= 0) + if tap_idx == WINDOW - 1: + value = tl.where(take, current.to(tl.float32), 0.0) + else: + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) + physical_block = tl.load( + block_table_ptr + sequence * stride_block_table_r + logical_block, + mask=take, + other=0, + ).to(tl.int64) + source_ptr = ( + cache_ptr + + physical_block * stride_cache_block + + head * stride_cache_head + + (safe_position % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) + value = tl.where(take, cached.to(tl.float32), 0.0) + if tap_idx == 0: + weight = weight0 + elif tap_idx == 1: + weight = weight1 + elif tap_idx == 2: + weight = weight2 + else: + weight = weight3 + conv += value * weight + + # Preserve both bf16 rounding points of the original sublayer. + short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) + output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( + tl.bfloat16 + ) + + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + packed = _pack_bf16_pairs(output) + row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair + + for destination in tl.static_range(WORLD): + base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) + tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _gather_norm_kernel( + output_peer_ptrs, + output_peer_offset_u32, + norm_weight_ptr, + normed_ptr, + residual_out_ptr, + eps, + stride_output_t, + C: tl.constexpr, + C_P2: tl.constexpr, + RANK: tl.constexpr, + HAS_NORM: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume a complete row; one CTA owns all outputs for one token.""" + token = tl.program_id(0).to(tl.int64) + pair = tl.arange(0, C_P2 // 2) + pair_mask = pair < C // 2 + output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + output_u32 += output_peer_offset_u32 + offsets = token * (C // 2) + pair + channel = tl.arange(0, C_P2) + channel_mask = channel < C + # Gamma is independent of the preceding sconv publication. Keep the + # acquire immediately before polling the Lamport output slots. + if HAS_NORM: + weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(output_u32, offsets, pair_mask) + row = _unpack_bf16_pairs(packed) + tl.store( + residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask + ) + if HAS_NORM: + row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) + inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) + tl.store( + normed_ptr + token * stride_output_t + channel, + (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), + mask=channel_mask, + ) + tl.store( + output_u32 + offsets, + tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=pair_mask, + ) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _validate_lamport_init_kernel( + input_peer_ptrs, + output_peer_ptrs, + bad_ptr, + num_pairs, + RANK: tl.constexpr, + BLOCK: tl.constexpr, +): + """Validate both complete local allocations through their fabric pointers.""" + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_pairs + ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) + local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) + value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) + value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) + bad = tl.max( + tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) + ) + if bad != 0: + tl.atomic_max(bad_ptr, 1) + + +class LamportRSConv: + """Persistent symmetric buffers for one TP group.""" + + def _initialize_lamport_buffers(self) -> None: + """Arm every slot and collectively verify the exact sentinel bits. + + Do not replace this with a zero-fill: Lamport readiness distinguishes + the bf16 bit pattern for -0.0 (0x8000) from every published payload. + The validation is deliberately collective so one rank cannot enter the + first polling kernel while another rank still has an unarmed buffer. + """ + if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): + raise RuntimeError("Lamport symmetric buffers must be contiguous") + if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: + raise RuntimeError("Lamport buffers must contain whole uint32 pairs") + + # Fill through int16 so each bf16 lane receives the exact -0.0 bits. + # This covers all three generations and all max-token slots, including + # slots that a smaller first invocation does not touch. + self.buf_in.view(torch.int16).fill_(-0x8000) + self.buf_out.view(torch.int16).fill_(-0x8000) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + # Scan the complete local allocations. Since every rank performs the + # scan and participates in MAX, success proves every symmetric backing + # allocation was armed before any rank is allowed to use generation 0. + bad = torch.logical_or( + self.buf_in.view(torch.int16).ne(-0x8000).any(), + self.buf_out.view(torch.int16).ne(-0x8000).any(), + ).to(dtype=torch.float32) + bad = self.tp.all_reduce(bad) + if int(bad.item()) != 0: + raise RuntimeError("Lamport sentinel initialization failed on a TP rank") + self.tp.barrier() + + def _initialize_mnnvl_buffers(self) -> None: + """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" + self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) + self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + bad = torch.zeros((), dtype=torch.int32, device=self.device) + num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 + _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( + self.input_peer_ptrs, + self.output_peer_ptrs, + bad, + num_pairs, + RANK=self.rank, + BLOCK=256, + num_warps=4, + ) + bad = self.tp.all_reduce(bad.to(torch.float32)) + if int(bad.item()) != 0: + raise RuntimeError("MNNVL Lamport sentinel initialization failed") + self.tp.barrier() + + def __init__( + self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS + ) -> None: + import torch.distributed._symmetric_memory as symm_mem + + tp = get_tp_group() + self.tp = tp + self.group = tp.device_group + self.world_size = tp.world_size + self.rank = tp.rank_in_group + self.device = torch.device(tp.device) + if self.world_size not in (2, 4, 8): + raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") + if hidden_size % (2 * self.world_size) != 0: + raise ValueError("hidden size must produce an even shard on every rank") + if window_size != 4: + raise ValueError(f"short-conv window size must be 4, got {window_size}") + if max_tokens < 1 or max_tokens > _MAX_TOKENS: + raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") + + is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) + self.hidden_size = hidden_size + self.window_size = window_size + self.max_tokens = max_tokens + self.shard_size = hidden_size // self.world_size + # Three generations follow FlashInfer's Lamport layout. A generation + # is reused only after two intervening collective calls. + self.num_buffers = 3 + self.input_generation_bytes = max_tokens * hidden_size * 2 + self.output_generation_bytes = max_tokens * hidden_size * 2 + self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 + if is_cross_node: + try: + from flashinfer.comm.mnnvl import ( + McastGPUBuffer, + TorchDistBackend, + is_mnnvl_fabric_supported, + ) + except ImportError as error: + raise RuntimeError( + "cross-node TP requires FlashInfer MNNVL support" + ) from error + + local_supported = int( + is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) + ) + unsupported = torch.tensor( + 1 - local_supported, dtype=torch.float32, device=self.device + ) + unsupported = tp.all_reduce(unsupported) + if int(unsupported.item()) != 0: + raise RuntimeError("cross-node TP is supported only on MNNVL fabric") + + comm_backend = TorchDistBackend(self.group) + allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 + self._mnnvl_input_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self._mnnvl_output_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() + self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() + self._initialize_mnnvl_buffers() + logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") + else: + self.buf_in = symm_mem.empty( + self.num_buffers, + max_tokens, + self.world_size, + self.shard_size, + dtype=torch.bfloat16, + device=self.device, + ) + self.buf_out = symm_mem.empty( + self.num_buffers, + max_tokens, + hidden_size, + dtype=torch.bfloat16, + device=self.device, + ) + group_name = self.group.group_name + input_handle = symm_mem.rendezvous(self.buf_in, group_name) + output_handle = symm_mem.rendezvous(self.buf_out, group_name) + self.input_peer_ptrs = input_handle.buffer_ptrs_dev + self.output_peer_ptrs = output_handle.buffer_ptrs_dev + self._input_handle = input_handle + self._output_handle = output_handle + self._initialize_lamport_buffers() + self.generation = 0 + + def usable(self, num_tokens: int) -> bool: + return 0 < num_tokens <= self.max_tokens + + def rs_sconv_ag_add_norm( + self, + input_tensor: torch.Tensor, + residual: torch.Tensor, + conv_weight: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float, + cache: torch.Tensor, + positions: torch.Tensor, + block_table: torch.Tensor, + seq_idx: torch.Tensor, + slot_mapping: torch.Tensor, + off_s: int, + ws: int, + block_size: int, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" + tokens, hidden_size = residual.shape + if not self.usable(tokens): + raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") + if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: + raise ValueError("residual must be bf16 [T, 6144]") + if ( + input_tensor.shape != residual.shape + or input_tensor.dtype != torch.bfloat16 + or input_tensor.stride(1) != 1 + ): + raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") + shard_size = hidden_size // self.world_size + if conv_weight.shape != (shard_size, self.window_size): + raise ValueError( + f"conv_weight must have shape [{shard_size}, {self.window_size}]" + ) + if ( + conv_weight.dtype != torch.bfloat16 + or conv_weight.stride(0) != self.window_size + ): + raise ValueError("conv_weight must be contiguous bf16") + if norm_weight is not None and ( + norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 + ): + raise ValueError("norm_weight must be bf16 [6144] or None") + if cache.dtype != torch.bfloat16 or cache.ndim != 4: + raise ValueError("cache must be a 4-D bf16 tensor") + if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: + raise ValueError("cache head layout is inconsistent with ws") + if off_s < 0 or off_s + ws > cache.shape[3]: + raise ValueError("cache channel offset is out of bounds") + + index = self.generation + input_offset = index * self.input_generation_bytes // 4 + output_offset = index * self.output_generation_bytes // 4 + normed = torch.empty_like(residual) if norm_weight is not None else None + residual_out = torch.empty_like(residual) + phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 + phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) + phase_grid = (tokens, phase_splits) + # Wide CTAs win before the grid saturates; smaller CTAs reduce register + # pressure once high-throughput batches provide enough parallelism. + if 128 <= tokens <= 2048: + phase_warps = 16 + elif tokens > 2048: + phase_warps = 8 + else: + phase_warps = 4 + gather_warps = 4 if tokens >= 256 else 8 + _publish_input_kernel[phase_grid]( + input_tensor, + self.input_peer_ptrs, + input_offset, + input_tensor.stride(0), + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _reduce_insert_kernel[phase_grid]( + self.input_peer_ptrs, + input_offset, + cache, + slot_mapping, + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _sconv_publish_kernel[phase_grid]( + self.output_peer_ptrs, + output_offset, + residual, + conv_weight, + cache, + positions, + seq_idx, + slot_mapping, + block_table, + residual.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_table.stride(0), + block_table.shape[1], + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + WINDOW=self.window_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _gather_norm_kernel[(tokens,)]( + self.output_peer_ptrs, + output_offset, + norm_weight if norm_weight is not None else residual, + normed if normed is not None else residual_out, + residual_out, + eps, + residual.stride(0), + C=hidden_size, + C_P2=triton.next_power_of_2(hidden_size), + RANK=self.rank, + HAS_NORM=norm_weight is not None, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=gather_warps, + ) + self.generation = (index + 1) % self.num_buffers + return normed, residual_out + + +_STATE: LamportRSConv | None = None +_STATE_FAILED = False + + +def initialize_lamport_rs_conv( + hidden_size: int, window_size: int, max_num_batched_tokens: int +) -> None: + """Collectively initialize the TP-group state during model construction.""" + global _STATE, _STATE_FAILED + if _STATE is not None: + if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: + raise RuntimeError("all Lamport users must share hidden and window sizes") + return + if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": + return + try: + max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) + _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) + except Exception: + _STATE_FAILED = True + logger.exception("fused collective unavailable; use the NCCL fallback") + + +def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: + """Return the state initialized with the model, or ``None`` for fallback.""" + if _STATE is not None and ( + _STATE.hidden_size != hidden_size or _STATE.window_size != window_size + ): + raise RuntimeError("all Lamport users must share hidden and window sizes") + return _STATE diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py new file mode 100644 index 000000000000..021cf5274bd2 --- /dev/null +++ b/vllm/models/inkling/amd/ops/mm_towers.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused CUDA kernels for the Inkling vision/audio towers. + +Both kernels keep the reference paths' fp32 accumulation and per-op bf16 +rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently +bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order +(real-checkpoint-weight cosine vs reference > 0.9999998). +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, tldevice, triton + +from .norm import _get_num_warps_from_block_size + + +@triton.jit +def _dmel_embed_sum_norm_kernel( + idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) + w_ptr, # [NB * VOCAB, D] bf16 embedding table + norm_w_ptr, # [D] (unused if HAS_NORM=False) + out_ptr, # [T, D] bf16 + eps, + D, + stride_idx_t, + NB: tl.constexpr, + VOCAB: tl.constexpr, + D_P2: tl.constexpr, + HAS_NORM: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, D_P2) + mask = offs < D + # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), + # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). + acc = tl.zeros([D_P2], dtype=tl.float32) + for b in tl.static_range(NB): + v = tl.load(idx_ptr + t * stride_idx_t + b) + row = (b * VOCAB + v).to(tl.int64) + acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + h = acc.to(tl.bfloat16) + if HAS_NORM: + # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded + # bf16 multiply with the bf16 weight. + x32 = h.to(tl.float32) + var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D + xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) + w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) + h = xn * w + tl.store(out_ptr + t * D + offs, h, mask=mask) + + +def dmel_embed_sum_norm( + dmel_idx: torch.Tensor, # [T, NB] int32 + weight: torch.Tensor, # [NB * VOCAB, D] bf16 + norm_weight: torch.Tensor | None, + eps: float, +) -> torch.Tensor: + """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no + [T, NB, D] intermediate).""" + T, nb = dmel_idx.shape + D = weight.shape[1] + assert weight.shape[0] % nb == 0 + vocab = weight.shape[0] // nb + out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) + if T == 0: + return out + d_p2 = triton.next_power_of_2(D) + _dmel_embed_sum_norm_kernel[(T,)]( + dmel_idx, + weight, + norm_weight if norm_weight is not None else weight, + out, + eps, + D, + dmel_idx.stride(0), + NB=nb, + VOCAB=vocab, + D_P2=d_p2, + HAS_NORM=norm_weight is not None, + # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% + # at large T (the 80-row gather chain is latency- not lane-bound). + num_warps=8, + ) + return out + + +@triton.jit +def _rmsnorm_gelu_kernel( + x_ptr, # [R, D] bf16 + w_ptr, # [D] + out_ptr, # [R, D] bf16 (or the folded layout when FOLD) + eps, + R, + D, + D_P2: tl.constexpr, + BLOCK_M: tl.constexpr, # rows per block (>1 for small D) + HAS_GELU: tl.constexpr, + FOLD: tl.constexpr, + # fold geometry: input rows index [N, T, H, W]; the store scatters each + # row to (out_row, slot) of fold_timespace_to_depth's output layout. + FT: tl.constexpr, + FH: tl.constexpr, + FW: tl.constexpr, + TF: tl.constexpr, + HF: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + rmask = rows < R + offs = tl.arange(0, D_P2) + mask = rmask[:, None] & (offs < D)[None, :] + x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + var = tl.sum(x32 * x32, axis=1) / D + xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) + w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) + h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm + if HAS_GELU: + # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, + # matching F.gelu's opmath on a bf16 tensor. + g32 = h.to(tl.float32) + h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) + if FOLD: + # Store directly into the next layer's folded layout (a pure + # permutation — replaces the separate fold copy pass). + t = (rows // (FH * FW)) % FT + hh = (rows // FW) % FH + ww = rows % FW + n = rows // (FT * FH * FW) + slot = ((t % TF) * HF + hh % HF) * HF + ww % HF + out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( + FW // HF + ) + ww // HF + base = (out_row * (TF * HF * HF) + slot) * D + tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) + else: + tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) + + +def rmsnorm_gelu( + x: torch.Tensor, # [..., D] bf16 contiguous + weight: torch.Tensor, + eps: float, + gelu: bool = True, + fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold +) -> torch.Tensor: + """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block + when D is small. With ``fold``, x must be [N, T, H, W, D] and the output + comes back as ``fold_timespace_to_depth(result, *fold)``.""" + D = x.shape[-1] + flat = x.reshape(-1, D) + assert flat.stride(1) == 1 and flat.stride(0) == D + R = flat.shape[0] + if fold is None: + out = torch.empty_like(flat) + ft = fh = fw = tf = hf = 1 + out_shape = x.shape + else: + tf, hf = fold + N, ft, fh, fw, _ = x.shape + out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + if R == 0: + return out.reshape(out_shape) + d_p2 = triton.next_power_of_2(D) + block_m = max(1, 4096 // d_p2) + _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( + flat, + weight, + out, + eps, + R, + D, + D_P2=d_p2, + BLOCK_M=block_m, + HAS_GELU=gelu, + FOLD=fold is not None, + FT=ft, + FH=fh, + FW=fw, + TF=tf, + HF=hf, + num_warps=_get_num_warps_from_block_size(d_p2 * block_m), + ) + return out.reshape(out_shape) diff --git a/vllm/models/inkling/amd/ops/norm.py b/vllm/models/inkling/amd/ops/norm.py new file mode 100644 index 000000000000..d34c1c3259b9 --- /dev/null +++ b/vllm/models/inkling/amd/ops/norm.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from functools import lru_cache + +import torch + +from vllm.triton_utils import tl, triton + +_MAX_FUSED_SIZE = 65536 + + +def _get_num_warps_from_block_size(block_size: int) -> int: + if block_size >= 32768: + return 32 + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +def _largest_power_of_2(n: int) -> int: + assert n > 0, f"{n=}" + return 1 << (n.bit_length() - 1) + + +@lru_cache(maxsize=128) +def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int: + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + return _largest_power_of_2(num_sms) * factor + + +@triton.jit +def _rmsnorm_fwd_kernel( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + pid_m, rstd) + + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +@triton.jit(do_not_specialize=["n_rows"]) +def _rmsnorm_fwd_kernel_block_m( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_rows, + n_cols, + block_size_m: tl.constexpr, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + num_blocks_m = tl.cdiv(n_rows, block_size_m) + blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0)) + block_id_start = pid_m * blocks_per_pid + block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m) + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + + for block_id in range(block_id_start, block_id_end): + offs_m = block_id * block_size_m + tl.arange(0, block_size_m) + mask_m = offs_m < n_rows + mask_mn = mask_m[:, None] & mask_n[None, :] + x = tl.load( + x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :], + mask=mask_mn, + other=0.0, + ).to(tl.float32) + row_var = tl.sum(x * x, axis=1) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + offs_m, rstd, mask=mask_m) + + y = x * rstd[:, None] * weight + tl.store( + y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :], + y, + mask=mask_mn, + ) + + +@triton.jit +def _add_rmsnorm_fwd_kernel( + res_ptr, # [T, N] residual (read) + delta_ptr, # [T, N] delta to add (read) + weight_ptr, + y_ptr, # [T, N] normed output + res_out_ptr, # [T, N] updated residual output + eps, + res_stride_0, + delta_stride_0, + y_stride_0, + res_out_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + # Round the sum to the residual dtype first (matches the eager + # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit). + s = (r + d).to(res_out_ptr.dtype.element_ty) + tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n) + x = s.to(tl.float32) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +def add_rmsnorm( + residual: torch.Tensor, + delta: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused ``res = residual + delta; y = rmsnorm(res)``. + + Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no + in-place update of the inputs). + """ + assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape) + n_rows, n_cols = residual.shape + assert weight.shape[0] == n_cols + y = torch.empty_like(residual) + res_out = torch.empty_like(residual) + if n_rows == 0: + return y, res_out + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // residual.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + num_warps = _get_num_warps_from_block_size(block_size_n) + _add_rmsnorm_fwd_kernel[(n_rows,)]( + residual, + delta, + weight, + y, + res_out, + eps, + residual.stride(0), + delta.stride(0), + y.stride(0), + res_out.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + return y, res_out + + +@triton.jit +def _embed_rmsnorm_kernel( + ids_ptr, # [T] token ids + table_ptr, # [V, N] embedding table + weight_ptr, # [N] (HAS_NORM only) + chain_weight_ptr, # [N] (HAS_CHAIN only) + out_ptr, # [T, N] rmsnorm(table[ids], weight) + chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only) + eps, + table_stride_0, + n_cols, + block_size_n: tl.constexpr, + HAS_NORM: tl.constexpr, + HAS_CHAIN: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + row = tl.load(ids_ptr + pid_m).to(tl.int64) + x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0) + if HAS_NORM: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + # Round to the output dtype so the chained norm is bit-exact vs the + # unfused pair (which stores bf16 in between). + x = (xf * rstd * w).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n) + if HAS_CHAIN: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + tl.store( + chain_out_ptr + pid_m * n_cols + offs_n, + (xf * rstd * w).to(chain_out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_rmsnorm( + input_ids: torch.Tensor, + embed_table: torch.Tensor, + weight: torch.Tensor | None, + eps: float, + chain_weight: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm. + + Requires the full vocab on-rank (replicated or tp_size == 1). + ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure + embedding-table gather. + ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the + first decoder layer's pre-attention norm) as a second output, still one + launch. Bit-exact vs the unfused module sequence.""" + ids = input_ids.view(-1) + (T,) = ids.shape + n = embed_table.shape[1] + out = torch.empty( + (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device + ) + chain_out = torch.empty_like(out) if chain_weight is not None else None + if T > 0: + block_size_n = triton.next_power_of_2(n) + _embed_rmsnorm_kernel[(T,)]( + ids, + embed_table, + weight if weight is not None else embed_table, + chain_weight if chain_weight is not None else embed_table, + out, + chain_out if chain_out is not None else out, + eps, + embed_table.stride(0), + n, + block_size_n, + HAS_NORM=weight is not None, + HAS_CHAIN=chain_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + if chain_out is not None: + return out, chain_out + return out + + +@triton.jit +def _embed_dual_rmsnorm_cat_kernel( + hidden_ptr, # [T, N] + emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER + ids_ptr, # [T] token ids (GATHER only) + w_hidden_ptr, # [N] + w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only) + w_embed_ptr, # [N] + out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))] + eps, + hidden_stride_0, + emb_stride_0, + n_cols, + block_size_n: tl.constexpr, + GATHER: tl.constexpr, + HAS_PRE_NORM: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + if which == 0: + x = tl.load( + hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0 + ).to(tl.float32) + w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + else: + row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m + x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + if HAS_PRE_NORM: + w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + # Round-trip through the output dtype so the chained norm is + # bit-exact vs the unfused pair (which stores bf16 in between). + x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32) + w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + tl.store( + out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n, + (x * rstd * w).to(out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_dual_rmsnorm_cat( + hidden: torch.Tensor, + hidden_weight: torch.Tensor, + embed_weight: torch.Tensor, + eps: float, + *, + embeds: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + embed_table: torch.Tensor | None = None, + pre_norm_weight: torch.Tensor | None = None, +) -> torch.Tensor: + """The MTP depth-layer input in one launch: + ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``. + + The embed side is either a fused row gather ``embed_table[input_ids]`` + (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged + multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the + backbone embed_norm in front of the depth embed_norm (bit-exact vs the + unfused sequence). The concat copies collapse into direct writes.""" + T, n = hidden.shape + if embeds is not None: + assert embeds.shape == hidden.shape + src, ids, src_stride = embeds, embeds, embeds.stride(0) + gather = False + else: + assert input_ids is not None and embed_table is not None + assert input_ids.shape == (T,) and embed_table.shape[1] == n + src, ids, src_stride = embed_table, input_ids, embed_table.stride(0) + gather = True + out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device) + if T == 0: + return out + block_size_n = triton.next_power_of_2(n) + _embed_dual_rmsnorm_cat_kernel[(T, 2)]( + hidden, + src, + ids, + hidden_weight, + pre_norm_weight if pre_norm_weight is not None else embed_weight, + embed_weight, + out, + eps, + hidden.stride(0), + src_stride, + n, + block_size_n, + GATHER=gather, + HAS_PRE_NORM=pre_norm_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + return out + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + assert x.ndim == 2, f"{x.shape=}" + assert weight.ndim == 1, f"{weight.shape=}" + n_rows, n_cols = x.shape + assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}" + y = torch.empty_like(x) + rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device) + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // x.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + block_size_m = max(1, 4096 // block_size_n) + num_warps = _get_num_warps_from_block_size(block_size_n) + + if block_size_m == 1: + _rmsnorm_fwd_kernel[(n_rows,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + else: + grid_size = _get_grid_size_for_mem_bw_kernel(x.device) + _rmsnorm_fwd_kernel_block_m[(grid_size,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_rows, + n_cols, + block_size_m, + block_size_n, + num_warps=num_warps, + ) + return y diff --git a/vllm/models/inkling/amd/ops/qkvr_prep.py b/vllm/models/inkling/amd/ops/qkvr_prep.py new file mode 100644 index 000000000000..a910e72170d1 --- /dev/null +++ b/vllm/models/inkling/amd/ops/qkvr_prep.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import aux_stream + +LOW_BLOCK_M = 32 +LOW_BLOCK_N = 64 +LOW_NUM_WARPS = 4 +THROUGHPUT_BLOCK_M = 32 +THROUGHPUT_BLOCK_N = 128 +THROUGHPUT_GROUP_M = 2 +THROUGHPUT_NUM_WARPS = 4 +SMALL_TOKEN_THRESHOLD = 128 +SMALL_NUM_WARPS = 2 +Q_BLOCK_ROWS = 8 +Q_NUM_WARPS = 2 +KV_BLOCK_ROWS = 4 +KV_NUM_WARPS = 2 + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_low_latency_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_throughput_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_M: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row_group = tl.program_id(0) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + row_offsets = tl.arange(0, BLOCK_M) + for group_offset in tl.static_range(GROUP_M): + row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[ + :, None + ] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool: + min_rows = 8192 if rel_extent == 512 else 2048 + return rows >= min_rows + + +def qkvr_rel_proj( + qkvr: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + log_scaling: torch.Tensor | None, + *, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, +) -> None: + rows = qkvr.shape[0] * num_q_heads + rel_extent = rel_proj.shape[1] + assert d_rel == 16 and rel_proj.shape[0] == 16 + r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim + log_scaling_ptr = log_scaling if log_scaling is not None else qkvr + common = dict( + R_OFFSET=r_offset, + NUM_Q_HEADS=num_q_heads, + REL_EXTENT=rel_extent, + APPLY_LOG_SCALING=log_scaling is not None, + ) + + if use_rel_proj_throughput(rows, rel_extent): + grid = ( + triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M), + triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N), + ) + _rel_proj_throughput_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=THROUGHPUT_BLOCK_M, + BLOCK_N=THROUGHPUT_BLOCK_N, + GROUP_M=THROUGHPUT_GROUP_M, + num_warps=THROUGHPUT_NUM_WARPS, + **common, + ) + return + + grid = ( + triton.cdiv(rows, LOW_BLOCK_M), + triton.cdiv(rel_extent, LOW_BLOCK_N), + ) + _rel_proj_low_latency_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=LOW_BLOCK_M, + BLOCK_N=LOW_BLOCK_N, + num_warps=LOW_NUM_WARPS, + **common, + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _qkvr_qkv_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + rel_proj_ptr, + rel_out_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + log_scaling_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, + D_REL: tl.constexpr, + REL_EXTENT: tl.constexpr, + REL_EXTENT_PADDED: tl.constexpr, +): + block = tl.program_id(0) + num_q_rows = tokens * NUM_Q_HEADS + dims = tl.arange(0, HEAD_DIM) + + if block < num_q_rows: + row = block + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps) + normalized = values * rstd * weight + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + normalized *= tl.load(log_scaling_ptr + token) + tl.store( + q_out_ptr + row * HEAD_DIM + dims, + normalized.to(q_out_ptr.dtype.element_ty), + ) + rel_cols = tl.arange(0, REL_EXTENT_PADDED) + rel_mask = rel_cols < REL_EXTENT + projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32) + rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL + for rel_dim in tl.static_range(D_REL): + rel_value = tl.load( + qkvr_ptr + token * stride_x_t + rel_offset + rel_dim + ).to(tl.float32) + proj = tl.load( + rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols, + mask=rel_mask, + other=0.0, + ).to(tl.float32) + projected += rel_value * proj + projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32) + if APPLY_LOG_SCALING: + projected *= tl.load(log_scaling_ptr + token) + tl.store( + rel_out_ptr + row * REL_EXTENT + rel_cols, + projected.to(rel_out_ptr.dtype.element_ty), + mask=rel_mask, + ) + else: + row = block - num_q_rows + if row < tokens * NUM_KV_HEADS: + token = row // NUM_KV_HEADS + head = row % NUM_KV_HEADS + position = tl.load(positions_ptr + token) + request = tl.load(seq_idx_ptr + token) + conv_slot = tl.load(conv_slot_mapping_ptr + token) + query_start = tl.load(query_start_ptr + token) + attention_slot = tl.load(attention_slot_mapping_ptr + token) + valid = conv_slot >= 0 + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims) + v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims) + + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base + (OFF_K + dims) * stride_cc_dim, + k_value, + mask=valid, + ) + tl.store( + cache_base + (OFF_V + dims) * stride_cc_dim, + v_value, + mask=valid, + ) + + acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row * stride_x_t + k_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row * stride_x_t + v_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum( + safe_position // conv_block_size, max_blocks - 1 + ) + physical_block = tl.load( + conv_block_table_ptr + + request * stride_block_table_req + + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base + (OFF_K + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base + (OFF_V + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight + acc_v += source_v * v_weight + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps) + k_normalized = (k_float * rstd * k_norm_weight).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + tl.store( + key_cache_ptr + + attention_block * stride_kc_block + + attention_offset * stride_kc_token + + head * stride_kc_head + + dims, + k_normalized, + mask=attention_slot >= 0, + ) + tl.store( + value_cache_ptr + + attention_block * stride_vc_block + + attention_offset * stride_vc_token + + head * stride_vc_head + + dims, + v_rounded, + mask=attention_slot >= 0, + ) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _q_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + log_scaling_ptr, + num_rows, + stride_x_t, + eps, + NUM_Q_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = rows < num_rows + tokens = rows // NUM_Q_HEADS + heads = rows % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + + tokens[:, None] * stride_x_t + + heads[:, None] * HEAD_DIM + + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps) + normalized = values * rstd[:, None] * weight[None, :] + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0) + normalized *= tau[:, None] + tl.store( + q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :], + normalized.to(q_out_ptr.dtype.element_ty), + mask=row_mask[:, None], + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _kv_kernel( + qkvr_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = token < tokens + head_id = tl.program_id(1) + head = tl.full([BLOCK_ROWS], head_id, tl.int64) + position = tl.load(positions_ptr + token, mask=row_mask, other=0) + request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0) + conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1) + query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0) + attention_slot = tl.load( + attention_slot_mapping_ptr + token, mask=row_mask, other=-1 + ) + valid = row_mask & (conv_slot >= 0) + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + v_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + k_value, + mask=valid[:, None], + ) + tl.store( + cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + v_value, + mask=valid[:, None], + ) + + acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1) + physical_block = tl.load( + conv_block_table_ptr + request * stride_block_table_req + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight[None, :] + acc_v += source_v * v_weight[None, :] + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps) + k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + attention_mask = row_mask & (attention_slot >= 0) + tl.store( + key_cache_ptr + + attention_block[:, None] * stride_kc_block + + attention_offset[:, None] * stride_kc_token + + head[:, None] * stride_kc_head + + dims[None, :], + k_normalized, + mask=attention_mask[:, None], + ) + tl.store( + value_cache_ptr + + attention_block[:, None] * stride_vc_block + + attention_offset[:, None] * stride_vc_token + + head[:, None] * stride_vc_head + + dims[None, :], + v_rounded, + mask=attention_mask[:, None], + ) + + +def _run_tiled_q( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + positions: torch.Tensor, + *, + eps: float, + num_q_heads: int, + head_dim: int, + log_scaling: torch.Tensor | None, +) -> None: + num_rows = qkvr.shape[0] * num_q_heads + _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)]( + qkvr, + q_norm_weight, + q_out, + log_scaling if log_scaling is not None else positions, + num_rows, + qkvr.stride(0), + eps, + NUM_Q_HEADS=num_q_heads, + HEAD_DIM=head_dim, + BLOCK_ROWS=Q_BLOCK_ROWS, + APPLY_LOG_SCALING=log_scaling is not None, + num_warps=Q_NUM_WARPS, + ) + + +def _run_tiled_kv( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, +) -> None: + tokens = qkvr.shape[0] + _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)]( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + BLOCK_ROWS=KV_BLOCK_ROWS, + num_warps=KV_NUM_WARPS, + ) + + +def _run_fused_small( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None, +) -> None: + tokens = qkvr.shape[0] + + num_q_rows = tokens * num_q_heads + grid = (num_q_rows + tokens * num_kv_heads,) + _qkvr_qkv_kernel[grid]( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + log_scaling if log_scaling is not None else positions, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + APPLY_LOG_SCALING=log_scaling is not None, + D_REL=16, + REL_EXTENT=rel_proj.shape[1], + REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]), + num_warps=SMALL_NUM_WARPS, + ) + + +def fused_qkvr_prep( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + rel_proj: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + conv_block_table: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + assert d_rel == 16 and rel_proj.shape[0] == 16 + assert head_dim == 128 + assert qkvr.is_contiguous() + assert k_weight.stride() == (k_weight.shape[1], 1) + assert v_weight.stride() == (v_weight.shape[1], 1) + assert rel_proj.stride() == (rel_proj.shape[1], 1) + assert conv_cache.stride(3) == 1 + assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1 + tokens = qkvr.shape[0] + q_out = torch.empty( + (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device + ) + rel_out = torch.empty( + (tokens, num_q_heads, rel_proj.shape[1]), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if tokens == 0: + return q_out, rel_out + + if tokens < SMALL_TOKEN_THRESHOLD: + _run_fused_small( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + log_scaling=log_scaling, + ) + return q_out, rel_out + + kv_stream = aux_stream() + assert kv_stream is not None + current_stream = torch.cuda.current_stream() + kv_stream.wait_stream(current_stream) + with torch.cuda.stream(kv_stream): + _run_tiled_kv( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + ) + _run_tiled_q( + qkvr, + q_norm_weight, + q_out, + positions, + eps=eps, + num_q_heads=num_q_heads, + head_dim=head_dim, + log_scaling=log_scaling, + ) + qkvr_rel_proj( + qkvr, + rel_proj, + rel_out, + log_scaling, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + d_rel=d_rel, + ) + current_stream.wait_stream(kv_stream) + return q_out, rel_out diff --git a/vllm/models/inkling/amd/ops/rel_attention_decode.py b/vllm/models/inkling/amd/ops/rel_attention_decode.py new file mode 100644 index 000000000000..810eb9cb5e6b --- /dev/null +++ b/vllm/models/inkling/amd/ops/rel_attention_decode.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Split-KV decode for Inkling relative attention on ROCm. + +Adapted from LightSeek TokenSpeed's portable Triton relative-MHA decode. +""" + +from __future__ import annotations + +import os + +import torch + +from vllm.triton_utils import tl, triton + +_MIN_BLOCK_KV = tl.constexpr(32) + + +def decode_split_count(max_kv_len: int, window_left: int) -> int: + """Return the number of parallel KV partitions for decode.""" + if window_left >= 0: + effective_len = max(1, min(max_kv_len, window_left + 1)) + return min(4, max(1, triton.cdiv(effective_len, 128))) + return min(32, max(1, triton.cdiv(max(1, max_kv_len), 2048))) + + +def use_split_kv_decode( + *, + max_query_len: int, + max_kv_len: int, + page_size: int, + window_left: int, +) -> bool: + """Select split-KV only where it outperforms the single-pass kernel.""" + if os.getenv("INKLING_SPLIT_KV", "1") != "1": + return False + if max_query_len != 1: + return False + if page_size >= 64: + return True + if window_left >= 0: + return page_size >= 32 + return max_kv_len >= 8192 + + +@triton.jit +def _split_kv_stage1( + q_ptr, + rel_ptr, + k_ptr, + v_ptr, + block_table_ptr, + cache_seqlens, + mid_out_ptr, + mid_lse_ptr, + softmax_scale, + stride_q_t, + stride_q_h, + stride_k_b, + stride_k_p, + stride_k_h, + stride_k_d, + stride_v_b, + stride_v_p, + stride_v_h, + stride_v_d, + stride_r_t, + stride_r_h, + stride_r_e, + stride_mo_t, + stride_mo_h, + stride_mo_s, + stride_ml_t, + stride_ml_h, + stride_ml_s, + stride_bt_b: tl.constexpr, + page_size: tl.constexpr, + window_left: tl.constexpr, + gqa_group_size: tl.constexpr, + num_q_heads: tl.constexpr, + head_dim: tl.constexpr, + rel_extent: tl.constexpr, + max_kv_splits: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + pid_s = tl.program_id(2) + + head_blocks_per_kv = tl.cdiv(gqa_group_size, BLOCK_H) + kv_head = pid_h // head_blocks_per_kv + valid_block_h: tl.constexpr = min(BLOCK_H, gqa_group_size) + off_h = pid_h * valid_block_h + tl.arange(0, BLOCK_H) + h_valid = (off_h < (pid_h + 1) * valid_block_h) & (off_h < num_q_heads) + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + + cache_len = tl.load(cache_seqlens + pid_t) + effective_len = ( + tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len + ) + kv_offset = cache_len - effective_len + split_len = ( + tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + split_start = split_len * pid_s + split_end = tl.minimum(split_start + split_len, effective_len) + + q_offsets = pid_t * stride_q_t + off_h[:, None] * stride_q_h + off_d[None, :] + row_max = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32) + row_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) + + if split_end > split_start: + q = tl.load( + q_ptr + q_offsets, + mask=h_valid[:, None] & d_valid[None, :], + other=0.0, + ) + for start_n in range(split_start, split_end, BLOCK_N): + off_n = start_n + tl.arange(0, BLOCK_N) + n_valid = off_n < split_end + token_idx = kv_offset + off_n + logical_page = token_idx // page_size + page_offset = token_idx % page_size + physical_page = tl.load( + block_table_ptr + pid_t * stride_bt_b + logical_page, + mask=n_valid, + other=0, + ) + k_offsets = ( + physical_page[None, :].to(tl.int64) * stride_k_b + + page_offset[None, :] * stride_k_p + + kv_head * stride_k_h + + off_d[:, None] * stride_k_d + ) + k = tl.load( + k_ptr + k_offsets, + mask=d_valid[:, None] & n_valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k.to(q.dtype)) * softmax_scale + + rel_dist = cache_len - 1 - token_idx + rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent) + rel_idx = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) + rel_offsets = ( + pid_t * stride_r_t + + off_h[:, None] * stride_r_h + + rel_idx[None, :] * stride_r_e + ) + rel_bias = tl.load( + rel_ptr + rel_offsets, + mask=h_valid[:, None] & rel_valid[None, :] & n_valid[None, :], + other=0.0, + ) + scores += rel_bias.to(tl.float32) + scores = tl.where( + h_valid[:, None] & n_valid[None, :], + scores, + float("-inf"), + ) + + v_offsets = ( + physical_page[:, None].to(tl.int64) * stride_v_b + + page_offset[:, None] * stride_v_p + + kv_head * stride_v_h + + off_d[None, :] * stride_v_d + ) + v = tl.load( + v_ptr + v_offsets, + mask=n_valid[:, None] & d_valid[None, :], + other=0.0, + ) + + next_max = tl.maximum(tl.max(scores, axis=1), row_max) + old_scale = tl.exp(row_max - next_max) + probs = tl.exp(scores - next_max[:, None]) + acc *= old_scale[:, None] + acc += tl.dot(probs.to(v.dtype), v) + row_sum = row_sum * old_scale + tl.sum(probs, axis=1) + row_max = next_max + + mid_out_offsets = ( + pid_t * stride_mo_t + + off_h[:, None] * stride_mo_h + + pid_s * stride_mo_s + + off_d[None, :] + ) + tl.store( + mid_out_ptr + mid_out_offsets, + acc / row_sum[:, None], + mask=h_valid[:, None] & d_valid[None, :], + ) + mid_lse_offsets = ( + pid_t * stride_ml_t + off_h * stride_ml_h + pid_s * stride_ml_s + ) + tl.store( + mid_lse_ptr + mid_lse_offsets, + row_max + tl.log(row_sum), + mask=h_valid, + ) + + +@triton.jit +def _split_kv_stage2( + mid_out_ptr, + mid_lse_ptr, + out_ptr, + cache_seqlens, + stride_mo_t, + stride_mo_h, + stride_mo_s, + stride_ml_t, + stride_ml_h, + stride_ml_s, + stride_o_t, + stride_o_h, + window_left: tl.constexpr, + head_dim: tl.constexpr, + max_kv_splits: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + + cache_len = tl.load(cache_seqlens + pid_t) + effective_len = ( + tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len + ) + split_len = ( + tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + row_max = -float("inf") + row_sum = 0.0 + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + + for split_id in range(max_kv_splits): + split_start = split_len * split_id + split_end = tl.minimum(split_start + split_len, effective_len) + if split_end > split_start: + value = tl.load( + mid_out_ptr + + pid_t * stride_mo_t + + pid_h * stride_mo_h + + split_id * stride_mo_s + + off_d, + mask=d_valid, + other=0.0, + ) + split_lse = tl.load( + mid_lse_ptr + + pid_t * stride_ml_t + + pid_h * stride_ml_h + + split_id * stride_ml_s + ) + next_max = tl.maximum(split_lse, row_max) + old_scale = tl.exp(row_max - next_max) + split_scale = tl.exp(split_lse - next_max) + acc = acc * old_scale + value * split_scale + row_sum = row_sum * old_scale + split_scale + row_max = next_max + + tl.store( + out_ptr + pid_t * stride_o_t + pid_h * stride_o_h + off_d, + acc / row_sum, + mask=d_valid, + ) + + +@torch.no_grad() +def inkling_rel_attention_split_kv_decode( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + softmax_scale: float, + window_left: int, + rel_extent: int, + rel_logits: torch.Tensor, + max_kv_len: int, + out: torch.Tensor, +) -> torch.Tensor: + """Run split-KV relative attention for single-token decode.""" + num_kv_heads = key_cache.shape[2] + gqa_group_size = q.shape[1] // num_kv_heads + max_kv_splits = decode_split_count(max_kv_len, window_left) + block_d = triton.next_power_of_2(q.shape[2]) + block_h = min(8, gqa_group_size) + + mid_out = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + q.shape[2], + dtype=torch.float32, + device=q.device, + ) + mid_lse = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + dtype=torch.float32, + device=q.device, + ) + stage1_grid = ( + q.shape[0], + triton.cdiv(q.shape[1], block_h), + max_kv_splits, + ) + _split_kv_stage1[stage1_grid]( + q, + rel_logits, + key_cache, + value_cache, + block_table, + cache_seqlens, + mid_out, + mid_lse, + softmax_scale, + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + key_cache.stride(3), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + value_cache.stride(3), + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + block_table.stride(0), + page_size=key_cache.shape[1], + window_left=window_left, + gqa_group_size=gqa_group_size, + num_q_heads=q.shape[1], + head_dim=q.shape[2], + rel_extent=rel_extent, + max_kv_splits=max_kv_splits, + BLOCK_D=block_d, + BLOCK_H=block_h, + BLOCK_N=key_cache.shape[1], + num_warps=4, + num_stages=1, + ) + stage2_grid = (q.shape[0], q.shape[1]) + _split_kv_stage2[stage2_grid]( + mid_out, + mid_lse, + out, + cache_seqlens, + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + out.stride(0), + out.stride(1), + window_left=window_left, + head_dim=q.shape[2], + max_kv_splits=max_kv_splits, + BLOCK_D=block_d, + num_warps=4, + num_stages=2, + ) + return out diff --git a/vllm/models/inkling/amd/ops/sconv.py b/vllm/models/inkling/amd/ops/sconv.py new file mode 100644 index 000000000000..7ebb2a8a0a87 --- /dev/null +++ b/vllm/models/inkling/amd/ops/sconv.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-convolution kernels backed by a paged sliding-window state cache. + +Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV +cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream +occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H`` +heads, so its flat per-token width is ``H * ws`` and that is the conv channel +dim. The cache stores the conv *input* at every absolute position. + +``fused_sconv`` is the single-launch path used by the model: per token it writes +the current input to its slot and convolves the ``W`` taps ending at its +absolute position. A tap landing inside the current forward is read from the +immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read +from the paged cache (window position ``src`` -> physical block via +``block_table[req, src // N]``). The just-written slot is never read back this +step, so there is no write/read hazard within or across programs -- which is +why this needs no decode-vs-prefill split and is valid for prefill / decode / +spec alike. + +All kernels address the cache purely by ``(slot, absolute_position)`` and +allocate nothing inside the captured region; their grids depend only on the +token count (``fused_sconv`` on a fixed token/channel tiling), so the same +path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs +without any data-dependent shape or branch. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_sconv_kernel( + x_ptr, # [T, H*WS] head-major current-token inputs (also residual) + cache_ptr, # [num_blocks, H, N, D] paged (page-strided view) + weight_ptr, # [H*WS, W] + out_ptr, # [T, H*WS] + pos_ptr, # [T] int64 absolute position per token + seq_idx_ptr, # [T] int32 token -> batch request + slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD + block_table_ptr, # [num_reqs, max_blocks] int32 block_table + qstart_ptr, # [T] int32 first x-row of the token's request + T, # num tokens + stride_x_t, + stride_c_blk, + stride_c_h, + stride_c_n, + stride_c_d, + stride_w_d, + stride_w_w, + stride_bt_r, + MAX_BLOCKS, + N, # block_size + W: tl.constexpr, + USE_SILU: tl.constexpr, + USE_RESIDUAL: tl.constexpr, + OFF_S: tl.constexpr, + WS: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + BLOCK_C: tl.constexpr, +): + # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel + # index packs all H heads head-major (head = c // WS, in-stream offset = + # c % WS), so one program spans heads -- no per-head launch and no + # next_power_of_2(WS) lane waste. + pid_t = tl.program_id(0) + pid_c = tl.program_id(1) + toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows + coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels + C = H * WS + t_mask = toff < T + c_mask = coff < C + head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer + cd = OFF_S + coff % WS # cache D-index of the channel's stream slot + + slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT] + valid = slot >= 0 + pos = tl.load(pos_ptr + toff, mask=t_mask, other=0) + req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0) + qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0) + + tc_mask = t_mask[:, None] & c_mask[None, :] + + # 1) Insert each token's input into its paged slot (skip PAD rows). + xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask) + safe_slot = tl.maximum(slot, 0) + dst = ( + cache_ptr + + (safe_slot // N)[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_slot % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + tl.store(dst, xv, mask=tc_mask & valid[:, None]) + + # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from + # `x` if it falls inside this forward (row >= the request's first row), else + # from the paged cache. Exactly one source is unmasked per tap, so we sum both. + acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32) + for iw in tl.static_range(W): + src = pos - (W - 1) + iw # [BT] absolute window position + row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token) + in_win = valid & (src >= 0) + intra = in_win & (row >= qstart) + cached = in_win & (row < qstart) + # intra-forward tap: read the immutable input x (never the slot we just + # wrote), so there is no write/read hazard. + safe_row = tl.maximum(row, 0) + xt = tl.load( + x_ptr + safe_row[:, None] * stride_x_t + coff[None, :], + mask=c_mask[None, :] & intra[:, None], + other=0.0, + ).to(tl.float32) + # pre-forward tap: read from the paged cache via the block table. Clamp + # addressing terms; the load is masked off when out of window. + safe_src = tl.maximum(src, 0) + safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1) + blk = tl.load( + block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0 + ).to(tl.int64) + cbase = ( + cache_ptr + + blk[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_src % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to( + tl.float32 + ) + wv = tl.load( + weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0 + ).to(tl.float32) + acc += (xt + cv) * wv[None, :] + + if USE_SILU: + acc = acc * tl.sigmoid(acc) + if USE_RESIDUAL: + acc += xv.to(tl.float32) + + tl.store( + out_ptr + toff[:, None] * stride_x_t + coff[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=tc_mask, + ) + + +def fused_sconv( + x: torch.Tensor, # [T, H*ws] head-major current-token inputs + weight: torch.Tensor, # [H*ws, W] + cache: torch.Tensor, # [num_blocks, H, N, D] paged + positions: torch.Tensor, # [T] int64 absolute position per token + block_table: torch.Tensor, # [num_reqs, max_blocks] int32 + seq_idx: torch.Tensor, # [T] int32 token -> batch request + slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip) + query_start: torch.Tensor, # [T] int32 first x-row of the token's request + off_s: int, + ws: int, + block_size: int, + activation: str | None = None, + use_residual: bool = True, +) -> torch.Tensor: + """Single-launch insert + depthwise causal conv1d over the paged cache. + + Reads same-forward taps from ``x`` and pre-forward taps from the cache, so + it is race-free in one launch for prefill / decode / spec and cudagraph-safe + under eager / piecewise / full capture. + """ + T = x.shape[0] + out = torch.empty_like(x) + if T == 0: + return out + assert x.is_contiguous() + assert cache.stride(3) == 1, "cache D-dim must be contiguous" + H = cache.shape[1] + W = weight.shape[1] + C = H * ws # flat conv channel dim (all heads, head-major) + # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to + # amortize the per-token addressing, while keeping the grid large for + # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element + # tile at 4 warps measured best on Blackwell; larger tiles spill registers. + BLOCK_C = min(triton.next_power_of_2(C), 256) + BT = 8 + grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C)) + _fused_sconv_kernel[grid]( + x, + cache, + weight, + out, + positions, + seq_idx, + slot_mapping, + block_table, + query_start, + T, + x.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + weight.stride(0), + weight.stride(1), + block_table.stride(0), + block_table.shape[1], + block_size, + W=W, + USE_SILU=activation in ("silu", "swish"), + USE_RESIDUAL=use_residual, + OFF_S=off_s, + WS=ws, + H=H, + BT=BT, + BLOCK_C=BLOCK_C, + num_warps=4, + ) + return out + + +@triton.jit +def _seq_metadata_kernel( + qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows + seq_idx_ptr, # [T] int32 out: token -> owning request + query_start_ptr, # [T] int32 out: first x-row of the token's request + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, # ceil(log2(num_reqs)): binary-search depth + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tok = offs.to(tl.int32) + # Largest j in [0, num_reqs) with qsl[j] <= tok. + lo = tl.zeros([BLOCK], tl.int32) + hi = tl.full([BLOCK], num_reqs - 1, tl.int32) + for _ in range(n_iters): + mid = (lo + hi + 1) // 2 + below = tl.load(qsl_ptr + mid) <= tok + lo = tl.where(below, mid, lo) + hi = tl.where(below, hi, mid - 1) + actual = offs < num_actual_tokens + padded = offs < num_padded_tokens + query_start = tl.load(qsl_ptr + lo) + tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded) + tl.store( + query_start_ptr + offs, + tl.where(actual, query_start, 0), + mask=padded, + ) + + +def sconv_seq_metadata( + query_start_loc: torch.Tensor, + num_reqs: int, + num_actual_tokens: int, + seq_idx_out: torch.Tensor, + query_start_out: torch.Tensor, + num_padded_tokens: int | None = None, +) -> None: + """Fill static per-token seq_idx / query_start buffers in one launch. + + Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the + sconv metadata build with a single kernel writing both persistent buffers. + Padded rows are filled with zero and must have ``slot_mapping == -1``. + """ + if num_padded_tokens is None: + num_padded_tokens = num_actual_tokens + if num_padded_tokens < num_actual_tokens: + raise ValueError("num_padded_tokens must cover all actual tokens") + if num_padded_tokens > seq_idx_out.shape[0]: + raise ValueError("seq_idx_out is too small for the padded token count") + if num_padded_tokens > query_start_out.shape[0]: + raise ValueError("query_start_out is too small for the padded token count") + + BLOCK = 256 + n_iters = (num_reqs - 1).bit_length() + grid = (triton.cdiv(num_padded_tokens, BLOCK),) + _seq_metadata_kernel[grid]( + query_start_loc, + seq_idx_out, + query_start_out, + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, + BLOCK=BLOCK, + ) diff --git a/vllm/models/inkling/amd/ops/silu_and_mul.py b/vllm/models/inkling/amd/ops/silu_and_mul.py new file mode 100644 index 000000000000..42ec25f63235 --- /dev/null +++ b/vllm/models/inkling/amd/ops/silu_and_mul.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SwiGLU kernels for the Inkling MLP layers. + +``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved +fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert +variant with the per-expert dequant scale and per-token gamma fused in. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["M"]) +def _silu_and_mul_triton_kernel( + gateup_out_ptr, + down_inp_ptr, + M, + N: tl.constexpr, + GRID_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + EVEN_N: tl.constexpr, + INT64_INDEX: tl.constexpr, +): + start_pid = tl.program_id(0) + if INT64_INDEX: + start_pid = start_pid.to(tl.int64) + M = M.to(tl.int64) + + NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) + num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N + + for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES): + pid_m = pid // NUM_BLOCKS_N + pid_n = pid % NUM_BLOCKS_N + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + mask_n = offs_n < N + + # Interleaved fused gate/up: [g0, u0, g1, u1, ...]. + mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2 + tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}") + mask_2n = mask_offs_2n < N + mask_2n = tl.max_constancy(mask_2n, [16]) + + offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) + offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :] + + if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N: + gateup_out = tl.load( + gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0 + ) + else: + mask_m2n = mask_m[:, None] & mask_2n[None, :] + gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0) + + gate_out, up_out = tl.split( + tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2)) + ) + gate_out = gate_out.to(tl.float32) + up_out = up_out.to(tl.float32) + + down_inp = gate_out * tl.sigmoid(gate_out) * up_out + + mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :] + offs_mn = offs_m[:, None] * N + offs_n[None, :] + tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn) + + +def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor: + """SiLU-and-mul for the interleaved fused gate/up layout. + + Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). + """ + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) + assert gateup_output.ndim == 2, f"{gateup_output.shape=}" + + M = gateup_output.shape[0] + hidden_size = gateup_output.shape[1] + assert hidden_size % 2 == 0, f"{hidden_size=}" + N = hidden_size // 2 + + down_input = torch.empty( + (M, N), device=gateup_output.device, dtype=gateup_output.dtype + ) + if M == 0: + return down_input + + BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N))) + if M <= 1: + BLOCK_SIZE_M = 4 + elif M <= 256: + BLOCK_SIZE_M = 2 + elif M < 4096: + BLOCK_SIZE_M = 4 + else: + BLOCK_SIZE_M = 16 + BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N))) + max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N) + num_sms = torch.cuda.get_device_properties( + gateup_output.device + ).multi_processor_count + grid_size = min(num_sms * 4, max_grid_size) + + _silu_and_mul_triton_kernel[(grid_size,)]( + gateup_out_ptr=gateup_output, + down_inp_ptr=down_input, + M=M, + N=N, + GRID_SIZE=grid_size, + NUM_STAGES=1, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + EVEN_N=N % BLOCK_SIZE_N == 0, + INT64_INDEX=gateup_output.nbytes >= 2**31, + num_warps=8, + ) + + return down_input + + +@triton.jit(do_not_specialize=["T"]) +def _sink_epilogue_kernel( + raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block + alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale + gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided) + ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio) + out_ptr, # [T, S * F] output + T, + stride_raw_0, + stride_gamma_0, + F: tl.constexpr, + S: tl.constexpr, + BLOCK_F: tl.constexpr, +): + pid_t = tl.program_id(0).to(tl.int64) + pid_sf = tl.program_id(1) + if pid_t >= T: + return + s = pid_sf // (F // BLOCK_F) + offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F) + + base = pid_t * stride_raw_0 + s * 2 * F + gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32) + up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32) + alpha = tl.load(alpha_ptr + s) + weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s) + + gate *= alpha + up *= alpha + h = gate * tl.sigmoid(gate) * up * weight + tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h) + + +def sink_silu_mul_epilogue( + raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows) + alphas: torch.Tensor, # [S] fp32 + gammas: torch.Tensor, # [T, S] fp32 + ratios: torch.Tensor, # [S] fp32 + n_experts: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e). + + One kernel replaces the per-expert dequant column scale, the SwiGLU, and + the per-token gamma multiply between the two sink GEMMs. + """ + tokens = raw.shape[0] + f = raw.shape[1] // (2 * n_experts) + out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype) + if tokens == 0: + return out + # raw may be a column-slice of a padded GEMM output (rows strided). + assert raw.stride(1) == 1 and gammas.stride(1) == 1 + # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512. + block_f = min(512, f & (-f)) + _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))]( + raw, + alphas, + gammas, + ratios, + out, + tokens, + raw.stride(0), + gammas.stride(0), + F=f, + S=n_experts, + BLOCK_F=block_f, + ) + return out diff --git a/vllm/models/inkling/amd/sconv_swa_attn.py b/vllm/models/inkling/amd/sconv_swa_attn.py new file mode 100644 index 000000000000..7b8958d556f2 --- /dev/null +++ b/vllm/models/inkling/amd/sconv_swa_attn.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-conv state managed as a sliding-window KV cache. + +Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that +emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V, +attn-output, mlp-output), packed head-major into one block: + + H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size, + D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp) + +``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv +reads/writes this cache out-of-band via a custom backend; the (smaller) conv page +is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + +import torch +from torch import nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec + +from .ops.sconv import sconv_seq_metadata + +# Stream order within the per-head packed D (== contiguous sub-ranges). +_K, _V, _ATTN, _MLP = 0, 1, 2, 3 + + +@dataclass +class InklingSconvMetadata(AttentionMetadata): + block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req + slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip) + seq_idx: torch.Tensor # [T] int32 token -> batch request + query_start: torch.Tensor # [T] int32 first x-row of each token's request + + +class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + assert isinstance(kv_cache_spec, SlidingWindowSpec) + # Persistent per-token buffers for CUDA graph capture. + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.seq_idx_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + self.query_start_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> InklingSconvMetadata: + num_reqs = common_attn_metadata.num_reqs + num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + num_padded_tokens = common_attn_metadata.slot_mapping.shape[0] + assert num_padded_tokens >= num_actual_tokens + + # Per-token seq_idx (owning request) and query_start (first x-row of + # that request; the fused kernel uses it to tell same-forward taps, + # read from x, from pre-forward taps, read from cache) in one launch. + sconv_seq_metadata( + common_attn_metadata.query_start_loc, + num_reqs, + num_actual_tokens, + self.seq_idx_buffer, + self.query_start_buffer, + num_padded_tokens, + ) + + return InklingSconvMetadata( + block_table=common_attn_metadata.block_table_tensor, + slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens], + seq_idx=self.seq_idx_buffer[:num_padded_tokens], + query_start=self.query_start_buffer[:num_padded_tokens], + ) + + +class InklingSconvBackend(AttentionBackend): + """Custom dummy backend for the sconv sliding-window cache management.""" + + @staticmethod + def get_name() -> str: + return "INKLING_SCONV_SWA" + + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the + # padded conv page is read through a strided view. + return True + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + # HND, num-blocks-first, head-major: [num_blocks, H, N, D]. + return (num_blocks, num_kv_heads, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + # Identity: physical layout == logical [num_blocks, H, N, D]. + if include_num_layers_dimension: + return (0, 1, 2, 3, 4) + return (0, 1, 2, 3) + + @staticmethod + def get_impl_cls(): + raise NotImplementedError( + "InklingSconvBackend has no attention impl; the conv runs out-of-band." + ) + + @staticmethod + def get_builder_cls() -> type[InklingSconvMetadataBuilder]: + return InklingSconvMetadataBuilder + + +class InklingConvState(nn.Module, AttentionLayerBase): + """Per-decoder-layer owner emitting one sliding-window conv-state spec.""" + + def __init__( + self, + *, + num_kv_heads: int, + head_dim: int, + hidden_size: int, + kernel_size: int, + prefix: str, + ) -> None: + super().__init__() + self.prefix = prefix + # Bound to the manager-allocated paged cache by bind_kv_cache; a + # placeholder until then. Read out-of-band by InklingShortConv. + self.kv_cache = torch.tensor([]) + tp_size = get_tensor_model_parallel_world_size() + # Guardrails for the conv-state layout below; only these are exercised. + # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no + # replication/clamping), so the per-head width stays TP-invariant. + assert tp_size <= num_kv_heads, ( + f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), " + f"got {tp_size}" + ) + # Per-rank head count; D is TP-invariant (K/V heads and the hidden + # chunk both scale 1/TP together). The attn-/mlp-output sconv streams + # are hidden-sharded: each rank owns its H/tp chunk (the sublayer + # outputs are reduce-scattered / all-gathered around the conv). + self.num_kv_heads = num_kv_heads // tp_size + hidden_per_head = hidden_size // num_kv_heads + # Packed per-head width: K + V + attn-output chunk + mlp-output chunk, + # padded to a power of two so every layer's conv page is the same size + # and an exact multiple of the attention page (the page unifier then + # scales attention block sizes instead of padding). + raw_head_size = 2 * head_dim + 2 * hidden_per_head + self.head_size = 1 << (raw_head_size - 1).bit_length() + self.sliding_window = kernel_size + self.block_size = kernel_size + # Per-head D-sub-range (offset, width) for each stream. Streams share + # the cache; each writes/reads its own width across all H heads. + self.stream_ranges: tuple[tuple[int, int], ...] = ( + (0, head_dim), # _K + (head_dim, head_dim), # _V + (2 * head_dim, hidden_per_head), # _ATTN + (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP + ) + vllm_config = get_current_vllm_config() + self._dtype = vllm_config.model_config.dtype + assert self._dtype == torch.bfloat16, ( + f"sconv SWA cache supports bfloat16 only, got {self._dtype}" + ) + # Register in the forward context so the runner enumerates this owner as + # an attention-like layer (get_kv_cache_spec / get_attn_backend). + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward(self): ... + + @property + def cache_block_size(self) -> int: + """Return the block size used by cache metadata and physical indexing.""" + if self.kv_cache.numel() > 0: + return self.kv_cache.shape[2] + return self.block_size + + def get_attn_backend(self) -> type[AttentionBackend]: + return InklingSconvBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return SlidingWindowSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=0, # all 4 streams packed into head_size + dtype=self._dtype, + sliding_window=self.sliding_window, + ) diff --git a/vllm/models/inkling/amd/short_conv.py b/vllm/models/inkling/amd/short_conv.py new file mode 100644 index 000000000000..3a4db8f88318 --- /dev/null +++ b/vllm/models/inkling/amd/short_conv.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short convolution: depthwise causal conv1d (+ residual) over a paged +sliding-window conv-state cache. + +Each decoder layer owns one ``InklingConvState`` (``sconv_swa_attn.py``) holding the +manager-allocated paged cache for the layer's 4 sconv streams (K, V, attn-output, +mlp-output), packed head-major into one block. Each ``InklingShortConv`` is a +stateless weight + kernel launcher that, per forward (positions-addressed, the +same path for prefill / decode / mixed), inserts the current tokens' inputs +into their paged slot and convolves each token against the ``W`` taps ending +at its absolute position, reading pre-forward window positions out of the +paged cache via the block table. + +Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` / +``query_start``) is built once by ``InklingSconvMetadataBuilder`` and published under +the owner's prefix in the forward context; the absolute ``positions`` are +threaded in from the model. The insert + conv run in a single ``fused_sconv`` +launch (same path for prefill / decode / mixed / spec). All inputs are +fixed-address persistent buffers and the grid is fixed, so the conv replays +correctly under eager, PIECEWISE, and FULL cudagraphs. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.model_executor.utils import set_weight_attrs + +from .ops import fused_sconv +from .sconv_swa_attn import InklingConvState, InklingSconvMetadata + + +class InklingShortConv(nn.Module): + def __init__( + self, dim: int, kernel_size: int, owner: InklingConvState, stream_idx: int + ) -> None: + super().__init__() + self.dim = dim + self.kernel_size = kernel_size + self.owner = owner + self.stream_idx = stream_idx + self.tp_rank = get_tensor_model_parallel_rank() + + # Depthwise conv weight; checkpoint stores (dim, 1, W). + self.weight = Parameter(torch.empty(dim, 1, kernel_size), requires_grad=False) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None: + if loaded_weight.shape[0] != param.shape[0]: + shard = param.shape[0] + loaded_weight = loaded_weight.narrow(0, self.tp_rank * shard, shard) + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # x: (num_tokens, dim); positions: (num_tokens,) absolute positions. + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + # Memory-profiling / no metadata: identity (residual). + return x + m = attn_metadata.get(self.owner.prefix) + if m is None: + return x + assert isinstance(m, InklingSconvMetadata) + cache = self.owner.kv_cache + if cache.numel() == 0: + # Cache not yet bound (profiling before KV alloc): identity. + return x + + off_s, ws = self.owner.stream_ranges[self.stream_idx] + # The hybrid KV-cache planner can enlarge the logical conv block so + # that its physical page size matches the attention caches (for + # example, W=4 becomes 32 when --block-size=128). Metadata slot + # mappings are built with that enlarged size, so index the bound cache + # with its runtime token dimension rather than the kernel window size. + block_size = self.owner.cache_block_size + x = x.contiguous() + weight = self.weight.squeeze(1) # (dim, W) + + return fused_sconv( + x, + weight, + cache, + positions, + m.block_table, + m.seq_idx, + m.slot_mapping, + m.query_start, + off_s, + ws, + block_size, + activation=None, + use_residual=True, + ) diff --git a/vllm/models/inkling/common/towers.py b/vllm/models/inkling/common/towers.py index 5ac3dbc20e6a..1c8739d0e8bc 100644 --- a/vllm/models/inkling/common/towers.py +++ b/vllm/models/inkling/common/towers.py @@ -9,6 +9,7 @@ from __future__ import annotations +from itertools import combinations from typing import cast import numpy as np @@ -45,6 +46,20 @@ def _prime_factors(n: int) -> list[int]: return factors +def linear_sum_assignment( + cost_matrix: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Implement SciPy's assignment for Inkling's ordered L1 cost matrix.""" + rows = np.arange(cost_matrix.shape[0]) + cols = np.array( + min( + combinations(range(cost_matrix.shape[1]), len(rows)), + key=lambda candidate: cost_matrix[rows, candidate].sum(), + ) + ) + return rows, cols + + def plan_out_scales( temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3 ) -> list[tuple[int, int, int, int]]: @@ -97,8 +112,6 @@ def _round_up(x: int) -> int: if n_layers >= len(scales): idxs = np.argmin(cost_matrix, axis=1) else: - from scipy.optimize import linear_sum_assignment - idxs = linear_sum_assignment(cost_matrix)[1] assert len(idxs) >= 2 diff --git a/vllm/models/inkling/nvidia/mlp.py b/vllm/models/inkling/nvidia/mlp.py index ab7dec0b709c..51deb38e64ef 100644 --- a/vllm/models/inkling/nvidia/mlp.py +++ b/vllm/models/inkling/nvidia/mlp.py @@ -58,6 +58,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = silu_and_mul_triton(gate_up) x, _ = self.down_proj(x) if self.global_scale is not None: - x = x * self.global_scale + x.mul_(self.global_scale) # TP-partial output: the layer's reduce-scatter fallback consumes it. return x diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index a8506d92327b..04ba0ee085a3 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -378,11 +378,20 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA): "language_model.lm_head.": "lm_head.", }, orig_to_new_suffix={ - # NVFP4 scale + # ModelOpt NVFP4 scales ".w13_weight.scale": ".w13_weight_scale", ".w13_weight.scale2": ".w13_weight_scale_2", ".w2_weight.scale": ".w2_weight_scale", ".w2_weight.scale2": ".w2_weight_scale_2", + # Compressed tensors NVFP4 parameters + ".w13_weight.input_global_scale": ".w13_input_global_scale", + ".w13_weight.weight_global_scale": ".w13_weight_global_scale", + ".w13_weight.weight_packed": ".w13_weight_packed", + ".w13_weight.weight_scale": ".w13_weight_scale", + ".w2_weight.input_global_scale": ".w2_input_global_scale", + ".w2_weight.weight_global_scale": ".w2_weight_global_scale", + ".w2_weight.weight_packed": ".w2_weight_packed", + ".w2_weight.weight_scale": ".w2_weight_scale", }, ) packed_modules_mapping = { diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 6d0a550b03d1..36ffd07b927a 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -10,11 +10,10 @@ activates every sink) and always bf16 (the checkpoint excludes every ``shared_experts`` from quantization). -NVFP4 routed experts reuse vLLM's ModelOpt NVFP4 fused-MoE method; excluded -(bf16) layers fall back to the unquantized method. The checkpoint's fused -stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` / -``.input_amax`` aux tensors) are translated to the standard per-expert loads -in :meth:`InklingMoE.load_expert_weight`. +NVFP4 routed experts reuse vLLM's fused-MoE methods; excluded (bf16) layers +fall back to the unquantized method. Checkpoint fused stacked tensors are +translated to the standard per-expert loads in +:meth:`InklingMoE.load_expert_weight`. """ from __future__ import annotations @@ -35,7 +34,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.triton_utils import tl, tldevice, triton @@ -440,7 +439,7 @@ def __init__( # power-of-two EP sizes (n_routed is a power of two). num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size() - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=num_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -528,7 +527,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor | None: ) self._routed_sel = None - return out + sink_out + return out.add_(sink_out) # -- weight loading ---------------------------------------------------- @@ -575,11 +574,24 @@ def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: lids = [slots[g] for g in gids] tp_rank = experts.moe_config.moe_parallel_config.tp_rank - if key.endswith("_scale_2"): + if key.endswith(("_scale_2", "_global_scale")): # Per-expert scalars, vectorized over the local experts. The - # fused w13 param carries one slot per gate/up half. - vals = weight[gids].float().to(param.device) - param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals + # fused w13 params carry one slot per gate/up half. A single + # checkpoint value is shared by both halves. + vals = weight[gids].float().reshape(len(gids), -1) + target_width = math.prod(param.shape[1:]) + if vals.shape[1] == 1: + vals = vals.expand(-1, target_width) + elif vals.shape[1] != target_width: + raise ValueError( + f"cannot load {tuple(weight.shape)} into {tuple(param.shape)}" + ) + param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to( + param.device + ) + elif key == "w2_weight_scale" and weight.shape[-1] == 1: + # Per-output-channel scales are replicated across TP ranks. + param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype) elif key.startswith("w13"): # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the # fused param layout is [w1(gate); w3(up)]. The TP-local rows form diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index a2559d4cf057..197894beb915 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -2,9 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). -Implements the first MTP depth from the reference ``mtp_model.py`` shipped with -the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` -input projection, and a full Inkling transformer block with a dense bf16 MLP. +Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP +depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj`` +(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the +same short convolutions as the backbone; its attention is full or sliding-window +per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared +``chain_norm`` is applied after every depth; its output is both the logits input +and the previous hidden state fed to the next depth. The draft shares the target's token embedding table and LM head (``load_eagle_model`` wires those references) and applies the backbone @@ -25,6 +29,9 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import maybe_prefix from vllm.sequence import IntermediateTensors @@ -50,6 +57,16 @@ def _mtp_depth_from_name(name: str) -> int | None: return int(m.group(1)) if m else None +def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int: + num_layers = min(n_predict, num_spec) if num_spec else n_predict + if num_layers <= 0: + raise ValueError( + "Inkling MTP requires num_nextn_predict_layers and " + "num_speculative_tokens to select at least one depth layer." + ) + return num_layers + + class InklingMTPDepthLayer(nn.Module): """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" @@ -93,14 +110,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: vllm_config.speculative_config.draft_model_config.hf_config ) self.config = config - if vllm_config.speculative_config.num_speculative_tokens != 1: - raise ValueError( - "Inkling MTP currently supports exactly one speculative token" - ) + # The checkpoint ships num_nextn_predict_layers depth blocks, but only + # the first ``num_speculative_tokens`` are exercised (step i uses depth + # i). Build only those to save memory — each depth is a full Inkling block + # with its own (large) full-history sconv caches and KV cache. + n_predict = config.num_nextn_predict_layers + num_spec = vllm_config.speculative_config.num_speculative_tokens + self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec) self.chain_hidden_post_norm = config.chain_hidden_post_norm + + # Depth blocks whose attention is sliding-window (swa_* head config) + # rather than full; keyed by MTP depth via the checkpoint's + # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors + # InklingModel's local_ids split, but over MTP depths, not backbone + # layers. local_ids = set(config.local_layer_ids) + + # Keyed by depth index (str) to mirror the checkpoint layout. self.layers = nn.ModuleDict( - {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + { + str(idx): InklingMTPDepthLayer( + config, f"{prefix}.layers.{idx}", idx in local_ids + ) + for idx in range(self.num_mtp_layers) + } ) self.chain_norm = ( InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -198,9 +231,8 @@ def forward( # auto-enumerated as a draft attention layer); its per-token metadata is # built by the speculator's build_attn_metadata and read from the # forward context, so nothing extra is threaded here. - if spec_step_idx != 0: - raise ValueError("Inkling MTP only supports spec_step_idx=0") - layer = self.layers["0"] + depth = spec_step_idx % self.num_mtp_layers + layer = self.layers[str(depth)] combined = self.fused_input_cat( layer, previous_hidden_states, input_ids, inputs_embeds ) @@ -352,8 +384,8 @@ def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: # Only consume the MTP weights; everything else belongs to the target. if ".mtp." not in name: continue - # Only the first checkpoint depth is used for MTP=1. - if depth is not None and depth != 0: + # Skip depth blocks beyond the ones we built (num_speculative_tokens). + if depth is not None and depth >= module.model.num_mtp_layers: continue # model.mtp.chain_norm.weight -> model.chain_norm.weight # model.mtp.layers.{i}.X -> model.layers.{i}.X @@ -396,7 +428,7 @@ def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: for name in params if name.startswith("model.layers.") or name.startswith("model.chain_norm.") } - if missing := sorted(required - loaded): + if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled(): raise ValueError( "Inkling MTP checkpoint is missing required parameters: " + ", ".join(missing) diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py index f547f6506291..883545e3f2da 100644 --- a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -130,13 +130,20 @@ def inkling_fa4_rel_attention( cute_window = (None, None) if window_size == (-1, -1) else window_size rel_logits = rel_logits.contiguous() + flash_attn_varlen_func: Callable[..., Any] if _use_sheared_bias(): - from vllm.third_party.tml_fa4 import flash_attn_varlen_func + from vllm.third_party.tml_fa4 import ( + flash_attn_varlen_func as tml_flash_attn_varlen_func, + ) + flash_attn_varlen_func = tml_flash_attn_varlen_func bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} else: - from vllm.vllm_flash_attn.cute import flash_attn_varlen_func + from vllm.vllm_flash_attn.cute import ( + flash_attn_varlen_func as cute_flash_attn_varlen_func, + ) + flash_attn_varlen_func = cute_flash_attn_varlen_func bias_kwargs = { "score_mod": _get_score_mod(rel_extent), "aux_tensors": [rel_logits], diff --git a/vllm/models/kimi_k3/__init__.py b/vllm/models/kimi_k3/__init__.py new file mode 100644 index 000000000000..52c8a6538e66 --- /dev/null +++ b/vllm/models/kimi_k3/__init__.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi K3 model — hardware-isolated entry point. + +The implementation lives under ``nvidia/`` and ``amd/``; this module picks the +right one for the current platform and re-exports the public classes used by +the model registry. (Mirrors ``vllm.models.minimax_m3``.) +""" + +from typing import TYPE_CHECKING + +from vllm.platforms import current_platform + +# The NVIDIA branch is the static default that type-checkers see; the ROCm +# branch overrides it at runtime (kept type-compatible via type: ignore). +# TPU plugins import the shared ``common`` modules through this package, but +# register their own model classes. Do not eagerly import a GPU implementation. +if TYPE_CHECKING: + from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM + from .nvidia.mtp import KimiK3MTP +elif current_platform.device_type == "tpu": + pass +elif current_platform.is_rocm(): + from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment] + from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment] + from .amd.mtp import KimiK3MTP # type: ignore[assignment] +else: + from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM + from .nvidia.mtp import KimiK3MTP + +__all__ = [ + "KimiK3ForConditionalGeneration", + "KimiK3MTP", + "KimiLinearForCausalLM", +] diff --git a/vllm/models/kimi_k3/amd/__init__.py b/vllm/models/kimi_k3/amd/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py new file mode 100644 index 000000000000..181704ebf9f0 --- /dev/null +++ b/vllm/models/kimi_k3/amd/linear.py @@ -0,0 +1,1068 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +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.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + get_spec_layer_idx_from_weight_name, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.math_utils import cdiv + +logger = init_logger(__name__) + + +class KimiMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu and situ are supported." + ) + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +def _apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: ReplicatedLinear, + norm: RMSNorm, + num_valid_blocks: int, +) -> torch.Tensor: + if num_valid_blocks <= 0: + return prefix_sum + + return attn_res( + prefix_sum, + block_residual, + norm.weight, + proj.weight.squeeze(0), + num_valid_blocks, + norm.variance_epsilon, + ) + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = ( + routed_expert_hidden_size + if routed_expert_hidden_size is not None + else hidden_size + ) + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr( + config, "min_moe_intermediate_per_partition", 256 + ) + if self.tp_size > 1: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = ( + min_moe_intermediate_per_partition * self.tp_size + ) + activation_situ_beta = ( + config.activation_situ_beta if config.hidden_act == "situ" else None + ) + activation_situ_linear_beta = ( + config.activation_situ_linear_beta if config.hidden_act == "situ" else None + ) + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + # Preserve FP32 checkpoint values and match FP32 router logits. + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=torch.float32) + ) + + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if self.latent_moe_use_norm + else None + ) + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + self.experts = FusedMoEFactory( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + routed_input_transform=self.routed_expert_down_proj, + routed_output_transform=self.routed_output_transform, + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = self.experts.w13_weight_packed + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = self.experts.w2_weight_packed + w13_weight.data.zero_() + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = ( + moe_intermediate_size // self.tp_size + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiMLAAttention(nn.Module): + """ + Main reference: DeepseekV2 vllm Implementation + """ + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_nope: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + self.num_local_heads = num_heads // tp_size + self.scaling = self.qk_head_dim**-0.5 + self.use_nope = use_nope + assert self.use_nope is True + assert num_heads % tp_size == 0 + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm( + self.q_lora_rank, + eps=config.rms_norm_eps, + ) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm( + self.kv_lora_rank, + eps=config.rms_norm_eps, + ) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.use_output_gate = config.mla_use_output_gate + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = ColumnParallelLinear( + self.hidden_size, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + mla_modules = MLAModules( # type: ignore[call-arg] + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=None, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=None, + is_sparse=False, + topk_indices_buffer=None, + g_proj=getattr(self, "g_proj", None), + ) + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + output: torch.Tensor, + ) -> None: + output[:] = self.mla_attn(positions, hidden_states) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + if config.is_kda_layer(layer_idx): + self.self_attn = KimiGatedDeltaNetAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope is not None + self.self_attn = KimiMLAAttention( + layer_idx=layer_idx, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + quant_config=quant_config, + cache_config=cache_config, + model_config=model_config, + prefix=f"{prefix}.self_attn", + config=config, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_nope=mla_use_nope, + ) + + if ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ): + self.block_sparse_moe = KimiMoE( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_residuals = attn_res_block_size is not None + if attn_res_block_size is not None: + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + attn_output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=attn_output, + ) + return attn_output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.use_attn_residuals: + assert residual is not None + return self.forward_attn_residual(positions, hidden_states, residual) + + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self._run_self_attn(positions, hidden_states) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + def forward_attn_residual( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + prefix_sum = hidden_states + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.prev_valid_blocks, + ) + + if self.is_block_write_layer: + block_residual[:, self.block_write_idx, :].copy_(prefix_sum) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self._run_self_attn(positions, hidden_states) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + mlp_valid_blocks = self.prev_valid_blocks + ( + 1 if self.is_block_write_layer else 0 + ) + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + mlp_valid_blocks, + ) + + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + prefix_sum = prefix_sum + hidden_states + return prefix_sum, block_residual + + +class KimiLinearModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + self.config = config + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + vllm_config, + prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.attn_res_block_size is not None: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if config.attn_res_block_size is not None: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.config.attn_res_block_size is not None: + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.config.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def _maybe_add_hidden_state( + self, + aux_hidden_states: list[torch.Tensor], + layer_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> list[torch.Tensor]: + if self.config.attn_res_block_size is not None: + # attn-res `residual` is a block-state bank, not an additive + # residual; None makes the mixin capture the prefix sum directly. + residual = None + return super()._maybe_add_hidden_state( + aux_hidden_states, layer_idx, hidden_states, residual + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, residual + ) + + if self.config.attn_res_block_size is None: + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + # NOTE: the final norm is applied in compute_logits instead of here, + # so the MTP draft model receives the pre-norm hidden states. + if residual is not None: + hidden_states = hidden_states + residual + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + attn_res_block_num = cdiv(self.end_layer, self.config.attn_res_block_size) + block_residual = hidden_states.new_empty( + hidden_states.size(0), attn_res_block_num, hidden_states.size(1) + ) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + # AMD attn-res layer already returns prefix_sum + MLP delta as + # hidden_states; the override drops the block bank in residual. + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states = _apply_attn_res( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + attn_res_block_num, + ) + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[ + tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]] + ], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + if self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + +class KimiLinearForCausalLM( + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + return hidden_states + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + return loader.load_weights(weights) diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py new file mode 100644 index 000000000000..ec49891ddb2c --- /dev/null +++ b/vllm/models/kimi_k3/amd/model.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +from collections.abc import Iterable +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from vllm.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) +from .linear import KimiLinearForCausalLM + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel( + config.vision_config.num_attention_heads + ) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(vllm_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to( + device=self.device, dtype=model_config.dtype + ) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to( + device=self.device, dtype=model_config.dtype + ) + + self.quant_config = quant_config + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + def _maybe_ignore_quant_config( + self, quant_config: QuantizationConfig | None + ) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input( + self, **kwargs: object + ) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError( + "pixel_values must be a tensor or a list of tensors, " + f"got {type(pixel_values)}" + ) + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] + ) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), ( + f"expect grid_thws to be a tensor, got {type(grid_thws)}" + ) + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( + f"unexpected shape for grid_thws: {grid_thws.shape}" + ) + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input( + self, media_input: KimiK25MediaPixelInputs + ) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + hidden_states = self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs + ) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( + batch_size + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/kimi_k3/amd/mtp.py b/vllm/models/kimi_k3/amd/mtp.py new file mode 100644 index 000000000000..8fd79d5a1492 --- /dev/null +++ b/vllm/models/kimi_k3/amd/mtp.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.logger import init_logger +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.quantization import QuantizationConfig +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.utils import get_pp_missing_layer_names, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .linear import KimiDecoderLayer, get_spec_layer_idx_from_weight_name + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_config.quant_config + + 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.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + # The MTP block is a standard KimiDecoderLayer, but it must NOT use the + # attn-residual (block-residual) scheme even when the base model does: + # the draft starts from the target's hidden state, without the main + # model's accumulated per-block residual tensor. We disable it by + # shallow-copying the config with ``attn_res_block_size=None``. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + self.mtp_block = KimiDecoderLayer(block_config, vllm_config, prefix=prefix) + + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + + hidden_states, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm( + hidden_states, residual + ) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer( + config, vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + 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, + ) -> tuple[torch.Tensor, 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, + ) -> 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)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor( + 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, + ) -> tuple[torch.Tensor, 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 | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + expert_params_mapping = ( + fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + if self.config.is_moe + else [] + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and ( + ".layers" not in name + ): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_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_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/vllm/models/kimi_k3/amd/ops/__init__.py b/vllm/models/kimi_k3/amd/ops/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py new file mode 100644 index 000000000000..c00a3422ca1a --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + eps: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + input_qk_weight = tl.load(norm_weight_ptr + d_offsets, mask=d_mask, other=0.0).to( + tl.float32 + ) * tl.load(qk_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(tl.float32) + + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + num_sources = num_blocks + 1 + + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + block_values = tl.load( + block_ptrs, + mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]), + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + values = tl.where(is_prefix[:, None], prefix[None, :], block_values) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + output = mixed / denominator + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert 0 < num_blocks <= blocks.shape[1] + assert blocks.shape[0] == num_tokens + assert norm_weight.numel() == hidden_size + assert qk_weight.numel() == hidden_size + assert prefix.stride(-1) == 1 + assert blocks.stride(-1) == 1 + assert norm_weight.stride(-1) == 1 + assert qk_weight.stride(-1) == 1 + + output = prefix.new_empty(prefix.shape) + if num_tokens == 0: + return output + + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + blocks, + norm_weight, + qk_weight, + output, + prefix.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + eps, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + ) + return output diff --git a/vllm/models/kimi_k3/amd/ops/third_party/__init__.py b/vllm/models/kimi_k3/amd/ops/third_party/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py new file mode 100644 index 000000000000..77ed09d489ee --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# AMD/ROCm vendored copy of the Kimi-K3 KDA triton kernels. +# +# Provenance: mirror of vllm/models/kimi_k3/nvidia/ops/third_party/kda (tracker +# mke-tracker @ 7adebfcf9; FLA vendored per PRs #39/#86). Split per-vendor so +# AMD can carry gfx950-specific kernel changes without touching the NVIDIA copy. +# +# fla-org/flash-linear-attention#869 (the unmerged ROCm fixes our earlier +# amd_fla shim carried against FLA 0.5.0) is covered here by the *newer* vendored +# FLA rather than the literal patch: +# - transpose-state-layout workaround: N/A (kernels rewritten; no +# transpose_state_layout path remains), +# - AMD autotune configs: present (is_amd num_warps/num_stages branches), +# - OOB-mask correctness fix: present (all tl.load use mask=..., other=0). +# Validated on gfx950: no core-dump, gsm8k 94.1%. +# +# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync +# with the NVIDIA copy on FLA updates; any divergence should be an intentional, +# documented gfx950-specific change (a #869-style AMD-only fix). + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py new file mode 100644 index 000000000000..7d0ea0204ce6 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py @@ -0,0 +1,935 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.op import exp2, log +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}") + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py new file mode 100644 index 000000000000..ca3bc8f75e33 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.op import exp2, gather +from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported +from vllm.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_g0 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_k1 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_g1 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_k2 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_g2 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_k3 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_g3 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b1 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) + ) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Aqk21 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b2 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) + ) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Aqk31 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Aqk32 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b3 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) + ) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.0) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2 * BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) + for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk10 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Akk20 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk21 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk30 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Akk31 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Akk32 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BK", "NC", "BT", "HV"], +) +@triton.jit(do_not_specialize=["B", "T"]) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather( + b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0 + ) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) + ) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.0) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 000000000000..cc448d95c0eb --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp2 +from vllm.triton_utils import tl, triton + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BH": BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=["T", "N"]) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H * K + k += bos * H * K + g += bos * HV * K + Aqk += bos * HV * BT + Akk += bos * HV * BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr( + g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr( + g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store( + Aqk + i_t * HV * BT + o_hv * BT + j % BT, + b_Aqk.to(Aqk.dtype.element_ty), + mask=m_hv, + ) + tl.store( + Akk + i_t * HV * BC + o_hv * BC + j - i_ts, + b_Akk.to(Akk.dtype.element_ty), + mask=m_hv, + ) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): + return (B * T, triton.cdiv(HV, meta["BH"])) + + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 000000000000..2f512df62643 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,621 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp, log +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, +): + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + BV = 32 if use_gate_in_kernel else 8 + num_warps = 4 if use_gate_in_kernel else 1 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=2, + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + # gfx950: always fuse the gate. + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_state_token: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + state_idx = tl.load(state_indices + i_n).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1 or state_indices.stride(0) != 1: + raise ValueError("`state_indices` must be contiguous and one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + ) + return out, initial_state diff --git a/vllm/models/kimi_k3/common/__init__.py b/vllm/models/kimi_k3/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/common/mm_preprocess.py b/vllm/models/kimi_k3/common/mm_preprocess.py new file mode 100644 index 000000000000..fa70ccdce0a0 --- /dev/null +++ b/vllm/models/kimi_k3/common/mm_preprocess.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared Kimi-K3 multimodal preprocessing.""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import torch +from transformers import BatchFeature + +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + InputProcessingContext, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config +from vllm.transformers_utils.processor import cached_get_image_processor +from vllm.transformers_utils.processors.kimi_k3 import KimiK3Processor + +logger = init_logger(__name__) + + +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, +): + # Apply the patch limits. + 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, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, 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: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + + 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, + } + + +class KimiK3ProcessingInfo(BaseProcessingInfo): + """Processing information for the image-only Kimi-K3 model. + + K3 uses the standard ``image`` modality (unlike K2.5's unified + ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around + the checkpoint's image processor and resolves the ``<|media_pad|>`` token + id the same way K2.5 does. + """ + + def __init__(self, ctx: InputProcessingContext) -> None: + super().__init__(ctx) + + self.hf_config = hf_config = self.get_hf_config() + + tokenizer = self.get_tokenizer() + 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, + ) + + # Resolve token ID from the tokenizer because transformers v5 + # may remap token IDs vs config.json. + config_token_id = hf_config.media_placeholder_token_id + resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") + unk_token_id = getattr(tokenizer, "unk_token_id", None) + is_valid_resolved = isinstance(resolved_token_id, int) and ( + unk_token_id is None or resolved_token_id != unk_token_id + ) + if is_valid_resolved and resolved_token_id != config_token_id: + logger.warning_once( + "Kimi-K3 config.media_placeholder_token_id (%d) disagrees " + "with tokenizer mapping for <|media_pad|> (%d). " + "Using tokenizer value.", + config_token_id, + resolved_token_id, + ) + media_token_id = resolved_token_id + # Patch config so downstream code also sees the correct ID. + hf_config.media_placeholder_token_id = resolved_token_id + else: + media_token_id = config_token_id + + self.media_token_id = media_token_id + self.media_token = tokenizer.decode(media_token_id) + + self.image_processor = image_processor + self.hf_processor = KimiK3Processor( + tokenizer=tokenizer, + image_processor=image_processor, + ) + self.media_tokens_calculator = image_processor.media_tokens_calculator + + def get_hf_processor(self, **kwargs: object) -> KimiK3Processor: + return self.hf_processor + + def get_hf_config(self) -> KimiK3Config: + return self.ctx.get_hf_config(KimiK3Config) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # None means unlimited + return {"image": None} + + @classmethod + def get_max_image_size( + cls, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, + ) -> ImageSize: + max_side = patch_limit_on_one_side * patch_size + best_score = (-1, -1) + best_size = (max_side, max_side) + + for width_patches in range(patch_limit_on_one_side + 1): + width = min((width_patches + 1) * patch_size - 1, max_side) + for height_patches in range(width_patches, patch_limit_on_one_side + 1): + height = min((height_patches + 1) * patch_size - 1, max_side) + resize_config = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit, + patch_limit_on_one_side, + fixed_output_tokens, + ) + padded_width = resize_config["new_width"] + resize_config["pad_width"] + padded_height = ( + resize_config["new_height"] + resize_config["pad_height"] + ) + num_patches = padded_width // patch_size * (padded_height // patch_size) + score = (resize_config["num_tokens"], num_patches) + if score > best_score: + best_score = score + best_size = (width, height) + return ImageSize(width=best_size[0], height=best_size[1]) + + +class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]): + """Builds image-based dummy inputs for K3 profiling. + + The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly + the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy + mm data is a plain list of PIL images under the ``image`` key. + """ + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + return self.info.get_hf_config().image_placeholder * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + media_proc_cfg = self.info.image_processor.media_proc_cfg + max_size = self.info.get_max_image_size( + media_proc_cfg["patch_size"], + media_proc_cfg["merge_kernel_size"], + media_proc_cfg["in_patch_limit"], + media_proc_cfg["patch_limit_on_one_side"], + media_proc_cfg["fixed_output_tokens"], + ) + num_images = mm_counts.get("image", 0) + image_overrides = cast( + ImageDummyOptions | None, + mm_options.get("image") if mm_options else None, + ) + return { + "image": self._get_dummy_images( + width=max_size.width, + height=max_size.height, + num_images=num_images, + overrides=image_overrides, + ) + } + + +class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]): + """Image-only multi-modal processor for Kimi-K3.""" + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Override so the base always routes through the text+mm path + # (`KimiK3Processor.__call__`). Otherwise the mm-only fast path calls + # the checkpoint image processor directly with bare PIL images, but it + # requires `{"type": "image", "image": PIL}` media dicts that only our + # wrapper builds. + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + """Slice the flattened patch tensor back into per-image items. + + ``pixel_values`` holds all patches from every image concatenated; each + image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one + ``[N_t, N_h, N_w]`` row per image. + """ + grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3))) + grid_sizes = grid_thws.prod(-1) + + return dict( + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes), + grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Expand each K3 image placeholder into a resolution-aware update. + + K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per + image. This replaces that token with + ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``, + embedding the per-image resolution in the prompt and marking only the + ``<|media_pad|>`` positions as embedding slots (the number of pads is + the feature size returned by ``media_tokens_calculator``). + """ + media_token_id = self.info.media_token_id + media_token = self.info.media_token + image_placeholder = self.info.get_hf_config().image_placeholder + + def get_replacement(item_idx: int) -> PromptUpdateDetails[str]: + images = mm_items.get_items("image", ImageProcessorItems) + image = images.get(item_idx) + if image is None: + raise ValueError(f"Missing image data at index {item_idx}") + + # The checkpoint image processor works on media dicts, so wrap the + # PIL image before asking it for the token count. + num_media_token = self.info.media_tokens_calculator( + {"type": "image", "image": image} + ) + pads = media_token * num_media_token + + # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the + # post-preprocess (smart-resized) ones. `image` comes from the + # untouched parsed `mm_items`; the checkpoint image processor + # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via + # `image.resize(...)` and never mutates the stored PIL. This matches + # the reference HF processor (`KimiK3Processor.preprocess_medias`), + # which also builds the prompt from the original `img.size`. The + # resize is reflected only in the pad count above. + width, height = images.get_image_size(item_idx) + full = ( + f"<|media_begin|>image {width}x{height}<|media_content|>" + f"{pads}<|media_end|>" + ) + + return PromptUpdateDetails.select_token_id(full, media_token_id) + + return [ + PromptReplacement( + modality="image", + target=image_placeholder, + replacement=get_replacement, + ), + ] diff --git a/vllm/models/kimi_k3/common/mtp.py b/vllm/models/kimi_k3/common/mtp.py new file mode 100644 index 000000000000..76c796c1b0bb --- /dev/null +++ b/vllm/models/kimi_k3/common/mtp.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Kimi-K3 MTP input preparation.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rms_norm(x, weight, eps, hidden_size: tl.constexpr): + x = x.to(tl.float32) + variance = tl.sum(x * x, axis=0) / hidden_size + return x * tl.rsqrt(variance + eps) * weight.to(tl.float32) + + +@triton.jit +def _fused_mtp_input_kernel( + positions_ptr, + inputs_embeds_ptr, + previous_hidden_states_ptr, + enorm_weight_ptr, + hnorm_weight_ptr, + output_ptr, + eps, + inputs_embeds_stride, + previous_hidden_states_stride, + output_stride, + hidden_size: tl.constexpr, + block_size: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + input_idx = tl.program_id(1) + offsets = tl.arange(0, block_size) + mask = offsets < hidden_size + + if input_idx == 0: + position = tl.load(positions_ptr + token_idx) + values = tl.load( + inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets, + mask=mask & (position != 0), + other=0.0, + ) + weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0) + else: + values = tl.load( + previous_hidden_states_ptr + + token_idx * previous_hidden_states_stride + + offsets, + mask=mask, + other=0.0, + ) + weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0) + + output = _rms_norm(values, weight, eps, hidden_size) + tl.store( + output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets, + output, + mask=mask, + ) + + +def fused_mtp_input( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Mask and normalize both MTP inputs into the projection layout.""" + num_tokens, hidden_size = inputs_embeds.shape + output = torch.empty( + num_tokens, + 2 * hidden_size, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + if num_tokens == 0: + return output + + _fused_mtp_input_kernel[(num_tokens, 2)]( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + output, + eps, + inputs_embeds.stride(0), + previous_hidden_states.stride(0), + output.stride(0), + hidden_size, + triton.next_power_of_2(hidden_size), + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/__init__.py b/vllm/models/kimi_k3/nvidia/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py new file mode 100644 index 000000000000..64714496a156 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -0,0 +1,521 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3 dense MLA draft model for DSpark speculative decoding.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import vllm._custom_ops as ops +from vllm.config import VllmConfig +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + get_draft_quant_config, + maybe_prefix, +) +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.model import KimiMLP +from vllm.utils.torch_utils import is_quantized_kv_cache + + +class K3DSparkDecoderLayer(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + config, + layer_idx: int, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + quant_config = get_draft_quant_config(vllm_config) + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=maybe_prefix( + prefix, f"layers.{start_layer_id + layer_idx}.self_attn" + ), + use_rope=True, + non_causal_multi_token_decode=True, + ) + # Both row-parallel outputs stay un-reduced; their all-reduces are fused + # into the RMSNorm that follows via fused_allreduce_rms_norm. + self.self_attn.o_proj.reduce_results = False + self.mlp = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + # First layer: hidden_states is the (already reduced) embedding. + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # The MLP output is reduced by the next layer's input_layernorm (or by + # the model's final_norm). + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class K3DSparkModel(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.quant_config = get_draft_quant_config(vllm_config) + + # The frozen target embedding is aliased after the draft checkpoint loads. + self.embed_tokens: nn.Module | None = None + + self.context_proj = ReplicatedLinear( + self.config.target_hidden_size * self.config.num_target_layers, + self.config.hidden_size, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "context_proj"), + ) + self.context_norm = RMSNorm( + self.config.hidden_size, eps=self.config.rms_norm_eps + ) + + self.layers = nn.ModuleList( + [ + K3DSparkDecoderLayer( + vllm_config=vllm_config, + config=self.config, + layer_idx=layer_idx, + start_layer_id=start_layer_id, + prefix=prefix, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.markov_head = DSparkMarkovHead( + self.config.vocab_size, + self.config.draft_vocab_size, + self.config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + self._context_kv_fusion_available: bool | None = None + self._max_num_context_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + assert self.embed_tokens is not None + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.context_norm(self.context_proj(hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project target-derived context into each draft layer's latent cache.""" + if self._context_kv_fusion_available is None: + self._build_fused_context_kv_buffers() + if self._context_kv_fusion_available: + self._precompute_fused_context_kv( + context_states, context_positions, context_slot_mapping + ) + return + + # Quantized fallback. Directly invoking the projection modules preserves + # their quantization methods, at the cost of also computing unused Q rows. + for layer_idx, layer in enumerate(self.layers): + attn = layer.self_attn + assert attn.fused_qkv_a_proj is not None + assert attn.q_lora_rank is not None + assert attn.rotary_emb is not None + qkv_lora = attn.fused_qkv_a_proj(context_states)[0] + kv_lora = qkv_lora[..., attn.q_lora_rank :] + kv_c, k_pe = kv_lora.split( + [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1 + ) + kv_c = attn.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors. + # The vLLM CUDA op supports rotating one tensor in place and + # consumes the same (possibly scaled fp32) cos/sin cache. + rotary_emb = attn.rotary_emb + ops.rotary_embedding( + context_positions, + k_pe, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn.impl.do_kv_cache_update( + kv_c, + k_pe, + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _build_fused_context_kv_buffers(self) -> None: + """Build a cross-layer KV-only A projection after checkpoint loading.""" + if self.quant_config is not None: + self._context_kv_fusion_available = False + return + + attentions = [layer.self_attn for layer in self.layers] + if not attentions or any( + attn.fused_qkv_a_proj is None + or not hasattr(attn.fused_qkv_a_proj, "weight") + for attn in attentions + ): + self._context_kv_fusion_available = False + return + + attn0 = attentions[0] + assert attn0.q_lora_rank is not None + kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim + kv_weights = [] + for attn in attentions: + assert attn.q_lora_rank is not None + assert ( + attn.q_lora_rank == attn0.q_lora_rank + and attn.kv_lora_rank == attn0.kv_lora_rank + and attn.qk_rope_head_dim == attn0.qk_rope_head_dim + and attn.kv_a_layernorm.variance_epsilon + == attn0.kv_a_layernorm.variance_epsilon + ), "All MLA DSpark layers must share their latent KV geometry." + kv_weights.append( + attn.fused_qkv_a_proj.weight.detach().narrow( + 0, attn.q_lora_rank, kv_width + ) + ) + + # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused + # A weights are replicated (`disable_tp=True`), so this is valid on + # every TP rank without communication. + self._fused_context_kv_weight = torch.cat(kv_weights, dim=0) + self._context_kv_norm_weights = torch.stack( + [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 + ).contiguous() + self._num_context_layers = len(attentions) + self._context_kv_width = kv_width + self._context_kv_lora_rank = attn0.kv_lora_rank + self._context_rope_dim = attn0.qk_rope_head_dim + self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._context_positions_repeated = torch.empty( + self._num_context_layers * self._max_num_context_tokens, + dtype=torch.int64, + device=self._fused_context_kv_weight.device, + ) + self._context_kv_fusion_available = True + + def _precompute_fused_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + num_ctx = context_states.shape[0] + num_layers = self._num_context_layers + + # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects + # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs). + all_kv = F.linear(context_states, self._fused_context_kv_weight) + all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width) + all_kv_c = all_kv[..., : self._context_kv_lora_rank] + all_k_pe = all_kv[..., self._context_kv_lora_rank :] + + # Layer-major layout lets the 2-D RMSNorm weights select a distinct row + # for each draft layer in one grouped kernel. + all_kv_c = all_kv_c.permute(1, 0, 2).contiguous() + all_kv_c_normed = torch.empty_like(all_kv_c) + ops.rms_norm( + all_kv_c_normed, + all_kv_c, + self._context_kv_norm_weights, + self._context_rms_norm_eps, + ) + + all_k_pe = all_k_pe.permute(1, 0, 2).contiguous() + all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim) + repeated_positions = self._context_positions_repeated[: num_layers * num_ctx] + repeated_positions.view(num_layers, num_ctx).copy_(context_positions) + # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; + # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. + rotary_emb = self.layers[0].self_attn.rotary_emb + assert rotary_emb is not None + ops.rotary_embedding( + repeated_positions, + all_k_pe_flat, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) + + if context_slot_mapping is None: + return + + cache_layers = [layer.self_attn for layer in self.layers] + if ( + not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype) + and self._has_uniform_block_layout(cache_layers) + and ( + isinstance(context_slot_mapping, torch.Tensor) + or all(s is not None for s in context_slot_mapping) + ) + ): + # Grouped context KV insert only supports unquantized (bf16) KV cache + # and assumes that all layers share the same block layout. + + if isinstance(context_slot_mapping, (list, tuple)): + per_layer_slot_mappings = [ + s for s in context_slot_mapping if s is not None + ] + if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1: + # All rows alias to the same slot mapping. + slot_mapping = ( + per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1) + ) + else: + slot_mapping = torch.stack(per_layer_slot_mappings, dim=0) + else: + # Broadcast the single shared context_slot_mapping tensor. + slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1) + + ref_cache = cache_layers[0].kv_cache + ops.concat_and_cache_mla_grouped( + all_kv_c_normed, + all_k_pe.squeeze(2), + self._get_context_kv_cache_ptrs(cache_layers), + slot_mapping, + ref_cache.size(1), + ref_cache.stride(0), + ref_cache.stride(1), + ) + return + + for layer_idx, layer in enumerate(self.layers): + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn = layer.self_attn + attn.impl.do_kv_cache_update( + all_kv_c_normed[layer_idx], + all_k_pe[layer_idx], + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _has_uniform_block_layout( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> bool: + if not hasattr(self, "_layers_share_kv_block_layout"): + ref_cache = cache_layers[0].kv_cache + self._layers_share_kv_block_layout = all( + cl.kv_cache.size(1) == ref_cache.size(1) + and cl.kv_cache.stride(0) == ref_cache.stride(0) + and cl.kv_cache.stride(1) == ref_cache.stride(1) + for cl in cache_layers + ) + return self._layers_share_kv_block_layout + + def _get_context_kv_cache_ptrs( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> torch.Tensor: + # The per-layer KV cache base pointers are stable after allocation, so + # build the pointer array once and return it on every call. + if not hasattr(self, "_context_cache_ptrs"): + ref_cache = cache_layers[0].kv_cache + cache_ptrs = torch.tensor( + [cl.kv_cache.data_ptr() for cl in cache_layers], + dtype=torch.int64, + device=ref_cache.device, + ) + self._context_cache_ptrs = cache_ptrs + return self._context_cache_ptrs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + hidden_states = inputs_embeds + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.final_norm + ) + return hidden_states + + +class K3DSparkForCausalLM(nn.Module): + has_own_embed_tokens = False + has_own_lm_head = False + draft_id_to_target_id = None + checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head") + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"": "model."}, + orig_to_new_stacked={ + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + ".q_a_proj": (".fused_qkv_a_proj", 0), + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + target_layer_num = vllm_config.model_config.get_num_layers( + vllm_config.parallel_config + ) + self.model = K3DSparkModel( + vllm_config=vllm_config, + start_layer_id=target_layer_num, + prefix=maybe_prefix(prefix, "model"), + ) + + # Assigned by load_dspark_model from the target. Keeping no placeholder + # avoids a transient full-vocabulary allocation for this 163k-vocab model. + self.lm_head: nn.Module | None = None + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.draft_vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.layer_name for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.lm_head is not None + return self.logits_processor(self.lm_head, hidden_states) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # confidence_head is training-only. The frozen target embedding and LM + # head are shared after this draft-specific checkpoint is loaded. + loader = AutoWeightsLoader( + self, + skip_substrs=list(self.checkpoint_skip_substrs), + ) + loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.model._build_fused_context_kv_buffers() + return loaded_weights diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py new file mode 100644 index 000000000000..3bc3ab5de0bc --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -0,0 +1,777 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch +from einops import rearrange +from torch import nn +from torch.nn.parameter import Parameter + +from vllm import _custom_ops as ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig +from vllm.distributed import divide, get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, +) +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.v1.attention.backend import AttentionBackend + +logger = init_logger(__name__) + +_KDA_GATE_LOGBOUND_MIN = -5.0 + + +def a_log_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load KDA A_log stored as either old 4D or current 1D weights.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + + if loaded_weight.dim() == 4: + assert loaded_weight.shape[:2] == (1, 1), ( + f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}" + ) + assert loaded_weight.shape[-1] == 1, ( + f"Expected old A_log last dim to be 1, got {loaded_weight.shape}" + ) + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + return default_weight_loader(param, loaded_weight) + + return loader + + +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks.""" + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_id: int, + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_id = replicated_shard_id + output_sizes = output_sizes.copy() + output_sizes[replicated_shard_id] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + +def is_fused_kda_decode_supported( + num_heads: int, + head_dim: int, + conv_width: int, + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, +) -> bool: + if ( + num_heads not in (12, 24, 48, 96) + or head_dim != 128 + or conv_width != 4 + or num_spec != 0 + or input_dtype != torch.bfloat16 + or conv_state_dtype != torch.bfloat16 + or is_conv_state_dim_first() + or not hasattr(torch.ops._C, "fused_kda_decode") + ): + return False + # SM90 is architecture-specific; SM10x and SM12x use family binaries. + return ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) + + +def is_flashkda_supported( + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + if not current_platform.is_cuda(): + return False + capability = current_platform.get_device_capability() + return ( + capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + + +def _flashkda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + import vllm._flashkda_C # noqa: F401 + + out = torch.empty(v.shape, dtype=v.dtype, device=v.device) + final_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size( + q.shape[0] * q.shape[1], + q.shape[2], + cu_seqlens.numel() - 1, + ), + dtype=torch.uint8, + device=q.device, + ) + # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because + # FlashKDA materializes its transposed [H, T] layout internally. + # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead + # of transposing it to contiguous [H, T] storage internally. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + q.shape[-1] ** -0.5, + out, + workspace, + A_log.contiguous(), + dt_bias.view(-1, q.shape[-1]).contiguous(), + lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + ) + return out, final_state + + +def resolve_kda_prefill_backend( + backend: str, + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> str: + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + supported = is_flashkda_supported(head_dim, dtype, lower_bound) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, " + "head_dim=128, and a bounded KDA gate." + ) + if supported and backend != "triton": + logger.info_once("Using FlashKDA KDA prefill backend.") + return "flashkda" + return "triton" + + +def _make_decode_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, + decode_conv1d_weight: torch.Tensor | None, +) -> Callable[..., None]: + sharded_dims = [dim // tp_size for dim in dims] + + def weight_loader( + param: torch.Tensor, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + if loaded_weight.dim() == 2: + loaded_weight = loaded_weight.unsqueeze(1) + shard_size = sharded_dims[loaded_shard_id] + source_start = tp_rank * shard_size + target_start = sum(sharded_dims[:loaded_shard_id]) + loaded_shard = loaded_weight[source_start : source_start + shard_size] + param.data[target_start : target_start + shard_size].copy_(loaded_shard) + if decode_conv1d_weight is not None and not param.is_meta: + decode_conv1d_weight[loaded_shard_id].copy_( + loaded_shard.squeeze(1).transpose(0, 1) + ) + + return weight_loader + + +def _make_decode_norm_weight_loader( + decode_norm_weight: torch.Tensor, +) -> Callable[..., None]: + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + default_weight_loader(param, loaded_weight) + if not param.is_meta: + decode_norm_weight.copy_(param.data) + + return weight_loader + + +class KimiK3DeltaAttention(GatedDeltaNetAttention): + def get_attn_backend(self) -> type[AttentionBackend]: + return KimiK3KDAAttentionBackend + + def get_state_dtype( + self, + ) -> tuple[torch.dtype, torch.dtype]: + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + return MambaStateDtypeCalculator.kda_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + + def get_state_shape( + self, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + return MambaStateShapeCalculator.kda_state_shape( + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) + + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__(config, vllm_config, prefix) + + kda_config = config.linear_attn_config # type: ignore[attr-defined] + assert kda_config is not None, "linear_attn_config must be set" + self.head_dim = kda_config["head_dim"] + self.num_heads = kda_config["num_heads"] + assert self.num_heads % self.tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.tp_size) + self.projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(self.projection_size, self.tp_size) + self.conv_size = kda_config["short_conv_kernel_size"] + assert kda_config.get("use_full_rank_gate", False), ( + "KimiK3DeltaAttention requires a full-rank gate" + ) + + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. + qkvg_output_sizes = [self.projection_size] * 4 + in_proj_output_sizes = qkvg_output_sizes + [ + self.head_dim, + self.num_heads, + ] + local_output_size = ( + 4 * self.local_projection_size + self.head_dim + self.local_num_heads + ) + self.in_proj_padding = -local_output_size % 16 + if self.in_proj_padding: + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( + self.hidden_size, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvgfab", + ) + if self.in_proj_padding: + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.dt_bias = nn.Parameter( + torch.empty(self.local_projection_size, dtype=torch.float32) + ) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + # One packed parameter and cache let decode run a single conv update. + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. + self.conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=3 * self.projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.conv1d", + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + # Keep a width-major copy for fused decode without changing the layout + # consumed by the prefill and fallback decode kernels. + conv_state_dtype, _ = self.get_state_dtype() + decode_conv1d_weight = None + if is_fused_kda_decode_supported( + self.local_num_heads, + self.head_dim, + self.conv_size, + self.num_spec, + vllm_config.model_config.dtype, + conv_state_dtype, + ): + logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") + decode_conv1d_weight = torch.empty( + 3, + self.conv_size, + self.local_projection_size, + dtype=self.conv1d.weight.dtype, + device=self.conv1d.weight.device, + ) + self.register_buffer( + "decode_conv1d_weight", decode_conv1d_weight, persistent=False + ) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_decode_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + decode_conv1d_weight, + ) + }, + ) + + self.A_log = nn.Parameter( + torch.empty(self.local_num_heads, dtype=torch.float32) + ) + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) + if self.gate_lower_bound is not None: + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, ( + "KDA gate lower bound must be in " + f"[{_KDA_GATE_LOGBOUND_MIN}, 0). " + f"Got {self.gate_lower_bound}." + ) + + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + self.kda_prefill_backend = resolve_kda_prefill_backend( + backend, + self.head_dim, + vllm_config.model_config.dtype, + self.gate_lower_bound, + ) + + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") + decode_norm_weight = None + if decode_conv1d_weight is not None: + decode_norm_weight = torch.empty( + self.head_dim, + dtype=torch.float32, + device=self.o_norm.weight.device, + ) + self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False) + if decode_norm_weight is not None: + # Upcast once while loading; direct BF16 norm weights slow the + # fully fused decode kernel. + if hasattr(self.o_norm.weight, "weight_loader"): + delattr(self.o_norm.weight, "weight_loader") + set_weight_attrs( + self.o_norm.weight, + {"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)}, + ) + self.o_proj = RowParallelLinear( + self.projection_size, + self.hidden_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.o_proj", + ) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden_states.size(0) + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + + g1 = self.f_b_proj(f_a)[0] + beta = beta.unsqueeze(0) + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) + g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, + ) + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") + return self.o_proj(core_attn_out)[0] + + @eager_break_during_capture + def _forward( + self, + mixed_qkv: torch.Tensor, + g1: torch.Tensor, + g2: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if attn_metadata_raw is None: + return + + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + + assert isinstance(attn_metadata_raw, dict) + attn_metadata_narrowed = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata) + m = attn_metadata_narrowed + has_initial_state = m.has_initial_state + non_spec_query_start_loc = m.non_spec_query_start_loc + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor + spec_token_indx = m.spec_token_indx + non_spec_token_indx = m.non_spec_token_indx + spec_state_indices_tensor = m.spec_state_indices_tensor + spec_query_start_loc = m.spec_query_start_loc + num_accepted_tokens = m.num_accepted_tokens + num_actual_tokens = m.num_actual_tokens + has_spec_decode = m.num_spec_decodes > 0 + mixed_qkv = mixed_qkv[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + beta = beta[:, :num_actual_tokens] + + conv_state, recurrent_state = self.kv_cache + # The convolution kernels consume (..., dim, width - 1). + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + + if ( + self.decode_conv1d_weight is not None + and self.decode_norm_weight is not None + and not has_spec_decode + and m.num_prefills == 0 + and m.num_decodes > 0 + ): + assert non_spec_state_indices_tensor is not None + ops.fused_kda_decode( + x=mixed_qkv, + weight=self.decode_conv1d_weight, + bias=self.conv1d.bias, + conv_state=conv_state, + raw_g=g1, + raw_beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=non_spec_state_indices_tensor[:num_actual_tokens], + state=recurrent_state, + out=core_attn_out[:, :num_actual_tokens], + lower_bound=self.gate_lower_bound, + output_gate=g2[:num_actual_tokens], + norm_weight=self.decode_norm_weight, + norm_eps=self.o_norm.eps, + ) + return + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split( + self.local_projection_size, dim=0 + ) + q_conv_state, k_conv_state, v_conv_state = conv_state.split( + self.local_projection_size, dim=-2 + ) + + # Separate multi-query speculative tokens from prefill/plain decode. + if has_spec_decode: + if m.num_prefills == 0 and m.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + g1_spec, beta_spec = g1, beta + mixed_qkv_ns = g1_ns = beta_ns = None + else: + assert spec_token_indx is not None + assert non_spec_token_indx is not None + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + else: + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta + + # Spec-decode multi-query path. + core_attn_out_spec = None + if has_spec_decode: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] + spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_conv_out = torch.empty_like(mixed_qkv_spec) + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=spec_conv_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_max_query_len, + validate_data=False, + out=spec_conv_out, + ) + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1) + ) + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1] + spec_out = ( + core_attn_out[:, : q_spec.shape[1]] + if m.num_prefills == 0 and m.num_decodes == 0 + else None + ) + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) + + # Prefill or plain-decode path. + core_attn_out_non_spec = None + if mixed_qkv_ns is not None: + assert g1_ns is not None and beta_ns is not None + if m.num_prefills > 0: + q_ns, k_ns, v_ns = mixed_qkv_ns.split( + self.local_projection_size, dim=-1 + ) + + # Separate convolution calls accept row-strided packed inputs + # and produce dense Q/K/V without an additional V copy. + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns, k_ns, v_ns = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in (q_ns, k_ns, v_ns) + ) + + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + initial_state = gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) + if self.kda_prefill_backend == "flashkda": + assert self.gate_lower_bound is not None + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + g_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + ) + recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + else: + # Pure non-speculative decode. + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : mixed_qkv_ns.size(0) + ] + packed_conv_out = torch.empty_like(mixed_qkv_ns) + mixed_qkv_ns = causal_conv1d_update( + mixed_qkv_ns, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=decode_conv_indices, + validate_data=True, + out=packed_conv_out, + ) + ( + core_attn_out_non_spec, + _, + ) = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + state_indices=decode_conv_indices, + ) + + # Restore the scheduler's original token order for mixed batches. + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + elif core_attn_out_non_spec is not None: + # TODO: prefill and decode kernels write directly to core_attn_out + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ + 0, :num_actual_tokens + ] + else: + assert core_attn_out_spec is not None + # Triton normalizes in place, so this is a self-copy with no device + # work. Keep it for the out-of-place native implementation. + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py new file mode 100644 index 000000000000..0bb4864d2918 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -0,0 +1,496 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 specialization of GDN attention metadata. + +The request classification and cudagraph staging intentionally mirror +``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its +prefill KDA kernel internally, so this builder omits the shared FLA chunk +metadata construction. +""" + +from dataclasses import dataclass +from functools import cache + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + split_decodes_and_prefills, +) +from vllm.v1.kv_cache_interface import MambaSpec + + +@cache +def _metadata_launch_pdl() -> bool: + return current_platform.is_arch_support_pdl() + + +@triton.jit(do_not_specialize=["num_requests"]) +def _get_aligned_state_indices_kernel( + block_table_ptr, + seq_lens_ptr, + state_indices_ptr, + block_table_stride_0: tl.constexpr, + block_table_stride_1: tl.constexpr, + seq_lens_stride: tl.constexpr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + num_requests, + CACHE_BLOCK_SIZE: tl.constexpr, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + valid_row = rows < num_requests + seq_lens = tl.load( + seq_lens_ptr + rows * seq_lens_stride, + mask=valid_row, + other=1, + ) + # Triton truncates signed division toward zero, unlike PyTorch floor + # division. Clamping makes both semantics equivalent for seq_lens <= 0. + first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + block_table_ptr + + rows[:, None] * block_table_stride_0 + + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + tl.store( + state_indices_ptr + + rows[:, None] * state_indices_stride_0 + + state_slots[None, :] * state_indices_stride_1, + state_indices, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + + +def _mamba_get_block_table_tensor( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_cache_spec: MambaSpec, + mamba_cache_mode: str, +) -> torch.Tensor: + if mamba_cache_mode in ("all", "none"): + return block_table + + assert block_table.is_cuda and seq_lens.is_cuda + num_requests = block_table.shape[0] + num_state_slots = 1 + kv_cache_spec.num_speculative_blocks + state_indices = torch.empty( + (num_requests, num_state_slots), + dtype=block_table.dtype, + device=block_table.device, + ) + BLOCK_ROWS = 32 + grid = (triton.cdiv(num_requests, BLOCK_ROWS),) + _get_aligned_state_indices_kernel[grid]( + block_table, + seq_lens, + state_indices, + block_table.stride(0), + block_table.stride(1), + seq_lens.stride(0), + state_indices.stride(0), + state_indices.stride(1), + num_requests, + CACHE_BLOCK_SIZE=kv_cache_spec.block_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + return state_indices + + +@triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"]) +def _stage_spec_decode_metadata_kernel( + state_indices_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + staged_state_indices_ptr, + staged_query_start_loc_ptr, + staged_num_accepted_tokens_ptr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + staged_state_indices_stride_0: tl.constexpr, + staged_state_indices_stride_1: tl.constexpr, + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + NULL_STATE_ID: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + real_request = rows < num_spec_decodes + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + state_indices_ptr + + rows[:, None] * state_indices_stride_0 + + state_slots[None, :] * state_indices_stride_1, + mask=real_request[:, None] & valid_state_slot[None, :], + other=NULL_STATE_ID, + ) + tl.store( + staged_state_indices_ptr + + rows[:, None] * staged_state_indices_stride_0 + + state_slots[None, :] * staged_state_indices_stride_1, + state_indices, + mask=(rows < batch_size)[:, None] & valid_state_slot[None, :], + ) + + query_row = tl.minimum(rows, num_spec_decodes) + query_start_loc = tl.load( + query_start_loc_ptr + query_row, + mask=rows <= batch_size, + ) + tl.store( + staged_query_start_loc_ptr + rows, + query_start_loc, + mask=rows <= batch_size, + ) + + num_accepted_tokens = tl.load( + num_accepted_tokens_ptr + rows, + mask=real_request, + other=1, + ) + tl.store( + staged_num_accepted_tokens_ptr + rows, + num_accepted_tokens, + mask=rows < batch_size, + ) + + +def stage_spec_decode_metadata( + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + staged_state_indices: torch.Tensor, + staged_query_start_loc: torch.Tensor, + staged_num_accepted_tokens: torch.Tensor, + *, + num_spec_decodes: int, +) -> None: + """Stage speculative-decode metadata into CUDA-graph buffers.""" + assert state_indices.is_cuda + assert state_indices.ndim == 2 + batch_size, num_state_slots = staged_state_indices.shape + BLOCK_ROWS = 32 + grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),) + _stage_spec_decode_metadata_kernel[grid]( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + state_indices.stride(0), + state_indices.stride(1), + staged_state_indices.stride(0), + staged_state_indices.stride(1), + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + NULL_STATE_ID=NULL_BLOCK_ID, + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + + +@dataclass +class KimiK3KDAMetadata(GDNAttentionMetadata): + pass + + +class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> KimiK3KDAMetadata: + m = common_attn_metadata + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + assert isinstance(self.kv_cache_spec, MambaSpec) + # Equivalent PyTorch "align" path: + # start = ((seq_lens - 1) // block_size).clamp_(min=0) + # offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32) + # indices = (start[:, None] + offsets).to(torch.int64) + # block_table_tensor = torch.gather(block_table, 1, indices) + block_table_tensor = _mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + # A nonnegative entry identifies a spec request. If no draft token + # was scheduled, process the whole batch as non-spec instead. + if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + + if num_spec_decodes == 0: + # The runner orders ordinary decodes before prefills. + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills(m, decode_threshold=1) + ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = None + non_spec_state_indices_tensor = block_table_tensor[:, 0] + spec_query_start_loc = None + non_spec_query_start_loc = query_start_loc + non_spec_query_start_loc_cpu = query_start_loc_cpu + num_accepted_tokens = None + else: + assert spec_sequence_masks_cpu is not None + assert num_accepted_tokens is not None + query_lens_cpu = query_start_loc_cpu.diff() + num_query_tokens = query_start_loc_cpu[-1].item() + + # Exclude zero-length cudagraph padding from request-indexed + # non-spec metadata. + active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0) + non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu] + num_non_spec_requests = non_spec_query_lens_cpu.numel() + num_non_spec_tokens = non_spec_query_lens_cpu.sum().item() + + # Query length alone cannot distinguish a true decode from a + # one-token prefill chunk. Packed decode is safe only when every + # active non-spec request is a true, one-token decode. + assert m.is_prefilling is not None + assert m.is_prefilling.device.type == "cpu" + non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu] + use_prefill = torch.any( + non_spec_is_prefilling | (non_spec_query_lens_cpu != 1) + ).item() + if use_prefill: + num_prefills = num_non_spec_requests + num_prefill_tokens = num_non_spec_tokens + num_decodes = 0 + num_decode_tokens = 0 + else: + num_prefills = 0 + num_prefill_tokens = 0 + num_decodes = num_non_spec_requests + num_decode_tokens = num_non_spec_tokens + + num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens + + if num_prefills == 0 and num_decodes == 0: + spec_token_indx = None + non_spec_token_indx = None + # Real requests precede trailing cudagraph padding. + spec_state_indices_tensor = block_table_tensor[ + :num_spec_decodes, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = None + # Padding trails real requests, so this prefix already contains + # the correct cumulative token counts. + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] + else: + query_lens = query_start_loc.diff() + spec_sequence_masks_gpu = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device + ) + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks_gpu, + query_lens, + output_size=num_query_tokens, + ) + # Stable partitioning preserves request-local token order in + # both subgroup tensors. + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + # Spec requests carry one state slot per speculative step; + # non-spec requests use only their current state slot. + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = block_table_tensor[ + active_non_spec_mask_cpu, 0 + ] + + spec_query_lens = query_lens[spec_sequence_masks_cpu] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + spec_query_lens, + dim=0, + out=spec_query_start_loc[1:], + ) + if num_prefills > 0: + non_spec_query_lens = query_lens[active_non_spec_mask_cpu] + non_spec_query_start_loc = torch.zeros( + non_spec_query_lens.size(0) + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + non_spec_query_lens, + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + non_spec_query_lens_cpu.size(0) + 1, + dtype=torch.int32, + ) + torch.cumsum( + non_spec_query_lens_cpu, + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + else: + # Packed decode consumes one row per request and does not + # use cumulative sequence lengths. + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + + # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares + # its own chunk indices. Only causal-convolution metadata is needed here. + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + if num_prefills > 0: + has_initial_state = m.compute_num_computed_tokens() > 0 + if spec_sequence_masks_cpu is not None: + has_initial_state = has_initial_state[active_non_spec_mask_cpu] + assert non_spec_query_start_loc_cpu is not None + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + non_spec_query_start_loc_cpu, + device=query_start_loc.device, + ) + ) + else: + has_initial_state = None + + # Prepare per-request tensors for cudagraph replay. num_actual_tokens + # may be token-padded, while state/query/acceptance metadata is indexed + # by request. + batch_size = m.num_reqs + if ( + self.use_full_cuda_graph + and num_spec_decodes > 0 + and num_prefills == 0 + and num_decodes == 0 + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_bs + ): + # Equivalent PyTorch staging: + # state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID) + # qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1]) + # accepted[:N].copy_(accepted_src); accepted[N:].fill_(1) + stage_spec_decode_metadata( + state_indices=spec_state_indices_tensor, + query_start_loc=spec_query_start_loc, + num_accepted_tokens=num_accepted_tokens, + staged_state_indices=self.spec_state_indices_tensor[:batch_size], + staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1], + staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size], + num_spec_decodes=num_spec_decodes, + ) + + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_spec_decodes == 0 + and num_decodes <= self.decode_cudagraph_max_bs + ): + self.non_spec_state_indices_tensor[:num_decodes].copy_( + non_spec_state_indices_tensor, non_blocking=True + ) + self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_( + NULL_BLOCK_ID + ) + non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[ + :batch_size + ] + + return KimiK3KDAMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + has_initial_state=has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=None, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + ) + + +class KimiK3KDAAttentionBackend(GDNAttentionBackend): + @staticmethod + def get_name() -> str: + return "KIMI_K3_KDA" + + @staticmethod + def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]: + return KimiK3KDAMetadataBuilder diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py new file mode 100644 index 000000000000..64ea39b48759 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103. + +Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module +name plays no role. Each measured shape maps to a :class:`ProjectionSpec` +holding the winning backend per token count. The static part of the decision is +resolved once per module at install time into a small ``{M: call}`` plan, so the +per-forward path is a single dict lookup. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform + +Backend = Literal["cute", "dsv3_fused_a"] +# A resolved per-token-count call: the backend plus its CuTe config (None for +# dsv3, which needs no config). +ResolvedCall = tuple[Backend, SkinnyGemmConfig | None] + + +@dataclass(frozen=True, slots=True) +class ProjectionSpec: + n: int + k: int + dsv3_tokens: frozenset[int] = frozenset() + cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + residual_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + name: str = "" # optional debug label; never used for dispatch + + def cute_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.cute_configs).get(num_tokens) + + def residual_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.residual_configs).get(num_tokens) + + +def _cute( + num_tokens: int, + block_size: int, + outputs_per_block: int, + k_unroll: int, + vector_width: int = 8, +) -> SkinnyGemmConfig: + return SkinnyGemmConfig( + num_tokens, + block_size, + outputs_per_block, + k_unroll, + vector_width, + ) + + +_M1_TO_16 = frozenset(range(1, 17)) +_M1 = frozenset({1}) + +# Keyed by local (N, K). Where two projections share a shape (only 1536x7168: +# shared_gate_up_proj and mla_g_proj) the entry is unified. +KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = { + (1536, 128): ProjectionSpec(1536, 128, _M1_TO_16, name="f_b_proj"), + (3072, 128): ProjectionSpec(3072, 128, _M1_TO_16, name="f_b_proj"), + # 1536x7168 is shared by shared_gate_up_proj and mla_g_proj. dsv3 M1..16 is + # only crash-safe once the mla_g aux-stream/PDL capture fix lands (subtask + # task_7388aba1); the fallback if it cannot be fixed is dsv3_tokens=_M1. + (1536, 7168): ProjectionSpec( + 1536, 7168, _M1_TO_16, name="shared_gate_up_proj/mla_g_proj" + ), + (3072, 7168): ProjectionSpec( + 3072, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 3, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="shared_gate_up_proj", + ), + (2112, 7168): ProjectionSpec(2112, 7168, _M1_TO_16, name="fused_qkv_a_proj"), + (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"), + (4608, 1536): ProjectionSpec(4608, 1536, _M1_TO_16, name="q_b_proj"), + (3584, 7168): ProjectionSpec( + 3584, + 7168, + frozenset(range(2, 9)), + ((1, _cute(1, 224, 2, 4)),), + name="routed_expert_down_proj", + ), + (6288, 7168): ProjectionSpec( + 6288, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 64, 3, 2)), + (3, _cute(3, 32, 3, 4)), + (4, _cute(4, 128, 6, 1)), + ), + name="in_proj_qkvgfab", + ), + (12448, 7168): ProjectionSpec( + 12448, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + ), + name="in_proj_qkvgfab", + ), + (7168, 768): ProjectionSpec(7168, 768, _M1_TO_16, name="shared_down_proj"), + (7168, 1536): ProjectionSpec( + 7168, 1536, cute_configs=((1, _cute(1, 96, 4, 2)),), name="o_proj" + ), + (7168, 3072): ProjectionSpec( + 7168, + 3072, + cute_configs=( + (1, _cute(1, 96, 2, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="o_proj", + ), + (7168, 3584): ProjectionSpec( + 7168, + 3584, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + ), + residual_configs=( + (1, _cute(1, 64, 4, 2)), + (2, _cute(2, 64, 7, 2)), + (3, _cute(3, 64, 2, 1)), + (4, _cute(4, 64, 2, 1)), + ), + name="routed_expert_up_proj", + ), + (7168, 4224): ProjectionSpec( + 7168, + 4224, + cute_configs=((1, _cute(1, 96, 4, 2, 4)),), + name="dense_down_proj", + ), + (7168, 8448): ProjectionSpec( + 7168, + 8448, + cute_configs=( + (1, _cute(1, 32, 4, 4)), + (2, _cute(2, 96, 4, 1)), + (3, _cute(3, 96, 4, 1)), + ), + name="dense_down_proj", + ), + (8448, 7168): ProjectionSpec( + 8448, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (16896, 7168): ProjectionSpec( + 16896, + 7168, + cute_configs=( + (1, _cute(1, 224, 6, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (20480, 7168): ProjectionSpec( + 20480, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + (40960, 7168): ProjectionSpec( + 40960, + 7168, + cute_configs=( + (1, _cute(1, 128, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # TP16. Measured on B300 over M=1..16 with the same >=5% threshold as the + # entries above. The replicated projections (2112x7168, 3584x7168, + # 7168x3584) keep their shapes at TP16 and reuse the entries above, and + # o_proj lands on 7168x768, which shared_down_proj already covers. + (3216, 7168): ProjectionSpec( + 3216, + 7168, + # Both gaps in this range are measured, not oversights: dsv3 is only + # 4% ahead at M6..M8, and at M16 cuBLAS switches to a faster kernel + # (11.42us vs dsv3's 11.83us) after trailing it by 6-8% at M9..M15. + frozenset(range(9, 16)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 4, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="in_proj_qkvgfab", + ), + (768, 7168): ProjectionSpec( + 768, + 7168, + frozenset(range(5, 17)), + cute_configs=( + (1, _cute(1, 224, 2, 4)), + (2, _cute(2, 224, 2, 2)), + (3, _cute(3, 224, 2, 2)), + (4, _cute(4, 224, 2, 2)), + ), + name="mla_g_proj/shared_gate_up_proj", + ), + (1152, 1536): ProjectionSpec( + 1152, + 1536, + frozenset(range(2, 17)), + ((1, _cute(1, 192, 3, 4)),), + name="q_b_proj", + ), + (768, 128): ProjectionSpec(768, 128, _M1_TO_16, name="f_b_proj"), + # dsv3 drops under 5% from M9 on for this shape. + (7168, 384): ProjectionSpec( + 7168, 384, frozenset(range(1, 9)), name="shared_down_proj" + ), + (4224, 7168): ProjectionSpec( + 4224, + 7168, + frozenset(range(4, 9)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 2, 1)), + (3, _cute(3, 64, 2, 2)), + ), + name="dense_gate_up_proj", + ), + (10240, 7168): ProjectionSpec( + 10240, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 32, 2, 4)), + (3, _cute(3, 64, 4, 1)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # 7168x2112 (TP16 dense down_proj) has no entry on purpose: K=2112 divides + # none of the fused-A tile_k values, and the CuTe kernel is left with + # vector_width=2, which measured slower than cuBLAS. +} + + +def _backend_for( + spec: ProjectionSpec, num_tokens: int, has_residual: bool +) -> Backend | None: + if has_residual: + return "cute" if spec.residual_config(num_tokens) is not None else None + if spec.cute_config(num_tokens) is not None: + return "cute" + if num_tokens in spec.dsv3_tokens: + return "dsv3_fused_a" + return None + + +def select_kimi_k3_backend( + num_tokens: int, + n: int, + k: int, + *, + has_residual: bool = False, +) -> Backend | None: + """Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back.""" + spec = KIMI_K3_PROJECTIONS.get((n, k)) + return _backend_for(spec, num_tokens, has_residual) if spec is not None else None + + +def _build_plan(spec: ProjectionSpec) -> dict[int, ResolvedCall]: + plan: dict[int, ResolvedCall] = {} + for num_tokens in range(1, 17): + backend = _backend_for(spec, num_tokens, has_residual=False) + if backend == "cute": + plan[num_tokens] = ("cute", spec.cute_config(num_tokens)) + elif backend == "dsv3_fused_a": + plan[num_tokens] = ("dsv3_fused_a", None) + return plan + + +def _build_residual_plan(spec: ProjectionSpec) -> dict[int, SkinnyGemmConfig]: + return {num_tokens: config for num_tokens, config in spec.residual_configs} + + +def _is_sm103() -> bool: + return current_platform.is_device_capability((10, 3)) + + +def _is_packed_row_major(tensor: torch.Tensor) -> bool: + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) + + +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + _is_packed_row_major(x) + and _is_packed_row_major(weight) + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.is_cuda + and weight.is_cuda + and x.device == weight.device + and x.shape[1] == weight.shape[1] + ) + + +def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool: + return ( + residual.dim() == 2 + and residual.dtype == torch.bfloat16 + and residual.is_cuda + and residual.device == x.device + and residual.is_contiguous() + and residual.shape == (x.shape[0], weight.shape[0]) + ) + + +def _run_plan( + plan: dict[int, ResolvedCall], x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor | None: + entry = plan.get(x.shape[0]) + if entry is None: + return None + backend, config = entry + if backend == "cute": + if not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, None) + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + return None + output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) + ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True) + return output + + +def _run_residual_plan( + residual_plan: dict[int, SkinnyGemmConfig], + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> torch.Tensor | None: + config = residual_plan.get(x.shape[0]) + if config is None or not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, residual) + + +def try_low_latency_gemm( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor | None = None, +) -> torch.Tensor | None: + """Run the shape-selected low-latency kernel, or None to fall back. + + Resolves the plan from the shape table on each call; production installs a + precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not + use this path. + """ + if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight): + return None + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + return None + if residual is None: + return _run_plan(_build_plan(spec), x, weight) + if not _residual_ok(x, weight, residual): + return None + return _run_residual_plan(_build_residual_plan(spec), x, weight, residual) + + +class _KimiK3LowLatencyApply: + """Mixin: try the precomputed plan, else defer to the base method.""" + + def __init__(self, plan: dict[int, ResolvedCall]) -> None: + self._plan = plan + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if ( + bias is None + and not envs.VLLM_BATCH_INVARIANT + and _runtime_ok(x, layer.weight) + ): + output = _run_plan(self._plan, x, layer.weight) + if output is not None: + return output + return super().apply(layer, x, bias) # type: ignore[misc] + + +class KimiK3LowLatencyLinearMethod(_KimiK3LowLatencyApply, UnquantizedLinearMethod): + def __init__( + self, + plan: dict[int, ResolvedCall], + residual_plan: dict[int, SkinnyGemmConfig], + ) -> None: + super().__init__(plan) + self._residual_plan = residual_plan + + def apply_with_residual( + self, + layer: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + ) -> torch.Tensor: + if ( + not envs.VLLM_BATCH_INVARIANT + and _runtime_ok(x, layer.weight) + and _residual_ok(x, layer.weight, residual) + ): + output = _run_residual_plan(self._residual_plan, x, layer.weight, residual) + if output is not None: + return output + return torch.addmm(residual, x, layer.weight.t()) + + +class KimiK3LowLatencyEmbeddingMethod( + _KimiK3LowLatencyApply, UnquantizedEmbeddingMethod +): + pass + + +def enable_kimi_k3_low_latency_gemm( + module: nn.Module, + dtype: torch.dtype, +) -> None: + """Install shape-selected low-latency GEMMs and register CuTe warmups. + + Modules are matched purely by type, an exactly-unquantized method, and a + local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`. + """ + if dtype != torch.bfloat16 or not _is_sm103(): + return + + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + for child in module.modules(): + is_linear = ( + isinstance(child, LinearBase) + and type(child.quant_method) is UnquantizedLinearMethod + ) + # ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is + # the parent type, so isinstance already excludes it. + is_head = ( + isinstance(child, ParallelLMHead) + and type(child.quant_method) is UnquantizedEmbeddingMethod + ) + if not (is_linear or is_head): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2: + continue + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + continue + if is_linear: + child.quant_method = KimiK3LowLatencyLinearMethod( + _build_plan(spec), _build_residual_plan(spec) + ) + else: + child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec)) + # Warm up only the configs measured for this module's local (N, K) so a + # TP8 deployment does not compile TP4 configs and vice versa. + warmup_configs.update(config for _, config in spec.cute_configs) + residual_warmup_configs.update(config for _, config in spec.residual_configs) + + if shape_dynamic_skinny_gemm.is_available(): + if warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs) + if residual_warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs( + dtype, residual_warmup_configs, has_residual=True + ) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py new file mode 100644 index 000000000000..3334f3bbbb49 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -0,0 +1,764 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clean Multi-head Latent Attention for Kimi-K3 (NVIDIA). + +This is a self-contained MLA layer that owns the full attention path: + + hidden_states + -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) + -> explicit prefill / decode split + prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens + (+ chunked-context merge); dispatched by cache dtype + (bf16 / plain fp8 / fp8_ds_mla) + decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel + -> impl.forward_mqa -> W_UV up-proj (MQA) + -> optional output gate + -> o_proj + +Unlike ``MultiHeadLatentAttentionWrapper`` (which delegates orchestration to +``MLAAttention.forward``), this class *is* the ``AttentionLayerBase``: it selects +the backend, builds the impl, registers itself in the forward context, owns the +KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the +``DeepseekV4Attention`` structure. + +K3 specifics: optional rotary embedding (disabled for the target model's NoPE +layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``). + +Out of scope (extension points, not wired here): context parallelism (DCP/PCP), +sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. +""" + +import math +from typing import TYPE_CHECKING, cast + +import torch +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.attention import ( + _init_kv_cache_quant, + set_default_quant_scales, + should_load_quant_weights, +) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from vllm.model_executor.utils import replace_parameter +from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import ( + is_quantized_kv_cache, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionType, + MLAAttentionImpl, +) +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend +from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.attention.selector import get_attn_backend +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata + +logger = init_logger(__name__) + +# Below this many tokens, overlap the g_proj GEMM on the aux stream with the +# attention front-end (the GEMM is small and launch-bound, so the overlap +# hides it); at or above it, run the gate on the main stream. +_GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512 + + +@torch.compile(backend=current_platform.simple_compile_backend) +def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply the sigmoid output gate to a precomputed ``g_proj`` projection.""" + return attn_out * gate.sigmoid() + + +class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): + """Kimi-K3 Multi-head Latent Attention with optional RoPE and output gate.""" + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_output_gate: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + use_rope: bool = False, + non_causal_multi_token_decode: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.non_causal_multi_token_decode = non_causal_multi_token_decode + # Latent "head" seen by the attention kernel / KV cache. + self.head_size = kv_lora_rank + qk_rope_head_dim + self.scale = self.qk_head_dim**-0.5 + self.rms_norm_eps = config.rms_norm_eps + self.layer_name = prefix + + self.rotary_emb: RotaryEmbedding | None = None + if use_rope: + rope_parameters = dict(config.rope_parameters) + if rope_parameters["rope_type"] != "default": + rope_parameters["rope_type"] = ( + "deepseek_yarn" + if rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + dtype=torch.float32, + ) + if rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = rope_parameters.get("mscale_all_dim", False) + scaling_factor = rope_parameters["factor"] + mscale = ( + 1.0 + if scaling_factor <= 1 + else 0.1 * float(mscale_all_dim) * math.log(scaling_factor) + 1.0 + ) + self.scale *= mscale * mscale + # The fused epilogues read the cos/sin table directly in fp32 and run + # the RoPE math in fp32, so there is no per-forward dtype cast (and no + # precision loss). deepseek_yarn builds cos_sin_cache in fp32 already; + # dtype=torch.float32 above forces it for the default rope too (the + # DSpark draft, which has no yarn scaling). + assert self.rotary_emb.cos_sin_cache.dtype == torch.float32, ( + "K3 fused MLA RoPE requires an fp32 cos/sin cache; got " + f"{self.rotary_emb.cos_sin_cache.dtype}." + ) + + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_heads = num_heads + self.num_local_heads = num_heads // tp_size + + # ---- Pre-attention projections (fusable front-end) ---- + # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the + # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank`` + # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``. + if self.q_lora_rank is not None: + # Fused q-down + kv-down projection. Replicated (disable_tp) because + # the low-rank latents are shared across TP ranks; TP splitting + # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj`` + # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively. + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + # Uncompressed query: full-rank q_proj (TP-split over heads) plus a + # replicated kv-down projection (shared latent across TP ranks). + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + + # ---- Post-attention projections ---- + self.use_output_gate = use_output_gate + self.g_proj = ( + ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + if use_output_gate + else None + ) + # Aux stream (created at the model level, DeepseekV4 convention) for + # overlapping the g_proj GEMM with the attention front-end. None on + # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self._gate_events = ( + [torch.cuda.Event(), torch.cuda.Event()] + if self.g_proj is not None and current_platform.is_cuda_alike() + else None + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # ---- Attention backend / impl / KV cache ---- + self.quant_config = quant_config + if cache_config is not None: + self.kv_cache_dtype = cache_config.cache_dtype + else: + self.kv_cache_dtype = "auto" + + dtype = torch.get_default_dtype() + self.attn_backend = get_attn_backend( + self.head_size, + dtype, + self.kv_cache_dtype, + use_mla=True, + use_sparse=False, + num_heads=self.num_local_heads, + ) + _init_kv_cache_quant(self, quant_config, prefix) + # Unit (1.0) scale for the fused fp8 prefill path: q/k/v are cast + # unscaled to match forward_mha (the prefill flash path does not + # dequantize); only the cache uses _k_scale. + self.register_buffer( + "_one_scale", torch.ones(1, dtype=torch.float32), persistent=False + ) + + impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) + self.impl = impl_cls( # type: ignore[assignment] + num_heads=self.num_local_heads, + head_size=self.head_size, + scale=self.scale, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=self.kv_cache_dtype, + logits_soft_cap=None, + attn_type=AttentionType.DECODER, + kv_sharing_target_layer_name=None, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + qk_head_dim=self.qk_head_dim, + v_head_dim=self.v_head_dim, + kv_b_proj=self.kv_b_proj, + indexer=None, + ) + self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) + + vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + assert ( + parallel_config.decode_context_parallel_size <= 1 + and parallel_config.prefill_context_parallel_size <= 1 + ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism." + self.prefill_backend = get_mla_prefill_backend(vllm_config)( + num_heads=self.num_local_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + vllm_config=vllm_config, + ) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # ------------------------------------------------------------------ + # AttentionLayerBase interface + # ------------------------------------------------------------------ + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + kv_cache_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + return MLAAttentionSpec( # type: ignore[call-arg] + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=kv_cache_dtype, + cache_dtype_str=self.kv_cache_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + non_causal_multi_token_decode=self.non_causal_multi_token_decode, + ) + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights. + + ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank`` + latent. For the MQA decode path we pre-split it so that queries are + projected into latent space by ``W_UK_T`` and the attention output is + projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. + """ + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + assert kv_b_proj_weight.shape == ( + self.kv_lora_rank, + self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim), + ), f"{kv_b_proj_weight.shape=}" + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, + self.num_local_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + W_UK, W_UV = kv_b_proj_weight.split( + [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + # (L, N, V) -> (N, L, V) + replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) + # (L, N, P) -> (N, P, L) + replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) + + quant_method = ( + self.quant_config.get_quant_method(self, prefix=self.layer_name) + if self.quant_config + else None + ) + if not should_load_quant_weights(quant_method): + set_default_quant_scales(self, register_buffer=False) + + # Precompute reciprocal scales once here (scales are final after load; + # K3 has no runtime calculate_kv_scales path) so the fp8 fused kernels + # in the decode/prefill hot path take a ready inverse instead of + # launching a per-step reciprocal kernel. + self.register_buffer( + "_q_scale_inv", self._q_scale.reciprocal().reshape(1), persistent=False + ) + self.register_buffer( + "_k_scale_inv", self._k_scale.reciprocal().reshape(1), persistent=False + ) + + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor) -> None: + """Project latent attention output back to ``v`` via ``W_UV`` (bmm).""" + # (B, N, L) -> (N, B, L) + x = x.view(-1, self.num_local_heads, self.kv_lora_rank).transpose(0, 1) + out = out.view(-1, self.num_local_heads, self.v_head_dim) + # (N, B, L) x (N, L, V) -> (N, B, V) written transposed into (B, N, V) + torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) + + def _attn_read_kv_cache(self) -> torch.Tensor: + """Latent cache as seen by the attention read kernels (decode / context). + + A plain per-tensor fp8 cache is stored as ``uint8``; view it as fp8 so + the backend reads it as E4M3 rather than fp4/E2M1 -- the latter doubles + the perceived head dim (``head_size * 2``) and fails the kernel's + ``head_dim_k == head_dim_q`` check. Mirrors ``MLAAttention.forward``; + the fp8_ds_mla layout keeps its native uint8 view. + """ + cache = self.kv_cache + if ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + ): + return cache.view(current_platform.fp8_dtype()) + return cache + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + def _forward_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention. + + Returns the pre-gate attention output ``[num_tokens, + num_local_heads * v_head_dim]``. On a profile/dummy run + it returns a zeroed buffer. + """ + if self.q_lora_rank is not None: + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] + q_c, kv_c, k_pe = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + q_c, kv_c_normed = fused_q_kv_rmsnorm( + q_c, + kv_c, + self.q_a_layernorm.weight.data, + self.kv_a_layernorm.weight.data, + self.rms_norm_eps, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + # Uncompressed query: project directly (no q-LoRA, no q norm) and + # normalize only the kv latent. + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] + kv_c, k_pe = kv_lora.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv_c_normed = self.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + + attn_out = torch.empty( + (hidden_states.shape[0], self.num_local_heads * self.v_head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._attention(positions, q, kv_c_normed, k_pe, attn_out) + return attn_out + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Both branches produce (attn_out, gate); they differ only in whether + # the g_proj GEMM is overlapped on the aux stream. + g_proj = self.g_proj + events = self._gate_events + if ( + g_proj is not None + and events is not None + and self.aux_stream is not None + and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD + ): + attn_out, gate = maybe_execute_in_parallel( + lambda: self._forward_attn(positions, hidden_states), + lambda: g_proj(hidden_states)[0], + events[0], + events[1], + self.aux_stream, + ) + else: + attn_out = self._forward_attn(positions, hidden_states) + gate = g_proj(hidden_states)[0] if g_proj is not None else None + + if gate is not None: + attn_out = _gate_sigmoid_mul(attn_out, gate) + + # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a + # fresh private tensor, so return it directly rather than copying into a + # caller buffer -- the previous ``output[:] = ...`` convention forced an + # extra [num_tokens, hidden] copy per layer. + return self.o_proj(attn_out)[0] + + @eager_break_during_capture + def _attention( + self, + positions: torch.Tensor, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_by_layer = forward_context.attn_metadata + if attn_metadata_by_layer is None: + attn_out.zero_() + return + assert isinstance(attn_metadata_by_layer, dict) + attn_metadata = cast( + "MLACommonMetadata", attn_metadata_by_layer[self.layer_name] + ) + + num_actual_toks = attn_metadata.num_actual_tokens + slot_mapping_by_layer = forward_context.slot_mapping + assert isinstance(slot_mapping_by_layer, dict) + slot_mapping = slot_mapping_by_layer[self.layer_name] + + q = q[:num_actual_toks] + kv_c_normed = kv_c_normed[:num_actual_toks] + k_pe = k_pe[:num_actual_toks] + positions = positions[:num_actual_toks] + attn_out = attn_out[:num_actual_toks] + + cos_sin_cache = None + rope_positions = None + if self.rotary_emb is not None: + # Pass the fp32 cos/sin table straight to the fused epilogue (it reads + # fp32 and does the RoPE math in fp32) -- no per-forward dtype cast. + cos_sin_cache = self.rotary_emb.cos_sin_cache + rope_positions = positions + + # Decode tokens are laid out first, prefill tokens after. The fused + # prefill covers every supported config (bf16 / plain-fp8 / + # fp8_ds_mla), so there is no dense-MHA (forward_mha) fallback. + num_mqa_tokens = attn_metadata.num_decode_tokens + num_mha_tokens = q.size(0) - num_mqa_tokens + + # Both the prefill and decode fused epilogues write their own cache + # slice, so there is no separate do_kv_cache_update. + + # ---- Prefill: fused key-concat + cache-insert + attention ---- + if num_mha_tokens > 0: + self._forward_prefill_fused( + q[num_mqa_tokens:], + kv_c_normed[num_mqa_tokens:], + k_pe[num_mqa_tokens:], + rope_positions[num_mqa_tokens:] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[num_mqa_tokens:num_actual_toks], + attn_metadata, + attn_out[num_mqa_tokens:], + ) + + # ---- Decode: latent multi-query attention ---- + if num_mqa_tokens > 0: + mqa_q_nope, mqa_q_pe = q[:num_mqa_tokens].split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + # BMM1: absorb q_nope into latent space. (N,B,P) x (N,P,L) -> (B,N,L) + ql_nope = torch.bmm(mqa_q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) + # Fused: concat mqa_q = [ql_nope | q_pe] and insert the decode-token + # latent into the paged cache (one launch, right before forward_mqa). + mqa_q = self._decode_concat_cache( + ql_nope, + mqa_q_pe, + kv_c_normed[:num_mqa_tokens], + k_pe[:num_mqa_tokens], + rope_positions[:num_mqa_tokens] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[:num_mqa_tokens], + ) + latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q, self._attn_read_kv_cache(), attn_metadata, self + ) + self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens]) + + def _decode_concat_cache( + self, + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + ) -> torch.Tensor: + """Fused decode query-concat + latent cache insert, dispatched by cache + dtype (same policy as prefill: fp8 cache -> fp8 query).""" + if self.kv_cache_dtype == "fp8_ds_mla": + cache = self.kv_cache + if cache.dtype != torch.uint8: + cache = cache.view(torch.uint8) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + ds_mla=True, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + if is_quantized_kv_cache(self.kv_cache_dtype): + assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined] + "Kimi-K3 fp8 KV cache decode requires a backend that accepts an " + "fp8 (quantized) query input." + ) + cache = self.kv_cache + if cache.dtype != torch.float8_e4m3fn: + cache = cache.view(torch.float8_e4m3fn) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + q_scale_inv=self._q_scale_inv, + cache_scale_inv=self._k_scale_inv, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + self.kv_cache, + slot_mapping, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + + def _forward_prefill_fused( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + attn_metadata, + out: torch.Tensor, + ) -> None: + """Prefill using the fused key-concat + cache-insert kernel. + + Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one + fused kernel launch, dispatched by cache dtype. The chunked context + gather + online-softmax merge are delegated to the impl. + + Supported configs (K3 fp8 policy): + - bf16 cache -> bf16 prefill query + - plain fp8 cache -> fp8 prefill query (unscaled q/k/v; cache _k_scale) + - fp8_ds_mla cache -> bf16 prefill query (656B per-tile self-scaled) + """ + prefill = attn_metadata.prefill + has_context = prefill.chunked_context is not None + fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype() + + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + if self.kv_cache_dtype == "fp8_ds_mla": + # fp8_ds_mla cache (656B, per-tile self-scaled); bf16 attention. + assert not fp8_prefill, ( + "Kimi-K3 fp8_ds_mla uses a bf16 prefill query; fp8 prefill " + "query is not supported with fp8_ds_mla." + ) + kv_cache = self.kv_cache + if kv_cache.dtype != torch.uint8: + kv_cache = kv_cache.view(torch.uint8) + k = fused_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + elif is_quantized_kv_cache(self.kv_cache_dtype): + assert fp8_prefill, ( + "Kimi-K3 fp8 KV cache requires an fp8 prefill query; enable " + "--attention-config '{\"use_prefill_query_quantization\": true}'." + ) + # Plain per-tensor fp8: quant q/k/v (unscaled, matching forward_mha's + # unscaled `.to(fp8)`) and insert the fp8 latent (scaled by _k_scale). + kv_cache = self.kv_cache + if kv_cache.dtype != torch.float8_e4m3fn: + kv_cache = kv_cache.view(torch.float8_e4m3fn) + q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + kv_cache, + slot_mapping, + self._one_scale, + self._one_scale, + self._one_scale, + self._k_scale_inv, + positions, + cos_sin_cache, + ) + else: + # Concat full K = [k_nope | k_pe] and insert [kv_c_normed | k_pe] + # into the paged cache for these prefill tokens, in one launch. + k = fused_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + self.kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + + # When there is no chunked context, backends that honor `out` write the + # attention result straight into it, avoiding a slice+flatten+copy. + writes_out = not has_context and prefill.prefill_backend.supports_out() + output_prefill = prefill.prefill_backend.run_prefill_new_tokens( + q=q, + k=k, + v=v, + return_softmax_lse=has_context, + out=( + out.view(-1, self.num_local_heads, self.v_head_dim) + if writes_out + else None + ), + ) + + if has_context: + context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] + q, self._attn_read_kv_cache(), attn_metadata, self._k_scale + ) + suffix_output, suffix_lse = output_prefill + out = out.view(-1, self.num_local_heads, self.v_head_dim) + merge_attn_states( + output=out, + prefix_output=context_output[..., : self.v_head_dim], + prefix_lse=context_lse, + suffix_output=suffix_output[..., : self.v_head_dim], + suffix_lse=suffix_lse, + prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context, + ) + elif not writes_out: + out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2)) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py new file mode 100644 index 000000000000..68015827b851 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -0,0 +1,1939 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +import math +from collections.abc import Iterable +from typing import Any, cast + +import torch +from torch import nn + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + fused_grouped_topk, +) +from vllm.model_executor.layers.fused_moe.runner.latent_moe_runner import ( + LatentMoERunner, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention as KimiLinearGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +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.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsEagle3, + SupportsEncoderCudaGraph, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from vllm.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts +from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention +from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( + enable_kimi_k3_low_latency_gemm, +) +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.math_utils import cdiv +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream +from vllm.v1.worker.ubatching import dbo_current_ubatch_id + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) + +logger = init_logger(__name__) + +# Token-count cutoff for overlapping the MoE router gate with the routed-expert +# down projection on a separate CUDA stream (latent MoE). At or below this many +# tokens the launch-bound decode path benefits from multi-stream overlap; above +# it the GEMMs saturate the device and the cross-stream sync is pure overhead, +# so it falls back to sequential. +_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256 + + +def shard_sequence_parallel_mlp( + hidden_size: int, + intermediate_size: int, + use_sequence_parallel: bool, + eligible: bool, +) -> bool: + """Whether to TP-shard a sequence-parallel MLP instead of replicating it. + + Opt-in via ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT``; see :class:`KimiMLP` for + the trade-off and :mod:`vllm.envs` for when it is worth enabling. + """ + enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT + if not (use_sequence_parallel and eligible and enabled): + return False + tp_size = get_tensor_model_parallel_world_size() + return ( + tp_size > 1 and intermediate_size % tp_size == 0 and hidden_size % tp_size == 0 + ) + + +class KimiMLP(nn.Module): + """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel. + + Under sequence parallelism each rank owns a distinct slice of the tokens, so + by default both projections are replicated (``disable_tp``) and the block + needs no collective. That makes every rank stream the entire weight to serve + its own token shard. + + With ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT`` the weights are TP-sharded + instead. A rank then holds only a slice of the intermediate dim, so it + cannot finish its own tokens alone: ``forward`` all-gathers the full token + set, computes this rank's partial, and reduce-scatters. The reduce-scatter + sums across TP and restores the sequence sharding in one collective, so the + block still ends with one collective per direction. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + use_sequence_parallel: bool = False, + can_shard_sequence_parallel: bool = False, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.shard_sequence_parallel = shard_sequence_parallel_mlp( + hidden_size, + intermediate_size, + use_sequence_parallel, + can_shard_sequence_parallel, + ) + replicate = use_sequence_parallel and not self.shard_sequence_parallel + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=replicate, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + # Sharded sequence parallel reduces via the reduce-scatter in + # forward(), which also restores the sequence sharding. + reduce_results=False if self.shard_sequence_parallel else reduce_results, + disable_tp=replicate, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu and situ are supported." + ) + + def forward(self, x): + if self.shard_sequence_parallel: + # Each rank holds a weight shard but only its own tokens, so it + # cannot finish those tokens alone: gather the full token set, + # compute this rank's partial for all of them, then reduce-scatter, + # which sums across TP and restores the sequence sharding. + x = sp_all_gather(x) + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + if self.shard_sequence_parallel: + x = sp_reduce_scatter(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): + """Kimi K3 adapter for the DeepGEMM MegaMoE kernel.""" + + _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {} + _synchronized_ep_groups: set[tuple[int, int]] = set() + + def __init__( + self, + *args, + activation: str, + activation_beta: float | None, + activation_linear_beta: float | None, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.activation = activation + self.activation_beta = activation_beta + self.activation_linear_beta = activation_linear_beta + + def synchronize_first_launch(self) -> None: + ep_group = get_ep_group() + device = torch.accelerator.current_device_index() + key = (id(ep_group.cpu_group), device) + if key in self._synchronized_ep_groups: + return + torch.accelerator.synchronize() + torch.distributed.barrier(group=ep_group.cpu_group) + self._synchronized_ep_groups.add(key) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, 32), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, 32), + self.num_local_experts, + ) + self._transformed_l1_weights, self._transformed_l2_weights = ( + deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + activation=self.activation, + ) + ) + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + def get_symm_buffer(self): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.activation, + ) + symm_buffer = self._kimi_symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + activation=self.activation, + ) + self._kimi_symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + self.synchronize_first_launch() + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but its symmetric buffer supports {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] + + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, + ) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + ) + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + activation=self.activation, + activation_beta=self.activation_beta, + activation_linear_beta=self.activation_linear_beta, + fast_math=fast_math, + ) + return y + + +def make_kimi_k3_mega_moe_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + mapping = [] + for expert_id in range(num_experts): + for shard_id in ("w1", "w2", "w3"): + param_prefix = "w13" if shard_id in ("w1", "w3") else "w2" + for suffix in ("weight_packed", "weight_scale"): + param_suffix = "weight" if suffix == "weight_packed" else suffix + mapping.append( + ( + f"experts.{param_prefix}_{param_suffix}", + f"experts.{expert_id}.{shard_id}.{suffix}", + expert_id, + shard_id, + ) + ) + return mapping + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + use_sequence_parallel: bool = False, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = ( + routed_expert_hidden_size + if routed_expert_hidden_size is not None + else hidden_size + ) + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_renormalize = moe_renormalize + self.use_grouped_topk = config.use_grouped_topk + self.num_expert_group = config.num_expert_group + self.topk_group = config.topk_group + self.moe_router_activation_func = config.moe_router_activation_func + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "Kimi K3 MegaMoE requires expert parallel. Enable it with " + "--enable-expert-parallel." + ) + if self.use_mega_moe and config.hidden_act != "situ": + raise ValueError("Kimi K3 MegaMoE requires SITU activation.") + if self.use_mega_moe and not self.use_latent_moe: + raise ValueError("Kimi K3 MegaMoE requires latent MoE projections.") + if self.use_mega_moe and not self.use_grouped_topk: + raise ValueError("Kimi K3 MegaMoE requires grouped top-k routing.") + if self.use_mega_moe and (self.num_expert_group != 1 or self.topk_group != 1): + raise NotImplementedError( + "Kimi K3 MegaMoE currently requires one expert group." + ) + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr( + config, "min_moe_intermediate_per_partition", 256 + ) + if self.tp_size > 1 and not vllm_config.parallel_config.enable_expert_parallel: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = ( + min_moe_intermediate_per_partition * self.tp_size + ) + activation_situ_beta = ( + config.activation_situ_beta if config.hidden_act == "situ" else None + ) + activation_situ_linear_beta = ( + config.activation_situ_linear_beta if config.hidden_act == "situ" else None + ) + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=torch.float32) + ) + + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + use_sequence_parallel=use_sequence_parallel, + # Only the MegaMoE path calls the shared experts directly; the + # FusedMoE path below hands them to the runner, which fuses + # their reduction and assumes the replicated layout. + can_shard_sequence_parallel=self.use_mega_moe, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if self.latent_moe_use_norm + else None + ) + # Replicated up-proj: the full weight lives on every rank and + # produces the full hidden dim locally. This lets LatentMoERunner + # fuse the latent and shared reductions into a single all-reduce + # (concat the two partials, reduce once), then run the up-proj and + # shared add locally with no further collective. + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + # Auxiliary CUDA stream to overlap the router gate with the routed + # down projection on decode-sized batches (gated by + # _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD). + self._down_proj_stream: torch.cuda.Stream | None = aux_stream() + self._down_proj_events = (torch.cuda.Event(), torch.cuda.Event()) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + if self.use_mega_moe: + ep_group = get_ep_group() + ep_size = ep_group.world_size + ep_rank = ep_group.rank_in_group + if num_experts % ep_size != 0: + raise ValueError( + f"Kimi K3 num_experts={num_experts} must be divisible by " + f"EP size {ep_size}." + ) + num_local_experts = num_experts // ep_size + self.experts = KimiK3MegaMoEExperts( + vllm_config, + num_experts=num_experts, + num_local_experts=num_local_experts, + experts_start_idx=ep_rank * num_local_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + prefix=f"{prefix}.experts", + activation="situ", + activation_beta=activation_situ_beta, + activation_linear_beta=activation_situ_linear_beta, + ) + else: + # The tail-fusion kernels are tcgen05-based, so they require an + # SM100 NVIDIA device; the runner falls back to the default latent + # MoE path everywhere else. + enable_tail_fusion = ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) + self.experts = FusedMoEFactory( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + # Down projection runs outside MoERunner so it can overlap the + # router gate on the aux stream (see forward()); the original + # hidden states are passed to forward() as shared_experts_input + # so shared experts still see the untransformed input. + routed_input_transform=None, + routed_output_transform=self.routed_output_transform, + is_sequence_parallel=use_sequence_parallel, + runner_cls=LatentMoERunner if self.use_latent_moe else None, + runner_args=( + {"enable_k3_latent_moe_tail_fusion": enable_tail_fusion} + if self.use_latent_moe + else None + ), + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = getattr(self.experts, "w13_weight_packed", None) + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = getattr(self.experts, "w2_weight_packed", None) + if w13_weight is not None: + w13_weight.data.zero_() + if w2_weight is not None: + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = ( + moe_intermediate_size // self.tp_size + ) + + def _maybe_overlap_router_and_down_proj( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Compute the routed-expert down projection alongside the router, + overlapping them on separate CUDA streams when latent MoE is enabled. + + The router gate and the down projection both read ``hidden_states``, so + the gate runs on the default stream and the down projection on the aux + stream, joined via ``maybe_execute_in_parallel``. For MegaMoE the + grouped top-k selection consumes only the gate logits, so it also runs + on the default stream and overlaps the down projection. + + Returns: + ``(routed_hidden_states, router_output, topk_ids)``. + ``routed_hidden_states`` is the down-projected latent (or the + original ``hidden_states`` when latent MoE is disabled). For MegaMoE + ``router_output`` holds the grouped top-k weights and ``topk_ids`` + the selected experts; otherwise ``router_output`` holds the raw gate + logits and ``topk_ids`` is ``None``. + """ + + def _router( + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + router_logits, _ = self.gate(hidden_states) + if not self.use_mega_moe: + return router_logits, None + return fused_grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.experts.top_k, + renormalize=self.moe_renormalize, + e_score_correction_bias=self.gate.e_score_correction_bias.data, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + scoring_func=self.moe_router_activation_func, + routed_scaling_factor=self.routed_scaling_factor, + ) + + down_proj = self.routed_expert_down_proj + if down_proj is None: + router_output, topk_ids = _router(hidden_states) + return hidden_states, router_output, topk_ids + num_tokens = hidden_states.shape[0] + (router_output, topk_ids), (routed_hidden_states, _) = ( + maybe_execute_in_parallel( + lambda: _router(hidden_states), + lambda: down_proj(hidden_states), + self._down_proj_events[0], + self._down_proj_events[1], + self._down_proj_stream + if num_tokens <= _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD + else None, + ) + ) + return routed_hidden_states, router_output, topk_ids + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + # Overlap the gate with the routed down projection; the returned hidden + # states are already down-projected. Keep the original ``hidden_states`` + # for the shared experts. + routed_hidden_states, router_output, topk_ids = ( + self._maybe_overlap_router_and_down_proj(hidden_states) + ) + if self.use_mega_moe: + assert self.routed_output_transform is not None + assert topk_ids is not None + final_hidden_states = self.experts( + routed_hidden_states, + router_output, + topk_ids, + activation_clamp=None, + ) + final_hidden_states = self.routed_output_transform(final_hidden_states) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts( + hidden_states + ) + else: + # Routed experts consume the down-projected latent; shared experts + # (inside MoERunner) get the original hidden states via + # shared_experts_input. + final_hidden_states = self.experts( + hidden_states=routed_hidden_states, + router_logits=router_output, + shared_experts_input=hidden_states, + ) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + self.is_moe_layer = ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + if config.is_kda_layer(layer_idx): + kda_config = config.linear_attn_config + assert kda_config is not None + # This class also serves standalone Kimi-Linear through the model + # registry. Only Kimi-K3's full-rank gate uses the private KDA path. + if kda_config.get("use_full_rank_gate", False): + self.self_attn = KimiK3DeltaAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = False + else: + self.self_attn = KimiLinearGatedDeltaNetAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = True + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope, "Kimi-K3 MLA (MultiHeadLatentAttention) is NoPE-only" + # q_lora_rank may be None (Kimi-Linear): the MLA layer then uses an + # uncompressed q_proj instead of the fused q-LoRA front-end. + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_output_gate=bool(config.mla_use_output_gate), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + aux_stream=aux_stream, + ) + self._self_attn_writes_output = False + + if self.use_sequence_parallel: + self.self_attn.o_proj.reduce_results = False + + if self.is_moe_layer: + self.block_sparse_moe = KimiMoE( + config=config, + vllm_config=vllm_config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + use_sequence_parallel=self.use_sequence_parallel, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + use_sequence_parallel=self.use_sequence_parallel, + can_shard_sequence_parallel=True, + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_res = attn_res_block_size is not None + if self.use_attn_res: + assert attn_res_block_size is not None + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if self._self_attn_writes_output: + output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=output, + ) + return output + return self.self_attn( + hidden_states=hidden_states, + positions=positions, + ) + + def _pre_attn_norm( + self, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + assert hidden_states is not None + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + assert residual is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.self_attention_res_norm.weight, + self.self_attention_res_proj.weight.squeeze(0), + self.input_layernorm.weight, + num_blocks=self.prev_valid_blocks, + block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), + eps=self.self_attention_res_norm.variance_epsilon, + output_norm_eps=self.input_layernorm.variance_epsilon, + ) + return hidden_states, prefix_sum, residual + + def _post_attn_norm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + if self.is_block_write_layer: + prefix_sum = hidden_states + prefix_delta = None + else: + prefix_delta = hidden_states + mlp_valid_blocks = self.prev_valid_blocks + self.is_block_write_layer + hidden_states = attn_res( + prefix_sum, + prefix_delta, + residual, + self.mlp_res_norm.weight, + self.mlp_res_proj.weight.squeeze(0), + self.post_attention_layernorm.weight, + num_blocks=mlp_valid_blocks, + block_write_idx=-1, + eps=self.mlp_res_norm.variance_epsilon, + output_norm_eps=self.post_attention_layernorm.variance_epsilon, + ) + return hidden_states, prefix_sum, residual + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + hidden_states, prefix_sum, residual = self._pre_attn_norm( + hidden_states, residual, prefix_sum + ) + assert hidden_states is not None + + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states) + # Remove SP padding before attention. + hidden_states = hidden_states[: positions.shape[0]] + + # Attention. + hidden_states = self._run_self_attn(positions, hidden_states) + + if self.use_sequence_parallel: + # Add SP padding if needed, and then perform reduce scatter. + hidden_states = sp_reduce_scatter(hidden_states) + + hidden_states, prefix_sum, residual = self._post_attn_norm( + hidden_states, residual, prefix_sum + ) + + # MoE/MLP. + hidden_states = self.mlp(hidden_states) + return hidden_states, prefix_sum, residual + + +class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"], + "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"], + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + self.config = config + self.attn_res_block_size: int | None = config.attn_res_block_size + self.use_attn_res = self.attn_res_block_size is not None + parallel_config = vllm_config.parallel_config + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + # Aux stream for overlapping the MLA g_proj output-gate GEMM with the + # attention front-end (DeepseekV4 convention: created at the model + # level and threaded into each attention layer). + aux_stream = torch.cuda.Stream() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + vllm_config, + prefix, + aux_stream=aux_stream, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + self.num_attn_res_blocks = ( + cdiv(self.end_layer, self.attn_res_block_size) + if self.attn_res_block_size is not None + else 0 + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.use_attn_res: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if self.use_attn_res: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.use_attn_res: + assert self.attn_res_block_size is not None + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + assert hidden_states is not None + + aux_hidden_states: list[torch.Tensor] = [] + if self.start_layer in self.aux_hidden_state_layers: + if self.use_attn_res or residual is None: + aux_hidden_states.append(hidden_states) + else: + aux_hidden_states.append(hidden_states + residual) + + full_num_tokens = positions.shape[0] + if self.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + assert residual is None, "Currently, SP is not supported with PP" + + prefix_sum = None + if self.use_attn_res: + block_residual = hidden_states.new_empty( + hidden_states.size(0), + self.num_attn_res_blocks, + hidden_states.size(1), + ) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + prefix_sum = hidden_states + hidden_states = None + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + prefix_sum=prefix_sum, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + if self.use_attn_res: + assert prefix_sum is not None + aux_hidden_state = prefix_sum + hidden_states + else: + assert residual is not None + aux_hidden_state = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded aux hidden states. + # TODO: Optimize this. + aux_hidden_state = sp_all_gather(aux_hidden_state) + aux_hidden_state = aux_hidden_state[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + + assert hidden_states is not None + assert residual is not None + if not get_pp_group().is_last_rank: + assert not self.use_sequence_parallel, ( + "Currently, SP is not supported with PP" + ) + if prefix_sum is not None: + hidden_states = hidden_states + prefix_sum + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if self.use_attn_res: + assert prefix_sum is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.output_attn_res_norm.weight, + self.output_attn_res_proj.weight.squeeze(0), + None, + num_blocks=self.num_attn_res_blocks, + block_write_idx=-1, + eps=self.output_attn_res_norm.variance_epsilon, + output_norm_eps=0.0, + ) + else: + hidden_states = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded hidden states. + hidden_states = sp_all_gather(hidden_states) + hidden_states = hidden_states[:full_num_tokens] + + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[ + tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]] + ], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + use_mega_moe = any( + module.use_mega_moe + for module in self.modules() + if isinstance(module, KimiMoE) + ) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping( + self.config.num_experts + ) + elif self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any( + n.endswith("w13_weight_packed") for n in params_dict + ) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + def finalize_mega_moe_weights(self) -> None: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + +class KimiLinearForCausalLM( + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + enable_kimi_k3_low_latency_gemm(self, self.model_config.dtype) + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + return self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + loaded = loader.load_weights(weights) + self.model.finalize_mega_moe_weights() + # The fused MultiHeadLatentAttention's process_weights_after_loading + # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load + # hook for any AttentionLayerBase, so no manual trigger is needed here. + return loaded + + +def get_spec_layer_idx_from_weight_name( + config: KimiLinearConfig, weight_name: str +) -> int | None: + if hasattr(config, "num_nextn_predict_layers") and ( + config.num_nextn_predict_layers > 0 + ): + layer_idx = config.num_hidden_layers + for i in range(config.num_nextn_predict_layers): + # Match regardless of the surrounding prefix. The name may arrive as + # ``model.layers.{i}.``, a bare ``layers.{i}.`` (after AutoWeightsLoader + # has stripped the ``model.`` prefix in the main model), or with the + # multimodal ``language_model.model.layers.{i}.`` prefix. + if f"layers.{layer_idx + i}." in weight_name: + return layer_idx + i + return None + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + self.model_config = model_config + quant_config = vllm_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel( + config.vision_config.num_attention_heads + ) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(vllm_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to( + device=self.device, dtype=model_config.dtype + ) + + vision_attn = self.vision_tower.encoder.blocks[0].attn + if vision_attn.is_flash_attn_backend and vision_attn._fa_version == 4: + from vllm.models.kimi_k3.nvidia.ops.vision_fa4_warmup import ( + KimiK3VisionFA4WarmupConfig, + register_kimi_k3_vision_fa4_warmup, + ) + + merge_height, merge_width = config.vision_config.merge_kernel_size + mm_config = model_config.get_multimodal_config() + assert mm_config is not None + register_kimi_k3_vision_fa4_warmup( + KimiK3VisionFA4WarmupConfig( + num_heads=vision_attn.num_heads, + head_dim=vision_attn.head_size, + dtype=vision_attn.dtype, + max_batch_size=( + vllm_config.scheduler_config.max_num_seqs + * mm_config.get_limit_per_prompt("image") + ), + max_seqlen=( + vllm_config.scheduler_config.max_num_encoder_input_tokens + * merge_height + * merge_width + ), + ) + ) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to( + device=self.device, dtype=model_config.dtype + ) + + self.quant_config = quant_config + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + # -- SupportsEncoderCudaGraph protocol methods -- + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "pixel_values", + "pos_embeds", + "rope_freqs_cis", + "cu_seqlens", + "max_seqlen", + "sequence_lengths", + "merge_gather_idx", + ], + out_hidden_size=self.hidden_size, + ) + + def get_encoder_cudagraph_budget_range( + self, vllm_config: VllmConfig + ) -> tuple[int, int]: + min_budget = 64 + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return min_budget, max_budget + + @staticmethod + def _get_grid_thws(mm_kwargs: dict[str, Any]) -> list[list[int]]: + grid_thws = mm_kwargs["grid_thws"] + if not isinstance(grid_thws, list): + grid_thws = grid_thws.tolist() + return grid_thws + + @staticmethod + def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor: + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + pixel_values = torch.cat(pixel_values) + if pixel_values.ndim in (3, 5): + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], + *pixel_values.shape[2:], + ) + return pixel_values + + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + kh, kw = self.config.vision_config.merge_kernel_size + return [ + EncoderItemSpec( + input_size=t * h * w, + output_tokens=(h // kh) * (w // kw), + ) + for t, h, w in self._get_grid_thws(mm_kwargs) + ] + + def select_encoder_cudagraph_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + grid_thws = self._get_grid_thws(mm_kwargs) + pixel_values = self._get_pixel_values(mm_kwargs) + source_grid = mm_kwargs["grid_thws"] + + if not indices: + empty_grid = ( + source_grid[:0] if isinstance(source_grid, torch.Tensor) else [] + ) + return {"pixel_values": pixel_values[:0], "grid_thws": empty_grid} + + patch_counts = [t * h * w for t, h, w in grid_thws] + offsets = [0] + for count in patch_counts: + offsets.append(offsets[-1] + count) + selected_pixel_values = torch.cat( + [pixel_values[offsets[i] : offsets[i + 1]] for i in indices] + ) + grid_device = ( + source_grid.device if isinstance(source_grid, torch.Tensor) else None + ) + selected_grid = torch.tensor( + [grid_thws[i] for i in indices], + dtype=torch.long, + device=grid_device, + ) + return {"pixel_values": selected_pixel_values, "grid_thws": selected_grid} + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + kh, kw = self.config.vision_config.merge_kernel_size + per_item_output = (token_budget + max_batch_size - 1) // max_batch_size + rope = self.vision_tower.encoder.rope_2d + max_output_width = rope.max_width // kw + max_output_height = rope.max_height // kh + output_width = min(math.ceil(math.sqrt(per_item_output)), max_output_width) + output_height = (per_item_output + output_width - 1) // output_width + if output_height > max_output_height: + output_height = max_output_height + output_width = (per_item_output + output_height - 1) // output_height + if output_width > max_output_width: + raise ValueError( + f"Encoder CUDA graph budget {token_budget} exceeds K3 RoPE " + f"capacity for max_batch_size={max_batch_size}" + ) + grid_thws = [ + [1, output_height * kh, output_width * kw] for _ in range(max_batch_size) + ] + + patch_size: int | tuple[int, int] = self.config.vision_config.patch_size + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + total_patches = sum(t * h * w for t, h, w in grid_thws) + pixel_values = torch.randn( + total_patches, + 3, + patch_size[0], + patch_size[1], + device=device, + dtype=dtype, + ) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + grid_thws, + max_batch_size=max_batch_size, + max_seqlen_override=max( + token_budget * kh * kw, + max(t * h * w for t, h, w in grid_thws), + ), + device=device, + ) + return EncoderCudaGraphCaptureInputs( + values=metadata | {"pixel_values": pixel_values} + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + pixel_values = self._get_pixel_values(mm_kwargs) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + self._get_grid_thws(mm_kwargs), + max_batch_size=max_batch_size, + device=pixel_values.device, + ) + return EncoderCudaGraphReplayBuffers( + values=metadata | {"pixel_values": pixel_values} + ) + + def _project_encoder_features(self, image_features: torch.Tensor) -> torch.Tensor: + projector_dtype = next(self.mm_projector.parameters()).dtype + if image_features.dtype != projector_dtype: + image_features = image_features.to(projector_dtype) + output = self.mm_projector(image_features) + return output.reshape(-1, output.shape[-1]) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + pixel_values = values.pop("pixel_values") + image_features = self.vision_tower(pixel_values, None, encoder_metadata=values) + return self._project_encoder_features(image_features) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + image_features = self.vision_tower( + self._get_pixel_values(mm_kwargs).to( + next(self.vision_tower.parameters()).dtype + ), + self._get_grid_thws(mm_kwargs), + ) + return self._project_encoder_features(torch.cat(image_features)) + + def _maybe_ignore_quant_config( + self, quant_config: QuantizationConfig | None + ) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input( + self, **kwargs: object + ) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError( + "pixel_values must be a tensor or a list of tensors, " + f"got {type(pixel_values)}" + ) + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] + ) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), ( + f"expect grid_thws to be a tensor, got {type(grid_thws)}" + ) + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( + f"unexpected shape for grid_thws: {grid_thws.shape}" + ) + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input( + self, media_input: KimiK25MediaPixelInputs + ) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if intermediate_tensors is not None: + inputs_embeds = None + return self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs + ) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( + batch_size + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/kimi_k3/nvidia/mtp.py b/vllm/models/kimi_k3/nvidia/mtp.py new file mode 100644 index 000000000000..45cbe3767e90 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/mtp.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +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.quantization import QuantizationConfig +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.utils import get_pp_missing_layer_names, maybe_prefix +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .low_latency_gemm import enable_kimi_k3_low_latency_gemm +from .model import ( + KimiDecoderLayer, + KimiMoE, + get_spec_layer_idx_from_weight_name, + make_kimi_k3_mega_moe_expert_params_mapping, +) + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_config.quant_config + + 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.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + # The MTP block starts without the base model's AttnRes state. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + # Own aux stream for the MLA g_proj output-gate overlap (DeepseekV4 + # convention: the MTP block creates its own stream). + aux_stream = torch.cuda.Stream() + self.mtp_block = KimiDecoderLayer( + block_config, vllm_config, prefix=prefix, aux_stream=aux_stream + ) + + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + if self.mtp_block.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + + hidden_states, _, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + if self.mtp_block.use_sequence_parallel: + assert residual is not None + hidden_states = hidden_states + residual + hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] + logits_hidden_states = self.shared_head.norm(hidden_states) + return logits_hidden_states, hidden_states + + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm( + hidden_states, residual + ) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer( + config, vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + 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, + ) -> tuple[torch.Tensor, 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, + ) -> 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)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + enable_kimi_k3_low_latency_gemm(self, vllm_config.model_config.dtype) + + 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, + ) -> tuple[torch.Tensor, 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 | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + use_mega_moe = any( + module.use_mega_moe + for module in self.modules() + if isinstance(module, KimiMoE) + ) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping( + self.config.num_experts + ) + elif self.config.is_moe: + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any( + n.endswith("w13_weight_packed") for n in params_dict + ) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and ( + ".layers" not in name + ): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_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_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + if use_mega_moe: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/vllm/models/kimi_k3/nvidia/ops/__init__.py b/vllm/models/kimi_k3/nvidia/ops/__init__.py new file mode 100644 index 000000000000..bbaee887ba22 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .attn_res import attn_res + +__all__ = ["attn_res"] diff --git a/vllm/models/kimi_k3/nvidia/ops/attn_res.py b/vllm/models/kimi_k3/nvidia/ops/attn_res.py new file mode 100644 index 000000000000..01078d6c01d0 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/attn_res.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + + +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_attn_res_triton_warmup_profiles( + max_blocks: int, +) -> tuple[tuple[int, bool, int, bool], ...]: + """Return the small-batch profiles that bypass the native kernel.""" + profiles = [ + (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1) + ] + profiles.extend( + (block_write_idx, True, block_write_idx, True) + for block_write_idx in range(2, max_blocks) + ) + profiles.append((max_blocks, True, -1, False)) + return tuple(profiles) + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + delta_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_norm_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_delta_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + block_write_idx: tl.constexpr, + eps: tl.constexpr, + output_norm_eps: tl.constexpr, + HAS_DELTA: tl.constexpr, + WRITE_BLOCK: tl.constexpr, + APPLY_OUTPUT_NORM: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, + launch_pdl: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + updated_prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + if HAS_DELTA: + delta = tl.load( + delta_ptr + row_idx * stride_delta_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + updated_prefix += delta + # Match the BF16 prefix-add result before using it as a residual source. + updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32) + tl.store( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + updated_prefix, + mask=d_mask, + ) + if WRITE_BLOCK: + tl.store( + blocks_ptr + + row_idx * stride_block_m + + block_write_idx * stride_block_r + + d_offsets, + updated_prefix, + mask=d_mask, + ) + # With only the prefix source, the AttnRes softmax is exactly one. + if num_blocks == 0: + mixed = updated_prefix + else: + # Reloading avoids keeping the full prefix vector live across the loop. + if HAS_DELTA: + tl.debug_barrier() + input_qk_weight = tl.load( + norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) * tl.load( + qk_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + + num_sources = num_blocks + 1 + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + prefix_ptrs = ( + prefix_ptr + + row_idx * stride_prefix_m + + source_offsets[:, None] * 0 + + d_offsets[None, :] + ) + value_ptrs = tl.where(is_prefix[:, None], prefix_ptrs, block_ptrs) + values = tl.load( + value_ptrs, + mask=source_mask[:, None] & d_mask[None, :], + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + mixed /= denominator + output = mixed + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + if APPLY_OUTPUT_NORM: + output_reciprocal_std = tl.rsqrt( + tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size) + + output_norm_eps + ) + output_norm_weight = tl.load( + output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + output = mixed * output_reciprocal_std * output_norm_weight + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, + block_write_idx: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert prefix.stride(-1) == 1 + assert delta is None or delta.stride(-1) == 1 + assert blocks.stride(-1) == 1 + assert norm_weight.stride(-1) == 1 + assert qk_weight.stride(-1) == 1 + assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 + # The in-tree NVIDIA kernel covers the common fused-add + output-norm path; + # Triton handles block boundaries and final pre-norm output. + if ( + hidden_size == 7168 + and delta is not None + and output_norm_weight is not None + and num_blocks > 0 + and block_write_idx < 0 + and current_platform.is_device_capability_family(100) + ): + return ops.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + eps, + output_norm_eps, + ) + output = prefix.new_empty(prefix.shape) + # Tuned on GB300: source tiling helps decode, while one-source tiles scale + # better for prefill. + # Keep get_attn_res_triton_warmup_profiles in sync with these fallbacks. + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + prefix.stride(0), + 0 if delta is None else delta.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + block_write_idx, + eps, + output_norm_eps, + HAS_DELTA=delta is not None, + WRITE_BLOCK=block_write_idx >= 0, + APPLY_OUTPUT_NORM=output_norm_weight is not None, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py new file mode 100644 index 000000000000..7eda574b83e1 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CuTe DSL kernels for KimiK3LatentMoETailOp.""" + +from .allreduce_rmsnorm_reduce_scatter_early_exit import CollectiveKernel +from .fused_add_multicast_gemm import AdaptiveUpProjectionKernel +from .lamport_copy import LamportCopyKernel + +__all__ = [ + "AdaptiveUpProjectionKernel", + "CollectiveKernel", + "LamportCopyKernel", +] diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py new file mode 100644 index 000000000000..4f33f8920645 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -0,0 +1,1012 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Specialized from FlashInfer's oneshotAllreduceFusionKernel. + +"""Routed AllReduce/RMSNorm with CTA-specialized ReduceScatter early exit.""" + +from __future__ import annotations + +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from .primitives import ( + NUM_LAMPORT_BUFFERS, + PACKED_BYTES, + VEC_BF16, + bf16x8_to_packed_u32x4, + block_sum_specialized, + fragment_is_dirty, + load_global_u32x4, + load_volatile_u32, + map_shared_to_peer, + packed_u32x4_to_bf16x8, + red_async_release_gpu_add_u32, + sanitize_negative_zero, + store_global_u32x4, + store_lamport_sentinel_128, + store_shared_cluster_f32, + to_cute, + to_cute_dynamic_m, +) + + +def _mapping( + tp_size: int, + latent_dim: int, + hidden_dim: int, +) -> tuple[int, int, int, int]: + shard_dim = hidden_dim // tp_size + cluster_ctas = latent_dim // shard_dim + threads = shard_dim // VEC_BF16 + shared_roles = tp_size // cluster_ctas + return shard_dim, cluster_ctas, threads, shared_roles + + +def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None: + """Validate constraints imposed by the fused collective mapping.""" + + if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0: + raise ValueError("collective dimensions must be positive") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by tp_size") + shard, cluster, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + if shard % VEC_BF16: + raise ValueError("hidden_dim / tp_size must be divisible by 8") + if latent_dim % shard: + raise ValueError("latent_dim must be an integer multiple of shard_dim") + if cluster > 16 or cluster & (cluster - 1): + raise ValueError("collective cluster width must be a power of two <= 16") + if tp_size % cluster: + raise ValueError("tp_size must be divisible by collective cluster width") + if not 32 <= threads <= 1024: + raise ValueError("collective threads per CTA must be in [32, 1024]") + if threads < cluster: + raise ValueError("collective threads must cover every cluster CTA") + + +def _select_routed_schedule( + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, +) -> tuple[int, int]: + """Match upstream MNNVL's one-cluster-per-token occupancy policy.""" + + _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + sm_count = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + while max_m * cluster_ctas > sm_count and cluster_ctas > 1 and threads <= 512: + cluster_ctas //= 2 + threads *= 2 + return threads, cluster_ctas + + +class AllReduceRMSNormWithReduceScatterEarlyExit: + """One routed role plus one ReduceScatter role per destination group.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool = False, + include_reduce_scatter: bool = True, + include_routed: bool = True, + ): + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + if not 0 <= rank < tp_size: + raise ValueError(f"rank must be in [0,{tp_size}), got {rank}") + if not include_routed and not include_reduce_scatter: + raise ValueError("at least one collective role must be enabled") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + ( + self.shard_dim, + mapped_cluster, + mapped_threads, + shared_roles, + ) = _mapping(tp_size, latent_dim, hidden_dim) + if include_reduce_scatter: + self.cluster_ctas = mapped_cluster + self.threads = mapped_threads + else: + self.threads, self.cluster_ctas = _select_routed_schedule( + tp_size, latent_dim, hidden_dim, max_m + ) + # Diagnostic specialization: a one-role grid executes only the routed + # AllReduce/RMSNorm path. Keep this compile-time so the production + # fused path is unchanged when include_reduce_scatter=True. + if include_routed and include_reduce_scatter: + self.roles = 1 + shared_roles + elif include_routed: + self.roles = 1 + else: + self.roles = shared_roles + self.warps = (self.threads + 31) // 32 + self.last_warp_lanes = self.threads - (self.warps - 1) * 32 + self.last_warp_mask = (1 << self.last_warp_lanes) - 1 + # The upstream MNNVL oneshot protocol assigns one cluster to each + # token. Keep that exact ownership in the routed-only diagnostic; + # reusing a cluster for multiple token waves allows the Lamport + # generation metadata to change between waves. + self.token_ctas = ( + min(max_m, max_token_ctas) if include_reduce_scatter else max_m + ) + self.fp32_internal = fp32_internal + self.include_reduce_scatter = include_reduce_scatter + self.include_routed = include_routed + + @cute.jit + def __call__( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + stream: cuda.CUstream, + ): + self.kernel( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + ).launch( + grid=(self.token_ctas, self.cluster_ctas, self.roles), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_ctas, 1), + smem=(self.warps + self.cluster_ctas) * 4, + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + token_cta, cta_y, role = cute.arch.block_idx() + logical_role = role + if cutlass.const_expr(not self.include_routed): + # A shared-only grid starts at z=0, while _token_device reserves + # logical role 0 for the routed collective. + logical_role = role + Int32(1) + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + + cute.arch.griddepcontrol_wait() + token = token_cta + while token < m: + self._token_device( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + token, + token_cta, + cta_y, + logical_role, + cluster_rank, + tidx, + ) + token = token + self.token_ctas + + @cute.jit + def _token_device( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + token: Int32, + token_cta: Int32, + cta_y: Int32, + role: Int32, + cluster_rank: Int32, + tidx: Int32, + ): + if role == 0: + # ---------------- routed AllReduce + RMSNorm ---------------- + packed_idx = cluster_rank * self.threads + tidx + element_offset = ( + Int64(token) * self.latent_dim + Int64(packed_idx) * VEC_BF16 + ) + current_index = cute.arch.load((latent_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((latent_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load( + (latent_flags.iterator + 2).llvm_ptr, Uint32 + ) + dirty_num_stages = cute.arch.load( + (latent_flags.iterator + 3).llvm_ptr, Uint32 + ) + bytes_to_clear = cute.arch.load( + (latent_flags.iterator + 4).llvm_ptr, Uint32 + ) + current_elements = Int64(current_index) * ( + Int64(bytes_per_buffer) // Int64(2) + ) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + local_ptr = cute.make_ptr( + BFloat16, + (latent_source.iterator + element_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero( + load_global_u32x4(local_ptr, volatile=False) + ) + multicast_offset = ( + Int64(current_index) * Int64(bytes_per_buffer) + + ( + (Int64(token) * self.tp_size + self.rank) * self.latent_dim + + Int64(packed_idx) * VEC_BF16 + ) + * 2 + ) + store_global_u32x4( + latent_multicast_ptr + multicast_offset, + local_packed, + volatile=False, + ) + + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(latent_flags.iterator + 8, Uint32(1)) + + global_tid = ( + Int64(token) * self.cluster_ctas + Int64(cta_y) + ) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.cluster_ctas * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + ( + routed_workspace.iterator + + dirty_elements + + clear_idx * VEC_BF16 + ).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + rank_words = cute.make_rmem_tensor( + cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32 + ) + for word in cutlass.range_constexpr(4): + rank_words[self.rank, word] = local_packed[word] + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + if cutlass.const_expr(source_rank != self.rank): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) + * self.latent_dim + + Int64(packed_idx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (routed_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8( + rank_words[source_rank, None].load() + ).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + + # Preserve the original early PDL point before RMSNorm. + cute.arch.griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.fp32_internal): + # High-precision fused mode: retain the rank reduction in + # FP32 through the RMS square and row reduction. + norm_input = accum.load() + norm_square = norm_input * norm_input + else: + norm_input_bf16 = accum.load().to(BFloat16) + norm_input = norm_input_bf16.to(Float32) + # Upstream-compatible mode: FlashInfer evaluates BF16 * + # BF16 first, then promotes the rounded square to FP32. + norm_square = (norm_input_bf16 * norm_input_bf16).to(Float32) + thread_sum = norm_square.reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_tensor( + Float32, cute.make_layout((self.warps,)), byte_alignment=4 + ) + cluster_sums = smem.allocate_tensor( + Float32, + cute.make_layout((self.cluster_ctas,)), + byte_alignment=4, + ) + block_sum = block_sum_specialized( + thread_sum, + warp_sums, + tidx, + self.warps, + self.last_warp_lanes, + self.last_warp_mask, + ) + if tidx < self.cluster_ctas: + local_slot = cluster_sums.iterator + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, block_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_ctas): + full_sum = full_sum + cluster_sums[peer] + inv_rms = cute.math.rsqrt( + full_sum / Float32(self.latent_dim) + epsilon, fastmath=True + ) + gamma_ptr = cute.make_ptr( + BFloat16, + (gamma.iterator + Int64(packed_idx) * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_values = packed_u32x4_to_bf16x8( + load_global_u32x4(gamma_ptr, volatile=False) + ) + result = (norm_input * inv_rms * gamma_values.to(Float32)).to(BFloat16) + store_global_u32x4( + Int64((latent_output.iterator + element_offset).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + # The x=0 CTA rotates only after reaching its final token wave. + # Waiting for all M arrivals then guarantees every token-wave CTA + # loaded the current generation before the metadata is advanced. + if ( + token_cta == 0 + and token + self.token_ctas >= m + and cta_y == 0 + and tidx == 0 + ): + access_counter = latent_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + while arrived < Uint32(m): + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2) + cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (latent_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((latent_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((latent_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (latent_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + else: + # ---------------- shared ReduceScatter ---------------- + shared_group = role - 1 + destination = shared_group * self.cluster_ctas + cluster_rank + current_index = cute.arch.load((shared_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((shared_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load( + (shared_flags.iterator + 2).llvm_ptr, Uint32 + ) + dirty_num_stages = cute.arch.load( + (shared_flags.iterator + 3).llvm_ptr, Uint32 + ) + bytes_to_clear = cute.arch.load( + (shared_flags.iterator + 4).llvm_ptr, Uint32 + ) + current_elements = Int64(current_index) * ( + Int64(bytes_per_buffer) // Int64(2) + ) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + source_element = ( + Int64(token) * self.hidden_dim + + Int64(destination) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + source_ptr = cute.make_ptr( + BFloat16, + (shared_source.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero( + load_global_u32x4(source_ptr, volatile=False) + ) + peer_base = cute.arch.load( + (shared_peer_ptrs.iterator + destination).llvm_ptr, + Int64, + ) + destination_element = current_elements + ( + (Int64(token) * self.tp_size + self.rank) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + store_global_u32x4( + peer_base + destination_element * 2, + local_packed, + volatile=False, + ) + + # One arrival per shared destination group and token. + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(shared_flags.iterator + 8, Uint32(1)) + + global_tid = ( + Int64(token) * self.tp_size + Int64(destination) + ) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.tp_size * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + ( + shared_workspace.iterator + + dirty_elements + + clear_idx * VEC_BF16 + ).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + if destination == self.rank: + rank_words = cute.make_rmem_tensor( + cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32 + ) + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (shared_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8( + rank_words[source_rank, None].load() + ).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + result = accum.load().to(BFloat16) + output_element = ( + Int64(token) * self.hidden_dim + + self.rank * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + store_global_u32x4( + Int64((shared_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + if destination == self.rank: + cute.arch.barrier() + + cute.arch.griddepcontrol_launch_dependents() + + if ( + token_cta == 0 + and token + self.token_ctas >= m + and shared_group == 0 + and cluster_rank == 0 + and tidx == 0 + ): + access_counter = shared_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + target = Uint32(m) * Uint32(self.tp_size // self.cluster_ctas) + while arrived < target: + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2) + cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (shared_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((shared_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((shared_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (shared_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + +_COMPILED: dict[tuple[object, ...], Any] = {} + + +def _routed_workspace_cute(workspace: torch.Tensor): + return to_cute(workspace.view(torch.bfloat16), 16) + + +def _compile_key( + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool, + include_routed: bool, +): + return ( + torch.accelerator.current_device_index(), + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + + +def _runtime_args( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, +): + return ( + to_cute_dynamic_m(latent_source, mode=0, assumed_align=16), + to_cute(gamma, 16), + to_cute(latent_output, 16), + _routed_workspace_cute(routed_workspace), + to_cute(routed_flags, 16), + Int64(routed_multicast_ptr), + to_cute_dynamic_m(shared_source, mode=0, assumed_align=16), + to_cute(shared_output, 16), + to_cute(shared_workspace, 16), + to_cute(shared_flags, 16), + to_cute(shared_peer_ptrs, 16), + Int32(latent_source.shape[0]), + Float32(rms_eps), + cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream), + ) + + +def compile_kernel( + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + """Compile the rank/M specialization without retaining caller tensors.""" + + key = _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + if key in _COMPILED: + return + device = latent_output.device + latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device) + gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device) + shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device) + kernel = AllReduceRMSNormWithReduceScatterEarlyExit( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[key] = cute.compile( + kernel, + *_runtime_args( + latent, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ), + ) + + +def launch( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=latent_output, + routed_workspace=routed_workspace, + routed_flags=routed_flags, + routed_multicast_ptr=routed_multicast_ptr, + shared_output=shared_output, + shared_workspace=shared_workspace, + shared_flags=shared_flags, + shared_peer_ptrs=shared_peer_ptrs, + rms_eps=rms_eps, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[ + _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + ]( + *_runtime_args( + latent_source, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ) + ) + + +class CollectiveKernel: + """Own and launch the routed AllReduce/RMSNorm plus shared ReduceScatter.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + rms_eps: float, + fp32_internal: bool, + ) -> None: + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.max_token_ctas = max_token_ctas + self.rms_eps = float(rms_eps) + self.fp32_internal = fp32_internal + device = torch.device("cuda", torch.accelerator.current_device_index()) + + bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2 + routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer + self._routed_workspace = symm_mem.empty( + routed_bytes // 4, + dtype=torch.float32, + device=device, + ) + self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group) + self._routed_workspace.fill_(-0.0) + actual_bytes_per_buffer = ( + self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16 + ) + if actual_bytes_per_buffer < bytes_per_routed_buffer: + raise RuntimeError("routed symmetric workspace is too small") + self._routed_flags = torch.tensor( + [0, 2, actual_bytes_per_buffer, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=device, + ) + routed_multicast_ptr = self._routed_symm_mem.multicast_ptr + if routed_multicast_ptr is None or routed_multicast_ptr == 0: + raise RuntimeError("routed NVLS multicast mapping is unavailable") + self._routed_multicast_ptr = int(routed_multicast_ptr) + + self._latent_output = torch.empty( + (max_m, latent_dim), dtype=torch.bfloat16, device=device + ) + self._shared_output = torch.empty( + (max_m, hidden_dim), dtype=torch.bfloat16, device=device + ) + shard_start = rank * self.shard_dim + shard_end = shard_start + self.shard_dim + self._shared_shard = self._shared_output[:, shard_start:shard_end] + + self._shared_workspace = symm_mem.empty( + (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim), + dtype=torch.bfloat16, + device=device, + ) + self._shared_symm_mem = symm_mem.rendezvous(self._shared_workspace, group) + self._shared_workspace.view(torch.int32).fill_(-0x80000000) + self._shared_flags = torch.zeros(12, dtype=torch.int32, device=device) + self._shared_flags[1] = 1 + self._shared_flags[2] = max_m * tp_size * self.shard_dim * 2 + peer_ptrs = [ + self._shared_symm_mem.get_buffer( + peer, + self._shared_workspace.shape, + torch.bfloat16, + ).data_ptr() + for peer in range(tp_size) + ] + if any(pointer == 0 for pointer in peer_ptrs): + raise RuntimeError("shared LSA peer mapping is unavailable") + self._shared_peer_ptrs = torch.tensor( + peer_ptrs, dtype=torch.int64, device=device + ) + + torch.accelerator.synchronize(device) + dist.barrier(group=group, device_ids=[device.index]) + for owner in range(tp_size): + if rank == owner: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=self._latent_output, + routed_workspace=self._routed_workspace, + routed_flags=self._routed_flags, + routed_multicast_ptr=self._routed_multicast_ptr, + shared_output=self._shared_output, + shared_workspace=self._shared_workspace, + shared_flags=self._shared_flags, + shared_peer_ptrs=self._shared_peer_ptrs, + rms_eps=self.rms_eps, + fp32_internal=fp32_internal, + ) + dist.barrier(group=group, device_ids=[device.index]) + + def __call__( + self, + latent_source: torch.Tensor, + shared_source: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if latent_source.ndim != 2 or shared_source.ndim != 2: + raise ValueError("latent_source and shared_source must be rank-2") + m = latent_source.shape[0] + device = self._routed_workspace.device + expected = ( + (latent_source, (m, self.latent_dim), "latent_source"), + (shared_source, (m, self.hidden_dim), "shared_source"), + (gamma, (self.latent_dim,), "gamma"), + ) + for tensor, shape, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != torch.bfloat16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + with torch.accelerator.device_index(device.index): + launch( + latent_source, + gamma, + self._latent_output, + self._routed_workspace, + self._routed_flags, + self._routed_multicast_ptr, + shared_source, + self._shared_output, + self._shared_workspace, + self._shared_flags, + self._shared_peer_ptrs, + self.rms_eps, + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + max_token_ctas=self.max_token_ctas, + fp32_internal=self.fp32_internal, + ) + return ( + self._latent_output[:m], + self._shared_shard, + ) + + @property + def latent_output(self) -> torch.Tensor: + return self._latent_output + + @property + def shared_output(self) -> torch.Tensor: + return self._shared_output diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py new file mode 100644 index 000000000000..6a1ba9d77b44 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -0,0 +1,1291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Blackwell GEMM with a fused shared-shard add and multicast epilogue. + +Modified from the original CUTLASS CuTe DSL SM100 persistent GEMM tutorial. +The PDL wait before A loading orders the up-projection and shared-expert +inputs after the producer collective. +""" + +import math +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.common import CacheEvictionPriority +from cutlass.cute.runtime import from_dlpack +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .fused_add_multicast_skinny_gemm import ( + FusedAddMulticastSkinnyGemmKernel, +) +from .primitives import CUDAGraphCompatibleWrapper + + +def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False): + converted = from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=16 + ) + if dynamic_m: + converted = converted.mark_compact_shape_dynamic( + mode=1, + stride_order=tensor.dim_order(), + ) + return converted + + +def validate_configuration( + *, + latent_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, +) -> None: + """Validate constraints imposed by this BF16 SM100 GEMM.""" + + if latent_dim <= 0 or shard_dim <= 0: + raise ValueError("GEMM K and N must be positive") + if latent_dim % 8 or shard_dim % 8: + raise ValueError("GEMM K and N must be divisible by 8 BF16 values") + if mma_tiler_mn[0] not in (64, 128): + raise ValueError("MMA M tile must be 64 or 128") + if mma_tiler_mn[1] not in range(32, 257, 32): + raise ValueError("MMA N tile must be a multiple of 32 in [32, 256]") + if ( + len(cluster_shape_mn) != 2 + or any(value <= 0 or value & (value - 1) for value in cluster_shape_mn) + or math.prod(cluster_shape_mn) > 16 + ): + raise ValueError("cluster dimensions must be powers of two with product <= 16") + if not 0 <= b_prime_stages <= math.ceil(latent_dim / 128): + raise ValueError("b_prime_stages exceeds the GEMM K-tile count") + + +@cute.jit +def _epilogue_tma_store_add_shared( + gemm_kernel, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + tCtAcc_base: cute.Tensor, + sC: cute.Tensor, + tCgC_base: cute.Tensor, + tCgShared_base: cute.Tensor, + tCcC_base: cute.Tensor, + mC_mnl: cute.Tensor, + mShared_mnl: cute.Tensor, + epi_tile: cute.Tile, + num_tiles_executed: cutlass.Int32, + mma_tile_coord_mnl, + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + c_pipeline: pipeline.PipelineTmaStore, +) -> pipeline.PipelineState: + """BF16 GEMM rounding + BF16 shared shard addition, then swizzled S2G TMA.""" + sm100 = utils.gemm.sm100 + tCgC = sm100.transform_partitioned_tensor_layout(tCgC_base) + tCgShared = sm100.transform_partitioned_tensor_layout(tCgShared_base) + tCcC = sm100.transform_partitioned_tensor_layout(tCcC_base) + tCtAcc = sm100.transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = sm100.epilogue_tmem_copy_and_partition( + gemm_kernel, + epi_tidx, + tCtAcc, + tCgC, + epi_tile, + False, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tTR_rShared = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = sm100.epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gShared_partitioned = thr_copy_t2r.partition_D( + cute.flat_divide(tCgShared, epi_tile) + ) + tTR_cC_partitioned = thr_copy_t2r.partition_D(cute.flat_divide(tCcC, epi_tile)) + tTR_gShared = tTR_gShared_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + tTR_cC = tTR_cC_partitioned[(None, None, None, None, None, *mma_tile_coord_mnl)] + + exemplar = tTR_gShared_partitioned[(None, None, None, 0, 0, 0, 0, 0)] + mcl_r = cute.max_common_layout(tTR_rShared.layout, exemplar.layout) + shared_copy_bits = min( + exemplar.iterator.alignment * 8, + cute.size(mcl_r) * gemm_kernel.c_dtype.width, + 128, + ) + shared_g2r_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + gemm_kernel.c_dtype, + num_bits_per_copy=shared_copy_bits, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gShared = cute.group_modes(tTR_gShared, 3, cute.rank(tTR_gShared)) + tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + previous_subtile_count = num_tiles_executed * subtile_cnt + cute.arch.griddepcontrol_wait() + for subtile_idx in range(subtile_cnt): + tTR_gShared_subtile = tTR_gShared[(None, None, None, subtile_idx)] + tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] + pred_shape = (1, *tTR_cC_subtile.shape[1:]) + pred = cute.make_rmem_tensor(pred_shape, cutlass.Boolean) + for m_idx in range(tTR_cC_subtile.shape[1]): + for n_idx in range(tTR_cC_subtile.shape[2]): + pred[(0, m_idx, n_idx)] = cute.elem_less( + tTR_cC_subtile[(0, m_idx, n_idx)], mC_mnl.shape + ) + tTR_rShared.store(cute.zeros_like(tTR_rShared, dtype=gemm_kernel.c_dtype)) + cute.copy( + shared_g2r_atom, + tTR_gShared_subtile, + tTR_rShared, + pred=pred, + ) + + # Load the shared addend before waiting for the accumulator. + if subtile_idx == 0: + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + gemm_vec = tiled_copy_r2s.retile(tTR_rAcc).load().to(gemm_kernel.c_dtype) + shared_vec = tiled_copy_r2s.retile(tTR_rShared).load() + fused_vec = (gemm_vec.to(cutlass.Float32) + shared_vec.to(cutlass.Float32)).to( + gemm_kernel.c_dtype + ) + # The symmetric output is an in-band Lamport mailbox whose empty + # marker contains BF16 -0. Normalize either signed zero to +0 so a + # legitimate result can never be mistaken for an unwritten fragment. + fused_vec = cute.where( + fused_vec == cute.zeros_like(fused_vec), + cute.zeros_like(fused_vec), + fused_vec, + ) + tRS_rC.store(fused_vec) + + c_buffer = (previous_subtile_count + subtile_idx) % gemm_kernel.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + epilog_sync_barrier.arrive_and_wait() + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype, + b_dtype, + c_dtype, + smem_capacity: int, + c_smem_layout, +) -> tuple[int, int, int]: + """Choose accumulator, mainloop, and epilogue stage counts.""" + num_acc_stage = 2 + num_c_stage = 2 + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + num_ab_stage = ( + smem_capacity - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + num_c_stage += ( + smem_capacity + - ab_bytes_per_stage * num_ab_stage + - (mbar_helpers_bytes + c_bytes) + ) // c_bytes_per_stage + return num_acc_stage, num_ab_stage, num_c_stage + + +class FusedAddMulticastGemm: + """Persistent Blackwell GEMM with a shared-add epilogue. + + B priming may overlap the producer collective. The PDL wait before A + loading orders both inputs before the shared-add epilogue. + """ + + def __init__( + self, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int = 2, + ): + self.acc_dtype = cutlass.Float32 + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + # B primes the combined A+B pipeline before the PDL wait. + self.b_prime_stages = b_prime_stages + self.cta_group = tcgen05.CtaGroup.ONE + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id) + ) + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + + def _create_tiled_mma(self): + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Derive layouts and stage counts from the compiled tensor shapes.""" + tiled_mma = self._create_tiled_mma() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + False, + self.c_layout, + self.c_dtype, + ) + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + utils.get_smem_capacity_in_bytes(), + c_smem_layout, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, "sm_100" + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + c_multicast_i64: cutlass.Int64, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Launch the persistent GEMM.""" + # Preserve C's logical strided layout but point the TMA descriptor at + # this rank's shard inside the LSA multicast mapping. One TMA store is + # therefore replicated into the same shard on all eight ranks. + c = cute.make_tensor( + cute.make_ptr( + c.element_type, + c_multicast_i64, + cute.AddressSpace.gmem, + assumed_align=16, + ), + c.layout, + ) + + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + tiled_mma = self._create_tiled_mma() + + self._setup_attributes() + if cutlass.const_expr(self.b_prime_stages > self.num_ab_stage): + raise ValueError( + "b_prime_stages exceeds the compiled A/B pipeline stage count" + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + tile_sched_params, + shared_shard, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=True, + ) + return + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + self._gemm_device( + tiled_mma, + tma_atom_a, + mA_mkl, + tma_atom_b, + mB_nkl, + tma_atom_c, + mC_mnl, + cluster_layout_vmnk, + a_smem_layout_staged, + b_smem_layout_staged, + c_smem_layout_staged, + epi_tile, + tile_sched_params, + shared_shard, + ) + + @cute.jit + def _gemm_device( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + ab_producer, ab_consumer = ab_pipeline.make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + # Shared shard is physically [M, shard_dim]. Give it the same logical MNL + # view as C so its epilogue partition is coordinate-identical. + mShared_mnl = cute.make_tensor( + shared_shard.iterator, + cute.append(shared_shard.layout, cute.make_layout((1,), stride=(0,))), + ) + gShared_mnl = cute.local_tile( + mShared_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgC = thr_mma.partition_C(gC_mnl) + tCgShared = thr_mma.partition_C(gShared_mnl) + + # Predicate the partial M tile when M is smaller than the MMA tile. + idC = cute.make_identity_tensor(mC_mnl.shape) + cC_mnl = cute.local_tile( + idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + tCcC = thr_mma.partition_C(cC_mnl) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + gemm_grid_z = cute.arch.grid_dim()[2] + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + ( + cute.arch.grid_dim()[0], + cute.arch.grid_dim()[1], + gemm_grid_z, + ), + ) + work_tile = tile_sched.initial_work_tile_info() + + if warp_idx == self.tma_warp_id: + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Prime a short prefix of the existing combined A+B ring with + # B only. Its barrier still expects A+B bytes, so the MMA + # consumer cannot observe a half-filled stage. + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < self.b_prime_stages: + peek_ab_empty_status = ab_producer.try_acquire() + + cute.arch.griddepcontrol_wait() + + # Supply A to the very same stages after the producer AR has + # programmatically released this dependent kernel. + a_fill_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + a_barrier = ab_pipeline.producer_get_barrier(a_fill_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_tile)], + tAsA[(None, a_fill_state.index)], + tma_bar_ptr=a_barrier, + mcast_mask=a_full_mcast_mask, + ) + a_fill_state.advance() + + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range( + self.b_prime_stages, k_tile_cnt, 1, unroll=1 + ): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + ab_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in range(k_tile_cnt): + if is_leader_cta: + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + handle.release() + + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + acc_pipeline.producer_tail(acc_producer_state) + + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + acc_consumer_state = _epilogue_tma_store_add_shared( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + tCgShared, + tCcC, + mC_mnl, + mShared_mnl, + epi_tile, + num_tiles_executed, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + c_pipeline.producer_tail() + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # Allow the Lamport copy to become resident before this grid + # fully retires. Its griddepcontrol.wait still enforces complete + # producer-grid ordering before mailbox inspection. + cute.arch.griddepcontrol_launch_dependents() + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Build the static persistent schedule.""" + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: tuple[int, int, int], + num_acc_stage: int, + arch: str, + ) -> int: + """Return the required tensor-memory column count.""" + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) + + return num_tmem_alloc_cols + + +@cute.jit +def launch_kernel( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, n, k) + c: cute.Tensor, # (l, m, n) + shared_shard: cute.Tensor, # (m, shard_dim), private TP shard + rows: cutlass.Int64, + c_multicast_i64: cutlass.Int64, + full_hidden_dim: cutlass.Constexpr, + shard_dim: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, +): + """Launch the fused-add multicast GEMM using PyTorch BMM tensor order.""" + # C is passed as the fixed-capacity [1,max_m,H] symmetric mailbox. Only M + # is runtime-variable; construct this rank's logical [1,M,S] view here so + # H and S remain compile-time constants and no dynamic strided host view is + # needed on every call. __call__ later replaces the local base pointer with + # the already rank-offset multicast address. + c = cute.make_tensor( + c.iterator, + cute.make_layout( + (1, rows, shard_dim), + stride=(0, full_hidden_dim, 1), + ), + ) + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,n,k) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[1, 2, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op( + a, + b, + c, + shared_shard, + c_multicast_i64, + max_active_clusters, + stream, + ) + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + mnkl: tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + full_hidden_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int] = (64, 32), + cluster_shape_mn: tuple[int, int] = (1, 8), + max_active_clusters: cutlass.Constexpr = None, + b_prime_stages: int = 2, +): + key = ( + torch.accelerator.current_device_index(), + mnkl, + full_hidden_dim, + shard_dim, + mma_tiler_mn, + cluster_shape_mn, + max_active_clusters, + b_prime_stages, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + gemm = FusedAddMulticastGemm( + mma_tiler_mn, + cluster_shape_mn, + b_prime_stages, + ) + validate_configuration( + latent_dim=mnkl[2], + shard_dim=mnkl[1], + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if any(tensor.element_type is not cutlass.BFloat16 for tensor in (a, b, c)): + raise ValueError("up-projection tensors must be BF16") + + # The producer writes [M, full_hidden_dim]. GEMM receives a rank-local + # [M, shard_dim] view: contiguous within a row, with full_hidden_dim as + # its leading dimension. + shared_shard_compile = make_fake_tensor( + shared_shard.element_type, + (mnkl[0], shard_dim), + stride=(full_hidden_dim, 1), + assumed_align=16, + ) + stream = make_fake_stream() + compiled = cute.compile( + launch_kernel, + gemm, + a, + b, + c, + shared_shard_compile, + cutlass.Int64(mnkl[0]), + cutlass.Int64(0), + full_hidden_dim, + shard_dim, + max_active_clusters, + stream, + ) + _COMPILED[key] = compiled + return compiled + + +class AdaptiveUpProjectionKernel: + """Dispatch static-M Skinny or dynamic-M WGMMA into one mailbox.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + skinny_max_m: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, + ) -> None: + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + if not 0 <= skinny_max_m <= min(8, max_m): + raise ValueError("skinny_max_m must be in [0, min(8, max_m)]") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.skinny_max_m = skinny_max_m + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + self.b_prime_stages = b_prime_stages + device = torch.device("cuda", torch.accelerator.current_device_index()) + self._device = device + self._dynamic: Any | None = None + self._skinny_by_m: dict[int, FusedAddMulticastSkinnyGemmKernel] = {} + validate_configuration( + latent_dim=latent_dim, + shard_dim=self.shard_dim, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if skinny_max_m and self.latent_dim % (224 * 8): + raise ValueError( + "Skinny up-projection requires latent_dim divisible by 1792." + ) + + self._mailbox = symm_mem.empty( + (1, max_m, hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + self._mailbox_symm_mem = symm_mem.rendezvous(self._mailbox, group) + self._mailbox.view(torch.int32).fill_(-0x80000000) + multicast_ptr = self._mailbox_symm_mem.multicast_ptr + if multicast_ptr is None or multicast_ptr == 0: + raise RuntimeError("mailbox NVLS multicast mapping is unavailable") + self._mailbox_multicast_ptr = int(multicast_ptr) + + cluster_size = math.prod(cluster_shape_mn) + self._max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_size + ) + self._mailbox_c = _as_cute(self._mailbox) + + def compile_dynamic(self) -> None: + if self._dynamic is not None: + return + device = self._device + with torch.accelerator.device_index(device.index): + compile_latent = torch.empty( + (1, self.max_m, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_weight = torch.empty( + (self.shard_dim, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_shared = torch.empty( + (self.max_m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + )[ + :, + self.rank * self.shard_dim : (self.rank + 1) * self.shard_dim, + ] + compile_latent_c = _as_cute(compile_latent, dynamic_m=True) + compile_weight_c = _as_cute(compile_weight.unsqueeze(0)) + compile_shared_c = _as_cute(compile_shared) + + self._dynamic = compile_kernel( + (self.max_m, self.shard_dim, self.latent_dim, 1), + compile_latent_c, + compile_weight_c, + self._mailbox_c, + compile_shared_c, + self.hidden_dim, + self.shard_dim, + self.mma_tiler_mn, + self.cluster_shape_mn, + self._max_active_clusters, + self.b_prime_stages, + ) + + def compile_skinny(self, m: int) -> None: + if not 1 <= m <= self.skinny_max_m: + raise ValueError( + f"Skinny up-projection requires M in [1, {self.skinny_max_m}]." + ) + if m in self._skinny_by_m: + return + with torch.accelerator.device_index(self._device.index): + self._skinny_by_m[m] = FusedAddMulticastSkinnyGemmKernel( + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + num_rows=m, + ) + + def ensure_compiled(self, m: int) -> None: + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + if m <= self.skinny_max_m: + self.compile_skinny(m) + else: + self.compile_dynamic() + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> torch.Tensor: + if latent.ndim != 2: + raise ValueError("latent must be rank-2") + m = latent.shape[0] + device = self._mailbox.device + expected = ( + (latent, (m, self.latent_dim), "latent"), + ( + weight, + (self.shard_dim, self.latent_dim), + "weight", + ), + ( + shared_shard, + (self.max_m, self.shard_dim), + "shared_shard", + ), + ) + for tensor, shape, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != torch.bfloat16 + or tensor.device != device + ): + raise ValueError(f"{name} must be CUDA torch.bfloat16 {list(shape)}") + if ( + not latent.is_contiguous() + or not weight.is_contiguous() + or shared_shard.stride() != (self.hidden_dim, 1) + ): + raise ValueError("up-projection inputs have unsupported strides") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + if m <= self.skinny_max_m: + skinny = self._skinny_by_m.get(m) + if skinny is None: + raise RuntimeError( + f"Skinny up-projection M={m} was not compiled before launch." + ) + return skinny( + latent, + weight, + shared_shard, + self._mailbox, + self._mailbox_multicast_ptr, + ) + + if self._dynamic is None: + raise RuntimeError("Dynamic up-projection was not compiled before launch.") + with torch.accelerator.device_index(device.index): + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + self._dynamic( + _as_cute(latent.unsqueeze(0), dynamic_m=True), + _as_cute(weight.unsqueeze(0)), + self._mailbox_c, + _as_cute(shared_shard), + cutlass.Int64(m), + cutlass.Int64( + self._mailbox_multicast_ptr + self.rank * self.shard_dim * 2 + ), + stream, + ) + return self._mailbox diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py new file mode 100644 index 000000000000..58d8f1ca18b3 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Static-M SIMT skinny GEMM with a shared-add multicast epilogue.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int64, const_expr +from cutlass.cute.runtime import from_dlpack + +from .primitives import ( + CUDAGraphCompatibleWrapper, + bf16x2_to_u32, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + sanitize_negative_zero, + sanitize_negative_zero_u32, + sanitize_negative_zero_u32x2, + store_global_u32, + store_global_u32x2, + store_global_u32x4, +) + + +@dataclass(frozen=True) +class SkinnyConfig: + block_size: int = 224 + outputs_per_block: int = 2 + k_unroll: int = 2 + vector_width: int = 8 + prefetch_b_before_pdl: bool = True + + +def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig: + if shard_dim == 448: + if num_rows >= 6: + return SkinnyConfig( + block_size=224, + outputs_per_block=2, + k_unroll=1, + ) + outputs_per_block = 2 if num_rows <= 3 else 4 + return SkinnyConfig( + block_size=448, + outputs_per_block=outputs_per_block, + k_unroll=1, + ) + if num_rows == 1: + return SkinnyConfig(outputs_per_block=8, k_unroll=1) + return SkinnyConfig(outputs_per_block=4) + + +def _as_cute(tensor: torch.Tensor): + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), + assumed_align=16, + ) + + +class FusedAddMulticastSkinnyGemm: + """SIMT GEMM adapted from the existing Skinny GEMM.""" + + def __init__( + self, + *, + num_rows: int, + hidden_dim: int, + config: SkinnyConfig, + ) -> None: + if config.block_size % 32: + raise ValueError("skinny block_size must be a multiple of 32") + self.num_rows = num_rows + self.hidden_dim = hidden_dim + self.block_size = config.block_size + self.outputs_per_block = config.outputs_per_block + self.k_unroll = config.k_unroll + self.vector_width = config.vector_width + self.prefetch_b_before_pdl = config.prefetch_b_before_pdl + self.num_warps = config.block_size // 32 + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + stream: cuda.CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + gA, + gB, + gShared, + output_multicast_ptr, + k, + copy_a, + copy_b, + ).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=(self.num_rows * self.outputs_per_block * self.num_warps * 4), + stream=stream, + use_pdl=True, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + num_rows: cutlass.Constexpr = self.num_rows + + acc = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, outputs_per_block), + stride=(outputs_per_block, 1), + ), + Float32, + ) + acc.fill(0.0) + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + + if const_expr(self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + + cute.arch.griddepcontrol_wait() + + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, 0], a_regs[mi, None]) + if const_expr(not self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ + ni, vi + ].to(Float32) + + for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy( + copy_a, + tA[mi, None, k_tile], + a_regs[mi, None], + ) + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ + ni, vi + ].to(Float32) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor( + Float32, + smem_layout, + byte_alignment=16, + ) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + fused = cute.make_rmem_tensor( + cute.make_layout((outputs_per_block,)), + BFloat16, + ) + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + ) + gemm_value = Float32(total).to(BFloat16) + fused[ni] = ( + gemm_value.to(Float32) + gShared[mi, n_base + ni].to(Float32) + ).to(BFloat16) + output_offset = Int64((mi * self.hidden_dim + n_base) * 2) + if const_expr(outputs_per_block == 2): + packed = sanitize_negative_zero_u32(bf16x2_to_u32(fused.load())) + store_global_u32( + output_multicast_ptr + output_offset, + packed, + ) + elif const_expr(outputs_per_block == 4): + packed = sanitize_negative_zero_u32x2( + bf16x4_to_packed_u32x2(fused.load()) + ) + store_global_u32x2( + output_multicast_ptr + output_offset, + packed, + ) + else: + packed = sanitize_negative_zero( + bf16x8_to_packed_u32x4(fused.load()) + ) + store_global_u32x4( + output_multicast_ptr + output_offset, + packed, + ) + + cute.arch.griddepcontrol_launch_dependents() + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + *, + num_rows: int, + latent_dim: int, + hidden_dim: int, + shard_dim: int, + config: SkinnyConfig, +): + key = ( + torch.accelerator.current_device_index(), + num_rows, + latent_dim, + hidden_dim, + shard_dim, + config, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + if shard_dim % config.outputs_per_block: + raise ValueError("shard_dim must be divisible by outputs_per_block") + if latent_dim % (config.block_size * config.vector_width): + raise ValueError("latent_dim must be divisible by block_size * vector_width") + + a = make_fake_tensor( + BFloat16, + (num_rows, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + b = make_fake_tensor( + BFloat16, + (shard_dim, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + shared = make_fake_tensor( + BFloat16, + (num_rows, shard_dim), + stride=(hidden_dim, 1), + assumed_align=16, + ) + compiled = cute.compile( + FusedAddMulticastSkinnyGemm( + num_rows=num_rows, + hidden_dim=hidden_dim, + config=config, + ), + a, + b, + shared, + Int64(0), + make_fake_stream(), + options="--ptxas-options -maxrregcount=128", + ) + _COMPILED[key] = compiled + return compiled + + +class FusedAddMulticastSkinnyGemmKernel: + """Buffer-free compiled launcher for one static M.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + num_rows: int, + ) -> None: + if not 1 <= num_rows <= 8: + raise ValueError("skinny backend requires static M in [1, 8]") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + self.rank = rank + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.num_rows = num_rows + self._skinny = compile_kernel( + num_rows=num_rows, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + shard_dim=self.shard_dim, + config=config_for_m(num_rows, self.shard_dim), + ) + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + mailbox_multicast_ptr: int, + ) -> torch.Tensor: + device = mailbox.device + expected = ( + ( + latent, + (self.num_rows, self.latent_dim), + torch.bfloat16, + "latent", + ), + ( + weight, + (self.shard_dim, self.latent_dim), + torch.bfloat16, + "weight", + ), + ( + shared_shard, + (mailbox.shape[1], self.shard_dim), + torch.bfloat16, + "shared_shard", + ), + ( + mailbox, + (1, mailbox.shape[1], self.hidden_dim), + torch.bfloat16, + "mailbox", + ), + ) + for tensor, shape, dtype, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != dtype + or tensor.device != device + ): + raise ValueError(f"{name} must be CUDA {dtype} {list(shape)}") + if ( + not latent.is_contiguous() + or not weight.is_contiguous() + or shared_shard.stride() != (self.hidden_dim, 1) + or not mailbox.is_contiguous() + ): + raise ValueError("skinny up-projection inputs have unsupported strides") + if mailbox.shape[1] < self.num_rows: + raise ValueError("mailbox capacity is smaller than runtime M") + + with torch.accelerator.device_index(device.index): + self._skinny( + _as_cute(latent), + _as_cute(weight), + _as_cute(shared_shard[: self.num_rows]), + Int64(mailbox_multicast_ptr + self.rank * self.shard_dim * 2), + cuda.CUstream(torch.cuda.current_stream(device).cuda_stream), + ) + return mailbox diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py new file mode 100644 index 000000000000..e178e88c4133 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Lamport mailbox-to-local copy and reset.""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from .primitives import ( + VEC_BF16, + fragment_is_dirty, + load_global_u32x4, + store_global_u32x4, + store_lamport_sentinel_128, + to_cute, + to_cute_dynamic_m, +) + + +class LamportCopy: + """Consume the local physical copy of an NVLS-multicast mailbox.""" + + def __init__(self, hidden_dim: int, ctas: int, threads: int): + self.hidden_dim = hidden_dim + self.ctas = ctas + self.threads = threads + + @cute.jit + def __call__( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + stream: cuda.CUstream, + ): + self.kernel(symmetric_mailbox, local_output, m).launch( + grid=(self.ctas, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + ): + # The CTA may be scheduled early, but mailbox inspection must not pass + # the producer GEMM's programmatic completion point. + cute.arch.griddepcontrol_wait() + + tidx, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + thread = cutlass.Int64(block * self.threads + tidx) + stride = cutlass.Int64(self.ctas * self.threads) + fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16) + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source, volatile=True) + while fragment_is_dirty(packed): + packed = load_global_u32x4(source, volatile=True) + + destination = cutlass.Int64((local_output.iterator + element).toint()) + store_global_u32x4(destination, packed, volatile=False) + fragment = fragment + stride + + # The returned ordinary tensor is complete. A same-stream successor + # may overlap the mailbox cleanup below. + cute.arch.griddepcontrol_launch_dependents() + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(source) + fragment = fragment + stride + + +@functools.cache +def compile_kernel( + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + device_index: int, +): + if hidden_dim <= 0 or hidden_dim % VEC_BF16: + raise ValueError("hidden_dim must be a positive multiple of 8") + if max_m <= 0: + raise ValueError("max_m must be positive") + if ctas <= 0: + raise ValueError("Lamport copy ctas must be positive") + if not 1 <= threads <= 1024: + raise ValueError("Lamport copy threads must be in [1, 1024]") + with torch.accelerator.device_index(device_index): + mailbox = make_fake_compact_tensor( + cutlass.BFloat16, (max_m * hidden_dim,), assumed_align=16 + ) + output = make_fake_compact_tensor( + cutlass.BFloat16, + (cute.sym_int32(divisibility=VEC_BF16),), + assumed_align=16, + ) + return cute.compile( + LamportCopy(hidden_dim, ctas, threads), + mailbox, + output, + cutlass.Int32(max_m), + make_fake_stream(), + ) + + +def launch( + symmetric_mailbox: torch.Tensor, + local_output: torch.Tensor, + *, + m: int, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, +) -> None: + if not 1 <= m <= max_m: + raise ValueError(f"runtime M={m} must be in [1, {max_m}]") + if ( + symmetric_mailbox.ndim != 3 + or symmetric_mailbox.dtype != torch.bfloat16 + or not symmetric_mailbox.is_cuda + or not symmetric_mailbox.is_contiguous() + or symmetric_mailbox.shape[0] != 1 + or symmetric_mailbox.shape[2] != hidden_dim + ): + raise ValueError( + f"symmetric_mailbox must be contiguous CUDA BF16 [1,M,{hidden_dim}]" + ) + if symmetric_mailbox.shape[1] != max_m: + raise ValueError("symmetric mailbox must retain its full max_m capacity") + if ( + local_output.shape != (1, m, hidden_dim) + or local_output.dtype != torch.bfloat16 + or local_output.device != symmetric_mailbox.device + or not local_output.is_contiguous() + ): + raise ValueError("local_output must be contiguous CUDA BF16 [1,M,H]") + + device_index = symmetric_mailbox.device.index + stream = cuda.CUstream( + torch.cuda.current_stream(symmetric_mailbox.device).cuda_stream + ) + compile_kernel(hidden_dim, max_m, ctas, threads, device_index)( + to_cute(symmetric_mailbox.flatten(), 16), + to_cute_dynamic_m( + local_output.flatten(), + mode=0, + assumed_align=16, + ), + cutlass.Int32(m), + stream, + ) + + +class LamportCopyKernel: + """Copy a borrowed symmetric mailbox into a fresh local tensor.""" + + def __init__( + self, + *, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + ) -> None: + self.hidden_dim = hidden_dim + self.max_m = max_m + self.ctas = ctas + self.threads = threads + compile_kernel( + hidden_dim, + max_m, + ctas, + threads, + torch.accelerator.current_device_index(), + ) + + def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor: + if not symmetric_mailbox.is_cuda: + raise ValueError("symmetric_mailbox must be a CUDA tensor") + device = symmetric_mailbox.device + with torch.accelerator.device_index(device.index): + output = torch.empty( + (1, m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + launch( + symmetric_mailbox, + output, + m=m, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + ctas=self.ctas, + threads=self.threads, + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py new file mode 100644 index 000000000000..1bcb6d400cc0 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared CuTe DSL primitives; this module does not define a CUDA kernel.""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +VEC_BF16 = 8 +PACKED_BYTES = 16 +NUM_LAMPORT_BUFFERS = 3 +NEG_ZERO_F32_BITS = 0x80000000 +NEG_ZERO_BF16_BITS = 0x8000 + + +class CUDAGraphCompatibleWrapper: + """DLPack view that does not synchronize with the producer stream.""" + + def __init__(self, tensor: torch.Tensor): + self.tensor = tensor + + def __dlpack__(self, stream=None): + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor: + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align + ) + + +def to_cute_dynamic_m( + tensor: torch.Tensor, + *, + mode: int, + assumed_align: int = 16, +) -> cute.Tensor: + """Expose exactly one compact tensor mode as a runtime shape. + + Model dimensions remain part of the compiled tensor type. Only the token + mode is symbolic, so changing M within an Op's capacity reuses the same + compiled kernel. + """ + + return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( + mode=mode, + stride_order=tensor.dim_order(), + ) + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + """Load one 128-bit fragment as four u32 registers. + + The volatile form is the Lamport polling load. Marking the asm + side-effecting prevents loop-invariant motion and common-subexpression + elimination across polling iterations. + """ + + address = pointer.toint(loc=loc, ip=ip) + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=volatile, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def store_global_u32x4( + address: Int64, + packed, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +) -> None: + """Store four packed words to an ordinary or NVLS multicast global VA.""" + + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] + opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + f"{opcode} [$0], {{$1, $2, $3, $4}};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32x2( + address: Int64, + packed, + *, + loc=None, + ip=None, +) -> None: + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32( + address: Int64, + word: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + word.ir_value(loc=loc, ip=ip), + ], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: + """Reset one Lamport fragment to four FP32 negative-zero bit patterns.""" + + address = pointer.toint(loc=loc, ip=ip) + value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value, value, value, value], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_release_gpu_add_u32( + pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None +) -> None: + """The exact SM100 arrival primitive used by upstream LamportFlags.""" + + address = pointer.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "red.async.release.global.gpu.add.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def map_shared_to_peer( + smem_ptr: cute.Pointer, + peer_rank: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Map a local shared-memory slot to the same slot in a peer CTA.""" + + smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return Int32( + llvm.inline_asm( + T.i32(), + [smem_address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, + value: Float32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + remote_address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc) + values = llvm.bitcast( + target_type, + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, BFloat16) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([4], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([2], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def bf16x2_to_u32(values, *, loc=None, ip=None): + packed = llvm.bitcast( + T.i32(), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return Uint32(packed) + + +@cute.jit +def sanitize_negative_zero_u32(word): + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for i in cutlass.range_constexpr(2): + result[i] = sanitize_negative_zero_u32(packed[i]) + return result.load() + + +@cute.jit +def sanitize_negative_zero(packed): + """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" + + result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for i in cutlass.range_constexpr(4): + word = packed[i] + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + result[i] = word + return result.load() + + +@cute.jit +def fragment_is_dirty(packed): + """Bit-exact upstream sentinel check: one comparison per 32-bit word.""" + + dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS) + for i in cutlass.range_constexpr(1, 4): + dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS)) + return dirty + + +@cute.jit +def warp_sum_specialized( + value: Float32, + warp_idx: Int32, + lane: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Warp sum supporting a compile-time partial final warp.""" + + if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): + for offset in cutlass.range_constexpr(16, 0, -1): + # range_constexpr does not provide powers-of-two stepping. + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + other = cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=last_warp_mask, + mask_and_clamp=31, + ) + if (lane ^ Int32(offset)) < Int32(last_warp_lanes): + value = value + other + else: + for offset in cutlass.range_constexpr(16, 0, -1): + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + return value + + +@cute.jit +def block_sum_specialized( + value: Float32, + warp_sums: cute.Tensor, + tidx: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Upstream-equivalent FP32 block reduction.""" + + lane = cute.arch.lane_idx() + warp_idx = cute.arch.warp_idx() + value = warp_sum_specialized( + value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask + ) + if lane == 0: + warp_sums[warp_idx] = value + cute.arch.barrier() + + block_sum = Float32(0.0) + if warp_idx == 0: + if lane < Int32(warps): + block_sum = warp_sums[lane] + block_sum = cute.arch.warp_reduction_sum(block_sum) + if lane == 0: + warp_sums[0] = block_sum + cute.arch.barrier() + return warp_sums[0] diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py new file mode 100644 index 000000000000..816da37153f5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MLA prefill and decode epilogues for Kimi-K3. + +Thin wrappers over the CUDA ops in +``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. + +- ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, + concat the full per-head key ``[k_nope | k_pe]`` into ``k_out``, and insert + the latent ``[kv_c_normed | k_pe]`` into the paged cache. +- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize + ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the + cache shares ``k_scale``, as in ``concat_and_cache_mla``). + +The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE +inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use +Programmatic Dependent Launch to overlap the tail of the producing GEMMs on +sm_90+. +""" + +import torch + + +def fused_mla_key_concat_kv_cache_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Apply optional RoPE, concat K, and insert the paged latent (bf16). + + Returns the full key ``[Tp, H, qk_nope_head_dim + qk_rope_head_dim]``; + optionally rotates ``q`` and writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty( + (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device + ) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_key_concat_ds_mla_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, 656] uint8 (fp8_ds_mla) + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout. + + The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of + 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no + scale argument. Returns the bf16 full key; optionally rotates ``q`` and + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty( + (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device + ) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_qkv_quant_kv_cache_fp8_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim] + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + v: torch.Tensor, # [Tp, H, v_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] fp8 + slot_mapping: torch.Tensor, # [Tp] int64 + q_scale_inv: torch.Tensor, # scalar fp32, 1 / q scale (attention query) + k_scale_inv: torch.Tensor, # scalar fp32, 1 / k scale (attention key) + v_scale_inv: torch.Tensor, # scalar fp32, 1 / v scale (attention value) + cache_scale_inv: torch.Tensor, # scalar fp32, 1 / kv scale (cache latent) + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize q/k/v to fp8 and insert the fp8 latent into the paged cache. + + The attention key ``k_fp8`` and the cache latent use *separate* scales + (``k_scale_inv`` vs ``cache_scale_inv``): the cache must be quantized with + ``_k_scale`` (read back by decode / context), while the prefill attention + q/k/v currently stay unscaled (the prefill flash path does not dequantize). + + Returns ``(q_fp8, k_fp8, v_fp8)``; writes the fp8 ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, _ = q.shape + qk_head_dim = q.shape[2] + v_head_dim = v.shape[2] + fp8 = torch.float8_e4m3fn + q_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + k_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + v_fp8 = torch.empty((tp, num_heads, v_head_dim), dtype=fp8, device=q.device) + if tp == 0: + return q_fp8, k_fp8, v_fp8 + torch.ops._C.fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + q_fp8, + k_fp8, + v_fp8, + kv_cache, + slot_mapping, + q_scale_inv, + k_scale_inv, + v_scale_inv, + cache_scale_inv, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return q_fp8, k_fp8, v_fp8 + + +def fused_mla_decode_q_concat_kv_cache_insert( + ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q) + q_pe: torch.Tensor, # [B, H, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [B, kv_lora_rank] + k_pe: torch.Tensor, # [B, qk_rope_head_dim] or [B, 1, qk_rope_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, entry] + slot_mapping: torch.Tensor, # [B] int64 + *, + ds_mla: bool = False, + q_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / q scale + cache_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / kv scale + positions: torch.Tensor | None = None, # [B] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat the latent decode query ``mqa_q = [ql_nope | q_pe]`` and insert the + latent ``[kv_c_normed | k_pe]`` into the paged cache, in one launch (runs + right before ``forward_mqa``). + + Dispatched by cache format: + - bf16 -> bf16 mqa_q, bf16 cache + - plain fp8 -> fp8 mqa_q (q_scale_inv), fp8 cache (cache_scale_inv) + - fp8_ds_mla -> bf16 mqa_q, 656B block-scaled cache + + Returns ``mqa_q`` of shape ``[B, H, kv_lora_rank + qk_rope_head_dim]``; + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + b, num_heads, kv_lora_rank = ql_nope.shape + entry = kv_lora_rank + q_pe.shape[-1] + fp8_q = q_scale_inv is not None + out_dtype = torch.float8_e4m3fn if fp8_q else ql_nope.dtype + mqa_q = torch.empty((b, num_heads, entry), dtype=out_dtype, device=ql_nope.device) + if b == 0: + return mqa_q + + if ds_mla: + cache = ( + kv_cache if kv_cache.dtype == torch.uint8 else kv_cache.view(torch.uint8) + ) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + cache.shape[1], + positions, + cos_sin_cache, + ) + elif fp8_q: + assert cache_scale_inv is not None, "fp8 decode requires cache_scale_inv" + cache = ( + kv_cache + if kv_cache.dtype == torch.float8_e4m3fn + else kv_cache.view(torch.float8_e4m3fn) + ) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + q_scale_inv, + cache_scale_inv, + cache.shape[1], + positions, + cos_sin_cache, + ) + else: + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return mqa_q diff --git a/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py new file mode 100644 index 000000000000..fe97c8fd3087 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from functools import partial +from typing import ClassVar + +import torch +import torch.distributed as dist + +from vllm.distributed import get_tp_group +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) + +_MAX_NUM_TOKENS = 16 +_SKINNY_MAX_NUM_TOKENS = 5 +_MMA_TILER_MN = (64, 32) +_GEMM_CLUSTER_MN = (1, 8) +_B_PRIME_STAGES = 2 +_COLLECTIVE_TOKEN_CTAS = 8 +_LAMPORT_COPY_CTAS = 32 +_LAMPORT_COPY_THREADS = 224 +_SUPPORTED_TP_SIZES = (8, 16) + + +@dataclass(frozen=True) +class KimiK3LatentMoETailContract: + tp_group_id: int + tp_size: int + device: torch.device + dtype: torch.dtype + hidden_size: int + latent_size: int + max_num_tokens: int + rms_eps: float + + +class KimiK3LatentMoETailOp: + """Process-wide cached K3 latent-MoE tail implementation.""" + + _instances: ClassVar[ + dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"] + ] = {} + + @classmethod + def _contract_and_group( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]: + tp = get_tp_group() + group = tp.device_group + device = torch.device(device) + tp_device = torch.device(tp.device) + if device != tp_device: + raise ValueError( + f"Input device {device} does not match TP device {tp_device}." + ) + return ( + KimiK3LatentMoETailContract( + tp_group_id=id(group), + tp_size=dist.get_world_size(group), + device=device, + dtype=dtype, + hidden_size=hidden_size, + latent_size=latent_size, + max_num_tokens=_MAX_NUM_TOKENS, + rms_eps=float(rms_eps), + ), + group, + ) + + @classmethod + def initialize( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> "KimiK3LatentMoETailOp": + contract, group = cls._contract_and_group( + hidden_size=hidden_size, + latent_size=latent_size, + dtype=dtype, + device=device, + rms_eps=rms_eps, + ) + op = cls._instances.get(contract) + if op is None: + op = cls(contract, group) + cls._instances[contract] = op + return op + + def __init__( + self, + contract: KimiK3LatentMoETailContract, + group: dist.ProcessGroup, + ) -> None: + if contract.tp_size not in _SUPPORTED_TP_SIZES: + raise ValueError( + "K3 latent-MoE tail fusion requires TP in " + f"{_SUPPORTED_TP_SIZES}, got {contract.tp_size}." + ) + if contract.device.type != "cuda": + raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.") + if torch.cuda.get_device_capability(contract.device)[0] != 10: + raise ValueError("K3 latent-MoE tail fusion requires SM100.") + if contract.dtype != torch.bfloat16: + raise ValueError("K3 latent-MoE tail fusion requires bfloat16.") + if (contract.hidden_size, contract.latent_size) != (7168, 3584): + raise ValueError( + "K3 latent-MoE tail fusion requires hidden_size=7168 and " + "latent_size=3584." + ) + + self.contract = contract + self.rank = dist.get_rank(group) + from .cute_dsl.latent_moe_tail import ( + AdaptiveUpProjectionKernel, + CollectiveKernel, + LamportCopyKernel, + ) + + with torch.accelerator.device_index(contract.device.index): + self._collective = CollectiveKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + max_token_ctas=_COLLECTIVE_TOKEN_CTAS, + rms_eps=contract.rms_eps, + fp32_internal=False, + ) + self._up_projection = AdaptiveUpProjectionKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + skinny_max_m=_SKINNY_MAX_NUM_TOKENS, + mma_tiler_mn=_MMA_TILER_MN, + cluster_shape_mn=_GEMM_CLUSTER_MN, + b_prime_stages=_B_PRIME_STAGES, + ) + self._lamport_copy = LamportCopyKernel( + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + ctas=_LAMPORT_COPY_CTAS, + threads=_LAMPORT_COPY_THREADS, + ) + register_cutedsl_warmup_provider(self) + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + contract = self.contract + skinny_units = tuple( + CuTeDSLCompileUnit( + name="K3 latent MoE tail Skinny up-projection", + key=( + "k3-latent-moe-tail-skinny-up-projection", + contract, + m, + ), + compile=partial(self._up_projection.compile_skinny, m), + ) + for m in range(1, self._up_projection.skinny_max_m + 1) + ) + return skinny_units + ( + CuTeDSLCompileUnit( + name="K3 latent MoE tail dynamic up-projection", + key=( + "k3-latent-moe-tail-dynamic-up-projection", + contract, + _MMA_TILER_MN, + _GEMM_CLUSTER_MN, + _B_PRIME_STAGES, + ), + compile=self._up_projection.compile_dynamic, + ), + ) + + def __call__( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> torch.Tensor: + self._validate_inputs( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + self._up_projection.ensure_compiled(routed_output.shape[0]) + latent, shared_shard = self._collective( + routed_output, + shared_output, + rms_weight, + ) + local_hidden_size = self.contract.hidden_size // self.contract.tp_size + local_up_weight = up_weight.narrow( + 0, + self.rank * local_hidden_size, + local_hidden_size, + ) + mailbox = self._up_projection( + latent, + local_up_weight, + shared_shard, + ) + return self._lamport_copy( + mailbox, + m=routed_output.shape[0], + ).squeeze(0) + + def _validate_inputs( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> None: + contract = self.contract + if routed_output.ndim != 2: + raise ValueError("routed_output must be a 2D tensor.") + num_tokens = routed_output.shape[0] + if routed_output.shape != (num_tokens, contract.latent_size): + raise ValueError( + f"routed_output must have shape [M, {contract.latent_size}]." + ) + if shared_output.shape != (num_tokens, contract.hidden_size): + raise ValueError( + f"shared_output must have shape [M, {contract.hidden_size}]." + ) + if rms_weight.shape != (contract.latent_size,): + raise ValueError(f"rms_weight must have shape [{contract.latent_size}].") + if up_weight.shape != (contract.hidden_size, contract.latent_size): + raise ValueError( + "up_weight must have shape " + f"[{contract.hidden_size}, {contract.latent_size}]." + ) + if not 1 <= num_tokens <= contract.max_num_tokens: + raise ValueError( + "K3 latent-MoE tail fusion requires between 1 and " + f"{contract.max_num_tokens} tokens." + ) + + tensors = (routed_output, shared_output, rms_weight, up_weight) + if any(tensor.device != contract.device for tensor in tensors): + raise ValueError("All inputs must be on the contract device.") + if any(tensor.dtype != contract.dtype for tensor in tensors): + raise ValueError("All inputs must use the contract dtype.") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("All inputs must be contiguous.") diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py b/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py new file mode 100644 index 000000000000..92936016ab64 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py new file mode 100644 index 000000000000..01bc60912891 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -0,0 +1,938 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.op import exp2, log +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError( + f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}" + ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + assert g.stride() == (HD, 1) + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py new file mode 100644 index 000000000000..087197d10a56 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,559 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.op import exp2, gather +from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported +from vllm.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT', 'HV'], +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() + and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 000000000000..ecd00a51f9dd --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp2 +from vllm.triton_utils import tl, triton + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H*K + k += bos * H*K + g += bos * HV*K + Aqk += bos * HV*BT + Akk += bos * HV*BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv) + tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 000000000000..b4d35d85571f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.op import exp, log +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T", "stride_beta_token"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_fused_recurrent_kda_fwd_warmup_profiles( + num_heads: int, +) -> tuple[int, ...]: + """Return representative sequence counts for gated launch variants.""" + # The region above 192 head-sequences reuses the second launch variant. + return ( + 1, + 48 // num_heads + 1, + 96 // num_heads + 1, + ) + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + if use_gate_in_kernel: + # Tuned on GB300 for Kimi-K3 shapes. Keep the warmup profiles above in + # sync with these boundaries. + head_sequences = H * N + if head_sequences <= 48: + BV, num_stages = 4, 4 + elif head_sequences <= 96: + BV, num_stages = 8, 3 + elif head_sequences <= 192: + BV, num_stages = 16, 3 + else: + BV, num_stages = 8, 3 + num_warps = 1 + else: + BV, num_warps, num_stages = 8, 1, 2 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit( + do_not_specialize=["stride_beta_token", "stride_state_indices"] +) +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_state_token: tl.constexpr, + stride_state_indices, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1: + raise ValueError("`state_indices` must be one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + stride_state_indices=state_indices.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state diff --git a/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py b/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py new file mode 100644 index 000000000000..3d429a9b58e4 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup compilation of Kimi-K3 vision FA4 kernels.""" + +from __future__ import annotations + +import math +from collections.abc import Iterator +from dataclasses import dataclass +from functools import partial + +import torch + +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) +from vllm.platforms import current_platform + +_FA4_TILE_SIZE = 128 +_FA4_MAX_SPLITS = 128 + + +@dataclass(frozen=True) +class KimiK3VisionFA4WarmupConfig: + num_heads: int + head_dim: int + dtype: torch.dtype + max_batch_size: int + max_seqlen: int + + +@dataclass(frozen=True) +class _FA4WarmupProbe: + batch_size: int + max_seqlen: int + dispatch_key: tuple[int, bool, int | None] + + +def _combine_log_max_splits(head_dim: int, num_splits: int) -> int: + k_block_size = 64 if head_dim <= 64 else 128 + tile_m = 8 if k_block_size == 128 else 16 + min_log_max_splits = 5 if tile_m == 8 else 4 + return max(math.ceil(math.log2(num_splits)), min_log_max_splits) + + +def _get_dispatch( + config: KimiK3VisionFA4WarmupConfig, + *, + batch_size: int, + max_seqlen: int, + num_sms: int, +) -> tuple[int, bool, int | None]: + # These are the shape-dependent fields in FA4's forward and combine + # compile keys for noncausal varlen MHA. Compile units still pass + # num_splits=0 so the production selector makes the final decision. + q_stage = 2 if max_seqlen > _FA4_TILE_SIZE else 1 + effective_q_tile = q_stage * _FA4_TILE_SIZE + num_m_blocks = math.ceil(max_seqlen / effective_q_tile) + num_n_blocks = math.ceil(max_seqlen / _FA4_TILE_SIZE) + + if num_n_blocks <= 4: + num_splits = 1 + else: + total_m_blocks = batch_size * config.num_heads * num_m_blocks + num_splits = min( + num_sms // total_m_blocks, + _FA4_MAX_SPLITS, + num_n_blocks, + ) + + is_split_kv = num_splits > 1 + combine_key = ( + _combine_log_max_splits(config.head_dim, num_splits) if is_split_kv else None + ) + return q_stage, is_split_kv, combine_key + + +def _get_warmup_probes( + config: KimiK3VisionFA4WarmupConfig, + *, + num_sms: int, +) -> tuple[_FA4WarmupProbe, ...]: + if config.max_batch_size <= 0 or config.max_seqlen <= 0: + return () + + probes: dict[tuple[int, bool, int | None], _FA4WarmupProbe] = {} + + def add_probe(batch_size: int, max_seqlen: int) -> None: + dispatch_key = _get_dispatch( + config, + batch_size=batch_size, + max_seqlen=max_seqlen, + num_sms=num_sms, + ) + if dispatch_key not in probes: + probes[dispatch_key] = _FA4WarmupProbe( + batch_size=batch_size, + max_seqlen=max_seqlen, + dispatch_key=dispatch_key, + ) + + add_probe(batch_size=1, max_seqlen=1) + if config.max_seqlen > _FA4_TILE_SIZE: + add_probe(batch_size=1, max_seqlen=_FA4_TILE_SIZE + 1) + + max_n_blocks = math.ceil(config.max_seqlen / _FA4_TILE_SIZE) + for num_n_blocks in range(5, max_n_blocks + 1): + max_seqlen = min( + (num_n_blocks - 1) * _FA4_TILE_SIZE + 1, + config.max_seqlen, + ) + num_m_blocks = math.ceil(max_seqlen / (2 * _FA4_TILE_SIZE)) + max_split_batch_size = min( + config.max_batch_size, + num_sms // (2 * config.num_heads * num_m_blocks), + ) + if max_split_batch_size == 0: + break + for batch_size in range(1, max_split_batch_size + 1): + add_probe(batch_size, max_seqlen) + + return tuple(probes.values()) + + +def _iter_compile_units( + config: KimiK3VisionFA4WarmupConfig, +) -> Iterator[CuTeDSLCompileUnit]: + device_id = current_platform.current_device() + num_sms = current_platform.num_compute_units(device_id) + for probe in _get_warmup_probes(config, num_sms=num_sms): + yield CuTeDSLCompileUnit( + name="kimi_k3_vision_fa4", + key=("kimi_k3_vision_fa4", config, probe.dispatch_key), + compile=partial(_compile, config, probe), + ) + + +def _compile( + config: KimiK3VisionFA4WarmupConfig, + probe: _FA4WarmupProbe, +) -> None: + from vllm.vllm_flash_attn import flash_attn_varlen_func + + device = current_platform.current_device() + total_tokens = probe.batch_size * probe.max_seqlen + qkv = torch.empty( + 3, + total_tokens, + config.num_heads, + config.head_dim, + device=device, + dtype=config.dtype, + ) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + probe.max_seqlen, + device=qkv.device, + dtype=torch.int32, + ) + flash_attn_varlen_func( + qkv[0], + qkv[1], + qkv[2], + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=probe.max_seqlen, + max_seqlen_k=probe.max_seqlen, + dropout_p=0.0, + causal=False, + softmax_scale=config.head_dim**-0.5, + fa_version=4, + num_splits=0, + ) + + +class _WarmupProvider: + def __init__(self) -> None: + self.configs: set[KimiK3VisionFA4WarmupConfig] = set() + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + return tuple( + unit for config in self.configs for unit in _iter_compile_units(config) + ) + + +_PROVIDER = _WarmupProvider() + + +def register_kimi_k3_vision_fa4_warmup( + config: KimiK3VisionFA4WarmupConfig, +) -> None: + _PROVIDER.configs.add(config) + register_cutedsl_warmup_provider(_PROVIDER) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 324104e055f8..696067666628 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -41,7 +41,7 @@ fused_allreduce_gemma_rms_norm, ) from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -391,7 +391,7 @@ def __init__( # always-on shared expert; aiter applies the routed scaling internally. # Every other path (vLLM top-k bias router, or no fusion) applies the # routed scaling to the MoE output here. - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, diff --git a/vllm/models/minimax_m3/amd/mtp.py b/vllm/models/minimax_m3/amd/mtp.py index f62face1d2e1..cfb26d7948df 100644 --- a/vllm/models/minimax_m3/amd/mtp.py +++ b/vllm/models/minimax_m3/amd/mtp.py @@ -35,6 +35,9 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -322,7 +325,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Validate that weights were loaded for each MTP layer. for layer_idx in range(self.model.num_mtp_layers): - if layer_idx not in loaded_mtp_layers: + if ( + layer_idx not in loaded_mtp_layers + and is_mtp_completeness_check_enabled() + ): raise ValueError( f"Failed to load MTP layer {layer_idx} weights from checkpoint." ) diff --git a/vllm/models/minimax_m3/amd/sparse_attention_msa.py b/vllm/models/minimax_m3/amd/sparse_attention_msa.py index b491b8fcb56f..5845d659bbd5 100644 --- a/vllm/models/minimax_m3/amd/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/amd/sparse_attention_msa.py @@ -23,6 +23,8 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, + *, + query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: from vllm.models.minimax_m3.amd.ops.sparse_pa import ( minimax_m3_sparse_attn_decode_aiter, diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index c66e7ce5267e..767aa2d19bcc 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -424,6 +424,9 @@ def forward( # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the # kernels' out= writes out[:, :total_q]. None -> allocate fresh. buf = self.topk_indices_buffer + buf_htk = ( + buf if buf is None or current_platform.is_rocm() else buf.transpose(0, 1) + ) decode_topk: torch.Tensor | None = None prefill_topk: torch.Tensor | None = None if index_md.num_decodes > 0: @@ -441,7 +444,7 @@ def forward( self.num_kv_heads, d.decode_query_len, d.max_decode_query_len, - out=buf, + out=buf_htk, ) if index_md.num_prefills > 0: p = index_md.prefill @@ -465,7 +468,7 @@ def forward( self.topk_blocks, self.init_blocks, self.local_blocks, - out=buf[:, nd:, :] if buf is not None else None, + out=buf_htk[:, nd:, :] if buf_htk is not None else None, ) return decode_topk, prefill_topk @@ -477,10 +480,10 @@ def select_indexer_impl_cls( ) -> type[MiniMaxM3IndexerImpl]: """Pick the indexer impl off the platform, top-k count, and cache dtype. - On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the - main MSA attend), the fmha_sm100 score path + Triton top-k is used for both - bf16 and fp8 index caches. Everything else falls back to the Triton indexer - (bf16 only). + On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's + ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is + used for both bf16 and fp8 index caches. Everything else falls back to the + Triton indexer (bf16 only). """ if indexer_kv_dtype in ("mxfp4", "nvfp4"): raise NotImplementedError( @@ -492,7 +495,7 @@ def select_indexer_impl_cls( ) use_msa = ( is_sm100 - and topk_blocks in (4, 8, 16, 32) + and topk_blocks == 16 and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3") ) if use_msa: @@ -502,7 +505,7 @@ def select_indexer_impl_cls( ) logger.info_once( - "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) " + "MiniMax M3 indexer: selected MSA (fmha_sm100 score + top-k) " "[topk_blocks=%d, indexer_kv_dtype=%s]", topk_blocks, indexer_kv_dtype, diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py index 208adfffea51..e8e0ca063dea 100644 --- a/vllm/models/minimax_m3/common/mm_preprocess.py +++ b/vllm/models/minimax_m3/common/mm_preprocess.py @@ -3,8 +3,9 @@ import math from collections.abc import Mapping, Sequence -from typing import cast +from typing import Any, Literal, cast +import numpy.typing as npt import torch from transformers import BatchFeature from transformers.video_utils import VideoMetadata @@ -469,10 +470,39 @@ def get_video_replacement(item_idx: int): ] -# TODO(Isotr0py): Tie with MinimaxVideoProcessor -# after https://github.com/vllm-project/vllm/pull/44126 -@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl") +@VIDEO_LOADER_REGISTRY.register( + name="minimax_m3_vl", + video_processor="MiniMaxM3VLVideoProcessor", +) class MiniMaxM3VideoBackend(VideoBackend): + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 1, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: Literal[ + "opencv", + "pyav", + "torchcodec", + "pynvvideocodec", + "deepstream", + ] = "opencv", + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + return super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + @classmethod def compute_frames_index_to_sample( cls, @@ -483,7 +513,6 @@ def compute_frames_index_to_sample( total_frames = source.total_frames_num video_fps = source.original_fps fps = target.fps - if total_frames <= 0 or video_fps <= 0 or fps <= 0: return [0] if total_frames > 0 else [] @@ -503,12 +532,6 @@ def compute_frames_index_to_sample( break indices.append(target_frame) prev_kept_ts = target_frame / video_fps - - last_frame_idx = total_frames - 1 - last_ts = last_frame_idx / video_fps - if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps: - indices.append(last_frame_idx) - if not indices: indices = [0] return indices diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index f887f14b643b..e8154459ddbf 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -330,6 +330,7 @@ def __init__( *, topk_blocks: int, sparse_block_size: int, + msa_decode_backend: str = "triton", ) -> None: self.num_heads = num_heads self.head_size = head_size @@ -355,6 +356,8 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, + *, + query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: """Attend the queries to the indexer-selected blocks. Per kernel. @@ -364,6 +367,9 @@ def forward( """ raise NotImplementedError + def should_use_msa_decode(self, layer_name: str) -> bool: + return False + class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl): """Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode.""" @@ -374,6 +380,8 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, + *, + query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): @@ -384,7 +392,14 @@ def forward( nd = main_md.num_decode_tokens num_tokens = main_md.num_actual_tokens # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:]. - topk = layer.topk_indices_buffer # type: ignore[attr-defined] + topk_buffer = layer.topk_indices_buffer # type: ignore[attr-defined] + assert topk_buffer is not None + + topk = ( + topk_buffer + if current_platform.is_rocm() + else topk_buffer[:num_tokens].transpose(0, 1) + ) assert topk is not None hd = self.head_size q = query[:num_tokens].view(-1, self.num_heads, hd) @@ -435,18 +450,18 @@ def forward( return output -def select_main_impl_cls( +def select_main_backend_and_impl_cls( *, topk_blocks: int, kv_cache_dtype: str, num_kv_heads: int, -) -> type[MiniMaxM3SparseImpl]: - """Pick the main attend impl off the main KV-cache dtype. +) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]: + """Pick the main attention backend and implementation. Blackwell (SM100) uses the MSA attend for supported top-k block counts when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA with shuffle KV cache layout; Other platforms and FP8 E5M2 fall - back to Triton. The MSA modules are imported lazily avoid import errors + back to Triton. The MSA modules are imported lazily to avoid import errors on unsupported platforms. """ use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(num_kv_heads) @@ -470,11 +485,26 @@ def select_main_impl_cls( MiniMaxM3SparseAiterPAImpl, ) - return MiniMaxM3SparseAiterPAImpl + return MiniMaxM3SparseBackend, MiniMaxM3SparseAiterPAImpl if use_msa: from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( + MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl, ) - return MiniMaxM3SparseMSAImpl - return MiniMaxM3SparseTritonImpl + return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl + return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl + + +def select_main_impl_cls( + *, + topk_blocks: int, + kv_cache_dtype: str, + num_kv_heads: int, +) -> type[MiniMaxM3SparseImpl]: + """Backward-compatible implementation-only selector.""" + return select_main_backend_and_impl_cls( + topk_blocks=topk_blocks, + kv_cache_dtype=kv_cache_dtype, + num_kv_heads=num_kv_heads, + )[1] diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index e29514ab9a32..03e28f344345 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -31,7 +31,7 @@ fused_allreduce_gemma_rms_norm, ) from vllm.model_executor.layers.fused_moe import ( - FusedMoE, + FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) @@ -79,7 +79,7 @@ from vllm.models.minimax_m3.common.sparse_attention import ( MiniMaxM3SparseBackend, MiniMaxM3SparseImpl, - select_main_impl_cls, + select_main_backend_and_impl_cls, ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY @@ -244,7 +244,7 @@ def __init__( prefix=f"{prefix}.shared_experts", ) - self.experts = FusedMoE( + self.experts = FusedMoEFactory( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -505,15 +505,15 @@ def __init__( # the attend impl reads them back (so nothing crosses the eager break as a # Python value, which would freeze at capture). self.topk_indices_buffer = topk_indices_buffer - self.attn_backend = MiniMaxM3SparseBackend # Indexer (top-k selection) and main attention are separate impls, each # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase # (broader than the AttentionImpl that AttentionLayerBase annotates). - self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment] + self.attn_backend, impl_cls = select_main_backend_and_impl_cls( topk_blocks=sparse_cfg["sparse_topk_blocks"], kv_cache_dtype=self.kv_cache_dtype, num_kv_heads=self.num_kv_heads, - )( + ) + self.impl: MiniMaxM3SparseImpl = impl_cls( # type: ignore[assignment] self.num_heads, self.head_dim, self.scaling, @@ -521,6 +521,9 @@ def __init__( kv_cache_dtype=self.kv_cache_dtype, topk_blocks=sparse_cfg["sparse_topk_blocks"], sparse_block_size=sparse_cfg["sparse_block_size"], + msa_decode_backend=( + vllm_config.attention_config.minimax_m3_msa_decode_backend + ), ) # Self-contained nn.Module: owns its side cache, selects its impl in init. self.indexer = MiniMaxM3Indexer( @@ -594,6 +597,16 @@ def forward( main_slot_mapping = fwd_slot_mapping[self.layer_name] index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) + use_msa_decode = self.impl.should_use_msa_decode(self.layer_name) + query_fp8 = ( + torch.empty( + (num_tokens, self.q_size), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ) + if use_msa_decode + else None + ) # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); # the fused kernel emits fp8 directly when this buffer is e4m3. index_q = qkv.new_empty( @@ -621,10 +634,12 @@ def forward( q, index_q, self.kv_cache_dtype, + q_fp8_out=query_fp8, + q_fp8_scale=self._q_scale_float, ) output = torch.empty_like(q) - attn_output = self._run_attention(q, index_q, output) + attn_output = self._run_attention(q, query_fp8, index_q, output) output, _ = self.o_proj(attn_output) return output @@ -632,6 +647,7 @@ def forward( def _run_attention( self, query: torch.Tensor, + query_fp8: torch.Tensor | None, index_query: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: @@ -639,7 +655,13 @@ def _run_attention( # metadata and can't be captured into a cudagraph. The indexer writes its # top-k into the shared ``topk_indices_buffer``; the attend reads it back. self.indexer(index_query) - return self.impl.forward(self, query, self.kv_cache, output) + return self.impl.forward( + self, + query, + self.kv_cache, + output, + query_fp8=query_fp8, + ) class MiniMaxM3DecoderLayer(nn.Module): diff --git a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py new file mode 100644 index 000000000000..04313af1070d --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax CUTLASS sparse decode using per-query-token page indices.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from vllm.config.attention import MiniMaxM3MSADecodeBackend +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +_MAX_NUM_Q_HEADS = 64 +_MAX_NUM_KV_HEADS = 4 +_HEAD_DIM = 128 +_PAGE_SIZE = 128 +_TOPK = 16 +# fmha_sm100 plans one row per query head. Keep every cached plan within the +# fixed planner allocation used by the MSA decode kernel. +_MAX_QUERY_HEAD_ROWS = 65536 +_MAX_DECODE_QUERY_LEN = 32 +# Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4. +_MIN_CUTLASS_BATCH_SIZE = 16 + + +@dataclass +class MSACutlassDecodeMetadata: + plan: Any + page_table: torch.Tensor + + +@triton.jit +def _update_runtime_metadata_kernel( + seq_lens_ptr, + kv_segment_lens_ptr, + qo_offset_ptr, + num_rows: tl.constexpr, + decode_query_len: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_rows + request = offsets // decode_query_len + local_query = offsets % decode_query_len + seq_len = tl.load(seq_lens_ptr + request, mask=mask) + tl.store(kv_segment_lens_ptr + offsets, seq_len, mask=mask) + tl.store( + qo_offset_ptr + offsets, + seq_len - decode_query_len + local_query, + mask=mask, + ) + + +@dataclass +class MSACutlassDecodePlanCache: + """Reusable plans whose mutable tensors retain cudagraph-stable addresses.""" + + plans: dict[tuple[int, ...], Any] = field(init=False, default_factory=dict) + + def _build_plan( + self, + *, + batch: int, + decode_query_len: int, + page_table_stride: int, + initial_seq_lens_cpu: torch.Tensor, + device: torch.device, + num_q_heads: int, + num_kv_heads: int, + page_size: int, + topk_blocks: int, + ) -> Any: + from vllm.third_party.fmha_sm100.api import fmha_sm100_plan + + qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32) + kv_lens_cpu = initial_seq_lens_cpu + plan = fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_q_heads, + num_kv_heads=num_kv_heads, + qo_offset=kv_lens_cpu - qo_lens_cpu, + page_size=page_size, + output_maxscore=False, + kv_block_num=topk_blocks, + causal=True, + sparse_kernel_mode="decode", + use_fp8_kvcache=True, + split_prefill_decode=False, + device=device, + ) + + plan_info = plan[3] + row_starts = ( + torch.arange(batch, dtype=torch.int32, device=device) + .mul_(page_table_stride) + .repeat_interleave(decode_query_len) + ) + page_indptr = torch.cat( + ( + row_starts, + torch.tensor( + [batch * page_table_stride], + dtype=torch.int32, + device=device, + ), + ) + ) + plan_info["kv_page_indptr"].copy_(page_indptr) + return plan + + def prepare( + self, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + decode_query_len: int, + *, + num_q_heads: int, + num_kv_heads: int, + page_size: int, + topk_blocks: int, + ) -> MSACutlassDecodeMetadata: + batch = int(seq_lens.shape[0]) + if ( + block_table.device.type != "cuda" + or block_table.dtype != torch.int32 + or not block_table.is_contiguous() + or block_table.shape[0] != batch + ): + raise ValueError( + "MSA sparse decode requires a contiguous CUDA int32 block " + "table with one row per request" + ) + if ( + seq_lens.dtype != torch.int32 + or not seq_lens.is_contiguous() + or seq_lens.device != block_table.device + ): + raise ValueError( + "MSA sparse decode requires contiguous CUDA int32 sequence " + "lengths on the block table device" + ) + if ( + seq_lens_cpu.device.type != "cpu" + or seq_lens_cpu.dtype != torch.int32 + or not seq_lens_cpu.is_contiguous() + or seq_lens_cpu.shape != seq_lens.shape + ): + raise ValueError( + "MSA sparse decode requires contiguous CPU int32 sequence " + "lengths matching the device sequence lengths" + ) + + page_table_stride = int(block_table.stride(0)) + key = ( + batch, + decode_query_len, + page_table_stride, + num_q_heads, + num_kv_heads, + page_size, + topk_blocks, + ) + plan = self.plans.get(key) + if plan is None: + plan = self._build_plan( + batch=batch, + decode_query_len=decode_query_len, + page_table_stride=page_table_stride, + initial_seq_lens_cpu=seq_lens_cpu, + device=seq_lens.device, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=page_size, + topk_blocks=topk_blocks, + ) + self.plans[key] = plan + + plan_info = plan[3] + num_rows = batch * decode_query_len + _update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)]( + seq_lens, + plan_info["kv_segment_lens"], + plan_info["qo_offset"], + num_rows=num_rows, + decode_query_len=decode_query_len, + BLOCK_SIZE=128, + ) + return MSACutlassDecodeMetadata( + plan=plan, + page_table=block_table.view(-1), + ) + + +def _supported_head_geometry(num_q_heads: int, num_kv_heads: int) -> bool: + return ( + 0 < num_q_heads <= _MAX_NUM_Q_HEADS + and 0 < num_kv_heads <= _MAX_NUM_KV_HEADS + and num_q_heads % num_kv_heads == 0 + ) + + +def supports_cutlass_sparse_decode( + *, + decode_backend: MiniMaxM3MSADecodeBackend, + num_q_heads: int, + num_kv_heads: int, + kv_cache_dtype: str, + page_size: int, + topk_blocks: int, +) -> bool: + """Return whether static model geometry supports CUTLASS sparse decode.""" + return ( + decode_backend == "cutlass" + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and kv_cache_dtype in ("fp8", "fp8_e4m3") + and _supported_head_geometry(num_q_heads, num_kv_heads) + and page_size == _PAGE_SIZE + and topk_blocks == _TOPK + ) + + +def should_prepare_decode_metadata( + batch_size: int, + decode_query_len: int, + *, + decode_backend: MiniMaxM3MSADecodeBackend, + num_q_heads: int, + num_kv_heads: int, + kv_cache_dtype: str, + page_size: int, + topk_blocks: int, +) -> bool: + """Return whether a graph shape can use the CUTLASS decode path.""" + total_q = batch_size * decode_query_len + return ( + supports_cutlass_sparse_decode( + decode_backend=decode_backend, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + kv_cache_dtype=kv_cache_dtype, + page_size=page_size, + topk_blocks=topk_blocks, + ) + and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN + and batch_size >= _MIN_CUTLASS_BATCH_SIZE + and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS + ) + + +@torch.no_grad() +def prepare_decode_metadata( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + decode_query_len: int, + *, + num_q_heads: int, + num_kv_heads: int, + page_size: int, + topk_blocks: int, + plan_cache: MSACutlassDecodePlanCache | None = None, +) -> MSACutlassDecodeMetadata: + """Prepare graph-stable runtime metadata for one sparse decode step.""" + cache = plan_cache or MSACutlassDecodePlanCache() + return cache.prepare( + block_table, + seq_lens, + seq_lens_cpu, + decode_query_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=page_size, + topk_blocks=topk_blocks, + ) + + +@torch.no_grad() +def msa_cutlass_sparse_decode( + query_fp8: torch.Tensor, + kv_cache: torch.Tensor, + topk: torch.Tensor, + output: torch.Tensor, + metadata: MSACutlassDecodeMetadata, + *, + scale: float, + q_scale_float: float, + k_scale_float: float, + v_scale_float: float, +) -> None: + """Run CUTLASS sparse decode with metadata prepared by the MSA builder.""" + key, value = kv_cache.split(_HEAD_DIM, dim=-1) + + from vllm.third_party.fmha_sm100.api import fmha_sm100 + + fmha_sm100( + query_fp8, + key, + value, + metadata.plan, + kv_indices=metadata.page_table, + kv_block_indexes=topk, + out=output, + output_maxscore=False, + output_o=True, + sm_scale=scale, + q_scale=q_scale_float, + k_scale=k_scale_float, + v_scale=v_scale_float, + o_scale=1.0, + ) diff --git a/vllm/models/minimax_m3/nvidia/mtp.py b/vllm/models/minimax_m3/nvidia/mtp.py index e2c7f8821d96..832c872a1006 100644 --- a/vllm/models/minimax_m3/nvidia/mtp.py +++ b/vllm/models/minimax_m3/nvidia/mtp.py @@ -19,6 +19,9 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -304,7 +307,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Validate that weights were loaded for each MTP layer. for layer_idx in range(self.model.num_mtp_layers): - if layer_idx not in loaded_mtp_layers: + if ( + layer_idx not in loaded_mtp_layers + and is_mtp_completeness_check_enabled() + ): raise ValueError( f"Failed to load MTP layer {layer_idx} weights from checkpoint." ) diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index 8666eaa28e70..d29e1346bd27 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -1,28 +1,198 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MSA (SM100/Blackwell) block-sparse attend for MiniMax M3. +"""MSA (SM100/Blackwell) block-sparse attention for MiniMax M3. -Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``); -decode falls back to the Triton split-K kernel (no MSA decode yet). ``fmha_sm100`` -imports are function-local, so this module is import-safe on AMD/non-SM100. +Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``). +Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100`` +path for regular decode and speculative verification. """ +from dataclasses import dataclass + import torch +from vllm.config import VllmConfig +from vllm.config.attention import MiniMaxM3MSADecodeBackend from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.models.minimax_m3.common.ops.sparse_attn import ( SPARSE_BLOCK_SIZE, minimax_m3_sparse_attn_decode, ) from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseBackend, + MiniMaxM3SparseDecodeMetadata, MiniMaxM3SparseImpl, MiniMaxM3SparseMetadata, + MiniMaxM3SparseMetadataBuilder, +) +from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( + MSACutlassDecodeMetadata, + MSACutlassDecodePlanCache, + msa_cutlass_sparse_decode, + prepare_decode_metadata, + should_prepare_decode_metadata, + supports_cutlass_sparse_decode, +) +from vllm.v1.attention.backend import ( + AttentionLayer, + CommonAttentionMetadata, ) -from vllm.v1.attention.backend import AttentionLayer +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend): + """MiniMax M3 backend with NVIDIA MSA-specific decode metadata.""" + + @staticmethod + def get_builder_cls() -> type["MiniMaxM3SparseMSAMetadataBuilder"]: + return MiniMaxM3SparseMSAMetadataBuilder + + +class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend): + """Attention-backend alias selecting CUTLASS MSA sparse decode.""" + + @staticmethod + def get_name() -> str: + return "CUTLASS_MSA" + + +class MiniMaxM3SparseTritonBackend(MiniMaxM3SparseMSABackend): + """Attention-backend alias selecting Triton MSA sparse decode.""" + + @staticmethod + def get_name() -> str: + return "TRITON_MSA" + + +@dataclass +class MiniMaxM3SparseMSADecodeMetadata(MiniMaxM3SparseDecodeMetadata): + msa_cutlass: MSACutlassDecodeMetadata | None = None + + +class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder): + """Prepare MSA plans only for decode shapes supported by ``fmha_sm100``.""" + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + config = vllm_config.model_config.hf_text_config + tp_size = vllm_config.parallel_config.tensor_parallel_size + self.num_q_heads = config.num_attention_heads // tp_size + self.num_kv_heads = kv_cache_spec.num_kv_heads + self.topk_blocks = config.sparse_attention_config["sparse_topk_blocks"] + # AttentionSpec stores every FP8 mode as uint8, so retain the configured + # format to distinguish E4M3 (supported) from E5M2 before planning. + self.kv_cache_dtype = vllm_config.cache_config.cache_dtype + self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend + self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache() + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> MiniMaxM3SparseMetadata: + metadata = super().build( + common_prefix_len, + common_attn_metadata, + fast_build, + ) + decode = metadata.decode + if decode is None: + return metadata + + msa_cutlass = None + if should_prepare_decode_metadata( + metadata.num_decodes, + decode.decode_query_len, + decode_backend=self.decode_backend, + num_q_heads=self.num_q_heads, + num_kv_heads=self.num_kv_heads, + kv_cache_dtype=self.kv_cache_dtype, + page_size=SPARSE_BLOCK_SIZE, + topk_blocks=self.topk_blocks, + ): + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + assert seq_lens_cpu is not None + msa_cutlass = prepare_decode_metadata( + decode.block_table, + decode.seq_lens, + seq_lens_cpu[: metadata.num_decodes], + decode.decode_query_len, + num_q_heads=self.num_q_heads, + num_kv_heads=self.num_kv_heads, + page_size=SPARSE_BLOCK_SIZE, + topk_blocks=self.topk_blocks, + plan_cache=self.msa_cutlass_plan_cache, + ) + metadata.decode = MiniMaxM3SparseMSADecodeMetadata( + seq_lens=decode.seq_lens, + block_table=decode.block_table, + decode_query_len=decode.decode_query_len, + msa_cutlass=msa_cutlass, + ) + return metadata class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): - """MSA block-sparse attend (``fmha_sm100``); Triton split-K decode.""" + """MSA block-sparse attention with guarded CUTLASS sparse decode.""" + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + kv_cache_dtype: str = "auto", + *, + topk_blocks: int, + sparse_block_size: int, + msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton", + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + kv_cache_dtype, + topk_blocks=topk_blocks, + sparse_block_size=sparse_block_size, + ) + self.use_cutlass_decode = supports_cutlass_sparse_decode( + decode_backend=msa_decode_backend, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + kv_cache_dtype=self.kv_cache_dtype, + page_size=self.block_size, + topk_blocks=self.topk_blocks, + ) + logger.info_once( + "MiniMax M3 MSA sparse decode selected %s", + "CUTLASS" if self.use_cutlass_decode else "Triton", + ) + + def should_use_msa_decode(self, layer_name: str) -> bool: + if not self.use_cutlass_decode: + return False + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return False + main_md = attn_metadata[layer_name] + if not isinstance(main_md, MiniMaxM3SparseMetadata): + return False + decode = main_md.decode + return ( + isinstance(decode, MiniMaxM3SparseMSADecodeMetadata) + and decode.msa_cutlass is not None + ) def forward( self, @@ -30,6 +200,8 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, + *, + query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): @@ -52,23 +224,42 @@ def forward( k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None - # Decode [:nd]: Triton split-K placeholder (no MSA decode yet). + # Decode [:nd]: CUTLASS for planned shapes, otherwise Triton. if main_md.num_decodes > 0: d = main_md.decode assert d is not None - minimax_m3_sparse_attn_decode( - q[:nd], - kv_cache, - topk[:nd].transpose(0, 1), - d.block_table, - d.seq_lens, - self.num_kv_heads, - self.scale, - out[:nd], - d.decode_query_len, - k_scale=k_scale, - v_scale=v_scale, + msa_metadata = ( + d.msa_cutlass + if isinstance(d, MiniMaxM3SparseMSADecodeMetadata) + else None ) + if self.use_cutlass_decode and msa_metadata is not None: + assert query_fp8 is not None + msa_cutlass_sparse_decode( + query_fp8[:nd].view(-1, self.num_heads, hd), + kv_cache, + topk[:nd], + out[:nd], + msa_metadata, + scale=self.scale, + q_scale_float=getattr(layer, "_q_scale_float", 1.0), + k_scale_float=getattr(layer, "_k_scale_float", 1.0), + v_scale_float=getattr(layer, "_v_scale_float", 1.0), + ) + else: + minimax_m3_sparse_attn_decode( + q[:nd], + kv_cache, + topk[:nd].transpose(0, 1), + d.block_table, + d.seq_lens, + self.num_kv_heads, + self.scale, + out[:nd], + d.decode_query_len, + k_scale=k_scale, + v_scale=v_scale, + ) # Prefill [nd:]: MSA sparse FMHA over the selected blocks. if main_md.num_prefills > 0: diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py index e9470c6e027b..50834efe5ca7 100644 --- a/vllm/multimodal/audio.py +++ b/vllm/multimodal/audio.py @@ -336,8 +336,8 @@ def split_audio( for splitting. Args: - audio_data: Audio array to split. Can be 1D (mono) or multi-dimensional. - Splits along the last dimension (time axis). + audio_data: 1D mono audio array to split. ASR models consume mono, so + callers must downmix before chunking. sample_rate: Sample rate of the audio in Hz. max_clip_duration_s: Maximum duration of each chunk in seconds. overlap_duration_s: Overlap duration in seconds between consecutive chunks. @@ -345,8 +345,10 @@ def split_audio( min_energy_window_size: Window size in samples for finding low-energy regions. Returns: - List of audio chunks. Each chunk is a numpy array with the same shape - as the input except for the last (time) dimension. + List of 1D audio chunks. + + Raises: + AssertionError: If ``audio_data`` is not 1D. Example: >>> audio = np.random.randn(1040000) # 65 seconds at 16kHz @@ -360,6 +362,11 @@ def split_audio( >>> len(chunks) 3 """ + if audio_data.ndim > 1: + raise ValueError( + f"split_audio expects mono audio, got shape {audio_data.shape}" + ) + chunk_size = int(sample_rate * max_clip_duration_s) overlap_size = int(sample_rate * overlap_duration_s) chunks = [] @@ -402,7 +409,7 @@ def find_split_point( RMS energy in sliding windows. Args: - wav: Audio array. Can be 1D or multi-dimensional. + wav: 1D mono audio array. start_idx: Start index of search region (inclusive). end_idx: End index of search region (exclusive). min_energy_window: Window size in samples for energy calculation. diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index c1ff600869bb..3042091611f9 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/vllm/multimodal/encoder_budget.py @@ -4,6 +4,7 @@ from vllm.config import ModelConfig, VllmConfig from vllm.logger import init_logger +from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.multimodal.registry import MultiModalRegistry from vllm.utils.torch_utils import set_default_torch_num_threads @@ -48,6 +49,8 @@ def __init__( self, vllm_config: VllmConfig, mm_registry: MultiModalRegistry, + *, + enable_cache: bool = True, ) -> None: super().__init__() @@ -58,7 +61,11 @@ def __init__( self.max_num_reqs = scheduler_config.max_num_seqs with set_default_torch_num_threads(): # Avoid hang during startup - cache = mm_registry.processor_only_cache_from_config(vllm_config) + cache = ( + mm_registry.processor_only_cache_from_config(vllm_config) + if enable_cache + else None + ) processor = mm_registry.create_processor(model_config, cache=cache) self.cache = cache @@ -191,3 +198,23 @@ def get_encoder_budget(self) -> int: def reset_cache(self) -> None: if self.cache is not None: self.cache.clear_cache() + + +def get_dummy_encoder_profile_inputs( + mm_registry: MultiModalRegistry, + budget: MultiModalBudget, +) -> list[tuple[str, MultiModalKwargsItem]]: + if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item: + return [] + + modality = budget.get_modality_with_max_tokens() + max_items_per_batch = budget.mm_max_items_per_batch[modality] + dummy_mm_inputs = mm_registry.get_dummy_mm_inputs( + budget.model_config, + mm_counts={modality: 1}, + processor=budget.processor, + ) + dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] + assert dummy_mm_item is not None, "Dummy item should be generated" + + return [(modality, dummy_mm_item)] * max_items_per_batch diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py index 15b912064a85..58cb98888021 100644 --- a/vllm/multimodal/gpu_ipc_memory.py +++ b/vllm/multimodal/gpu_ipc_memory.py @@ -17,9 +17,14 @@ """ import threading +from typing import TYPE_CHECKING from vllm.logger import init_logger from vllm.utils.mem_constants import GiB_bytes +from vllm.utils.mem_utils import format_gib + +if TYPE_CHECKING: + from vllm.config.multimodal import MultiModalConfig logger = init_logger(__name__) @@ -145,3 +150,108 @@ def maybe_init_mm_gpu_ipc_pool( api_process_count, ) return pool + + +def reserve_mm_ipc_gpu_memory( + available_kv_cache_memory_bytes: int, + mm_config: "MultiModalConfig | None", + api_process_count: int = 1, +) -> int: + """Return KV-cache memory remaining after frontend multimodal reservations. + + The reservation covers: + + * The total ``mm_ipc_gpu_memory_gb`` budget for transient decoded-frame + buffers. This budget is divided among API processes, so it is not + multiplied by ``api_process_count``. + * For GPU video backends, a fixed upper bound for each API process's + retained decoder surfaces and CUDA context. The PyNvVideoCodec surface + reservation scales with its configured ``hw_decoders`` value, and the + entire decoder reservation scales with ``api_process_count`` because + these resources are not shared between processes. + + Args: + available_kv_cache_memory_bytes: KV-cache capacity before reserving + memory for frontend multimodal processing. + mm_config: Multimodal configuration, or ``None`` when multimodal + processing is disabled. + api_process_count: Number of frontend API processes sharing the GPU. + Values below one are treated as one. + + Returns: + KV-cache capacity after subtracting the frontend reservation. + + Raises: + ValueError: If the reservation leaves no memory for the KV cache. + """ + if mm_config is None: + return available_kv_cache_memory_bytes + + from vllm import envs + from vllm.multimodal.video import ( + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, + validate_pynvvideocodec_hw_decoders, + ) + + raw_frame_reserved_bytes = int(mm_config.mm_ipc_gpu_memory_gb * GiB_bytes) + # Each API server process runs its own decoder surfaces and NVDEC/CUVID CUDA + # context on the GPU, outside the worker memory pool. Reserve that footprint + # per process so gpu_memory_utilization bounds total GPU usage across them. + num_api_servers = max(1, api_process_count) + 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") + uses_pynvvideocodec = ( + video_loader_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + or codec_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + ) + hw_decoders = ( + validate_pynvvideocodec_hw_decoders( + video_kwargs.get("hw_decoders", PYNVVIDEOCODEC_DEFAULT_HW_DECODERS) + ) + if uses_pynvvideocodec + else 1 + ) + per_server_decoder_bytes = ( + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * hw_decoders + + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES + ) + decoder_reserved_bytes = ( + num_api_servers * per_server_decoder_bytes + if mm_config.use_gpu_video_backend() + else 0 + ) + reserved_bytes = raw_frame_reserved_bytes + decoder_reserved_bytes + if reserved_bytes <= 0: + return available_kv_cache_memory_bytes + + remaining = available_kv_cache_memory_bytes - reserved_bytes + if remaining <= 0: + raise ValueError( + f"frontend multimodal GPU decoding reserves " + f"{format_gib(reserved_bytes)} GiB " + f"({format_gib(raw_frame_reserved_bytes)} GiB raw-frame budget, " + f"{format_gib(decoder_reserved_bytes)} GiB decoder cache budget), " + f"but only {format_gib(available_kv_cache_memory_bytes)} GiB is " + "available for the KV cache. Reduce mm_ipc_gpu_memory_gb or " + "hw_decoders, use a different video backend, or increase " + "gpu_memory_utilization." + ) + logger.info_once( + "Reserving %s GiB of GPU memory for frontend multimodal decoding " + "(%s GiB raw-frame semaphore budget, %s GiB decoder+CUDA-context " + "across %d API server(s) @ %s GiB/server); " + "KV cache memory reduced to %s GiB.", + format_gib(reserved_bytes), + format_gib(raw_frame_reserved_bytes), + format_gib(decoder_reserved_bytes), + num_api_servers, + format_gib(per_server_decoder_bytes), + format_gib(remaining), + ) + return remaining diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 6caf9c114277..3d86b33d5138 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -11,7 +11,7 @@ import torch from PIL import Image -import vllm.envs as envs +from vllm.config.multimodal import MMHasherAlgorithm from vllm.logger import init_logger from .media import MediaWithBytes @@ -20,7 +20,9 @@ @functools.lru_cache(maxsize=3) -def _get_hasher_factory(algorithm: str) -> Callable[[], "hashlib._Hash"]: +def _get_hasher_factory( + algorithm: MMHasherAlgorithm, +) -> Callable[[], "hashlib._Hash"]: """ Get the hasher factory based on the configured algorithm. @@ -32,7 +34,6 @@ def _get_hasher_factory(algorithm: str) -> Callable[[], "hashlib._Hash"]: See: https://github.com/vllm-project/vllm/issues/18334 """ - algorithm = algorithm.lower() if algorithm == "blake3": from blake3 import blake3 @@ -43,7 +44,7 @@ def _get_hasher_factory(algorithm: str) -> Callable[[], "hashlib._Hash"]: elif algorithm == "sha512": return hashlib.sha512 else: - # This should never happen due to env_with_choices validation + # This should never happen due to config validation raise ValueError(f"Unsupported hash algorithm: {algorithm}") @@ -81,8 +82,19 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: ): return (exif[Image.ExifTags.Base.ImageID].bytes,) + if obj.io_config: + return cls.iter_item_to_bytes( + "image", + {"io_config": obj.io_config, "data": obj.original_bytes}, + ) return cls.iter_item_to_bytes("image", obj.original_bytes) + if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray): + frames = obj.media + if frames.nbytes < len(obj.original_bytes): + return cls.iter_item_to_bytes("video", frames) + return cls.iter_item_to_bytes("video", obj.original_bytes) + if isinstance(obj, torch.Tensor): tensor_obj: torch.Tensor = obj.cpu() tensor_dtype = tensor_obj.dtype @@ -151,8 +163,13 @@ def iter_item_to_bytes( yield from cls.serialize_item(obj) @classmethod - def hash_kwargs(cls, **kwargs: object) -> str: - hasher_factory = _get_hasher_factory(envs.VLLM_MM_HASHER_ALGORITHM) + def hash_kwargs( + cls, + algorithm: MMHasherAlgorithm, + /, + **kwargs: object, + ) -> str: + hasher_factory = _get_hasher_factory(algorithm) hasher = hasher_factory() for k, v in sorted(kwargs.items(), key=lambda kv: kv[0]): diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index c55bbe4623cb..ab04dc71e8cd 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -66,7 +66,10 @@ """ VideoItem: TypeAlias = Union[ - HfVideoItem, "torch.Tensor", tuple[HfVideoItem, dict[str, Any]] + HfVideoItem, + "torch.Tensor", + tuple[HfVideoItem, dict[str, Any]], + MediaWithBytes[tuple[HfVideoItem, dict[str, Any]]], ] """ A `transformers.video_utils.VideoInput` representing a single video item. @@ -226,40 +229,57 @@ def __eq__(self, other: object) -> bool: """ -def nested_tensors_equal(a: NestedTensors, b: NestedTensors) -> bool: +def nested_tensors_equal( + a: NestedTensors, + b: NestedTensors, + check_dtype: bool = True, +) -> bool: """ Equality check between [`NestedTensors`][vllm.multimodal.inputs.NestedTensors] objects. + + If `check_dtype` is `True`, the tensors must have the same dtype. """ + check_dtype_func = ( + lambda a, b, check_dtype: a.dtype == b.dtype if check_dtype else True + ) if isinstance(a, torch.Tensor): - return isinstance(b, torch.Tensor) and torch.equal(a, b) + return ( + isinstance(b, torch.Tensor) + and torch.equal(a, b) + and check_dtype_func(a, b, check_dtype) + ) elif isinstance(b, torch.Tensor): - return isinstance(a, torch.Tensor) and torch.equal(b, a) + return ( + isinstance(a, torch.Tensor) + and torch.equal(b, a) + and check_dtype_func(b, a, check_dtype) + ) if isinstance(a, list): return ( isinstance(b, list) and len(a) == len(b) - and all(nested_tensors_equal(a_, b_) for a_, b_ in zip(a, b)) + and all(nested_tensors_equal(a_, b_, check_dtype) for a_, b_ in zip(a, b)) ) if isinstance(b, list): return ( isinstance(a, list) and len(b) == len(a) - and all(nested_tensors_equal(b_, a_) for b_, a_ in zip(b, a)) + and all(nested_tensors_equal(b_, a_, check_dtype) for b_, a_ in zip(b, a)) ) if isinstance(a, tuple): return ( isinstance(b, tuple) and len(a) == len(b) - and all(nested_tensors_equal(a_, b_) for a_, b_ in zip(a, b)) + and all(nested_tensors_equal(a_, b_, check_dtype) for a_, b_ in zip(a, b)) ) if isinstance(b, tuple): return ( isinstance(a, tuple) and len(b) == len(a) - and all(nested_tensors_equal(b_, a_) for b_, a_ in zip(b, a)) + and all(nested_tensors_equal(b_, a_, check_dtype) for b_, a_ in zip(b, a)) ) # Both a and b are scalars @@ -454,6 +474,8 @@ def reduce_data( device = "cpu" if pin_memory and self.keep_on_cpu: pin_memory = False + if device == "cpu" or device == torch.device("cpu"): + pin_memory = False batch = [elem.data for elem in elems] out = self._reduce_data(batch, pin_memory=pin_memory) diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index 91e7a4947170..2927914e6459 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, cast import numpy as np @@ -22,16 +23,28 @@ class MediaWithBytes(Generic[_T]): The wrapper delegates attribute access to the underlying media object, making it behave transparently like the wrapped type (e.g., PIL.Image). - NOTE: Currently, this wrapper is used only for the image modality. + NOTE: Currently, this wrapper is used only for the image and video + modalities. """ media: _T original_bytes: bytes = field(repr=False) + io_config: dict[str, Any] | None = None + """Decode settings that altered the media relative to `original_bytes` + (e.g. `image_mode` conversion), so they participate in cache hashing.""" def __array__(self, *args, **kwargs) -> np.ndarray: """Allow np.array(obj) to return np.array(obj.media).""" return np.array(self.media, *args, **kwargs) + def __iter__(self) -> Iterator[Any]: + """Allow unpacking obj to unpack obj.media (e.g. video tuples).""" + return iter(cast(Iterable[Any], self.media)) + + def __getitem__(self, index: Any) -> Any: + """Allow obj[i] to index obj.media (e.g. video tuples).""" + return cast(Any, self.media)[index] + def __getstate__(self): return self.__dict__.copy() diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index a440e69ef1ca..fed41657e98c 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -29,7 +29,7 @@ from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO +from .base import MediaIO, MediaWithBytes from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VideoMediaIO @@ -302,14 +302,18 @@ def _load_data_url( media_io: MediaIO[_M], ) -> _M: # type: ignore[type-var] # Format per RFC 2397: - # data:[][;base64], - data_spec, data = url[5:].split(",", 1) - media_type, data_type = data_spec.split(";", 1) - - if data_type != "base64": + # data:[][;=]*[;base64], + data_spec, sep, data = url[5:].partition(",") + if not sep: + msg = f"Invalid data URL {url[:32]!r}: missing ',' separator." + raise ValueError(msg) + + media_type, sep, encoding = data_spec.rpartition(";") + if not sep or encoding != "base64": msg = "Only base64 data URLs are supported for now." raise NotImplementedError(msg) + media_type = media_type.partition(";")[0] return media_io.load_base64(media_type, data) def _load_file_url( @@ -476,15 +480,17 @@ def fetch_image( self, image_url: str, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", ) -> Image.Image: """ Load a PIL image from an HTTP or base64 data URL. - By default, the image is converted into RGB format. + By default, the image is converted into RGB format. Set + `media_io_kwargs={"image": {"image_mode": None}}` to keep the + original image mode (e.g. preserving the alpha channel). """ image_io = ImageMediaIO( - image_mode=image_mode, **self.media_io_kwargs.get("image", {}) + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) ) try: @@ -501,15 +507,17 @@ async def fetch_image_async( self, image_url: str, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", ) -> Image.Image: """ Asynchronously load a PIL image from an HTTP or base64 data URL. - By default, the image is converted into RGB format. + By default, the image is converted into RGB format. Set + `media_io_kwargs={"image": {"image_mode": None}}` to keep the + original image mode (e.g. preserving the alpha channel). """ image_io = ImageMediaIO( - image_mode=image_mode, **self.media_io_kwargs.get("image", {}) + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) ) try: @@ -526,14 +534,14 @@ def fetch_video( self, video_url: str, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", video_processor: str | None = None, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Load video from an HTTP or base64 data URL. """ image_io = ImageMediaIO( - image_mode=image_mode, **self.media_io_kwargs.get("image", {}) + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) ) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( @@ -552,16 +560,18 @@ async def fetch_video_async( self, video_url: str, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", video_processor: str | None = None, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Asynchronously load video from an HTTP or base64 data URL. - By default, the image is converted into RGB format. + By default, the image is converted into RGB format. Set + `media_io_kwargs={"image": {"image_mode": None}}` to keep the + original image mode (e.g. preserving the alpha channel). """ image_io = ImageMediaIO( - image_mode=image_mode, **self.media_io_kwargs.get("image", {}) + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) ) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( @@ -587,6 +597,20 @@ def fetch_image_embedding( return image_embedding_io.load_base64("", data) + async def fetch_image_embedding_async( + self, + data: str, + ) -> torch.Tensor: + """ + Asynchronously load image embedding from a URL. + """ + image_embedding_io = ImageEmbeddingMediaIO() + loop = asyncio.get_running_loop() + + return await loop.run_in_executor( + global_thread_pool, image_embedding_io.load_base64, "", data + ) + def fetch_audio_embedding( self, data: str, @@ -597,3 +621,17 @@ def fetch_audio_embedding( audio_embedding_io = AudioEmbeddingMediaIO() return audio_embedding_io.load_base64("", data) + + async def fetch_audio_embedding_async( + self, + data: str, + ) -> torch.Tensor: + """ + Asynchronously load audio embedding from a URL. + """ + audio_embedding_io = AudioEmbeddingMediaIO() + loop = asyncio.get_running_loop() + + return await loop.run_in_executor( + global_thread_pool, audio_embedding_io.load_base64, "", data + ) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index 8a49574612e5..d9fc2132ed55 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -25,9 +25,11 @@ class ImageMediaIO(MediaIO[Image.Image]): error handling. """ - def __init__(self, image_mode: str = "RGB", **kwargs) -> None: + def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None: super().__init__() + # Target mode for loaded images; `None` keeps the original mode + # (i.e. no conversion, alpha channel is preserved as-is). self.image_mode = image_mode # `kwargs` contains custom arguments from # --media-io-kwargs for this modality, merged with @@ -62,7 +64,7 @@ def _convert_image_mode( """Convert image mode with custom background color.""" if isinstance(image, MediaWithBytes): image = image.media - if image.mode == self.image_mode: + if self.image_mode is None or image.mode == self.image_mode: return image elif image.mode == "RGBA" and self.image_mode == "RGB": return rgba_to_rgb(image, self.rgba_background_color) @@ -84,10 +86,17 @@ def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]: ) image = normalize_image(image) image.load() - image = self._convert_image_mode(image) + converted = self._convert_image_mode(image) except (OSError, Image.UnidentifiedImageError) as e: raise ValueError(f"Failed to load image: {e}") from e - return MediaWithBytes(image, data) + + io_config = None + if converted is not image: + io_config = { + "image_mode": self.image_mode, + "rgba_background_color": self.rgba_background_color, + } + return MediaWithBytes(converted, data, io_config) def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]: return self.load_bytes(pybase64.b64decode(data, validate=True)) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 45ea4c2fdf45..124dcd7f7e50 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -13,13 +13,13 @@ from vllm.logger import init_logger from ..video import VIDEO_LOADER_REGISTRY -from .base import MediaIO +from .base import MediaIO, MediaWithBytes from .image import ImageMediaIO logger = init_logger(__name__) -class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): +class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): """Configuration values can be user-provided either by --media-io-kwargs or by the runtime API field "media_io_kwargs". Ensure proper validation and error handling. @@ -32,6 +32,10 @@ def merge_kwargs( runtime_kwargs: dict[str, Any] | None, ) -> dict[str, Any]: if runtime_kwargs: + # Decoder GPU memory is reserved from the startup value. + runtime_kwargs = dict(runtime_kwargs) + runtime_kwargs.pop("hw_decoders", None) + # Block request-level selection of GPU video backends that # were not configured (and VRAM-reserved) at startup. for key in ("video_backend", "backend"): @@ -87,14 +91,17 @@ def __init__( self.kwargs = kwargs self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend) - def load_bytes(self, data: bytes) -> tuple[npt.NDArray, dict[str, Any]]: - return self.video_loader.load_bytes( + def load_bytes( + self, data: bytes + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) + return MediaWithBytes(video, data) def load_base64( self, media_type: str, data: str - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: if media_type.lower() == "video/jpeg": load_frame = partial( self.image_io.load_base64, @@ -156,11 +163,13 @@ def load_base64( "frames_indices": frames_indices, "do_sample_frames": self.kwargs.get("do_sample_frames", False), } - return frames, metadata + return MediaWithBytes((frames, metadata), data.encode()) return self.load_bytes(pybase64.b64decode(data)) - def load_file(self, filepath: Path) -> tuple[npt.NDArray, dict[str, Any]]: + def load_file( + self, filepath: Path + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: with filepath.open("rb") as f: data = f.read() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 9d1e005da7f4..729a5bad18a1 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -366,6 +366,20 @@ def __init__( self.metadata = metadata + def _unwrap(self, item: Any) -> Any: + if isinstance(item, tuple): + frames, metadata = item + return super()._unwrap(frames), metadata + return super()._unwrap(item) + + def get_item_for_hash(self, index: int) -> Any: + item = self.data[index] + if isinstance(item, MediaWithBytes) and isinstance(self.metadata, list): + metadata = self.metadata[index] + if metadata is not None: + return item, metadata + return item + def get_num_frames(self, item_idx: int) -> int: video = self.get(item_idx) if video is None: @@ -552,7 +566,10 @@ def _get_audio_with_sr( def _get_video_with_metadata( self, video: VideoItem, - ) -> tuple[np.ndarray, dict[str, Any] | None]: + ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]: + if isinstance(video, MediaWithBytes): + new_video, metadata = self._get_video_with_metadata(video.media) + return MediaWithBytes(new_video, video.original_bytes), metadata if isinstance(video, tuple): return video if isinstance(video, list): @@ -653,7 +670,11 @@ def _parse_video_data( else: data_items = data # type: ignore[assignment] - new_videos = list[tuple[np.ndarray, dict[str, Any] | None]]() + new_videos = list[ + np.ndarray + | MediaWithBytes[np.ndarray] + | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]] + ]() metadata_lst: list[dict[str, Any] | None] = [] for data_item in data_items: video, metadata = self._get_video_with_metadata(data_item) @@ -664,12 +685,9 @@ def _parse_video_data( "Please check your video input in `multi_modal_data`" ) new_videos.append((video, metadata)) - metadata_lst.append(metadata) else: new_videos.append(video) - - if not self.video_needs_metadata: - metadata = None + metadata_lst.append(metadata) return VideoProcessorItems(new_videos, metadata=metadata_lst) diff --git a/vllm/multimodal/processing/inputs.py b/vllm/multimodal/processing/inputs.py index ae12f4e51b86..73a734d9c2fe 100644 --- a/vllm/multimodal/processing/inputs.py +++ b/vllm/multimodal/processing/inputs.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field +from vllm.config.multimodal import MMHasherAlgorithm from vllm.inputs import MultiModalHashes from ..hasher import MultiModalHasher @@ -22,7 +23,11 @@ class ProcessorInputs: hf_processor_mm_kwargs: Mapping[str, object] = field(default_factory=dict) tokenization_kwargs: Mapping[str, object] = field(default_factory=dict) - def get_mm_hashes(self, model_id: str) -> MultiModalHashes: + def get_mm_hashes( + self, + model_id: str, + hash_algorithm: MMHasherAlgorithm, + ) -> MultiModalHashes: mm_data_items = self.mm_data_items mm_uuid_items = self.mm_uuid_items or {} hf_processor_mm_kwargs = self.hf_processor_mm_kwargs @@ -49,6 +54,7 @@ def get_mm_hashes(self, model_id: str) -> MultiModalHashes: item = uuid_item if uuid_item is not None else item hashes.append( hasher.hash_kwargs( + hash_algorithm, model_id=model_id, **{modality: item}, **hf_processor_mm_kwargs, @@ -61,11 +67,12 @@ def get_mm_hashes(self, model_id: str) -> MultiModalHashes: else: mm_hashes[modality] = [ hasher.hash_kwargs( + hash_algorithm, model_id=model_id, **{modality: item}, **hf_processor_mm_kwargs, ) - for item in data_items + for item in data_items.get_all_items_for_hash() ] return mm_hashes diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 7b24cd3fcb5e..f43b45ff7e4d 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1422,7 +1422,10 @@ def _apply_hf_processor( # Use overrides if provided; fallback to data-dependent hashing. with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) mm_prompt_updates = self._get_mm_prompt_updates( inputs.mm_data_items, @@ -1454,7 +1457,10 @@ def _cached_apply_hf_processor( return self._apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) with timing_ctx.record("get_cache_missing_items"): mm_is_cached, mm_missing_data_items = self._get_cache_missing_items( diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index f34769d838bd..25c811fd424f 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -16,15 +16,21 @@ from vllm.inputs import MultiModalPlaceholders from vllm.utils.import_utils import LazyLoader -from .hasher import MultiModalHasher from .inputs import ( BatchedTensorInputs, MultiModalFeatureSpec, MultiModalFieldElem, MultiModalKwargsItem, MultiModalSharedField, + nested_tensors_equal, +) +from .media import ( + AudioMediaIO, + ImageMediaIO, + MediaConnector, + MediaWithBytes, + VideoMediaIO, ) -from .media import AudioMediaIO, ImageMediaIO, MediaConnector, VideoMediaIO if TYPE_CHECKING: import torch.types @@ -58,13 +64,14 @@ def encode_audio_url( def encode_image_base64( image: Image.Image, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", format: str = "PNG", ) -> str: """ Encode a pillow image to base64 format. By default, the image is converted into RGB format before being encoded. + Pass `image_mode=None` to keep the original image mode. """ image_io = ImageMediaIO(image_mode=image_mode) return image_io.encode_base64(image, image_format=format) @@ -73,13 +80,14 @@ def encode_image_base64( def encode_image_url( image: Image.Image, *, - image_mode: str = "RGB", + image_mode: str | None = "RGB", format: str = "PNG", ) -> str: """ Encode a pillow image as a data URL. By default, the image is converted into RGB format before being encoded. + Pass `image_mode=None` to keep the original image mode. """ image_b64 = encode_image_base64(image, image_mode=image_mode, format=format) mimetype = mimetypes.types_map.get("." + format.lower(), "image") @@ -157,11 +165,28 @@ def argsort_mm_positions( return [(modality, idx) for modality, idx, _ in sorted_flat_items] -def _get_group_hash(elem: MultiModalFieldElem): - if not isinstance(elem.field, MultiModalSharedField): - return None +def _can_batch_mm_items( + left: MultiModalKwargsItem, + right: MultiModalKwargsItem, +) -> bool: + if left.keys() != right.keys(): + return False + + for key, left_elem in left.items(): + right_elem = right[key] + left_field, right_field = left_elem.field, right_elem.field + is_shared_field = isinstance(left_field, MultiModalSharedField) and isinstance( + right_field, MultiModalSharedField + ) + if (type(left_field) is not type(right_field)) or ( + is_shared_field + and not nested_tensors_equal( + left_elem.data, right_elem.data, check_dtype=True + ) + ): + return False - return MultiModalHasher.hash_kwargs(data=elem.data) + return True def _batch_mm_items( @@ -209,26 +234,21 @@ def group_and_batch_mm_items( - `kwargs` is a dictionary of keyword arguments to pass to the model; - `num_items` is the corresponding number of items. """ - group_ids = [ - tuple( - (key, _get_group_hash(elem)) - for key, elem in sorted(item.items(), key=lambda kv: kv[0]) - ) - for item in items - ] - group_sizes = [sum(1 for _ in group) for _, group in groupby(group_ids)] - start_idx = 0 - for group_size in group_sizes: + for end_idx in range(1, len(items) + 1): + if end_idx < len(items) and _can_batch_mm_items( + items[end_idx - 1], items[end_idx] + ): + continue + group_data = _batch_mm_items( - items[start_idx : start_idx + group_size], + items[start_idx:end_idx], device=device, pin_memory=pin_memory, ) - yield group_size, group_data - - start_idx += group_size + yield end_idx - start_idx, group_data + start_idx = end_idx assert start_idx == len(items) @@ -330,7 +350,7 @@ def fetch_image( def fetch_video( video_url: str, video_io_kwargs: dict[str, Any] | None = None, -) -> tuple[npt.NDArray, dict[str, Any]]: +) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Args: video_url: URL of the video file to fetch. diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index ca111f0b3f86..3e0969f1ad9a 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -210,15 +210,25 @@ def create_hf_metadata( VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" -# Fixed upper bound reserved for persistent PyNvVideoCodec decoder surfaces. +# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces. PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 -PYNVVIDEOCODEC_MAX_RETAINED_DECODERS = 1 +PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2 # Per-API-server CUDA context and driver allocation, measured with # PyNvVideoCodec 2.0.4 on H100. PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes) +def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int: + if ( + isinstance(hw_decoders, bool) + or not isinstance(hw_decoders, int) + or hw_decoders < 1 + ): + raise ValueError("hw_decoders must be a positive integer") + return hw_decoders + + class PyNvVideoCodecDecoderSlot: """A retained PyNv decoder slot and its CUDA stream. @@ -628,6 +638,7 @@ class PyNvVideoCodecVideoBackendMixin: _decoder_slots: ClassVar[list[PyNvVideoCodecDecoderSlot]] = [] _active_decoder_slots: ClassVar[int] = 0 _decoder_slot_cond: ClassVar[threading.Condition] = threading.Condition() + _max_decoder_slots: ClassVar[int | None] = None _DEVICE_INDEX: ClassVar[int] = 0 @classmethod @@ -651,6 +662,18 @@ def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX)) + @classmethod + def _configure_decoder_slots(cls, hw_decoders: object) -> None: + hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) + with cls._decoder_slot_cond: + if cls._max_decoder_slots is None: + cls._max_decoder_slots = hw_decoders + elif cls._max_decoder_slots != hw_decoders: + raise RuntimeError( + "PyNvVideoCodec decoder count is already configured as " + f"{cls._max_decoder_slots}, got {hw_decoders}" + ) + @staticmethod @contextmanager def _torch_stream_context(stream): @@ -669,11 +692,14 @@ def _torch_stream_context(stream): def _borrow_decoder_slot(cls): create_slot = False with cls._decoder_slot_cond: + max_decoder_slots = cls._max_decoder_slots + if max_decoder_slots is None: + raise RuntimeError("PyNvVideoCodec decoder slots are not configured") while True: if cls._decoder_slots: slot = cls._decoder_slots.pop() break - if cls._active_decoder_slots < PYNVVIDEOCODEC_MAX_RETAINED_DECODERS: + if cls._active_decoder_slots < max_decoder_slots: cls._active_decoder_slots += 1 create_slot = True break @@ -999,6 +1025,7 @@ def load_bytes( ] = "opencv", num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", + hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -1025,6 +1052,8 @@ def load_bytes( 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. + hw_decoders: Maximum number of concurrent PyNvVideoCodec decoder + slots. Defaults to 2 and must be a positive integer. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -1088,6 +1117,7 @@ def load_bytes( "frame_recovery is not supported for " f"`{PYNVVIDEOCODEC_VIDEO_BACKEND}` backend" ) + cls._configure_decoder_slots(hw_decoders) frames, source, frame_idx, valid = cls.decode_frames_pynvvideocodec( data, target, @@ -1177,7 +1207,7 @@ def load_bytes( @VIDEO_LOADER_REGISTRY.register( "qwen3_vl", - video_processor="Qwen3VLVideoProcessor", + video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"), ) class Qwen3VLVideoBackend(VideoBackend): @classmethod diff --git a/vllm/multimodal/video_prune/__init__.py b/vllm/multimodal/video_prune/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/multimodal/video_prune/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/multimodal/evs.py b/vllm/multimodal/video_prune/evs.py similarity index 100% rename from vllm/multimodal/evs.py rename to vllm/multimodal/video_prune/evs.py diff --git a/vllm/multimodal/video_prune/vidcom2.py b/vllm/multimodal/video_prune/vidcom2.py new file mode 100644 index 000000000000..830e47bc055f --- /dev/null +++ b/vllm/multimodal/video_prune/vidcom2.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# VidCom2 (Video Compression Commander) video token pruning. +# Liu et al., EMNLP 2025 — https://arxiv.org/abs/2505.14454 +# Adapted from the reference implementation: +# https://github.com/xuyang-liu16/VidCom2 (Apache-2.0, +# Copyright (c) 2025 the VidCom2 authors). + +import torch +import torch.nn.functional as F + +# Multi-scale Gaussian bandwidths from the reference implementation. +_ALPHAS: tuple[float, ...] = tuple(2.0**k for k in range(-3, 2)) +_LOW_VAR_CHANNEL_RATIO: float = 0.5 +_SOFTMAX_TEMPERATURE: float = 0.01 + + +def compute_retained_tokens_count( + tokens_per_frame: int, num_frames: int, q: float +) -> int: + """Number of video tokens retained after VidCom2 pruning. + + The target is `(1 - q) * total_tokens`, i.e. a retention ratio of + `1 - q` averaged across frames. Because the per-frame budget is floored + at one token, the global minimum is `num_frames` (one token per frame). + """ + total_tokens = tokens_per_frame * num_frames + base_num = int(total_tokens * (1.0 - q)) + return max(num_frames, min(base_num, total_tokens)) + + +def compute_retention_mask( + video_embeds: torch.Tensor, + video_size_thw: torch.LongTensor | tuple[int, int, int], + spatial_merge_size: int, + q: float, +) -> torch.Tensor: + """Compute the VidCom2 retention mask for a single video. + + Args: + video_embeds: `(T*H*W/merge^2, hidden_size)` post-ViT token features. + video_size_thw: `(T, H, W)` grid dimensions. + spatial_merge_size: ViT spatial merge factor (e.g. 2). + q: Pruning fraction in `[0, 1)`; retention ratio is `1 - q`. + + Returns: + Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained + tokens. The True count equals `compute_retained_tokens_count` so + placeholders sized at prompt-processing time match exactly. + """ + T, H, W = map(int, video_size_thw) + rows = H // spatial_merge_size + cols = W // spatial_merge_size + tokens_per_frame = rows * cols + total_tokens = T * tokens_per_frame + + device = video_embeds.device + if tokens_per_frame == 0 or total_tokens == 0: + return torch.ones(0, dtype=torch.bool, device=device) + + target_retained = compute_retained_tokens_count( + tokens_per_frame=tokens_per_frame, num_frames=T, q=q + ) + target_retained = min(target_retained, total_tokens) + + # 1. Score in the lowest-variance half of channels. + variances = video_embeds.var(dim=0, unbiased=False) + k_channels = max(1, int(video_embeds.size(-1) * _LOW_VAR_CHANNEL_RATIO)) + _, low_var_idx = torch.topk(variances, k=k_channels, largest=False) + sel = video_embeds.index_select(-1, low_var_idx) + + # 2. Multi-scale Gaussian similarity to video and per-frame centers. + frames = sel.view(T, tokens_per_frame, sel.size(-1)) + frames = F.normalize(frames, dim=-1) + vid_center = frames.mean(dim=(0, 1), keepdim=True) # (1, 1, C) + frame_center = frames.mean(dim=1, keepdim=True) # (T, 1, C) + v_score = _multi_scale_gaussian(frames, vid_center) + f_score = _multi_scale_gaussian(frames, frame_center) + # Higher similarity = more redundant; lowest-similarity tokens are kept. + similarity = v_score + f_score # (T, tpf) + + # 3. Per-frame dynamic budget: distinctive frames get a larger share. + base = 1.0 - q + frame_scores = -v_score.mean(dim=-1) # (T,) + probs = F.softmax((frame_scores - frame_scores.max()) / _SOFTMAX_TEMPERATURE, dim=0) + scales = (base * (1.0 + probs - probs.mean())).clamp(max=1.0) + ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame) + + # 4. Retain the smallest-similarity tokens per frame. + mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device) + for i in range(T): + k_i = int(ks[i].item()) + if k_i <= 0: + continue + _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False) + mask_2d[i].scatter_(0, idx, True) + + # 5. Reconcile rounding/clamp drift to the exact target count by score. + flat_mask = mask_2d.view(-1) + flat_sim = similarity.view(-1) + current = int(flat_mask.sum().item()) + if current > target_retained: + drop_n = current - target_retained + retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1) + retained_sim = flat_sim[retained_idx] + _, worst = torch.topk(retained_sim, k=drop_n, largest=True, sorted=False) + flat_mask[retained_idx[worst]] = False + elif current < target_retained: + add_n = target_retained - current + available_idx = (~flat_mask).nonzero(as_tuple=False).squeeze(-1) + if available_idx.numel() > 0: + available_sim = flat_sim[available_idx] + add_n = min(add_n, available_idx.numel()) + _, best = torch.topk(available_sim, k=add_n, largest=False, sorted=False) + flat_mask[available_idx[best]] = True + + return flat_mask + + +def _multi_scale_gaussian(x: torch.Tensor, center: torch.Tensor) -> torch.Tensor: + """Sum Gaussian kernels over `_ALPHAS`; `(T, N, C) -> (T, N)` scores.""" + dist_sq = ((x - center) ** 2).sum(dim=-1) + return sum(torch.exp(-dist_sq / (2.0 * a)) for a in _ALPHAS) diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 37d508407cfd..f7119d84ef8f 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -122,7 +122,7 @@ def __init__( self._tool_parser: ToolParser | None = None if self.__class__.reasoning_parser_cls is not None: self._reasoning_parser = self.__class__.reasoning_parser_cls( - tokenizer, *args, **kwargs + tokenizer, *args, model_config=model_config, **kwargs ) if self.__class__.tool_parser_cls is not None: self._tool_parser = self.__class__.tool_parser_cls(tokenizer, tools) diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 895f7516a53e..048e714cb4ae 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -396,6 +396,9 @@ def _is_valid_tool_name(self, name: str) -> bool: return True return find_tool_name(self._tools, name) + def _accept_tool_name(self, name: str) -> bool: + return bool(name) and self._is_valid_tool_name(name) + # ── Private helpers ───────────────────────────────────────────── def _check_skip_tool_parsing( @@ -807,7 +810,7 @@ def _emit_name_delta( deltas: list[DeltaToolCall], name: str | None, ) -> None: - if not name or not self._is_valid_tool_name(name): + if name is None or not self._accept_tool_name(name): return slot = self._tool_slots[idx] slot.name = name @@ -866,8 +869,8 @@ def _handle_tool_end( slot = self._tool_slots[idx] if not slot.name_sent: - name = slot.name or self._try_extract_name(idx) - if name and self._is_valid_tool_name(name): + name = slot.name or self._try_extract_name(idx) or "" + if self._accept_tool_name(name): slot.name = name slot.name_sent = True slot.string_keys = self._streamable_string_keys( @@ -1042,7 +1045,7 @@ def _build_extracted_result( else: args_json = "{}" - if name and self._is_valid_tool_name(name): + if self._accept_tool_name(name): self._ensure_tool_id(slot, name) args_json = self._fix_arg_types(args_json, name) tool_calls.append( diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index d259ade545c2..57865e96b525 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -15,6 +15,7 @@ from vllm.parser.inkling import InklingParser from vllm.parser.kimi_k2 import KimiK2Parser from vllm.parser.minimax_m2 import MinimaxM2Parser +from vllm.parser.mistral import MistralParser from vllm.parser.nemotron_v3 import NemotronV3Parser from vllm.parser.qwen3 import Qwen3Parser from vllm.parser.seed_oss import SeedOssParser @@ -68,3 +69,8 @@ InklingParserReasoningAdapter, InklingParserToolAdapter, ) = make_adapters(InklingParser) + +( + MistralParserReasoningAdapter, + MistralParserToolAdapter, +) = make_adapters(MistralParser) diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py index 8cafdf8e625b..f5c395f76d3f 100644 --- a/vllm/parser/engine/streaming_parser_engine.py +++ b/vllm/parser/engine/streaming_parser_engine.py @@ -185,6 +185,7 @@ def reset(self, initial_state: ParserState | None = None) -> None: # implicit-reasoning-end (content returns None). self._scanner.reset() self._lexer.reset() + self._message_header_buffer = "" self._reset_args_state() def feed( @@ -270,6 +271,15 @@ def finish(self) -> list[SemanticEvent]: ) self.state = ParserState.CONTENT elif self.state == ParserState.MESSAGE_HEADER: + if self._message_header_buffer: + events.append( + SemanticEvent( + EventType.TEXT_CHUNK, + value=self._message_header_buffer, + tool_index=self.tool_index, + ) + ) + self._message_header_buffer = "" self.state = ParserState.CONTENT return events @@ -304,20 +314,14 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: transition = self.config.transitions.get(key) if transition is None: - if ( - self._has_drops - and terminal == DROP_TERMINAL - # Preserve drop tokens when skip_tool_parsing is active so - # the reasoning pass doesn't silently remove tokens that a - # later tool-call pass might need to see. - and not self.skip_tool_parsing - ): + if self._has_drops and terminal == DROP_TERMINAL: return [] return self._emit_for_state(value) if self.skip_tool_parsing and terminal in self._tool_terminals: if self.state == ParserState.MESSAGE_HEADER: self.state = ParserState.CONTENT + self._message_header_buffer = "" return [ SemanticEvent( EventType.TEXT_CHUNK, @@ -352,6 +356,9 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: return self._apply_transition(transition, value) def _emit_for_state(self, text: str) -> list[SemanticEvent]: + if self.state == ParserState.MESSAGE_HEADER: + self._message_header_buffer += text + return [] if self.state == ParserState.TOOL_ARGS: if self.config.tool_args_json: return self._feed_args_text(text) @@ -378,6 +385,8 @@ def _apply_transition( value: str, ) -> list[SemanticEvent]: events: list[SemanticEvent] = [] + previous_state = self.state + message_header = "" if ( self.state == ParserState.TOOL_ARGS @@ -393,15 +402,27 @@ def _apply_transition( ) self._args_buffer = "" + if previous_state == ParserState.MESSAGE_HEADER: + message_header = self._message_header_buffer + self._message_header_buffer = "" + self.state = transition.next_state for event_type in transition.events: if event_type == EventType.TOOL_CALL_START: self.tool_index += 1 + event_value = ( + message_header + if previous_state == ParserState.MESSAGE_HEADER + and event_type == EventType.TEXT_CHUNK + else value + ) + if event_type == EventType.TEXT_CHUNK and not event_value: + continue events.append( SemanticEvent( event_type, - value=value, + value=event_value, tool_index=self.tool_index, ) ) diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 1442246f139c..3cc493b65923 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -5,14 +5,31 @@ import json from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum, auto from typing import TYPE_CHECKING, NamedTuple from openai_harmony import HarmonyError, Message, Role +from xgrammar import StructuralTag +from xgrammar.openai_tool_call_schema import BuiltinToolParam, FunctionToolParam +from xgrammar.structural_tag import ( + AnyTextFormat, + ConstStringFormat, + Format, + GrammarFormat, + JSONSchemaFormat, + OptionalFormat, + OrFormat, + RegexFormat, + SequenceFormat, + TagFormat, + TriggeredTagsFormat, +) from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -28,7 +45,13 @@ from vllm.logger import init_logger from vllm.parser.abstract_parser import DelegatingParser from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.sampling_params import StructuredOutputsParams from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser +from vllm.tool_parsers.structural_tag_registry import ( + SimplifiedToolChoice, + get_function_parameters, + register_vllm_structural_tag, +) if TYPE_CHECKING: from openai_harmony import Message, StreamableParser @@ -346,6 +369,12 @@ def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: reasoning_token_count=reasoning_token_count, ) + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + request = _adjust_output_format(request) + return super().adjust_request(request) + @staticmethod def _normalize_recipient(recipient: str | None) -> str | None: """Remove constrained formats misparsed into recipients by older Harmony.""" @@ -356,3 +385,202 @@ def _normalize_recipient(recipient: str | None) -> str | None: if constrain_index == -1: return recipient return recipient[:constrain_index].rstrip() or None + + +# Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|> +# <|return|> is represented as "" since it's the default stop token, which xgrammar +# disallows under constraints, leading to bad or infinite generation. +# StreamableParser doesn't consider <|endoftext|> as a message end, so it's excluded +# TODO: Remove <|call|> once #50595 lands. +_END_TAG = ["<|end|>", "<|call|>", ""] +_FINAL_BEGIN = "<|channel|>final{constrain}<|message|>" +_TOOL_CALL_CHANNELS = [ + "<|channel|>commentary", + "<|channel|>analysis", + "<|channel|>final", +] +_FUNCTION_CALL_BEGINS = [ + "to=functions.{name} {channel} json<|message|>", + "to=functions.{name} {channel} <|constrain|>json<|message|>", + "{channel} to=functions.{name} json<|message|>", + "{channel} to=functions.{name} <|constrain|>json<|message|>", +] +_JSON_CONTENT = JSONSchemaFormat(json_schema={"type": "object"}) +_ANY_CONTENT = AnyTextFormat() + + +def _assemble_tag( + allow_analysis: bool, allow_commentary: bool, content: Format +) -> StructuralTag: + tags = [] + if allow_analysis: + analysis_tag = OptionalFormat( + content=SequenceFormat( + elements=[ + TagFormat( + begin="<|channel|>analysis<|message|>", + content=_ANY_CONTENT, + end="<|end|>", + ), + ConstStringFormat(value="<|start|>assistant"), + ] + ) + ) + tags.append(analysis_tag) + + if allow_commentary: + commentary_tag = OptionalFormat( + content=SequenceFormat( + elements=[ + TagFormat( + begin="<|channel|>commentary<|message|>", + content=_ANY_CONTENT, + end="<|end|>", + ), + ConstStringFormat(value="<|start|>assistant"), + ] + ) + ) + tags.append(commentary_tag) + + tags.append(content) + + return StructuralTag(format=SequenceFormat(elements=tags)) + + +@register_vllm_structural_tag("harmony") +def get_harmony_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + # reasoning always enabled for Harmony + del reasoning + + if builtin_tools: + # Fallback for built-in tools + tags = [ + TagFormat( + begin="to=", + content=AnyTextFormat(excludes=["<|start|>"]), + end=_END_TAG, + ) + ] + tags.extend( + TagFormat( + begin=channel + " to=", + content=AnyTextFormat(excludes=["<|start|>", "<|channel|>"]), + end=_END_TAG, + ) + for channel in _TOOL_CALL_CHANNELS + ) + else: + tags = [ + TagFormat( + begin=pattern.format(name=tool.function.name, channel=channel), + content=JSONSchemaFormat( + json_schema=get_function_parameters(tool.function) + ), + end=_END_TAG, + ) + for tool in tools + for pattern in _FUNCTION_CALL_BEGINS + for channel in _TOOL_CALL_CHANNELS + ] + + if tool_choice == "auto": + tags.append( + TagFormat( + begin=_FINAL_BEGIN.format(constrain=" <|constrain|>json"), + content=_ANY_CONTENT, + end=_END_TAG, + ) + ) + tags.append( + TagFormat( + begin=_FINAL_BEGIN.format(constrain=""), + content=_ANY_CONTENT, + end=_END_TAG, + ) + ) + + return _assemble_tag( + allow_analysis=True, allow_commentary=True, content=OrFormat(elements=tags) + ) + + +def _params_to_final_content(params: StructuredOutputsParams) -> Format | None: + """Map StructuredOutputsParams in a XGrammar Format.""" + if params.json_object: + return _JSON_CONTENT + if params.json is not None: + schema = params.json + if isinstance(schema, str): + schema = json.loads(schema) + return JSONSchemaFormat(json_schema=schema) + if params.regex is not None: + return RegexFormat(pattern=params.regex) + if params.choice is not None: + return OrFormat( + elements=[ConstStringFormat(value=choice) for choice in params.choice] + ) + if params.grammar is not None: + return GrammarFormat(grammar=params.grammar) + if params.structural_tag is not None: + s_tag = json.loads(params.structural_tag) + if "structures" in s_tag: + # LegacyStructuralTagResponseFormat + return TriggeredTagsFormat( + triggers=s_tag["triggers"], + tags=[ + TagFormat( + begin=structure["begin"], + content=JSONSchemaFormat(json_schema=structure["schema"]), + end=structure["end"], + ) + for structure in s_tag["structures"] + ], + ) + # StructuralTagResponseFormat + return StructuralTag.model_validate(s_tag).format + return None + + +def _adjust_output_format( + request: ChatCompletionRequest | ResponsesRequest, +) -> ChatCompletionRequest | ResponsesRequest: + """Canonicalize request constraints into a reasoning-aware StructuralTag.""" + params = request.extract_structured_outputs() + if params is None: + return request + + final_content = _params_to_final_content(params) + if final_content is None: + return request + + if isinstance(final_content, JSONSchemaFormat): + begin = _FINAL_BEGIN.format(constrain=" <|constrain|>json") + else: + begin = _FINAL_BEGIN.format(constrain="") + + structural_tag = _assemble_tag( + allow_analysis=True, + allow_commentary=False, + content=TagFormat(begin=begin, content=final_content, end=_END_TAG), + ) + + request.structured_outputs = replace( + params, + json=None, + regex=None, + choice=None, + grammar=None, + json_object=None, + structural_tag=json.dumps(structural_tag.model_dump()), + ) + if isinstance(request, ResponsesRequest): + request.text = None + else: + request.response_format = None + return request diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 84502be018e3..994382539ec1 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -268,7 +268,7 @@ def inkling_config() -> ParserEngineConfig: ) transitions[(ParserState.MESSAGE_HEADER, end)] = Transition( ParserState.CONTENT, - (), + (EventType.TEXT_CHUNK,), ) return ParserEngineConfig( diff --git a/vllm/parser/kimi_k3.py b/vllm/parser/kimi_k3.py new file mode 100644 index 000000000000..380e2cce5f91 --- /dev/null +++ b/vllm/parser/kimi_k3.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + FunctionCall, +) +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser + +if TYPE_CHECKING: + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + + +class KimiK3Parser(DelegatingParser): + """Compose the Kimi K3 reasoning and tool parsers for XTML output.""" + + # TODO: Switch Kimi K3 to the parser engine once its XTML reasoning/tool + # path is covered there. + def _extract_tool_calls( + self, + content: str | None, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + ) -> tuple[list[FunctionCall] | None, str | None]: + if self._tool_parser is None or not enable_auto_tools: + return super()._extract_tool_calls(content, request, enable_auto_tools) + + tool_call_info = self.extract_tool_calls(content or "", request=request) + if request.tool_choice == "none": + return [], tool_call_info.content + if not tool_call_info.tools_called: + return None, tool_call_info.content + + tool_calls = [ + FunctionCall( + id=tool_call.id, + name=tool_call.function.name, + arguments=tool_call.function.arguments, + ) + for tool_call in tool_call_info.tool_calls + ] + parsed_content = tool_call_info.content + if parsed_content and parsed_content.strip() == "": + parsed_content = None + return tool_calls, parsed_content + + def _extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest | ResponsesRequest, + tool_call_idx: int | None = None, + tool_call_id_type: str = "random", + function_name_returned: bool = False, + ) -> tuple[DeltaMessage | None, bool]: + if request.tool_choice != "none": + return super()._extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + tool_call_idx=tool_call_idx, + tool_call_id_type=tool_call_id_type, + function_name_returned=function_name_returned, + ) + + delta_message = self.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + if delta_message is not None: + delta_message.tool_calls = [] + return delta_message, False + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + state = self._stream_state + previous_content = state.previous_text if state.reasoning_ended else "" + delta_message = super().parse_delta( + delta_text, + delta_token_ids, + request, + prompt_token_ids, + finished=finished, + ) + + if ( + self._tool_parser is not None + or not isinstance(self._reasoning_parser, KimiK3ReasoningParser) + or not state.reasoning_ended + or delta_message is None + ): + return delta_message + + stripped = self._reasoning_parser.strip_content_streaming( + previous_text=previous_content, + current_text=state.previous_text, + ) + delta_message.content = stripped.content if stripped is not None else None + if ( + delta_message.role is None + and delta_message.content is None + and delta_message.reasoning is None + and not delta_message.tool_calls + ): + return None + return delta_message diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index 52f16136ee3b..edc7018620cc 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -3,80 +3,979 @@ from __future__ import annotations +import functools +import json from collections.abc import Sequence -from typing import TYPE_CHECKING +from enum import Enum, auto +from random import choices +from string import ascii_letters, digits +from typing import TYPE_CHECKING, Any, Literal -from vllm.entrypoints.openai.engine.protocol import DeltaMessage, FunctionCall -from vllm.parser.abstract_parser import DelegatingParser +import ijson +import regex as re +from mistral_common.protocol.instruct.tool_calls import ( + NamedToolChoice as MistralNamedToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + Tool as MistralTool, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoice as MistralToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoiceEnum as MistralToolChoiceEnum, +) +from mistral_common.tokens.tokenizers.base import SpecialTokens +from pydantic import Field + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.logger import init_logger +from vllm.parser.engine.events import EventType +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) +from vllm.sampling_params import StructuredOutputsParams +from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) - from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.parser.engine.events import SemanticEvent + from vllm.parser.engine.parser_engine import ToolCallSlot + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + +logger = init_logger(__name__) + +_ALPHANUMERIC = ascii_letters + digits +_DEFAULT_JSON_SCHEMA: dict[str, Any] = { + "anyOf": [{"type": "object"}, {"type": "array"}] +} + +# Special token strings as emitted/decoded by MistralTokenizer. +_TOOL_CALLS = SpecialTokens.tool_calls.value +_ARGS = SpecialTokens.args.value +_THINK_START_SPECIAL = SpecialTokens.begin_think.value +_THINK_END_SPECIAL = SpecialTokens.end_think.value + +# Plain-text reasoning markers used by v11 tokenizers. +_THINK_START_TEXT = "" +_THINK_END_TEXT = "" + +_OPEN_BRACE = "{" + + +class StreamingState(Enum): + """Streaming parsing state for pre-v11 tool call extraction.""" + + WAITING_FOR_TOOL_START = auto() + WAITING_FOR_TOOL_KEY = auto() + PARSING_NAME = auto() + PARSING_NAME_COMPLETED = auto() + WAITING_FOR_ARGUMENTS_START = auto() + PARSING_ARGUMENTS = auto() + PARSING_ARGUMENTS_COMPLETED = auto() + TOOL_COMPLETE = auto() + ALL_TOOLS_COMPLETE = auto() + + +class MistralToolCall(ToolCall): + """ToolCall with a Mistral-compatible random alphanumeric id.""" + + id: str = Field(default_factory=lambda: MistralToolCall.generate_random_id()) + + @staticmethod + def generate_random_id() -> str: + # Mistral Tool Call Ids must be alphanumeric with a length of 9. + # https://github.com/mistralai/mistral-common/blob/21ee9f6cee3441e9bb1e6ed2d10173f90bd9b94b/src/mistral_common/protocol/instruct/validator.py#L299 + return "".join(choices(_ALPHANUMERIC, k=9)) + + @staticmethod + def is_valid_id(id: str) -> bool: + return id.isalnum() and len(id) == 9 + + +def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool: + if is_mistral_tokenizer(model_tokenizer): + return model_tokenizer.version < 11 + vocab: dict[str, int] = getattr(model_tokenizer, "get_vocab", lambda: {})() + return _ARGS not in vocab + + +@functools.cache +def mistral_config( + *, + reasoning_encoding: Literal["special_token", "text", "none"], + name: str = "mistral", +) -> ParserEngineConfig: + """Return a :class:`ParserEngineConfig` for the Mistral output format. + + Args: + reasoning_encoding: Reasoning token format. + ``"special_token"`` – v13+ ``[THINK]``/``[/THINK]`` special + tokens placed in both ``terminals`` and ``token_id_terminals``. + ``"text"`` – v11 ````/```` plain text placed in + ``terminals`` only; no ``token_id_terminals`` for think tokens. + ``"none"`` – no reasoning support. + name: Name embedded in the returned config (used for debugging). + + Returns: + A frozen :class:`ParserEngineConfig` with ``initial_state=CONTENT``. + """ + if reasoning_encoding == "special_token": + think_start = _THINK_START_SPECIAL + think_end = _THINK_END_SPECIAL + reasoning_terminals: dict[str, str] = { + "THINK_START": think_start, + "THINK_END": think_end, + } + reasoning_token_id_terminals: dict[str, str] = { + "THINK_START": think_start, + "THINK_END": think_end, + } + elif reasoning_encoding == "text": + think_start = _THINK_START_TEXT + think_end = _THINK_END_TEXT + reasoning_terminals = { + "THINK_START": think_start, + "THINK_END": think_end, + } + # Text think markers have no token-id terminals. + reasoning_token_id_terminals = {} + else: + reasoning_terminals = {} + reasoning_token_id_terminals = {} + + if reasoning_encoding != "none": + reasoning_transitions: dict[tuple[ParserState, str], Transition] = { + (ParserState.CONTENT, "THINK_START"): Transition( + ParserState.REASONING, + (), + ), + # Absorb a duplicate/re-emitted THINK_START inside reasoning. + (ParserState.REASONING, "THINK_START"): Transition( + ParserState.REASONING, + (), + ), + (ParserState.REASONING, "THINK_END"): Transition( + ParserState.CONTENT, + (EventType.REASONING_END,), + ), + # Absorb stray THINK_END that arrives after reasoning ended. + (ParserState.CONTENT, "THINK_END"): Transition( + ParserState.CONTENT, + (), + ), + # [TOOL_CALLS] directly from reasoning implicitly ends it. + (ParserState.REASONING, "TOOL_CALLS"): Transition( + ParserState.TOOL_NAME, + (EventType.REASONING_END, EventType.TOOL_CALL_START), + ), + } + else: + reasoning_transitions = {} + + return ParserEngineConfig( + name=name, + initial_state=ParserState.CONTENT, + terminals={ + **reasoning_terminals, + "TOOL_CALLS": _TOOL_CALLS, + "ARGS": _ARGS, + "OPEN_BRACE": _OPEN_BRACE, + }, + token_id_terminals={ + **reasoning_token_id_terminals, + "TOOL_CALLS": _TOOL_CALLS, + "ARGS": _ARGS, + }, + transitions={ + **reasoning_transitions, + # A tool call from content implicitly ends reasoning when enabled. + (ParserState.CONTENT, "TOOL_CALLS"): Transition( + ParserState.TOOL_NAME, + (EventType.REASONING_END, EventType.TOOL_CALL_START) + if reasoning_encoding != "none" + else (EventType.TOOL_CALL_START,), + ), + # NAME→ARGS via explicit [ARGS] separator (v11+): consumed, no events. + (ParserState.TOOL_NAME, "ARGS"): Transition( + ParserState.TOOL_ARGS, + (), + ), + # NAME→ARGS via "{": carried as ARG_VALUE_CHUNK so the opening + # brace lands in the JSON argument buffer (fallback for name{args}). + (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition( + ParserState.TOOL_ARGS, + (EventType.ARG_VALUE_CHUNK,), + ), + # Parallel tool calls: next [TOOL_CALLS] ends current call. + (ParserState.TOOL_ARGS, "TOOL_CALLS"): Transition( + ParserState.TOOL_NAME, + (EventType.TOOL_CALL_END, EventType.TOOL_CALL_START), + ), + }, + stream_arg_deltas=True, + tool_args_json=True, + strip_trailing_reasoning_whitespace=True, + drop_whitespace_only_content_before_tools=True, + ) + + +class MistralParser(ParserEngine): + """Mistral parser: engine-based reasoning + ``[TOOL_CALLS]`` tool calls. + Reasoning encoding is auto-detected from the tokenizer: -class MistralParser(DelegatingParser): - def __init__(self, tokenizer, tools=None, *args, **kwargs): - super().__init__(tokenizer, tools, *args, **kwargs) - from vllm.tool_parsers.mistral_tool_parser import MistralToolParser + - ``"special_token"`` – ``[THINK]`` present in vocab (v13+). + - ``"text"`` – tokenizer supports grammar but has no ``[THINK]`` (v11). + - ``"none"`` – no grammar support; reasoning disabled. - if not isinstance(self._tool_parser, MistralToolParser): - raise ValueError( - "MistralParser requires --tool-call-parser mistral, " - f"got {self._tool_parser.__class__.__name__}." + Tool calls use the ``[TOOL_CALLS]func_name{...}`` format. The opening + ``{`` doubles as the NAME→ARGS separator and is included in the argument + buffer via an ``ARG_VALUE_CHUNK`` event on the transition (mirrors + kimi_k2, JSON args, `tool_args_json=True`). + + When the tokenizer does not support grammar (``_reasoning_encoding == + "none"``), the legacy ``[TOOL_CALLS]``-based extraction path is used + instead of the declarative engine, handling both pre-v11 JSON-array and + v11+ ``funcname{args}`` formats. + """ + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + vocab = tokenizer.get_vocab() + self._reasoning_encoding: Literal["special_token", "text", "none"] + if _THINK_START_SPECIAL in vocab: + self._reasoning_encoding = "special_token" + elif getattr(tokenizer, "supports_grammar", False): + self._reasoning_encoding = "text" + else: + self._reasoning_encoding = "none" + + kwargs.setdefault( + "parser_engine_config", + mistral_config(reasoning_encoding=self._reasoning_encoding), + ) + super().__init__(tokenizer, tools, **kwargs) + + self._tool_calls_token_id: int | None = self.vocab.get(_TOOL_CALLS) + + # Tool calls use the legacy parser for all tokenizer versions; + # reasoning is handled by the engine. + self.bot_token: str = _TOOL_CALLS + self.bot_token_id: int | None = self._tool_calls_token_id + if self.bot_token_id is None: + raise RuntimeError( + "Mistral parser could not locate the tool call token in the tokenizer!" + ) + + # Legacy tool-call streaming state. + self.prev_tool_call_arr: list[dict[str, Any]] = [] + self.current_tool_id: int = -1 + self.streaming_state: StreamingState = StreamingState.WAITING_FOR_TOOL_START + self.tool_call_started: bool = False + self.current_tool_name: str | None = None + self.current_tool_mistral_id: str | None = None + self.starting_new_tool: bool = False + self.streamed_args_for_tool: list[str] = [] + self._is_pre_v11: bool = _is_pre_v11_tokeniser(tokenizer) + self.parse_coro = None + if self._is_pre_v11: + self.parse_coro = ijson.parse_coro( + self.update_stream_state_pre_v11_tokenizer() ) + self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL) - def _maybe_force_auto_tool_parsing( + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest - ) -> None: - # When the Mistral grammar factory injected structured outputs, - # the model emits v11+ format ([TOOL_CALLS]name{args}) that the - # named/required parsers can't handle. Disable them so all - # tool_choice modes fall back to auto tool parsing via - # extract_tool_calls. - if getattr(request, "_grammar_from_tool_parser", False): - assert self._tool_parser is not None - self._tool_parser.supports_required_and_named = False - - def parse( + ) -> ChatCompletionRequest | ResponsesRequest: + if not isinstance(request, ResponsesRequest) and request._grammar_from_parser: + return request + so_non_supported_attributes = [ + "regex", + "choice", + "grammar", + # whitespace_pattern is not a constraint type but an option; + # Mistral grammar factory does not support it. + "whitespace_pattern", + "structural_tag", + ] + any_so_non_supported_active = request.structured_outputs is not None and any( + getattr(request.structured_outputs, attribute) is not None + for attribute in so_non_supported_attributes + ) + response_format_non_supported_active = ( + isinstance(request, ResponsesRequest) + or request.response_format is not None + and request.response_format.type == "structural_tag" + ) + + if ( + not is_mistral_tokenizer(self.model_tokenizer) + or isinstance(request, ResponsesRequest) + or not self.model_tokenizer.supports_grammar + or any_so_non_supported_active + or response_format_non_supported_active + ): + request = super().adjust_request(request) + if request.tools and request.tool_choice != "none": + # Keep special tokens so the [TOOL_CALLS] marker + # survives for tool detection. + request.skip_special_tokens = False + # Inject a guided JSON schema for pre-v11 required/named tool choice + # so the model emits a well-formed bare JSON array instead of + # rambling. tool_choice forces a tool call, so a json_object / + # json_schema response_format is cleared here (the tool schema is + # the sole structured-output constraint), mirroring the base + # ToolParser; a structural_tag response_format is left untouched. + req_tool_choice = request.tool_choice + is_required_or_named = req_tool_choice == "required" or isinstance( + req_tool_choice, ChatCompletionNamedToolChoiceParam + ) + response_format = getattr(request, "response_format", None) + response_format_overridable = response_format is None or ( + response_format.type in ("text", "json_object", "json_schema") + ) + # This runs only in the non-grammar (legacy) branch, so it covers + # both pre-v11 Mistral tokenizers and non-Mistral (e.g. HF-mode) + # tokenizers driving the Mistral tool parser. The remaining guards + # keep it off when the user supplied their own structured output. + if ( + not isinstance(request, ResponsesRequest) + and request.tools + and request.structured_outputs is None + and is_required_or_named + and response_format_overridable + ): + schema = self._build_guided_schema_pre_v11(request) + if schema is not None: + request.structured_outputs = StructuredOutputsParams(json=schema) + request.response_format = None + return request + + json_schema: dict[str, Any] | None = None + if request.structured_outputs is not None: + if request.structured_outputs.json_object is not None: + json_schema = _DEFAULT_JSON_SCHEMA + elif request.structured_outputs.json is not None: + if isinstance(request.structured_outputs.json, str): + json_schema = json.loads(request.structured_outputs.json) + else: + json_schema = request.structured_outputs.json + else: + raise ValueError( + "Unsupported request.structured_outputs for MistralParser. " + "Only `json` and `json_object` are supported." + ) + elif ( + request.response_format is not None + and request.response_format.type != "text" + ): + if request.response_format.type == "json_object": + json_schema = _DEFAULT_JSON_SCHEMA + elif request.response_format.type == "json_schema": + if request.response_format.json_schema is not None: + json_schema = request.response_format.json_schema.json_schema + else: + json_schema = _DEFAULT_JSON_SCHEMA + else: + raise ValueError( + "MistralParser only accepts `text`, `json_object` or " + f"`json_schema`, got {request.response_format=}" + ) + request.response_format = None + + grammar_factory = self.model_tokenizer.grammar_factory + + # Rendering grammar is cached in mistral-common given tools, template and mode. + template = grammar_factory.select_jinja_template() + + mistral_tools = ( + [MistralTool.from_openai(tool.model_dump()) for tool in request.tools] + if request.tools is not None + else None + ) + + tool_choice: MistralToolChoice + match request.tool_choice: + case "none" | "auto" | "required": + tool_choice = MistralToolChoiceEnum(request.tool_choice) + case None: + tool_choice = MistralToolChoiceEnum.auto + # _ == Named tool choice + case _: + tool_choice = MistralNamedToolChoice.model_validate( + { + "type": "function", + "function": {"name": request.tool_choice.function.name}, + } + ) + + match tool_choice, json_schema is not None: + case MistralToolChoiceEnum.none, True: + lark_grammar = grammar_factory.get_lark_for_json_schema( + template=template, json_schema=json_schema + ) + case _, _: + lark_grammar = grammar_factory.get_lark_from_jinja( + template=template, + mode=tool_choice, + tools=mistral_tools, + json_schema=json_schema, + parallel_tool_calls=request.parallel_tool_calls, + json_only=False, + ) + + request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar) + request._grammar_from_parser = True + return request + + def _build_guided_schema_pre_v11( + self, + request: ChatCompletionRequest, + ) -> dict[str, Any] | None: + """Build a guided JSON schema for pre-v11 required/named tool choice. + + The schema enforces the Mistral-native array format + ``[{"name": ..., "arguments": {...}}]`` so the model emits a parseable + bare JSON array instead of free-form text. + + Args: + request: The chat completion request carrying `tools` and + `tool_choice`. + + Returns: + A JSON Schema dict, or ``None`` if the named tool is not found. + """ + tool_choice = request.tool_choice + tools = request.tools or [] + extra: dict[str, Any] = {} + + if tool_choice == "required": + applicable_tools = tools + else: + # Named tool choice — restrict to the single requested tool. + assert isinstance(tool_choice, ChatCompletionNamedToolChoiceParam) + chosen_name = tool_choice.function.name + applicable_tools = [t for t in tools if t.function.name == chosen_name] + if not applicable_tools: + logger.warning( + "Named tool %r not found in tools list; " + "skipping guided schema injection.", + chosen_name, + ) + return None + extra["maxItems"] = 1 + + any_of = [ + { + "type": "object", + "properties": { + "name": {"type": "string", "enum": [tool.function.name]}, + "arguments": tool.function.parameters or {"type": "object"}, + }, + "required": ["name", "arguments"], + } + for tool in applicable_tools + ] + + return { + "type": "array", + "minItems": 1, + **extra, + "items": { + "type": "object", + "anyOf": any_of, + }, + } + + def _ensure_tool_id(self, slot: ToolCallSlot, name: str) -> None: + """Assign a Mistral-compatible 9-char alphanumeric id to `slot`.""" + if not slot.id: + slot.id = MistralToolCall.generate_random_id() + + def extract_tool_calls_from_content( + self, + content: str, + request: ChatCompletionRequest, + ) -> ExtractedToolCallInformation: + if self._is_pre_v11: + return self._legacy_extract_tool_calls(content, request) + return super().extract_tool_calls_from_content(content, request) + + def _legacy_extract_tool_calls( self, model_output: str, + request: ChatCompletionRequest | None, + ) -> ExtractedToolCallInformation: + """Pre-v11 non-streaming extraction. + + Handles ``[TOOL_CALLS][{...}]`` and guided bare-array formats. + """ + if request is None: + tool_choice = None + tools = None + else: + tool_choice = request.tool_choice + tools = request.tools + + # tool_choice="none" with tools: never produce tool calls. + if tool_choice == "none" and tools: + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + content: str | None = None + if self.bot_token in model_output: + content_and_raw_tool_calls = model_output.split(self.bot_token) + content = content_and_raw_tool_calls[0] + raw_tool_calls = content_and_raw_tool_calls[1:] + # pre-v11: content[BOT] [{tool_call1},{tool_call2}] + if len(raw_tool_calls) != 1: + raise ValueError( + "Only one BOT token should have been outputted, " + f"but got {model_output}." + ) + stringified_tool_calls = raw_tool_calls[0].strip() + elif tool_choice == "required" or isinstance( + tool_choice, ChatCompletionNamedToolChoiceParam + ): + # Guided bare-array output (no [TOOL_CALLS] marker). + stringified_tool_calls = model_output.strip() + else: + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + try: + # Use raw_decode to parse the first valid JSON value, + # ignoring trailing tokens the model may emit after + # the tool call array. + tool_calls, _ = json.JSONDecoder().raw_decode(stringified_tool_calls) + except json.JSONDecodeError: + try: + raw_tool_call = self.tool_call_regex.findall(stringified_tool_calls)[0] + tool_calls = json.loads(raw_tool_call) + tool_calls = [ + { + "name": tool_call["name"], + "arguments": json.dumps( + tool_call.get("arguments", {}), + ensure_ascii=False, + ), + } + for tool_call in tool_calls + ] + except (IndexError, json.JSONDecodeError): + logger.exception("Error in extracting tool call from response.") + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=stringified_tool_calls, + ) + else: + tool_calls = [ + { + "name": tool_call["name"], + "arguments": json.dumps( + tool_call.get("arguments", {}), + ensure_ascii=False, + ), + } + for tool_call in tool_calls + ] + + mistral_tool_calls: list[MistralToolCall] = [ + MistralToolCall( + type="function", + function=FunctionCall( + name=tool_call["name"], + arguments=tool_call.get("arguments", "{}"), + ), + ) + for tool_call in tool_calls + ] + + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=mistral_tool_calls, + content=content if content and content.strip() else None, + ) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], request: ChatCompletionRequest | ResponsesRequest, - enable_auto_tools: bool = False, - model_output_token_ids: Sequence[int] = (), - ) -> tuple[str | None, str | None, list[FunctionCall] | None]: - self._maybe_force_auto_tool_parsing(request) - reasoning, content, tool_calls = super().parse( - model_output, - request, - enable_auto_tools, - model_output_token_ids, + ) -> DeltaMessage | None: + if not self._is_pre_v11: + return super().extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + # Pre-v11: latch on [TOOL_CALLS] or on the first content of a guided + # required/named request (bare JSON array, no special token). + if self.bot_token_id in delta_token_ids or self.bot_token in delta_text: + self.tool_call_started = True + elif not self.tool_call_started and delta_text: + is_guided = request.tool_choice == "required" or isinstance( + request.tool_choice, ChatCompletionNamedToolChoiceParam + ) + if is_guided: + self.tool_call_started = True + if not self.tool_call_started: + return DeltaMessage(content=delta_text) + try: + return self._extract_tool_calls_streaming_pre_v11_tokenizer( + delta_text=delta_text, + delta_token_ids=delta_token_ids, + ) + except Exception: + logger.exception("Error trying to handle streaming tool call.") + return None + + @ijson.coroutine + def update_stream_state_pre_v11_tokenizer(self): + while True: + (prefix, event, value) = yield + + 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": + self.current_tool_name = value + self.streaming_state = StreamingState.PARSING_NAME_COMPLETED + if prefix == "item" and event == "map_key" and value == "arguments": + self.streaming_state = StreamingState.WAITING_FOR_ARGUMENTS_START + if prefix == "item.arguments" and event == "start_map": + self.streaming_state = StreamingState.PARSING_ARGUMENTS + if prefix == "item.arguments" and event == "end_map": + self.streaming_state = StreamingState.PARSING_ARGUMENTS_COMPLETED + if prefix == "item" and event == "end_map": + self.streaming_state = StreamingState.TOOL_COMPLETE + if prefix == "" and event == "end_array": + self.streaming_state = StreamingState.ALL_TOOLS_COMPLETE + + def _extract_tool_calls_streaming_pre_v11_tokenizer( + self, + delta_text: str, + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + """Extract tool calls for pre-v11 Mistral models. + + Handles ``[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5}}]``. + """ + assert self.parse_coro is not None + content = None + delta_tool_calls: list[DeltaToolCall] = [] + current_tool_call: DeltaToolCall = DeltaToolCall( + index=self.current_tool_id, type="function" ) - if tool_calls: - from vllm.tool_parsers.mistral_tool_parser import MistralToolCall + current_tool_call_modified = False + if self.bot_token_id in delta_token_ids or self.bot_token in delta_text: + # this is the first tool call + if not delta_text.startswith(self.bot_token): + content = delta_text.split(self.bot_token)[0] + delta_text = "".join(delta_text.split(self.bot_token)[1:]) + + # ijson gives no text index per event, so split the delta manually + # to know where each event is emitted from. + while len(delta_text) > 0: + streaming_state_before_parse = self.streaming_state + + if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_opening_curly_braces=1, + ) + elif self.streaming_state == StreamingState.WAITING_FOR_TOOL_KEY: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_colon=1, + stop_after_opening_curly_braces=1, + ) + elif self.streaming_state == StreamingState.PARSING_NAME: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_comma=1, + stop_after_closing_brackets=1, + ) + elif self.streaming_state == StreamingState.WAITING_FOR_ARGUMENTS_START: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_opening_curly_braces=1, + ) + elif self.streaming_state == StreamingState.PARSING_ARGUMENTS: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_closing_curly_braces=1, + ) + elif self.streaming_state in [ + StreamingState.PARSING_ARGUMENTS_COMPLETED, + StreamingState.PARSING_NAME_COMPLETED, + ]: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_closing_curly_braces=1, + stop_after_closing_brackets=1, + ) + elif self.streaming_state == StreamingState.TOOL_COMPLETE: + delta_to_be_parsed, delta_text = self._split_delta( + delta_text=delta_text, + stop_after_opening_curly_braces=1, + stop_after_closing_brackets=1, + ) + elif self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE: + content = delta_text + delta_text = "" + else: + delta_to_be_parsed = delta_text + delta_text = "" + + if self.streaming_state != StreamingState.ALL_TOOLS_COMPLETE: + self.parse_coro.send(delta_to_be_parsed.encode("utf-8")) + + # 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 + self.current_tool_mistral_id = None + self._track_streamed_args_pre_v11(current_tool_call) + delta_tool_calls.append(current_tool_call) + current_tool_call_modified = False + self.current_tool_id += 1 + self.streamed_args_for_tool.append("") + self.prev_tool_call_arr.append({}) + self.current_tool_mistral_id = MistralToolCall.generate_random_id() + current_tool_call = DeltaToolCall( + index=self.current_tool_id, + type="function", + ) + if current_tool_call.function is None: + current_tool_call.function = DeltaFunctionCall() + + if self.current_tool_name is not None: + current_tool_call_modified = True + current_tool_call.function.name = self.current_tool_name + self.prev_tool_call_arr[self.current_tool_id]["name"] = ( + self.current_tool_name + ) + self.current_tool_name = None + if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED: + self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY + if self.streaming_state in [ + StreamingState.PARSING_ARGUMENTS, + StreamingState.PARSING_ARGUMENTS_COMPLETED, + ]: + if self.streaming_state == StreamingState.PARSING_ARGUMENTS_COMPLETED: + self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY + current_tool_call_modified = True + if current_tool_call.function.arguments is None: + current_tool_call.function.arguments = delta_to_be_parsed + else: + current_tool_call.function.arguments += delta_to_be_parsed + if streaming_state_before_parse != StreamingState.PARSING_ARGUMENTS: + # It's the first chunk of arg. let's lstrip it + current_tool_call.function.arguments = ( + current_tool_call.function.arguments.lstrip() + ) - # Named/required tool_choice builds FunctionCalls without - # ID, backfill with Mistral-format IDs. - for tc in tool_calls: - if not tc.id: - tc.id = MistralToolCall.generate_random_id() - return reasoning, content, tool_calls + if current_tool_call_modified: + if self.current_tool_mistral_id is not None: + current_tool_call.id = self.current_tool_mistral_id + self.current_tool_mistral_id = None + self._track_streamed_args_pre_v11(current_tool_call) + delta_tool_calls.append(current_tool_call) - def parse_delta( + if content or len(delta_tool_calls) > 0: + delta_message = DeltaMessage() + if content: + delta_message.content = content + if len(delta_tool_calls) > 0: + delta_message.tool_calls = delta_tool_calls + return delta_message + else: + if self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE: + return DeltaMessage() + else: + return None + + def _track_streamed_args_pre_v11(self, tool_call: DeltaToolCall) -> None: + r"""Accumulate `tool_call` arguments into the streaming state.""" + if tool_call.function is not None and tool_call.function.arguments is not None: + self.streamed_args_for_tool[self.current_tool_id] += ( + tool_call.function.arguments + ) + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = ( + self.streamed_args_for_tool[self.current_tool_id] + ) + + def _split_delta( self, delta_text: str, - delta_token_ids: list[int], + stop_after_quotes: int = -1, + stop_after_opening_curly_braces: int = -1, + stop_after_closing_curly_braces: int = -1, + stop_after_closing_brackets: int = -1, + stop_after_colon: int = -1, + stop_after_comma: int = -1, + ) -> tuple[str, str]: + delta_to_be_parsed = "" + for i, c in enumerate(delta_text): + if c in ['"', "'"]: + delta_to_be_parsed += c + stop_after_quotes -= 1 + if stop_after_quotes == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + elif c == "{": + delta_to_be_parsed += c + stop_after_opening_curly_braces -= 1 + if stop_after_opening_curly_braces == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + elif c == "}": + delta_to_be_parsed += c + stop_after_closing_curly_braces -= 1 + if stop_after_closing_curly_braces == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + elif c == "]": + delta_to_be_parsed += c + stop_after_closing_brackets -= 1 + if stop_after_closing_brackets == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + elif c == ":": + delta_to_be_parsed += c + stop_after_colon -= 1 + if stop_after_colon == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + elif c == ",": + delta_to_be_parsed += c + stop_after_comma -= 1 + if stop_after_comma == 0: + return (delta_to_be_parsed, delta_text[i + 1 :]) + else: + delta_to_be_parsed += c + + return (delta_to_be_parsed, "") + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + if self._reasoning_encoding == "none": + return True + if super().is_reasoning_end(input_ids): + return True + # [TOOL_CALLS] acts as an implicit reasoning-end marker + if self.bot_token_id is not None: + reasoning_start_id = self._reasoning_start_token_id + for i in range(len(input_ids) - 1, -1, -1): + if ( + reasoning_start_id is not None + and input_ids[i] == reasoning_start_id + ): + return False + if input_ids[i] == self.bot_token_id: + return True + return False + + def extract_reasoning( + self, + model_output: str, request: ChatCompletionRequest | ResponsesRequest, - prompt_token_ids: list[int] | None = None, - *, - finished: bool, - ) -> DeltaMessage | None: - self._maybe_force_auto_tool_parsing(request) - return super().parse_delta( - delta_text, - delta_token_ids, - request, - prompt_token_ids, - finished=finished, + ) -> tuple[str | None, str | None]: + if self._reasoning_encoding == "none": + return None, model_output + return super().extract_reasoning(model_output, request) + + def _accept_tool_name(self, name: str) -> bool: + # Once `[ARGS]` or the opening `{` has moved the slot past its name, + # that name is final, empty included. Emitting the call with "" + # surfaces the malformed generation instead of dropping it. + return self._is_valid_tool_name(name) + + def _try_extract_name(self, idx: int) -> str | None: + # No JSON-embedded "name" key in this format; the slot name is final. + return self._tool_slots[idx].name + + def _extract_name_and_args(self, raw_body: str) -> tuple[str, str]: + # Never a {"name": ..., "arguments": ...} envelope -- a literal "name" + # argument key must not be mistaken for the tool name. + return "", self._extract_args_json(raw_body, "") + + def _handle_arg_chunk( + self, + event: SemanticEvent, + deltas: list[DeltaToolCall], + ) -> None: + """Emit the opening ``{`` as an arg delta when the name is first sent. + + When the TOOL_NAME→TOOL_ARGS transition fires, ``{`` arrives as an + ARG_VALUE_CHUNK before ``name_sent`` is True. The parent emits the + name delta but not the ``{`` arg delta. This override re-emits the + current chunk so streaming clients receive a valid JSON prefix. + """ + idx = event.tool_index + name_sent_before = ( + 0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent ) + super()._handle_arg_chunk(event, deltas) + if ( + event.value + and not name_sent_before + and 0 <= idx < len(self._tool_slots) + and self._tool_slots[idx].name_sent + ): + deltas.append( + DeltaToolCall( + index=idx, + function=DeltaFunctionCall(arguments=event.value), + ) + ) + + def _extract_args_json(self, raw_args: str, func_name: str) -> str: + """Return the first complete JSON value in ``raw_args``. + + v11+ tool calls are emitted as ``name{args}`` with no terminator, so a + model may append ordinary text after the closing brace. Parse the first + JSON value with ``raw_decode`` and drop any trailing output, mirroring + the pre-v11 path (fixes gh#48975). + """ + stripped = raw_args.strip() + if not stripped: + return "{}" + try: + _, end = json.JSONDecoder().raw_decode(stripped) + except json.JSONDecodeError: + return stripped + return stripped[:end] diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 1b5133f5a8f9..865f749a79c9 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -109,8 +109,6 @@ def get_parser( if reasoning_parser_cls is None and tool_parser_cls is None: return None - from vllm.utils.mistral import is_mistral_tool_parser - if is_harmony: from vllm.parser.harmony import HarmonyParser @@ -118,12 +116,17 @@ def get_parser( HarmonyParser.tool_parser_cls = tool_parser_cls return HarmonyParser - if is_mistral_tool_parser(tool_parser_cls): - from vllm.parser.mistral import MistralParser + if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": + from vllm.parser.kimi_k3 import KimiK3Parser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _KimiK3Parser(KimiK3Parser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls - MistralParser.reasoning_parser_cls = reasoning_parser_cls - MistralParser.tool_parser_cls = tool_parser_cls - return MistralParser + return _KimiK3Parser from vllm.parser.abstract_parser import DelegatingParser diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 0e805794f5fd..698e83a8aba9 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -461,11 +461,7 @@ def import_kernels(cls) -> None: @classmethod def pack_kv_cache( cls, - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_ids: list[int], + kv_cache: torch.Tensor, indices: torch.Tensor, ) -> None: """ @@ -476,15 +472,26 @@ def pack_kv_cache( from vllm._custom_ops import cpu_attn_reshape_and_cache from vllm.v1.attention.backends.cpu_attn import _get_attn_isa + num_blocks, num_kv_heads, block_size, fused_head_size = kv_cache.shape + head_size = fused_head_size // 2 + + # Fused path used by heterogeneous NIXL CPU_ATTN post-processing. + blocks_to_update = kv_cache.index_select(0, indices) + key = blocks_to_update[..., :head_size] + value = blocks_to_update[..., head_size:] + + key_cache, value_cache = kv_cache.view( + num_blocks, num_kv_heads, block_size * 2, head_size + ).chunk(2, dim=2) + dtype = key.dtype # For CPU_ATTN, the shape is [N, num_kv_heads, block_size, head_size] - _, _, block_size, head_size = key_cache.shape key = key.permute(0, 2, 1, 3).flatten(0, 1) value = value.permute(0, 2, 1, 3).flatten(0, 1) isa = _get_attn_isa(dtype, block_size, head_size) block_offsets = torch.arange(block_size, device="cpu", dtype=torch.long) - num_blocks = len(block_ids) + num_blocks = indices.numel() slot_mapping = ( block_offsets.reshape(1, block_size) + indices.reshape(num_blocks, 1) * block_size diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 07dce7d138b2..f0b93da2922c 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -880,6 +880,12 @@ def _align_hybrid_block_size( ), cache_config.block_size, ) + if model_config.use_mla: + # TRTLLM/FlashInfer MLA decode kernels require the physical + # number of kernel blocks to be aligned to 128 / kernel_block_size. + # For hybrid MLA/Mamba models, make the manager block size a + # multiple of 128 so split kernel blocks keep that invariant. + kernel_block_alignment_size = max(kernel_block_alignment_size, 128) if cache_config.mamba_cache_mode == "all": # With prefix caching, align to mamba chunk size for kernel perf diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 4b182aeb8852..d7e5bf02f210 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -194,11 +194,6 @@ def _get_gcn_arch() -> str: return _query_gcn_arch_from_amdsmi() except Exception as e: logger.debug("Failed to get GCN arch via amdsmi: %s", e) - logger.warning_once( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda. " - "This will initialize CUDA and may cause " - "issues if CUDA_VISIBLE_DEVICES is not set yet." - ) # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName @@ -218,6 +213,11 @@ def _get_gcn_arch() -> str: _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH +_ON_GFX1250 = "gfx1250" in _GCN_ARCH + +_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9", "gfx1250"]) +# RDNA = gfx11/gfx12 minus the CDNA-classified gfx1250. +_ON_RDNA = _ON_GFX1X and not _ON_CDNA def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -292,7 +292,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X + return _ON_GFX1X and not _ON_CDNA def on_gfx11() -> bool: @@ -308,7 +308,11 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return _ON_GFX12X + return _ON_GFX12X and not _ON_CDNA + + +def on_gfx1250() -> bool: + return _ON_GFX1250 def on_mi3xx() -> bool: @@ -331,13 +335,33 @@ def on_gfx950() -> bool: return _ON_GFX950 +def on_cdna() -> bool: + return _ON_CDNA + + +def on_rdna() -> bool: + return _ON_RDNA + + +def get_cdna_version() -> int: + if on_gfx90a(): + return 2 + if on_gfx942(): + return 3 + if on_gfx950(): + return 4 + if on_gfx1250(): + return 5 + return 0 + + # Enable HIP online tuning early, before hipBLASLt initializes. # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM. if ( envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - and on_mi3xx() + and get_cdna_version() > 2 ): os.environ["HIP_ONLINE_TUNING"] = "1" @@ -356,7 +380,7 @@ def use_rocm_custom_paged_attention( ) -> bool: # custom paged attn always supported on V0. On V1, requires sliding window # disabled due to observed numerical discrepancy. - if _ON_GFX9: + if on_cdna(): return ( (sliding_window == 0 or sliding_window == (-1, -1)) and (qtype == torch.half or qtype == torch.bfloat16) @@ -389,9 +413,21 @@ def flash_attn_triton_available() -> bool: try: from importlib.util import find_spec - if find_spec("flash_attn") is None: - return False - if find_spec("flash_attn.flash_attn_triton_amd") is None: + # Locate the Triton-AMD kernels. Older ROCm/flash-attention (pre + # 2026-03) shipped them as the flash_attn.flash_attn_triton_amd + # subpackage. The main_perf migration commit 3f94643 moved them + # into aiter at aiter.ops.triton._triton_kernels.flash_attn_triton_amd, + # so accept either location. + def _has_spec(name: str) -> bool: + try: + return find_spec(name) is not None + except (ImportError, ValueError): + return False + + if not ( + _has_spec("flash_attn.flash_attn_triton_amd") + or _has_spec("aiter.ops.triton._triton_kernels.flash_attn_triton_amd") + ): return False if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE") != "TRUE": logger.info_once( @@ -657,12 +693,12 @@ def get_vit_attn_backend( from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled() and on_gfx9(): + if rocm_aiter_ops.is_mha_enabled() and on_cdna(): logger.info_once("Using AITER Flash Attention backend for ViT model.") return AttentionBackendEnum.ROCM_AITER_FA if ( - on_gfx9() + on_cdna() and find_spec("flash_attn") is not None and (dtype == torch.float16 or dtype == torch.bfloat16) ): @@ -878,11 +914,11 @@ def get_device_communicator_cls(cls) -> str: @classmethod def supports_mx(cls) -> bool: - return any(gfx in _GCN_ARCH for gfx in ["gfx95"]) + return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"]) @classmethod def supports_fp8(cls) -> bool: - return on_gfx9() or on_gfx12x() + return on_cdna() or on_gfx12x() @classmethod def is_fp8_fnuz(cls) -> bool: diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 80f99acc5acb..1104f291891e 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -110,6 +110,23 @@ class XPUPlatform(Platform): ray_device_key: str = "GPU" dist_backend: str = "xccl" # xccl only device_control_env_var: str = "ZE_AFFINITY_MASK" + supported_quantization: list[str] = [ + "awq", + "gptq", + "auto_awq", + "auto_gptq", + "inc", + "fp8", + "deepseek_v4_fp8", + "mxfp4", + "mxfp8", + "fp8_per_tensor", + "fp8_per_block", + "online", + "gpt_oss_mxfp4", + "modelopt", + "compressed-tensors", + ] @classmethod def import_kernels(cls) -> None: @@ -282,6 +299,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: "XPU Graph is disabled by environment variable, " "please set VLLM_XPU_ENABLE_XPU_GRAPH=1 to enable it." ) + else: + logger.warning_once( + "XPU Graph support is experimental and has known limitations: " + "(1) only single-GPU execution is supported; " + "(2) FLASH_ATTN supports PIECEWISE mode only; use TRITON_ATTN " + "for FULL mode; " + "(3) XPU Graph may increase device memory usage, " + "potentially causing OOM errors or leaving less memory " + "for the KV cache and reducing performance." + ) # Disable fusion passes not yet supported on XPU. from vllm.config.compilation import CompilationMode @@ -294,7 +321,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: "fuse_act_padding": "Activation + padding fusion", "fuse_rope_kvcache": "RoPE + KV cache fusion", "fuse_rope_kvcache_cat_mla": "RoPE + KV cache + MLA fusion", - "enable_qk_norm_rope_fusion": "QK Norm + RoPE fusion", } if compilation_config.mode != CompilationMode.NONE: for flag, feature_name in fusion_passes_to_disable.items(): @@ -439,7 +465,7 @@ def get_default_ir_op_priority( # use fused kernels where available when no codegen cc = vllm_config.compilation_config using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE - default = ["native"] if using_inductor else ["xpu_kernels", "native"] + default = ["native"] if using_inductor else ["vllm_c", "native"] return IrOpPriorityConfig.with_default(default) diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 6cb130fdbbbb..5280e997c9c9 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -7,6 +7,7 @@ import msgspec from vllm.config import ModelConfig, PoolerConfig +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind from vllm.tasks import PoolingTask, check_removed_pooling_task @@ -145,7 +146,7 @@ def _verify_step_pooling( invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise VLLMValidationError( f"Task {self.task} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -170,21 +171,21 @@ def _set_default_parameters(self, model_config: ModelConfig): valid_range = f"[1, {embedding_size}]" dimensions_in_range = 1 <= dimensions <= embedding_size if not model_config.is_matryoshka: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} does not support Matryoshka " f"embeddings; dimensions must be unset " f"(received dimensions={dimensions})." ) if not dimensions_in_range: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} only supports dimensions in " f"range {valid_range}, got {dimensions}." ) mds = model_config.matryoshka_dimensions if mds is not None and dimensions not in mds: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} only supports Matryoshka " f"dimensions {str(mds)}, got {dimensions}." ) @@ -208,7 +209,7 @@ def _verify_valid_parameters(self): invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise VLLMValidationError( f"Task {self.task!r} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -231,7 +232,7 @@ def __repr__(self) -> str: def __post_init__(self) -> None: check_removed_pooling_task(self.task) if self.output_kind != RequestOutputKind.FINAL_ONLY: - raise ValueError( + raise VLLMValidationError( "For pooling output_kind has to be FINAL_ONLY, " f"got {self.output_kind!r}" ) diff --git a/vllm/ray/ray_env.py b/vllm/ray/ray_env.py index 5ecca742cb0b..4103b93172ce 100644 --- a/vllm/ray/ray_env.py +++ b/vllm/ray/ray_env.py @@ -35,6 +35,7 @@ # --------------------------------------------------------------------------- DEFAULT_ENV_VAR_PREFIXES: set[str] = { "VLLM_", + "FLASH_ATTENTION_", "LMCACHE_", "NCCL_", "UCX_", diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 84682cdc8bcb..b5111a9ec642 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -84,6 +84,10 @@ "kimi_k2_reasoning_parser", "KimiK2ReasoningParser", ), + "kimi_k3": ( + "kimi_k3_reasoning_parser", + "KimiK3ReasoningParser", + ), "mimo": ( "qwen3_engine_reasoning_parser", "Qwen3ParserReasoningAdapter", @@ -102,7 +106,7 @@ ), "mistral": ( "mistral_reasoning_parser", - "MistralReasoningParser", + "MistralParserReasoningAdapter", ), "nemotron_v3": ( "nemotron_v3_engine_reasoning_parser", diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index f0e7aed0b034..340f7e022c3c 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -20,6 +20,11 @@ ) from e +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation, + CitationSource, + CohereDeltaMessage, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) @@ -404,6 +409,83 @@ def _schema_dict_from_structured_outputs( ) +def _melody_sources_to_vllm( + raw_sources: Any, + position_to_source: Mapping[tuple[int, int], CitationSource] | None, +) -> list[CitationSource]: + """Convert melody's ``Source`` objects into resolved :class:`CitationSource`. + + Melody's ``Source`` shape is + ``{tool_call_index, tool_result_indices, document_ids}``. Without + :meth:`PyFilterOptions.with_message_history` the ``document_ids`` + list is empty, so we resolve the numeric address ourselves using + ``position_to_source``, which is + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source` + applied to the inbound request and forwarded through + ``chat_template_kwargs`` at parser-construction time. Each melody + source with ``tool_result_indices=[i, j, ...]`` fans out to one + :class:`CitationSource` per resolved position, populating the + wire-shape ``type`` / ``id`` / ``document`` / ``tool_output`` + fields; unresolved positions are silently skipped (the calling + citation is dropped at wire-coercion time if every source in it + ends up unresolved -- see ``_to_wire_citation`` in + :mod:`vllm.entrypoints.cohere.serving`). + + When ``position_to_source`` is ``None`` (parser wired outside of + :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion + request that happens to hit a cohere reasoning parser) sources + can't be resolved and get dropped; citations then never reach any + citation-aware handler downstream. + """ + if not position_to_source: + return [] + out: list[CitationSource] = [] + for s in raw_sources or []: + bucket = getattr(s, "tool_call_index", None) + indices = list(getattr(s, "tool_result_indices", None) or []) + if bucket is None or not indices: + continue + for idx in indices: + info = position_to_source.get((bucket, idx)) + if info is not None: + out.append(info) + return out + + +def _melody_citations_to_vllm( + raw_citations: Any, + position_to_source: Mapping[tuple[int, int], CitationSource] | None, +) -> list[Citation] | None: + """Convert melody's ``FilterCitation`` objects into :class:`Citation`. + + Resolves each source in-place via ``position_to_source`` (see + :func:`_melody_sources_to_vllm`). Citations whose sources all fail + to resolve are still emitted here with ``sources=[]``; the serving + layer drops them at wire-coercion time so the fail-closed policy + lives in exactly one place. + """ + if not raw_citations: + return None + out: list[Citation] = [] + for c in raw_citations: + out.append( + Citation( + start=getattr(c, "start_index", None), + end=getattr(c, "end_index", None), + text=getattr(c, "text", None), + sources=_melody_sources_to_vllm( + getattr(c, "sources", None), position_to_source + ), + type=( + "THINKING_CONTENT" + if getattr(c, "is_thinking", False) + else "TEXT_CONTENT" + ), + ) + ) + return out + + class BaseCohereCommandReasoningParser(ReasoningParser): def __init__( self, @@ -420,6 +502,32 @@ def __init__( self.unary_opts = unary_opts self.melody_unary = PyFilter(unary_opts) self.melody_streaming = PyFilter(streaming_opts) + # Citations extracted by the most recent ``extract_reasoning`` call. + # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its + # ``_finalize_response_message`` hook) read this back from the + # parser instance (which is constructed per-request) and attach + # the result to :class:`CohereChatMessage.citations`. ``None`` when + # the last parse produced no citations. + self.last_unary_citations: list[Citation] | None = None + # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved + # CitationSource`` map, forwarded from + # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via + # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm` + # for how it's consumed. Absent for parser instances that + # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI + # chat completion request routed through a cohere reasoning + # parser); citations from those requests are dropped since we + # have no way to attribute their sources. + ctk = kwargs.get("chat_template_kwargs") or {} + # Lazy import: ``renderers.cohere`` transitively imports + # ``cohere_melody`` and other heavy deps; hoisting the constant + # to module scope would tie parser importability to that graph. + from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY + + raw_map = ctk.get(POSITION_TO_SOURCE_KEY) + self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = ( + raw_map if isinstance(raw_map, Mapping) else None + ) @property def reasoning_start_str(self) -> str | None: @@ -439,9 +547,22 @@ def extract_reasoning_streaming( delta_token_ids: Sequence[int], ) -> DeltaMessage | None: r = self.melody_streaming.write_decoded(delta_text) - if r.content is None and r.reasoning is None and not r.tool_calls: + citations = _melody_citations_to_vllm( + getattr(r, "citations", None), self._position_to_source + ) + if ( + r.content is None + and r.reasoning is None + and not r.tool_calls + and not citations + ): return None - msg = DeltaMessage() + # Always emit CohereDeltaMessage so citations reach the wire via + # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope. + # When ``citations`` is unset, ``CohereDeltaMessage._serialize`` + # drops the field so the emitted shape matches plain + # :class:`DeltaMessage` + msg = CohereDeltaMessage() if r.content is not None: msg.content = r.content if r.reasoning is not None: @@ -456,12 +577,23 @@ def extract_reasoning_streaming( ) for tc in r.tool_calls ] + if citations: + msg.citations = citations return msg def extract_reasoning( self, model_output: str, request: ChatCompletionRequest | ResponsesRequest ) -> tuple[str | None, str | None]: result = self.melody_unary.process_full_text(model_output) + # Cache citations so citation-aware handlers can surface them + # on :class:`CohereChatMessage.citations`. The base + # :meth:`ReasoningParser.extract_reasoning` contract only + # returns ``(reasoning, content)``, so citations are passed + # back via parser-instance state -- safe because the parser is + # constructed per-request. + self.last_unary_citations = _melody_citations_to_vllm( + getattr(result, "citations", None), self._position_to_source + ) return result.reasoning, result.content def extract_content_ids(self, input_ids: list[int]) -> list[int]: @@ -549,12 +681,25 @@ def adjust_request( return request +# melody's streaming filter only buffers a partial ```` citation +# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is +# set: otherwise, the moment an opening ```` in the same delta, the filter emits the partial +# marker bytes verbatim as plain content. In vLLM's streaming path the +# parser is fed one token (1-4 chars) per call, so an unbuffered filter +# will leak ````-style markers into ``delta.content`` and never +# emit a ``FilterCitation`` for them. Enabling the flag flips the +# partial-match branch in melody's ``parse_citations`` (see +# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e. +# keep buffering -- which lets a full citation eventually resolve. +# Non-streaming (unary) parsing receives the whole output in one call so +# the flag is a no-op there and we leave ``unary_opts`` alone. class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd3(), + streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd3().no_tools(), **kwargs, ) @@ -565,7 +710,7 @@ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd4(), + streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd4().no_tools(), **kwargs, ) diff --git a/vllm/reasoning/gptoss_reasoning_parser.py b/vllm/reasoning/gptoss_reasoning_parser.py index d7bdca829126..b846768f4f35 100644 --- a/vllm/reasoning/gptoss_reasoning_parser.py +++ b/vllm/reasoning/gptoss_reasoning_parser.py @@ -1,65 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING from transformers import PreTrainedTokenizerBase -from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - -no_func_reasoning_tag = { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "tags": [ - { - "begin": "<|channel|>analysis<|message|>", - "content": {"type": "any_text"}, - "end": "<|end|>", - } - ], - "triggers": ["<|channel|>analysis"], - "stop_after_first": False, - }, -} - - -def from_builtin_tool_to_tag(tool: str) -> list[dict]: - tag = [ - { - "begin": f"<|channel|>commentary to={tool}", - "content": {"type": "any_text"}, - "end": "<|end|>", - }, - { - "begin": f"<|channel|>analysis to={tool}", - "content": {"type": "any_text"}, - "end": "<|end|>", - }, - ] - return tag - - -def tag_with_builtin_funcs(no_func_reasoning_tag, builtin_tool_list: list[str]) -> dict: - import copy - - new_tag = copy.deepcopy(no_func_reasoning_tag) - new_tag["format"]["triggers"].append("<|channel|>commentary to=") - - for tool in builtin_tool_list: - new_tag["format"]["tags"].extend(from_builtin_tool_to_tag(tool)) - return new_tag - class GptOssReasoningParser(ReasoningParser): """ @@ -71,64 +23,14 @@ class GptOssReasoningParser(ReasoningParser): def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): super().__init__(tokenizer, *args, **kwargs) - # The model can output some special tokens between "final" and "<|message|>" - # So we need to look for both sequences to determine the end of reasoning. - self.reasoning_end_token_ids_prefix = self.model_tokenizer.encode( - "<|channel|>final" - ) - self.reasoning_end_token_ids_suffix = self.model_tokenizer.encode("<|message|>") - # We also need to check for the <|end|> token to avoid false positives from - # previous messages in multi-turn conversations. - self.eom_token_id = self.vocab["<|end|>"] - self.reasoning_max_num_between_tokens = 20 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - end_token_ids_prefix = self.reasoning_end_token_ids_prefix - end_token_ids_suffix = self.reasoning_end_token_ids_suffix - assert len(end_token_ids_prefix) > 0, "reasoning_end_token_ids_prefix is empty" - assert len(end_token_ids_suffix) > 0, "reasoning_end_token_ids_suffix is empty" - # Check if the end sequence is present in the input_ids. - # We search from the end of input_ids to find the last match. - for i in range(len(input_ids) - len(end_token_ids_prefix), -1, -1): - if input_ids[i] == self.eom_token_id: - # We looped backwards far enough to find the end of a previous message, - # which means we have searched the entirety of the current message - # and can exit early without searching further back into prior - # messages of the conversation. - return False - if input_ids[i : i + len(end_token_ids_prefix)] == end_token_ids_prefix: - # We have found the prefix, now we look for the suffix after the prefix. - suffix_start = i + len(end_token_ids_prefix) - for j in range( - suffix_start, len(input_ids) - len(end_token_ids_suffix) + 1 - ): - if j - suffix_start >= self.reasoning_max_num_between_tokens: - break - if ( - input_ids[j : j + len(end_token_ids_suffix)] - == end_token_ids_suffix - ): - return True - return False + return True def is_reasoning_end_streaming( self, input_ids: Sequence[int], delta_ids: Iterable[int] ) -> bool: - # The pattern window covers the end-of-reasoning marker itself. - # We add len(delta_ids) so that under speculative decoding (where - # a single step can accept many tokens) the entire accepted chunk - # is always inside the scan region. - delta_ids = tuple(delta_ids) - pattern_len = ( - len(self.reasoning_end_token_ids_prefix) - + self.reasoning_max_num_between_tokens - + len(self.reasoning_end_token_ids_suffix) - ) - window = pattern_len + len(delta_ids) - n = len(input_ids) - if n <= window: - return self.is_reasoning_end(input_ids) - return self.is_reasoning_end(input_ids[n - window :]) + return True def extract_content_ids(self, input_ids: list[int]) -> list[int]: raise NotImplementedError( @@ -159,33 +61,3 @@ def extract_reasoning( "GptOssReasoningParser only provides boundary detection. " "Use HarmonyParser for output parsing." ) - - # This function prepares the structural tag to format reasoning output - def prepare_structured_tag( - self, original_tag: str | None, tool_server: ToolServer | None - ) -> str | None: - if original_tag is None: - if tool_server is None: - return json.dumps(no_func_reasoning_tag) - else: - builtin_tool_list: list[str] = [] - if tool_server.has_tool("browser"): - builtin_tool_list.append("browser") - if tool_server.has_tool("python"): - builtin_tool_list.append("python") - if tool_server.has_tool("container"): - builtin_tool_list.append("container") - - if len(builtin_tool_list) > 0: - logger.info("Builtin_tool_list: %s", builtin_tool_list) - func_tag = json.dumps( - tag_with_builtin_funcs(no_func_reasoning_tag, builtin_tool_list) - ) - else: - logger.info("Builtin_tool_list is empty") - func_tag = json.dumps(no_func_reasoning_tag) - - return func_tag - else: - # There is potential risk for appending the tag to the original tag - return original_tag diff --git a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py index 257dc0f95409..1da1938e72bd 100644 --- a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py +++ b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py @@ -130,18 +130,6 @@ def extract_reasoning( return None, model_output - def _is_strict_increasing_subsequence( - self, subsequence: Sequence[int], sequence: Sequence[int] - ) -> bool: - if not subsequence: - return False - - sub_idx = 0 - for num in sequence: - if sub_idx < len(subsequence) and num == subsequence[sub_idx]: - sub_idx += 1 - return sub_idx == len(subsequence) - def extract_reasoning_streaming( self, previous_text: str, diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py new file mode 100644 index 000000000000..43e4338f322f --- /dev/null +++ b/vllm/reasoning/kimi_k3_reasoning_parser.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning parser for the Kimi K3 (XTML) chat format. + +This strips the ``think`` channel out of generated text and hands the remainder +(``response`` + ``tools`` channels) downstream. Kimi K3 wraps the thinking +channel as an XTML element built from special tokens:: + + <|open|>think<|sep|> <|close|>think<|sep|> + +Two subtleties drive the implementation: + * Unlike Kimi-K2 (a single ```` token), each K3 marker is a 3-token + sequence, so the token-id helpers search for the marker *subsequence* + rather than a single id. + * In thinking mode the serving layer may feed ``<|open|>think<|sep|>`` as the + generation prefix, so the model's output can begin *inside* the think + channel with no open marker. The text paths therefore treat a missing open + marker as "reasoning starts at offset 0". + +When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or +``{"enable_thinking": False}``, i.e. instruct mode) the parser returns every +delta as normal content; there is simply no think channel to extract. +""" + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import regex as re +from transformers import PreTrainedTokenizerBase + +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.reasoning import ReasoningParser + +if TYPE_CHECKING: + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + + +def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int: + """Return start index of the last occurrence of needle in haystack, or -1.""" + n = len(needle) + if n == 0: + return -1 + for i in range(len(haystack) - n, -1, -1): + if list(haystack[i : i + n]) == list(needle): + return i + return -1 + + +class KimiK3ReasoningParser(ReasoningParser): + """Reasoning parser for the Kimi K3 (XTML) think channel.""" + + def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): + super().__init__(tokenizer) + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ReasoningParser " + "constructor during construction." + ) + + # thinking can be disabled via chat_template_kwargs -> identity fallthrough + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + thinking = chat_kwargs.get("thinking", None) + if thinking is None: + thinking = chat_kwargs.get("enable_thinking", True) + self._thinking_enabled = bool(thinking) + + # XTML markers as literal strings (skip_special_tokens=False at serve time) + self._think_open = "<|open|>think<|sep|>" + self._think_close = "<|close|>think<|sep|>" + + # Content-channel markers that must be stripped from the final output. + # The tool parser handles these when active, but when tool_choice is + # "none" (the default for requests without tools) the tool parser is + # bypassed and the reasoning parser must strip them itself. + self._response_open = "<|open|>response<|sep|>" + self._response_close = "<|close|>response<|sep|>" + self._message_close = "<|close|>message<|sep|>" + + # Tolerant matchers: defense-in-depth against vLLM's added-token spacing + # ("<|open|> think <|sep|>"). The `\s*` is a no-op on clean input, so the + # normal path stays byte-exact; it only helps if some serving path leaves + # spaces_between_special_tokens on. See adjust_request for the real fix. + open_marker = r"<\|open\|>" + close_marker = r"<\|close\|>" + sep_marker = r"<\|sep\|>" + self._think_open_re = re.compile(open_marker + r"\s*think\s*" + sep_marker) + self._think_close_re = re.compile(close_marker + r"\s*think\s*" + sep_marker) + self._response_open_re = re.compile( + open_marker + r"\s*response\s*" + sep_marker + ) + self._response_close_re = re.compile( + close_marker + r"\s*response\s*" + sep_marker + ) + self._message_close_re = re.compile( + close_marker + r"\s*message\s*" + sep_marker + ) + + # marker token-id subsequences (each marker is 3 tokens) + self._think_open_ids = tokenizer.encode( + self._think_open, add_special_tokens=False + ) + self._think_close_ids = tokenizer.encode( + self._think_close, add_special_tokens=False + ) + self._last_streaming_delta_token_ids: tuple[int, ...] | None = None + self._last_streaming_content_token_ids: list[int] | None = None + + @property + def reasoning_start_str(self) -> str | None: + return self._think_open + + @property + def reasoning_end_str(self) -> str | None: + return self._think_close + + def adjust_request( + self, + request: "ChatCompletionRequest | ResponsesRequest", + ) -> "ChatCompletionRequest | ResponsesRequest": + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + if not self._thinking_enabled: + return True + # Reasoning has ended only if the *most recent* think block is closed: + # the last close marker must come after the last open marker. A plain + # "a close marker exists anywhere" check false-positives in multi-turn / + # agent continuations, where the chat template keeps a prior turn's + # think channel (with its <|close|>think<|sep|>) in the prompt while the + # current turn is still reasoning (its <|open|>think<|sep|> is the newest + # marker). See _subseq_index (returns the last occurrence). + last_close = _subseq_index(input_ids, self._think_close_ids) + last_open = _subseq_index(input_ids, self._think_open_ids) + if last_open == -1: + # No open marker in scope (e.g. the open was consumed as the + # generation prefix): a close marker means reasoning has ended. + return last_close != -1 + return last_close > last_open + + def _extract_content_ids(self, input_ids: list[int]) -> list[int]: + if not self._thinking_enabled: + return input_ids + idx = _subseq_index(input_ids, self._think_close_ids) + if idx == -1: + return [] # still reasoning + return input_ids[idx + len(self._think_close_ids) :] + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + cached_delta_ids = self._last_streaming_delta_token_ids + cached_content_ids = self._last_streaming_content_token_ids + self._last_streaming_delta_token_ids = None + self._last_streaming_content_token_ids = None + if cached_delta_ids == tuple(input_ids) and cached_content_ids is not None: + return cached_content_ids + return self._extract_content_ids(input_ids) + + def _strip_content_wrapper(self, text: str) -> str: + """Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and + ``<|close|>message<|sep|>`` from *text*. + + When Kimi K3 tool parsing is active it needs the raw XTML + ``response`` + ``tools`` channels. Otherwise the reasoning parser cleans + up the response wrapper itself so API users do not see XTML markers. + """ + # Try structured unwrap first: extract body of response channel + m_ro = self._response_open_re.search(text) + m_rc = self._response_close_re.search(text, m_ro.end() if m_ro else 0) + if m_ro is not None and m_rc is not None: + text = text[m_ro.end() : m_rc.start()] + elif m_ro is not None: + # response opened but not closed (truncated) + text = text[m_ro.end() :] + else: + # No response wrapper — strip stray markers if present + text = self._response_open_re.sub("", text) + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + return text + + @staticmethod + def _should_preserve_tool_channels( + request: "ChatCompletionRequest | ResponsesRequest", + ) -> bool: + return bool(getattr(request, "tools", None)) and ( + getattr(request, "tool_choice", None) != "none" + ) + + def _content_after_reasoning( + self, + text: str, + request: "ChatCompletionRequest | ResponsesRequest", + ) -> str | None: + if self._should_preserve_tool_channels(request): + return text or None + return self._strip_content_wrapper(text) or None + + def extract_reasoning( + self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" + ) -> tuple[str | None, str | None]: + """Split full text into ``(reasoning, rest)`` for the non-streaming path. + + Handles three shapes: + * no think channel at all -> ``(None, model_output)`` (all content) + * open marker present -> reasoning starts after ``<|open|>think<|sep|>`` + * open marker absent but a close marker exists (gen-prefix consumed + the open) -> reasoning starts at offset 0 + ``rest`` is whatever follows the close marker, fed on to the tool parser. + """ + if not self._thinking_enabled: + return None, self._content_after_reasoning(model_output, request) + + m_open = self._think_open_re.search(model_output) + # reasoning content begins right after think-open (or at start if the + # open marker was already consumed as a generation prefix) + content_start = m_open.end() if m_open is not None else 0 + # if there is no think channel at all, everything is content + if m_open is None and self._think_close_re.search(model_output) is None: + return None, self._content_after_reasoning(model_output, request) + + m_close = self._think_close_re.search(model_output, content_start) + if m_close is not None: + reasoning = model_output[content_start : m_close.start()] + rest = model_output[m_close.end() :] + return (reasoning or None, self._content_after_reasoning(rest, request)) + # think not closed -> still reasoning, no content yet + return (model_output[content_start:] or None, None) + + def _reasoning_text_ready_to_emit(self, text: str) -> str: + """Return the reasoning prefix that is safe to stream now. + + Work from accumulated text, not from the current delta alone. That turns + split-open and split-close handling into the same prefix-diff problem. + + Example 1: split think-open marker. + chunks: ``<|open|>`` / ``think`` / ``<|sep|>reasoning`` + current text after chunk 1: ``<|open|>`` -> emit ``""`` + current text after chunk 2: ``<|open|>think`` -> emit ``""`` + current text after chunk 3: ``<|open|>think<|sep|>reasoning`` + -> emit ``reasoning`` + + Example 2: split think-close marker. + chunks: ``reasoning`` / ``<|close|>`` / ``think<|sep|>...`` + after ``<|close|>``, the suffix is only a partial close marker, so the + sendable reasoning is still ``reasoning`` and the delta is empty. Once + ``think<|sep|>`` arrives, the close branch hands the following response + or tools text to the downstream parser. + """ + m_open = self._think_open_re.search(text) + if m_open is not None: + text = text[m_open.end() :] + overlap = 0 + for marker in (self._think_open, self._think_close): + max_check = min(len(marker) - 1, len(text)) + for n in range(max_check, 0, -1): + if text.endswith(marker[:n]): + overlap = max(overlap, n) + break + return text[:-overlap] if overlap else text + + def _content_ready_to_emit(self, text: str) -> str: + """Return the content prefix that is safe to stream now. + + Mirrors ``_reasoning_text_ready_to_emit`` but for the post-reasoning + content phase. Strips the ``<|open|>response<|sep|>`` prefix, holds back + any partial marker suffix, and removes complete + ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` markers. + """ + # Strip response-open prefix + m_open = self._response_open_re.search(text) + if m_open is not None: + text = text[m_open.end() :] + + # Remove complete close/message markers + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + + # Hold back partial markers at the end + overlap = 0 + for marker in ( + self._response_open, + self._response_close, + self._message_close, + ): + max_check = min(len(marker) - 1, len(text)) + for n in range(max_check, 0, -1): + if text.endswith(marker[:n]): + overlap = max(overlap, n) + break + return text[:-overlap] if overlap else text + + def strip_content_streaming( + self, + previous_text: str, + current_text: str, + ) -> DeltaMessage | None: + """Strip XTML content wrappers from streaming deltas after reasoning. + + Called by KimiK3Parser when no tool parser is configured, so the + reasoning parser handles ``<|open|>response<|sep|>`` / + ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` stripping itself. + + Works from accumulated text (``previous_text`` / ``current_text`` + already contain only post-reasoning content). + """ + current_safe = self._content_ready_to_emit(current_text) + previous_safe = self._content_ready_to_emit(previous_text) + if current_safe.startswith(previous_safe): + delta = current_safe[len(previous_safe) :] + else: + delta = current_safe + return DeltaMessage(content=delta) if delta else None + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + self._last_streaming_delta_token_ids = None + self._last_streaming_content_token_ids = None + if not self._thinking_enabled: + return DeltaMessage(content=delta_text) + + # reasoning already ended -> downstream content + if self._think_close_re.search(previous_text): + return DeltaMessage(content=delta_text) + + # the close marker completes within this delta's accumulated text: + # split the buffer at the close marker into reasoning vs trailing content. + m_close = self._think_close_re.search(current_text) + if m_close is not None: + self._last_streaming_delta_token_ids = tuple(delta_token_ids) + self._last_streaming_content_token_ids = self._extract_content_ids( + list(current_token_ids) + ) + m_open = self._think_open_re.search(current_text) + r_start = m_open.end() if m_open is not None else 0 + reasoning = current_text[r_start : m_close.start()] + already_sent = self._reasoning_text_ready_to_emit(previous_text) + if reasoning.startswith(already_sent): + reasoning_delta = reasoning[len(already_sent) :] + else: + reasoning_delta = reasoning + content = current_text[m_close.end() :] + return DeltaMessage( + reasoning=reasoning_delta or None, + content=content or None, + ) + + current_reasoning = self._reasoning_text_ready_to_emit(current_text) + previous_reasoning = self._reasoning_text_ready_to_emit(previous_text) + if current_reasoning.startswith(previous_reasoning): + reasoning_delta = current_reasoning[len(previous_reasoning) :] + else: + reasoning_delta = current_reasoning + if not reasoning_delta: + return None + return DeltaMessage(reasoning=reasoning_delta) + + # Backward-compatible aliases for existing unit tests and downstream users + # that still call the pre-split method names. + extract_reasoning_content = extract_reasoning + extract_reasoning_content_streaming = extract_reasoning_streaming diff --git a/vllm/reasoning/mistral_reasoning_parser.py b/vllm/reasoning/mistral_reasoning_parser.py index c224c3c165c2..6adbb08cb5bf 100644 --- a/vllm/reasoning/mistral_reasoning_parser.py +++ b/vllm/reasoning/mistral_reasoning_parser.py @@ -1,162 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence -from functools import cached_property -from typing import TYPE_CHECKING +from vllm.parser.engine.registered_adapters import MistralParserReasoningAdapter -from vllm.reasoning import ReasoningParser -from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser -from vllm.tokenizers.mistral import MistralTokenizer - -if TYPE_CHECKING: - from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest - from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - - -class MistralReasoningParser(BaseThinkingReasoningParser): - """ - Reasoning parser for Mistral models. - - The Mistral models uses `[THINK]`...`[/THINK]` tokens to denote reasoning - text. This parser extracts the reasoning content from the model output. - - A valid reasoning trace should always start with a `[THINK]` token and end with - a `[/THINK]` token. - - If `[THINK]` token is not generated, then this parser only returns content. - """ - - def __init__(self, tokenizer: MistralTokenizer, *args, **kwargs): - if not isinstance(tokenizer, MistralTokenizer): - raise ValueError("The tokenizer must be an instance of MistralTokenizer.") - - ReasoningParser.__init__(self, tokenizer, *args, **kwargs) - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ReasoningParser " - "constructor during construction." - ) - - self.start_token_id = tokenizer.tokenizer.get_special_token(self.start_token) - self.end_token_id = tokenizer.tokenizer.get_special_token(self.end_token) - - if self.start_token_id is None or self.end_token_id is None: - raise RuntimeError( - "Mistral reasoning parser could not locate think start/end " - "tokens in the tokenizer!" - ) - - @cached_property - def start_token(self) -> str: - """The token that starts reasoning content.""" - from mistral_common.tokens.tokenizers.base import SpecialTokens - - return SpecialTokens.begin_think - - @cached_property - def end_token(self) -> str: - """The token that ends reasoning content.""" - from mistral_common.tokens.tokenizers.base import SpecialTokens - - return SpecialTokens.end_think - - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - has_eot_token = False - - for id in reversed(input_ids): - if id == self.start_token_id: - # Reasoning ends only if a BOT token is found before a EOT token. - return has_eot_token - elif id == self.end_token_id: - has_eot_token = True - return False - - def is_reasoning_end_streaming( - self, input_ids: Sequence[int], delta_ids: Iterable[int] - ) -> bool: - if self.end_token_id in delta_ids: - return True - # Grammar's think? is optional — if [THINK] was never generated, - # reasoning was skipped entirely. - return self.start_token_id not in input_ids - - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """ - Extract the content - """ - has_bot_token = False - has_eot_token = False - bot_token_index = -1 - eot_token_index = -1 - # One for loop instead of multiple lookups - for i, token_id in enumerate(input_ids): - # We filter that we have multiple BOT tokens which should not - # happen for a well prompted trained model - if token_id == self.start_token_id and not has_bot_token: - has_bot_token = True - bot_token_index = i - elif token_id == self.end_token_id: - has_eot_token = True - eot_token_index = i - break - - # 1. Only BOT has been outputted - if has_bot_token and not has_eot_token: - # Should be = [] if model is well prompted and trained. - return input_ids[:bot_token_index] - # 2. Neither BOT or EOT have been outputted - elif not has_bot_token and not has_eot_token: - return input_ids - # 3. Both BOT and EOT have been outputted. - elif has_bot_token and has_eot_token: - return input_ids[:bot_token_index] + input_ids[eot_token_index + 1 :] - # 4. Only EOT has been outputted => this should not have occurred for a model - # well prompted and trained. - else: - return input_ids[:eot_token_index] + input_ids[eot_token_index + 1 :] - - def extract_reasoning( - self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" - ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from the model output. - """ - if not model_output: - return (None, "") - - # Check if the start token is present in the model output, remove it - # if it is present. - prev_bot_token, bot_token, post_bot_token = model_output.partition( - self.start_token - ) - - has_bot_token = bool(bot_token) - # Valid EOT tokens should follow BOT token - has_valid_eot_token = has_bot_token and self.end_token in post_bot_token - - # 1. If there is BOT token followed by EOT token - if has_bot_token and has_valid_eot_token: - prev_eot_token, _, post_eot_token = post_bot_token.partition(self.end_token) - # If model is well prompted and trained prev_bot_token should be "" - content = prev_bot_token + post_eot_token - return prev_eot_token, content if content else None - # 2. Only BOT token - elif has_bot_token: - # If model is well prompted and trained prev_bot_token should be "" - return post_bot_token, prev_bot_token if prev_bot_token else None - # 3. EOT token has been outputted without BOT or neither has been outputted - else: - has_non_valid_eot_token = self.end_token in prev_bot_token - # 3.a EOT token has been outputted without BOT - # If model is well prompted and trained `has_non_valid_eot_token` should - # be `False` and the parser outputs all tokens as 'content' - if has_non_valid_eot_token: - prev_eot_token, _, post_eot_token = prev_bot_token.partition( - self.end_token - ) - return None, prev_eot_token + post_eot_token - # 3.b neither BOT or EOT have been outputted - else: - return None, prev_bot_token +__all__ = ["MistralParserReasoningAdapter"] diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index bb02c115caea..832f954eaf9c 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -79,16 +79,15 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: self.tokenizer = tokenizer - # Shared thread pool executor for blocking tokenizer and - # multimodal preprocessing operations. The multimodal processor - # receives a deep-copied tokenizer (see #36557) so it is safe to - # run tokenization and MM preprocessing concurrently. + # Thread pool executor for blocking tokenizer operations. The + # multimodal processor receives a deep-copied tokenizer (see #36557) + # so it is safe to run tokenization and MM preprocessing concurrently. pool_workers = config.model_config.renderer_num_workers self._executor = ThreadPoolExecutor(max_workers=pool_workers) - # Multimodal preprocessing is always offloaded to the thread pool - # to keep the asyncio event loop responsive under concurrent load. - self._mm_executor: Executor = self._executor + # Separate single-worker executor so tokenization never queues behind + # MM preprocessing; must stay single-worker per #38418 (P0/P1 order). + self._mm_executor: Executor = ThreadPoolExecutor(max_workers=1) # Offload tokenization to the thread pool. The sync # ``_tokenize_prompt`` already encapsulates the unified ``__call__`` @@ -103,7 +102,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: self._readonly_mm_processor: BaseMultiModalProcessor | None = None self._mm_cache_stats: MultiModalCacheStats | None = None self._clear_mm_cache_async = make_async( - self.clear_mm_cache, executor=self._executor + self.clear_mm_cache, executor=self._mm_executor ) self._process_multimodal_async = make_async( self._process_multimodal, executor=self._mm_executor @@ -244,45 +243,47 @@ def warmup(self, chat_params: ChatParams) -> None: """ from vllm.entrypoints.chat_utils import ChatTemplateResolutionError - try: - logger.debug("Warming up chat template processing...") - start_time = time.perf_counter() - - self.render_chat([[{"role": "user", "content": "warmup"}]], chat_params) - - elapsed = time.perf_counter() - start_time - logger.debug("Chat template warmup completed in %.3fs", elapsed) - except ChatTemplateResolutionError: - logger.debug("This model does not support chat template.") - except Exception: - logger.warning("Chat template warmup failed", exc_info=True) - - if self.mm_processor: + # prevent MM processor hangs + with set_default_torch_num_threads(1): try: - logger.debug("Warming up multi-modal processing...") - self._warmup_mm_processor( - self.mm_processor, - log_prefix="Multi-modal", - ) - except Exception: - logger.warning("Multi-modal warmup failed") - finally: - self.clear_mm_cache() + logger.debug("Warming up chat template processing...") + start_time = time.perf_counter() - if self._readonly_mm_processor is not None: - try: - logger.debug("Warming up readonly multi-modal processing...") - self._warmup_mm_processor( - self._readonly_mm_processor, - log_prefix="Readonly multi-modal", - ) + self.render_chat([[{"role": "user", "content": "warmup"}]], chat_params) + + elapsed = time.perf_counter() - start_time + logger.debug("Chat template warmup completed in %.3fs", elapsed) + except ChatTemplateResolutionError: + logger.debug("This model does not support chat template.") except Exception: - logger.warning("Readonly multi-modal warmup failed") - finally: - self._clear_processor_cache(self._readonly_mm_processor) + logger.warning("Chat template warmup failed", exc_info=True) + + if self.mm_processor: + try: + logger.debug("Warming up multi-modal processing...") + self._warmup_mm_processor( + self.mm_processor, + log_prefix="Multi-modal", + ) + except Exception: + logger.warning("Multi-modal warmup failed") + finally: + self.clear_mm_cache() + + if self._readonly_mm_processor is not None: + try: + logger.debug("Warming up readonly multi-modal processing...") + self._warmup_mm_processor( + self._readonly_mm_processor, + log_prefix="Readonly multi-modal", + ) + except Exception: + logger.warning("Readonly multi-modal warmup failed") + finally: + self._clear_processor_cache(self._readonly_mm_processor) async def clear_mm_cache_async(self) -> None: - """Serialize clear_mm_cache through the shared executor to avoid + """Serialize clear_mm_cache through the multimodal executor to avoid races with concurrent process_inputs on the mm_processor_cache.""" await self._clear_mm_cache_async() diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py new file mode 100644 index 000000000000..dbe5bf0ae5d9 --- /dev/null +++ b/vllm/renderers/cohere.py @@ -0,0 +1,704 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere prompt renderer. + +Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using +the ``cohere_melody`` Rust bindings instead of Jinja. Enabled by +passing ``--tokenizer-mode cohere`` on the engine; tokenization itself +still flows through the cached HuggingFace tokenizer. + +The renderer's client surface is ``chat_template_kwargs`` (vLLM's +generic passthrough for template-time inputs). It is fed by two +different code paths, which speak different vocabularies: + +1. :class:`vllm.entrypoints.cohere.serving.CohereServingChatV2` -- + translates a native `Cohere Chat v2 request + `__ into + ``chat_template_kwargs``. +2. :class:`OpenAIServingChat` -- forwards whatever the client puts in + ``request.chat_template_kwargs`` directly, without translation. + +Accordingly the renderer accepts keys from both vocabularies. + +**Cohere Chat v2 request fields** (populated automatically by +``CohereServingChatV2``; a direct ``chat_template_kwargs`` caller may +also pass them): + +* ``documents``: list of documents (v2 ``documents``) +* ``tools``: list of tools (v2 ``tools``) +* ``safety_mode``: v2 ``safety_mode`` (``CONTEXTUAL`` / ``STRICT`` / + ``NONE``) -- forwarded to melody's cmd3 ``safety_mode`` slot; + lowercased. +* ``citation_options``: v2 ``citation_options``. Its ``.mode`` + (``ENABLED`` / ``DISABLED`` / ``FAST`` / ``ACCURATE`` / ``OFF``) is + normalized into melody's cmd3 ``citation_quality`` (``on`` / ``off``) + and cmd4 ``grounding`` (``enabled`` / ``disabled`` / ``unknown``) + slots. cmd4 has no fast/accurate distinction at the prompt-template + layer -- both mean grounding-on. +* ``response_format``: v2 ``response_format`` (``json_object`` / + ``json_schema``) -- mapped onto melody's ``json_mode`` / + ``json_schema`` config slots. +* ``thinking``: v2 ``thinking``; its ``.type`` (``enabled`` / + ``disabled``) becomes melody's ``reasoning_type``. + +**Melody template-config knobs** (accepted for direct-passthrough +callers; not part of the Cohere Chat v2 surface): + +* ``cohere_format``: renderer selector, ``cmd3`` or ``cmd4`` (default + ``cmd4``). +* ``template_id``: pick one of melody's built-in template variants. +* ``available_tools``: melody's own name for tools; takes precedence + over ``tools`` when both are set. +* ``reasoning_type``: direct ``enabled`` / ``disabled`` toggle; + overrides derivation from ``thinking.type``. +* ``dev_instruction``: developer-instruction override. +* ``json_mode`` / ``json_schema``: direct structured-output toggles; + override derivation from ``response_format``. +* cmd3-only: ``citation_quality`` (``on`` / ``off``), ``skip_preamble``. +* cmd4-only: ``grounding`` (``enabled`` / ``disabled`` / ``unknown``; + overrides derivation from ``citation_options.mode``), + ``platform_instruction``. + +**Rejected inputs**: ``template_jinja`` and ``template`` are Cohere's +own inlets for raw Jinja template source. They are valid at Cohere's +API surface but not at vLLM's -- raw template source must flow through +the standard ``chat_template`` request field (guarded by +``--trust-request-chat-template``), which this renderer forwards to +melody as ``template_jinja``. Setting these keys explicitly in +``chat_template_kwargs`` raises ``ValueError`` so client +misconfiguration surfaces loudly. + +**Everything else** in ``chat_template_kwargs`` is forwarded verbatim +to melody as ``additional_template_fields`` and becomes accessible as +Jinja variables inside the template. This matches vLLM's documented +contract for ``chat_template_kwargs`` ("kwargs accessible by the +template"), so e.g. ``chat_template_kwargs={"reasoning_effort": +"low"}`` resolves ``{{ reasoning_effort }}`` inside cmd3 / cmd4 +templates. + +Citations produced by Cohere models are surfaced through the +Cohere-scoped ``CohereChatMessage.citations`` / +``CohereDeltaMessage.citations`` fields (see +:mod:`vllm.entrypoints.cohere.cohere_chat_message`), populated by the +``cohere2`` reasoning parser. The base OpenAI ``ChatMessage`` / +``DeltaMessage`` keep their declared schemas unchanged; the response +envelope declares them as ``SerializeAsAny[...]`` so the subclass +fields survive JSON serialization. +""" + +from __future__ import annotations + +import copy +import json +from enum import Enum +from typing import Any + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.logger import init_logger +from vllm.tokenizers.hf import HfTokenizer +from vllm.utils.async_utils import make_async + +from .base import BaseRenderer +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +logger = init_logger(__name__) + + +_DEFAULT_FORMAT = "cmd4" +_VALID_FORMATS = ("cmd3", "cmd4") + + +class MelodyContentType(str, Enum): + """Wire-format discriminator for melody content blocks. + + These strings are what the cmd3 / cmd4 Jinja templates check against + in ``message.content[0].type`` (e.g. the ``thinking`` branch in + ``cmd4-v1.jinja``). Keep new values in sync with melody's template + schema. + """ + + TEXT = "text" + THINKING = "thinking" + IMAGE = "image" + DOCUMENT = "document" + + +# Reserved ``chat_template_kwargs`` key that carries +# ``AssistantChatMessageV2.citations`` from a Cohere v2 request through +# to the renderer, keyed by request-message index. Populated by +# ``CohereServingChatV2._apply_cohere_template_kwargs`` and consumed by +# ``_conversation_to_melody_messages``. The value is a +# ``dict[int, list[dict]]`` of melody-shape ``FilterCitation`` dicts. +# The leading underscore signals "internal, not part of the public +# ``chat_template_kwargs`` surface" -- clients should not set this key +# directly. +MESSAGES_CITATIONS_KEY = "_messages_citations" + + +# Reserved ``chat_template_kwargs`` key carrying a request-side +# ``dict[tuple[int, int], CitationSource]`` map from melody's numeric +# ``(tool_call_index, tool_result_index)`` addressing to fully-resolved +# wire-shape sources (type, id, document/tool_output payload). +# Populated by :meth:`CohereServingChatV2._apply_cohere_template_kwargs` +# and consumed by :class:`BaseCohereCommandReasoningParser` at request +# start so the parser can attach real document ids on each melody source +# it emits -- eliminating the numeric-coord round-trip through the +# internal streaming pipeline. Leading underscore signals "internal, +# not part of the public ``chat_template_kwargs`` surface" -- clients +# should not set this key directly. +POSITION_TO_SOURCE_KEY = "_position_to_source" + + +# Keys this renderer interprets directly from ``chat_template_kwargs`` and +# maps onto typed melody render-config fields. Everything *not* in this +# set is forwarded verbatim to melody as ``additional_template_fields`` +# (i.e. as Jinja template variables), so callers can write +# ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}`` +# resolve inside the template -- matching vLLM's documented contract +# for ``chat_template_kwargs``. +_RENDERER_CONSUMED_KEYS = frozenset( + { + "cohere_format", + "template_id", + # ``template_jinja`` / ``template`` are Cohere-only inputs for + # raw Jinja source -- valid at Cohere's own API surface but not + # at vLLM's. In vLLM, raw template source must flow through the + # standard ``chat_template`` request field, so a non-``None`` + # value here is rejected in ``_build_render_config``. The keys + # stay in the consumed set so an explicit ``None`` (e.g. from + # ``.model_dump(exclude_none=False)``) is dropped rather than + # surfacing as a stray ``{{ template_jinja }}`` Jinja variable. + "template_jinja", + "template", + "use_jinja", + "documents", + "available_tools", + "tools", + "reasoning_type", + "thinking", + "dev_instruction", + "response_format", + "json_schema", + "json_mode", + "safety_mode", + "citation_quality", + "citation_options", + "skip_preamble", + "grounding", + "platform_instruction", + MESSAGES_CITATIONS_KEY, + POSITION_TO_SOURCE_KEY, + } +) + + +def _try_import_melody(): + try: + import cohere_melody # type: ignore + + return cohere_melody + except ImportError as e: # pragma: no cover - exercised at runtime + raise ImportError( + "The `cohere` tokenizer/renderer mode requires the " + "`cohere_melody` package. Install it via " + "`pip install cohere-melody` or build from " + "https://github.com/cohere-ai/melody." + ) from e + + +_MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"}) + + +# Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``. +# v2 surfaces five modes (``ENABLED`` / ``DISABLED`` / ``FAST`` / +# ``ACCURATE`` / ``OFF``) but cmd4's template doesn't differentiate +# fast vs accurate at the prompt layer -- both just turn grounding on, +# and ``ENABLED`` / ``DISABLED`` are direct aliases. Melody itself +# accepts ``unknown`` / ``enabled`` / ``disabled``. +_CMD4_GROUNDING_FROM_MODE = { + "fast": "enabled", + "accurate": "enabled", + "on": "enabled", + "enabled": "enabled", + "off": "disabled", + "disabled": "disabled", + "unknown": "unknown", +} + + +def _normalize_cmd4_grounding(value: Any) -> str: + """Coerce a user-facing grounding/citation mode to melody's vocab. + + Raises ``ValueError`` rather than letting an unrecognized value + slip through to ``render_cmd4`` and surface as a generic + ``Invalid config: grounding`` from melody. + """ + out = _CMD4_GROUNDING_FROM_MODE.get(value.lower()) + if out is None: + raise ValueError( + f"Unrecognized cmd4 grounding value: {value!r}. Expected one " + f"of ENABLED / DISABLED / FAST / ACCURATE / OFF (from " + f"citation_options.mode), or the direct melody values " + f"enabled / disabled / unknown." + ) + return out + + +def _role_to_melody(role: str) -> str: + """Map an OpenAI role to the role string melody expects. + + The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``, + ``assistant``/``chatbot``, and ``tool``: any other role is silently + dropped by the template's role-dispatch chain (no fallback branch), + which produces a malformed prompt without any error. We therefore + refuse unknown roles up front rather than letting them disappear. + + Aliases: + + * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role + name used by the templates). + * ``developer`` -> ``system`` (OpenAI's ``developer`` role is + documented as high-priority instructions, which maps onto the + ``system`` slot in Cohere's prompt format). + """ + role = role.lower() + if role == "assistant": + return "chatbot" + if role == "developer": + return "system" + if role in _MELODY_ROLES: + return role + raise ValueError( + f"Unsupported message role for the cohere renderer: {role!r}. " + "Expected one of: system, developer, user, assistant, chatbot, tool." + ) + + +def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]: + """Normalize an OpenAI tool call dict into melody's tool call shape. + + melody expects ``{id, name, parameters: }`` whereas OpenAI + delivers ``{id, type, function: {name, arguments: }}``. + """ + # Pydantic objects + if hasattr(tc, "model_dump"): + tc = tc.model_dump() + if not isinstance(tc, dict): + raise TypeError(f"Unexpected tool_call value: {tc!r}") + + fn = tc.get("function") or {} + name = fn.get("name") or tc.get("name") or "" + args = fn.get("arguments") + if args is None: + args = tc.get("arguments", {}) + # melody expects a JSON-encoded string + if not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + return { + "id": tc.get("id") or "", + "name": name, + "parameters": args, + } + + +def _content_blocks(content: Any) -> list[dict[str, Any]]: + """Convert OpenAI ``ConversationMessage.content`` into melody content blocks. + + The chat_utils ``content_format="openai"`` produces a list of + ``{type: text, text: ...}`` dicts already; we pass those through with + minimal coercion. Plain string content is wrapped in a single text block. + Image / multimodal placeholder dicts (``{"type": "image"}``) are + forwarded as-is so that templates can reference image placeholders that + the upstream tokenizer expands separately. + """ + if content is None: + return [] + if isinstance(content, str): + return [{"type": MelodyContentType.TEXT, "text": content}] + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": MelodyContentType.TEXT, "text": part}) + continue + if not isinstance(part, dict): + raise TypeError(f"Unexpected content part: {part!r}") + part_type = part.get("type", MelodyContentType.TEXT) + if part_type in ("text", "input_text", "output_text", "refusal"): + blocks.append( + {"type": MelodyContentType.TEXT, "text": part.get("text", "")} + ) + elif part_type == MelodyContentType.THINKING: + blocks.append( + { + "type": MelodyContentType.THINKING, + "thinking": part.get("thinking", ""), + } + ) + elif part_type == MelodyContentType.IMAGE: + blocks.append( + { + "type": MelodyContentType.IMAGE, + "image": { + "template_placeholder": part.get( + "template_placeholder", "" + ), + }, + } + ) + elif part_type == MelodyContentType.DOCUMENT: + doc = part.get("document") + if isinstance(doc, dict): + blocks.append({"type": MelodyContentType.DOCUMENT, "document": doc}) + else: + # Fall back to wrapping arbitrary string as text. + blocks.append({"type": MelodyContentType.TEXT, "text": json.dumps(doc)}) + elif part_type == "tool_reference": + # Tool references are rendered by name; emit as text so the + # renderer downstream is content-format agnostic. + blocks.append( + { + "type": MelodyContentType.TEXT, + "text": part.get("name") or part.get("text", ""), + } + ) + else: + # Unknown block type: render as text fallback. + text = part.get("text") or part.get(part_type) or "" + text_str = text if isinstance(text, str) else json.dumps(text) + blocks.append({"type": MelodyContentType.TEXT, "text": text_str}) + return blocks + + +def _document_to_melody(doc: Any) -> dict[str, Any]: + """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape.""" + match doc: + case str(): + return {"text": doc} + # Cohere v2 wraps documents in {id, data: {...}}; the ``**rest`` + # capture keeps the outer ``id`` (and any other outer keys) so we + # can lift them onto the payload without losing them. + case {"data": dict() as data, **rest}: + payload = dict(data) + if (outer_id := rest.get("id")) is not None and "id" not in payload: + payload["id"] = outer_id + return payload + case dict(): + return dict(doc) + case _: + raise TypeError(f"Unsupported document type: {type(doc).__name__}") + + +def _tool_to_melody(tool: Any) -> dict[str, Any]: + """Coerce a chat completions tool definition into melody's ``Tool`` shape. + + Accepts either the raw OpenAI tool wrapper ``{type:"function", function: + {name, description, parameters}}`` or a flat ``{name, description, + parameters}`` dict (which is what melody itself expects). + """ + if hasattr(tool, "model_dump"): + tool = tool.model_dump() + if not isinstance(tool, dict): + raise TypeError(f"Unsupported tool type: {type(tool).__name__}") + if "function" in tool and isinstance(tool["function"], dict): + fn = tool["function"] + else: + fn = tool + return { + "name": fn.get("name", ""), + "description": fn.get("description", "") or "", + "parameters": fn.get("parameters") or {}, + } + + +def _conversation_to_melody_messages( + conversation: list[ConversationMessage], + messages_citations: dict[int, list[dict[str, Any]]] | None = None, +) -> list[dict[str, Any]]: + """Convert vLLM ``ConversationMessage``s into melody's message dict shape. + + ``messages_citations`` is an optional index-keyed lookup of + per-message citations from the request, populated by + ``CohereServingChatV2`` under the + :data:`MESSAGES_CITATIONS_KEY` chat_template_kwargs entry. Values + must already be in melody's ``FilterCitation`` dict shape (see + ``CohereServingChatV2._sdk_citation_to_melody``). Index is the + position in the *input* ``conversation`` list, which currently maps + 1-to-1 with ``CohereChatV2Request.messages`` for text-only Cohere + requests. This mapping breaks silently if ``parse_chat_messages`` + ever expands a single input message into multiple entries (e.g. + multi-modal splitting), which is not the case today. + """ + out: list[dict[str, Any]] = [] + for i, msg in enumerate(conversation): + role = _role_to_melody(msg.get("role", "user")) + content_blocks = _content_blocks(msg.get("content")) + + # Treat reasoning content as a thinking block on assistant turns + # so multi-turn reasoning is preserved in the rendered prompt. + # + # Cohere models output thought as either a ``thinking`` + # block (reasoning models) or a ``tool_plan`` field (older non- + # reasoning Command models), but vLLM's ConversationMessage has a + # single unified ``reasoning`` field that drops that difference. + # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both + # the case of tool calls and regular thinking blocks, so this is fine + # from the renderer's perspective. + reasoning = msg.get("reasoning") or msg.get("reasoning_content") + if role == "chatbot" and reasoning: + content_blocks.insert( + 0, + {"type": MelodyContentType.THINKING, "thinking": reasoning}, + ) + + tool_calls = [_normalize_tool_call(tc) for tc in (msg.get("tool_calls") or [])] + + out_msg: dict[str, Any] = { + "role": role, + "content": content_blocks, + "tool_calls": tool_calls, + } + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + out_msg["tool_call_id"] = tool_call_id + if messages_citations and (cites := messages_citations.get(i)): + out_msg["citations"] = cites + out.append(out_msg) + return out + + +def _build_render_config( + conversation: list[ConversationMessage], + chat_template_kwargs: dict[str, Any], + chat_template: str | None = None, +) -> tuple[str, dict[str, Any]]: + """Build the ``render_cmd3`` / ``render_cmd4`` config dict. + + ``chat_template`` is the standard vLLM per-request template override + (from ``ChatCompletionRequest.chat_template`` / the server-side + ``--chat-template`` flag). When non-``None`` it is forwarded to melody + as ``template_jinja`` so a single vLLM-native path drives template + source through the trust guard. + + Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"`` + or ``"cmd4"``. + """ + fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT) + if fmt not in _VALID_FORMATS: + raise ValueError( + f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}" + ) + + # Reject Cohere-only inlets for raw Jinja source. These keys are + # valid at Cohere's own API surface but not at vLLM's: callers must + # use the standard vLLM ``chat_template`` request field so the + # ``--trust-request-chat-template`` guard applies uniformly. + for cohere_only_key in ("template_jinja", "template"): + if chat_template_kwargs.get(cohere_only_key) is not None: + raise ValueError( + f"chat_template_kwargs.{cohere_only_key!r} is not accepted " + "in vLLM (it is a Cohere-only field); pass the template " + "body via the standard 'chat_template' request field " + "instead (requires --trust-request-chat-template)." + ) + + config: dict[str, Any] = { + "messages": _conversation_to_melody_messages( + conversation, + chat_template_kwargs.get(MESSAGES_CITATIONS_KEY), + ), + } + + # ``template_id`` is a selector for one of melody's built-in template + # variants (not raw source), so it is safe to accept from the client. + if template_id := chat_template_kwargs.get("template_id"): + config["template_id"] = template_id + # Raw template source flows exclusively through the standard vLLM + # ``chat_template`` field so it is subject to the + # ``--trust-request-chat-template`` guard in OnlineRenderer. + if chat_template: + config["template_jinja"] = chat_template + # Only support Jinja with vllm + config["use_jinja"] = True + + # Documents + documents = chat_template_kwargs.get("documents") or [] + if documents: + config["documents"] = [_document_to_melody(d) for d in documents] + + # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools`` + tools = ( + chat_template_kwargs.get("available_tools") + or chat_template_kwargs.get("tools") + or [] + ) + if tools: + config["available_tools"] = [_tool_to_melody(t) for t in tools] + + # Reasoning toggle (cmd3 + cmd4) + if (rt := chat_template_kwargs.get("reasoning_type")) is not None: + config["reasoning_type"] = str(rt) + elif "thinking" in chat_template_kwargs: + # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand + thinking = chat_template_kwargs["thinking"] + t = thinking.get("type") if isinstance(thinking, dict) else thinking + if t in ("enabled", "disabled"): + config["reasoning_type"] = t + + if (di := chat_template_kwargs.get("dev_instruction")) is not None: + config["dev_instruction"] = str(di) + + # JSON / structured outputs + if (rf := chat_template_kwargs.get("response_format")) is not None: + rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf) + rf_type = rf.get("type") + if rf_type == "json_object": + config["json_mode"] = True + elif rf_type in ("json_schema", "json"): + schema = rf.get("schema") or rf.get("json_schema") + if isinstance(schema, dict) and "schema" in schema: + schema = schema["schema"] + if schema is not None: + config["json_schema"] = ( + schema if isinstance(schema, str) else json.dumps(schema) + ) + if (js := chat_template_kwargs.get("json_schema")) is not None: + config["json_schema"] = js if isinstance(js, str) else json.dumps(js) + if "json_mode" in chat_template_kwargs: + config["json_mode"] = bool(chat_template_kwargs["json_mode"]) + + if fmt == "cmd3": + if (sm := chat_template_kwargs.get("safety_mode")) is not None: + config["safety_mode"] = str(sm).lower() + # citation_quality: ``on`` / ``off`` + cq = chat_template_kwargs.get("citation_quality") + if cq is None and (co := chat_template_kwargs.get("citation_options")): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + cq = "on" if str(mode).lower() != "off" else "off" + if cq is not None: + config["citation_quality"] = str(cq).lower() + if "skip_preamble" in chat_template_kwargs: + config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"]) + else: # cmd4 + # cmd4 uses ``grounding`` rather than safety_mode/citation_quality. + # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``, + # so the Cohere v2 ``citation_options.mode`` values + # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a + # raw lowercased passthrough would raise from + # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the + # prompt-template layer; both request grounding-on). + if (gr := chat_template_kwargs.get("grounding")) is not None: + config["grounding"] = _normalize_cmd4_grounding(gr) + elif co := chat_template_kwargs.get("citation_options"): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + config["grounding"] = _normalize_cmd4_grounding(mode) + if (pi := chat_template_kwargs.get("platform_instruction")) is not None: + config["platform_instruction"] = str(pi) + + # Anything we didn't explicitly interpret above is forwarded to melody + # as a Jinja template variable. This matches vLLM's documented contract + # for ``chat_template_kwargs`` ("kwargs accessible by the template") + # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly + # without a nested ``additional_template_fields`` wrapper. + extra = { + k: v + for k, v in chat_template_kwargs.items() + if k not in _RENDERER_CONSUMED_KEYS + } + if extra: + config["additional_template_fields"] = extra + + return fmt, config + + +class CohereRenderer(BaseRenderer[HfTokenizer]): + """Renderer that templates Cohere prompts via the melody Rust bindings. + + Tokenization is delegated to the standard HF tokenizer; only the + chat-template step is replaced with ``cohere_melody.render_cmd3`` / + ``render_cmd4``. Enabled via ``--tokenizer-mode cohere``. + """ + + def __init__( + self, + config: VllmConfig, + tokenizer: HfTokenizer | None, + ) -> None: + # Match HfRenderer in not mutating the cached tokenizer instance + tokenizer = copy.copy(tokenizer) + super().__init__(config, tokenizer) + + # Lazy import to keep `cohere_melody` an optional dependency + self._melody = _try_import_melody() + # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the + # thread-pool wrapper once so the async path doesn't allocate a new + # adapter on every request. + self._render_async = make_async(self._render, executor=self._executor) + + def _render(self, fmt: str, config_dict: dict[str, Any]) -> str: + if fmt == "cmd3": + return self._melody.render_cmd3(config_dict) + return self._melody.render_cmd4(config_dict) + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config( + conversation, chat_template_kwargs, params.chat_template + ) + prompt_text = self._render(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config( + conversation, chat_template_kwargs, params.chat_template + ) + prompt_text = await self._render_async(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index a7a6693154b0..95210e4a47b4 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1122,6 +1122,8 @@ async def render_messages_async( and mm_uuids is not None and mm_data is not None ): + mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data) + # get video placeholder, replace it with runtime video-chunk prompts video_placeholder = getattr( model_config.hf_config, "video_placeholder", None @@ -1274,8 +1276,8 @@ def _apply_prompt_embeds_to_prompt( embeds_prompt["prompt_embeds"] = full_embeds embeds_prompt["prompt_is_token_ids"] = is_token_ids_mask - @staticmethod def _apply_prompt_embeds_to_engine_input( + self, engine_input: MultiModalInput, prompt_embeds_tensors: list[torch.Tensor], mm_updates: MultiModalPromptUpdates, @@ -1298,6 +1300,7 @@ def _apply_prompt_embeds_to_engine_input( pe_kwargs_items: list[MultiModalKwargsItem] = [] pe_hashes: list[str] = [] pe_placeholders: list[PlaceholderRange] = [] + mm_config = self.model_config.get_multimodal_config() for tensor, (start, length) in zip( prompt_embeds_tensors, positions, strict=True ): @@ -1311,7 +1314,11 @@ def _apply_prompt_embeds_to_engine_input( } ) ) - pe_hashes.append(MultiModalHasher.hash_kwargs(prompt_embeds=tensor)) + pe_hashes.append( + MultiModalHasher.hash_kwargs( + mm_config.mm_hasher_algorithm, prompt_embeds=tensor + ) + ) # `is_embed=None` matches the existing image_embeds-style # "no encoder, just splice the tensor directly" semantics. pe_placeholders.append( diff --git a/vllm/renderers/kimi_k3.py b/vllm/renderers/kimi_k3.py new file mode 100644 index 000000000000..e0f937f2e6ce --- /dev/null +++ b/vllm/renderers/kimi_k3.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any, cast + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.exceptions import VLLMValidationError +from vllm.multimodal.media.connector import merge_media_io_kwargs +from vllm.tokenizers.hf import HfTokenizer +from vllm.utils.async_utils import make_async + +from .base import BaseRenderer +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +# Keep the original image mode (including the alpha channel) for K3 instead of +# flattening images onto a background color. Server-level (--media-io-kwargs) +# and request-level media_io_kwargs still take precedence over this default. +_K3_MEDIA_IO_DEFAULTS: dict[str, dict[str, Any]] = {"image": {"image_mode": None}} +_K3_THINKING_EFFORTS = ("low", "high", "max") + + +def _merge_k3_media_io_kwargs( + media_io_kwargs: dict[str, dict[str, Any]] | None, +) -> dict[str, dict[str, Any]] | None: + return merge_media_io_kwargs(_K3_MEDIA_IO_DEFAULTS, media_io_kwargs) + + +def _dump_k3_template_value(value: Any) -> Any: + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + return model_dump(mode="json", exclude_none=True) + + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + return to_dict() + + return value + + +def _apply_k3_thinking_kwargs(kwargs: dict[str, Any]) -> None: + if (enable_thinking := kwargs.pop("enable_thinking", None)) is not None: + kwargs.setdefault("thinking", enable_thinking) + + reasoning_effort = kwargs.pop("reasoning_effort", None) + if reasoning_effort == "none": + kwargs.setdefault("thinking", False) + elif reasoning_effort is not None: + kwargs.setdefault("thinking_effort", reasoning_effort) + + thinking_effort = kwargs.get("thinking_effort") + if thinking_effort is not None and thinking_effort not in _K3_THINKING_EFFORTS: + supported = ", ".join(_K3_THINKING_EFFORTS) + raise VLLMValidationError( + f"Kimi K3 supports thinking_effort values: {supported}", + parameter="thinking_effort", + value=thinking_effort, + ) + + +def _normalize_k3_tool_messages( + conversation: list[ConversationMessage], +) -> list[dict[str, Any]]: + """Reorder tool-result messages to match assistant tool_call order. + + Supports matching by ``tool_call_id`` or by the synthetic + ``"{tool}:{zero_based_index}"`` alias. When any tool message in a + block cannot be resolved, the whole block is left in its original + order (graceful fallback). + + Returns a new list; caller-owned message dicts are not mutated. + """ + normalized: list[dict[str, Any]] = [] + i = 0 + + while i < len(conversation): + message = conversation[i] + normalized.append(dict(message)) + i += 1 + + if message.get("role") != "assistant": + continue + + tool_calls = message.get("tool_calls") + if not tool_calls: + continue + + # Build the lookup table from the assistant's tool_calls. + targets_by_id: dict[str, tuple[int, str]] = {} + for position, tool_call in enumerate(tool_calls): + function = tool_call.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + + call_target = (position, name) + aliases = [f"{name}:{position}"] + if tool_call_id := tool_call.get("id"): + aliases.insert(0, str(tool_call_id)) + for alias in aliases: + targets_by_id.setdefault(alias, call_target) + + # Collect consecutive tool messages. + block_start = i + while i < len(conversation) and conversation[i].get("role") == "tool": + i += 1 + if i == block_start: + continue + + # Try to resolve every tool message in the block. + resolved: list[tuple[int, int, dict[str, Any]]] = [] + for order, tool_message in enumerate(conversation[block_start:i]): + tool_call_id = tool_message.get("tool_call_id") + resolved_target = ( + targets_by_id.get(str(tool_call_id)) + if tool_call_id is not None + else None + ) + if resolved_target is None: + normalized.extend(dict(item) for item in conversation[block_start:i]) + break + + position, name = resolved_target + enriched = dict(tool_message) + enriched["tool"] = name + enriched["index"] = position + 1 + resolved.append((position, order, enriched)) + else: + resolved.sort(key=lambda item: (item[0], item[1])) + normalized.extend(item for _, _, item in resolved) + + return normalized + + +class KimiK3Renderer(BaseRenderer[HfTokenizer]): + """Render chat prompts with Kimi K3's Python XTML encoding. + + K3 ships no Jinja chat template; its tokenizer renders messages through + ``encoding_k3`` instead. We tokenize eagerly so the structural markers keep + their special-token ids while user- and tool-supplied text stays ordinary. + """ + + def __init__(self, config: VllmConfig, tokenizer: HfTokenizer | None) -> None: + super().__init__(config, tokenizer) + + self._apply_chat_template_async = make_async( + self._apply_chat_template, executor=self._executor + ) + + def _apply_chat_template( + self, + conversation: list[dict[str, Any]], + params: ChatParams, + ) -> list[int]: + # Tokenize eagerly: K3 encodes structural markers as special tokens and + # user/tool text as ordinary tokens, so we cannot defer to a plain + # re-tokenization of the rendered string downstream. + kwargs = params.get_apply_chat_template_kwargs() + _apply_k3_thinking_kwargs(kwargs) + if params.tool_choice not in (None, "auto"): + kwargs["tool_choice"] = _dump_k3_template_value(params.tool_choice) + if params.response_format is not None: + kwargs["response_format"] = _dump_k3_template_value(params.response_format) + kwargs["tokenize"] = True + return self.get_tokenizer().apply_chat_template(conversation, **kwargs) + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="string", + media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + rendered_conversation = _normalize_k3_tool_messages(conversation) + prompt = parse_dec_only_prompt( + self._apply_chat_template(rendered_conversation, params) + ) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return cast(list[ConversationMessage], rendered_conversation), prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="string", + media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + rendered_conversation = _normalize_k3_tool_messages(conversation) + token_ids = await self._apply_chat_template_async(rendered_conversation, params) + prompt = parse_dec_only_prompt(token_ids) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return cast(list[ConversationMessage], rendered_conversation), prompt diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 3ce4f74d9d82..a1285d8e178d 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -10,20 +10,31 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, ) from vllm.entrypoints.openai.completion.protocol import ( CompletionLogProbs, + CompletionRequest, CompletionResponseChoice, + CompletionResponseStreamChoice, + CompletionStreamResponse, +) +from vllm.entrypoints.openai.engine.protocol import DeltaMessage, ToolCall, UsageInfo +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderStreamState, + GenerateResponse, + GenerateStreamResponse, ) -from vllm.entrypoints.openai.engine.protocol import ToolCall -from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateResponse from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.detokenizer_utils import detokenize_incrementally from vllm.utils import random_uuid +from vllm.utils.async_utils import make_async logger = init_logger(__name__) @@ -73,10 +84,26 @@ def __init__( self.supports_browsing = False self.supports_code_interpreter = False + # Detokenization, logprob resolution and parsing are CPU-bound; + # offload them in one hop to keep the event loop responsive. + self._derender_chat_async = make_async( + self._derender_chat, executor=renderer._executor + ) + self._derender_completion_async = make_async( + self._derender_completion, executor=renderer._executor + ) + async def derender_chat( self, generate_response: GenerateResponse, chat_request: ChatCompletionRequest | None = None, + ) -> list[ChatCompletionResponseChoice]: + return await self._derender_chat_async(generate_response, chat_request) + + def _derender_chat( + self, + generate_response: GenerateResponse, + chat_request: ChatCompletionRequest | None = None, ) -> list[ChatCompletionResponseChoice]: tokenizer = self.renderer.get_tokenizer() choices: list[ChatCompletionResponseChoice] = [] @@ -151,9 +178,15 @@ async def derender_chat( tool_calls=tc_items, ) else: - # No parser: plain detokenization. + # No parser: plain detokenization honouring the request's + # skip_special_tokens (default True when no request was given). + skip_special = ( + chat_request.skip_special_tokens + if chat_request is not None + else True + ) decoded_text = tokenizer.decode( - choice.token_ids, skip_special_tokens=True + choice.token_ids, skip_special_tokens=skip_special ) message = ChatMessage(role="assistant", content=decoded_text) @@ -168,16 +201,200 @@ async def derender_chat( return choices + def _detokenize_delta( + self, + tokenizer: TokenizerLike, + delta_token_ids: list[int], + state: DerenderStreamState, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, + ) -> tuple[str, DerenderStreamState]: + """Incrementally detokenize ``delta_token_ids`` from prior stream state. + + Resumes decoding from the offsets carried in ``state`` rather than + replaying token history. ``state.prev_tokens`` holds the trailing decode + window (from ``prefix_offset`` onward) that ``detokenize_incrementally`` + still needs to reproduce any partially read multi-byte character + (tracked by ``read_offset``). The delta tokens are fed straight onto it. + + The window is bounded. ``detokenize_incrementally`` never reads before + ``prefix_offset``, so after processing we trim ``prev_tokens`` to that + tail and rebase the offsets to it. State transport therefore stays + O(window) per chunk instead of re-sending the full token history. + + Args: + tokenizer: The tokenizer to decode with. + delta_token_ids: New token IDs from this generate chunk. + state: Client carried detok state from the previous call. + skip_special_tokens: Passed through to the tokenizer. + spaces_between_special_tokens: Passed through to the tokenizer. + + Returns: + (new_text, updated_state) — the delta text for this chunk and the + state to pass to the next call. + """ + prev_tokens = list(state.prev_tokens) + prefix_offset = state.prefix_offset + read_offset = state.read_offset + + text_parts: list[str] = [] + for tok_id in delta_token_ids: + # prev_tokens is a (possibly empty) list, never None, so this + # always takes the non first iter path and only consumes + # all_input_ids[-1]. + new_toks, text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer=tokenizer, + all_input_ids=[tok_id], + prev_tokens=prev_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + ) + prev_tokens = prev_tokens + new_toks + text_parts.append(text) + + # Trim to the tail still readable by detokenize_incrementally + # (everything before prefix_offset is dead) and rebase the offsets so + # the carried window stays bounded regardless of generation length. + trimmed = prev_tokens[prefix_offset:] + updated_state = state.model_copy( + update={ + "prev_tokens": trimmed, + "prefix_offset": 0, + "read_offset": read_offset - prefix_offset, + } + ) + return "".join(text_parts), updated_state + + async def derender_chat_stream( + self, + model: str, + generate_chunk: GenerateStreamResponse, + state: DerenderStreamState | None = None, + chat_request: ChatCompletionRequest | None = None, + prompt_tokens: int | None = None, + ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState]: + """Process one GenerateStreamResponse chunk for streaming chat derender. + + TODO: parse path for reasoning and tool calls is implemented in future PR. + + Unlike OpenAI's API, which always emits ``role: "assistant"`` on the + very first chunk, this emits it on the first chunk with a non empty + ``choices`` list. A leading usage only chunk therefore defers the + role to the following content chunk instead of sending an empty + role only delta. + + Args: + model: Model name for the response object. + generate_chunk: One SSE chunk from ``/inference/v1/generate``. + state: Client carried detok state (``None`` for first call). + chat_request: Original ChatCompletionRequest from ``/render``. + prompt_tokens: Prompt token count for the usage chunk. + + Returns: + (chunk, updated_state) — the derendered SSE chunk and the state + the client must pass to the next call. + """ + if state is None: + state = DerenderStreamState() + + if self.parser is not None: + # TODO: Follow on PR will implement the parse path. Check on the + # parser alone (fail closed). A parser configured model must never + # fall through to plain detok on the streaming path, even when + # ``chat_request`` is omitted or reasoning/tool markup would leak + # into ``delta.content``. + raise NotImplementedError( + "Streaming chat derender is not yet supported for models with " + "a reasoning or tool parser configured. Use the non-streaming " + "derender endpoint (stream=false) for parsed output." + ) + + # A single DerenderStreamState is threaded through every choice in + # this chunk. Correct only when there is at most one choice per SSE + # event (n=1, one call per index), as the streaming derender + # protocol assumes. Multiple choices sharing one chunk would corrupt + # each other's detok window. + if len(generate_chunk.choices) > 1: + raise ValueError( + "derender_chat_stream expects at most one choice per chunk" + ) + + tokenizer = self.renderer.get_tokenizer() + skip_special = ( + chat_request.skip_special_tokens if chat_request is not None else True + ) + stream_choices: list[ChatCompletionResponseStreamChoice] = [] + updated_state = state + + for choice in generate_chunk.choices: + delta_tids = choice.token_ids or [] + new_text, updated_state = self._detokenize_delta( + tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special + ) + + include_role = not updated_state.role_sent + if include_role: + updated_state = updated_state.model_copy(update={"role_sent": True}) + + delta = DeltaMessage( + role="assistant" if include_role else None, + content=new_text if new_text else None, + ) + stream_choices.append( + ChatCompletionResponseStreamChoice( + index=choice.index, + delta=delta, + finish_reason=choice.finish_reason, + ) + ) + + usage: UsageInfo | None = None + if generate_chunk.usage is not None: + u = generate_chunk.usage + pt = prompt_tokens if prompt_tokens is not None else (u.prompt_tokens or 0) + ct = u.completion_tokens or 0 + usage = UsageInfo( + prompt_tokens=pt, + completion_tokens=ct, + total_tokens=pt + ct, + ) + + chunk = ChatCompletionStreamResponse( + id=generate_chunk.request_id, + model=model, + choices=stream_choices, + usage=usage, + ) + return chunk, updated_state + async def derender_completion( self, generate_responses: list[GenerateResponse], prompt_tokens: list[int] | None = None, + completion_request: CompletionRequest | None = None, + ) -> tuple[list[CompletionResponseChoice], int, int]: + return await self._derender_completion_async( + generate_responses, prompt_tokens, completion_request + ) + + def _derender_completion( + self, + generate_responses: list[GenerateResponse], + prompt_tokens: list[int] | None = None, + completion_request: CompletionRequest | None = None, ) -> tuple[list[CompletionResponseChoice], int, int]: n = len(generate_responses) prompt_tokens_list: list[int] = ( prompt_tokens if prompt_tokens is not None else [0] * n ) + skip_special = ( + completion_request.skip_special_tokens + if completion_request is not None + else True + ) tokenizer = self.renderer.get_tokenizer() choices: list[CompletionResponseChoice] = [] total_prompt_tokens = 0 @@ -193,7 +410,7 @@ async def derender_completion( ) decoded_text = tokenizer.decode( - choice.token_ids, skip_special_tokens=True + choice.token_ids, skip_special_tokens=skip_special ) completion_logprobs = None if choice.logprobs is not None: @@ -215,6 +432,88 @@ async def derender_completion( return choices, total_prompt_tokens, total_completion_tokens + async def derender_completion_stream( + self, + model: str, + generate_chunk: GenerateStreamResponse, + state: DerenderStreamState | None = None, + prompt_tokens: int | None = None, + completion_request: CompletionRequest | None = None, + ) -> tuple[CompletionStreamResponse, DerenderStreamState]: + """Process one GenerateStreamResponse chunk for streaming completions. + + Each call takes one SSE chunk from ``/inference/v1/generate`` plus the + client carried ``stream_state`` and returns a ``CompletionStreamResponse`` + chunk and the updated state. + + The generate stream emits one choice per SSE event, so this method + processes one output sequence at a time. For ``n > 1`` the client + maintains one ``DerenderStreamState`` per ``choice.index``. + + Args: + model: Model name for the response object. + generate_chunk: One SSE chunk from ``/inference/v1/generate``. + state: Client carried detok state (``None`` → first call). + prompt_tokens: Prompt token count for usage (from the render step). + completion_request: Original CompletionRequest from ``/render``; + supplies ``skip_special_tokens``. + + Returns: + (chunk, updated_state) — the derendered chunk and updated state. + """ + if state is None: + state = DerenderStreamState() + + # See the equivalent check in derender_chat_stream: a single + # DerenderStreamState is threaded through every choice in this + # chunk, so more than one choice per chunk would corrupt the + # detok window across choices. + if len(generate_chunk.choices) > 1: + raise ValueError( + "derender_completion_stream expects at most one choice per chunk" + ) + + tokenizer = self.renderer.get_tokenizer() + skip_special = ( + completion_request.skip_special_tokens + if completion_request is not None + else True + ) + stream_choices: list[CompletionResponseStreamChoice] = [] + updated_state = state + + for choice in generate_chunk.choices: + delta_tids = choice.token_ids or [] + new_text, updated_state = self._detokenize_delta( + tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special + ) + stream_choices.append( + CompletionResponseStreamChoice( + index=choice.index, + text=new_text, + finish_reason=choice.finish_reason, + ) + ) + + usage: UsageInfo | None = None + if generate_chunk.usage is not None: + u = generate_chunk.usage + pt = prompt_tokens if prompt_tokens is not None else (u.prompt_tokens or 0) + ct = u.completion_tokens or 0 + usage = UsageInfo( + prompt_tokens=pt, + completion_tokens=ct, + total_tokens=pt + ct, + ) + + chunk = CompletionStreamResponse( + id=generate_chunk.request_id, + model=model, + choices=stream_choices, + usage=usage, + ) + return chunk, updated_state + def _parse_token_id_placeholder(token: str) -> int | None: """Extract token ID from a 'token_id:N' placeholder string.""" diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index 15a4023fceeb..4d3b9911355a 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -36,7 +36,7 @@ ) from vllm.logger import init_logger from vllm.parser import Parser, ParserManager -from vllm.renderers import BaseRenderer, merge_kwargs +from vllm.renderers import BaseRenderer, ChatParams, merge_kwargs from vllm.renderers.inputs.preprocess import ( parse_model_prompt, prompt_to_seq, @@ -47,6 +47,19 @@ logger = init_logger(__name__) +def _reused_prompt_token_ids(request: Any) -> list[int] | None: + """Pop prompt token ids forwarded for decode-side reuse, if any. + + Disaggregated serving carries the prefill stage's ids in + ``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing + the key keeps the id list out of the engine's sampling metadata. + """ + kv = getattr(request, "kv_transfer_params", None) + if not isinstance(kv, dict): + return None + return kv.pop("prompt_token_ids", None) or None + + class OnlineRenderer: def __init__( self, @@ -92,6 +105,15 @@ def __init__( self.supports_browsing = False self.supports_code_interpreter = False + def warmup(self) -> None: + self.renderer.warmup( + ChatParams( + chat_template=self.chat_template, + chat_template_content_format=self.chat_template_content_format, + chat_template_kwargs=self.default_chat_template_kwargs, + ) + ) + async def render_chat( self, request: ChatCompletionRequest, @@ -102,6 +124,12 @@ async def render_chat( Called directly by render_chat_request and delegated to by OpenAIServingChat.render_chat_request after its engine-aware checks. + + Decode-side token reuse (ids forwarded in ``kv_transfer_params``) is + handled deeper, in ``preprocess_chat`` / ``_make_request_with_harmony``, + so it skips only templating and tokenization while tool-choice + validation and ``adjust_request`` still run and the output is + detokenized (text-out). """ tokenizer = self.renderer.tokenizer @@ -173,6 +201,15 @@ async def render_chat( ) else: # For GPT-OSS. + if self.parser is not None: + # HarmonyParser doesn't need chat_template_kwargs + # TODO: Unify adjust_request() call with non-harmony branch + self.parser( + self.renderer.get_tokenizer(), + request.tools, + model_config=self.model_config, + ).adjust_request(request=request) + should_include_tools = tool_dicts is not None conversation, engine_inputs = self._make_request_with_harmony( request, should_include_tools @@ -186,6 +223,13 @@ def _make_request_with_harmony( should_include_tools: bool = True, ): """Build Harmony (GPT-OSS) messages and engine prompt from a chat request.""" + reuse_ids = _reused_prompt_token_ids(request) + if reuse_ids: + # Decode-side token reuse: feed the forwarded ids straight to the + # engine. Harmony has no adjust_request hook to preserve. + engine_input = tokens_input(reuse_ids, cache_salt=request.cache_salt) + return [], [engine_input] + messages: list[OpenAIMessage] = [] # because of issues with pydantic we need to potentially @@ -368,17 +412,28 @@ async def preprocess_chat( default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) - (conversation,), (engine_input,) = await renderer.render_chat_async( - [messages], - chat_params, - tok_params, - prompt_extras={ - k: v - for k in ("mm_processor_kwargs", "cache_salt") - if (v := getattr(request, k, None)) is not None - }, - skip_mm_cache=skip_mm_cache, - ) + reuse_ids = _reused_prompt_token_ids(request) + if reuse_ids: + # Decode-side token reuse: feed the forwarded ids straight to the + # engine, skipping templating and tokenization. ``messages`` are not + # tokenized, so conversation is empty. The adjust_request tail below + # still runs. + conversation: list[ConversationMessage] = [] + engine_input = tokens_input( + reuse_ids, cache_salt=getattr(request, "cache_salt", None) + ) + else: + (conversation,), (engine_input,) = await renderer.render_chat_async( + [messages], + chat_params, + tok_params, + prompt_extras={ + k: v + for k in ("mm_processor_kwargs", "cache_salt") + if (v := getattr(request, k, None)) is not None + }, + skip_mm_cache=skip_mm_cache, + ) # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index a07a49230676..c0e2d2129e54 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -90,6 +90,12 @@ class ChatParams: return_assistant_tokens_mask: bool = False """Request a per-token assistant mask from apply_chat_template.""" + tool_choice: Any | None = None + """Request-level tool choice for renderers that need API metadata.""" + + response_format: Any | None = None + """Request-level response format for renderers that need API metadata.""" + def with_defaults( self, default_chat_template_kwargs: dict[str, Any] | None = None, @@ -119,6 +125,8 @@ def with_defaults( self.mm_processor_kwargs, ), return_assistant_tokens_mask=self.return_assistant_tokens_mask, + tool_choice=self.tool_choice, + response_format=self.response_format, ) def get_apply_chat_template_kwargs(self) -> dict[str, Any]: diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py index 395132c56e15..c7e1fc11adc8 100644 --- a/vllm/renderers/registry.py +++ b/vllm/renderers/registry.py @@ -20,10 +20,12 @@ _VLLM_RENDERERS = { + "cohere": ("cohere", "CohereRenderer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"), "hf": ("hf", "HfRenderer"), "kimi_audio": ("hf", "HfRenderer"), + "kimi_k3": ("kimi_k3", "KimiK3Renderer"), "mistral": ("mistral", "MistralRenderer"), "terratorch": ("terratorch", "TerratorchRenderer"), "inkling": ("inkling", "InklingRenderer"), diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 08580f6e8f67..5dedbde372e4 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -100,12 +100,12 @@ def __post_init__(self): ] ) if count > 1: - raise ValueError( + raise VLLMValidationError( "You can only use one kind of structured outputs constraint " f"but multiple are specified: {self.__dict__}" ) if count < 1: - raise ValueError( + raise VLLMValidationError( "You must use one kind of structured outputs constraint " f"but none are specified: {self.__dict__}" ) @@ -166,13 +166,13 @@ def __post_init__(self): or self.min_pattern_size < 0 or self.min_pattern_size > self.max_pattern_size ): - raise ValueError( + raise VLLMValidationError( "max_pattern_size, min_pattern_size must be >=0, " "with min_pattern_size <= max_pattern_size. " "Set both to 0 to disable repetitive pattern detection." ) if self.max_pattern_size > 0 and self.min_count < 2: - raise ValueError( + raise VLLMValidationError( "min_count must be >= 2 to detect repetitive patterns " "in engine output. If you do not wish to detect repetitive " "patterns, set max_pattern_size to 0." @@ -299,6 +299,11 @@ class SamplingParams( include_stop_str_in_output: bool = False """Whether to include the stop strings in output text.""" output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE + stream_interval: int | None = None + """Number of newly generated tokens to batch into each streamed + `RequestOutput`. Raises the interval above the engine-level + `--stream-interval`. Values below engine setting are clamped up to it. + The first and final outputs are always emitted immediately.""" skip_clone: bool = False """Internal flag indicating that this SamplingParams instance is safe to reuse without cloning. When True, clone() will return self without @@ -377,6 +382,7 @@ def from_optional( skip_special_tokens: bool = True, spaces_between_special_tokens: bool = True, output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, + stream_interval: int | None = None, structured_outputs: StructuredOutputsParams | None = None, logit_bias: dict[int, float] | dict[str, float] | None = None, allowed_token_ids: list[int] | None = None, @@ -439,6 +445,7 @@ def from_optional( skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, output_kind=output_kind, + stream_interval=stream_interval, structured_outputs=structured_outputs, logit_bias=logit_bias, allowed_token_ids=allowed_token_ids, @@ -507,31 +514,33 @@ def __post_init__(self) -> None: def _verify_args(self) -> None: if not isinstance(self.n, int): - raise ValueError(f"n must be an int, but is of type {type(self.n)}") + raise VLLMValidationError( + f"n must be an int, but is of type {type(self.n)}" + ) if self.n < 1: - raise ValueError(f"n must be at least 1, got {self.n}.") + raise VLLMValidationError(f"n must be at least 1, got {self.n}.") max_n = envs.VLLM_MAX_N_SEQUENCES if self.n > max_n: - raise ValueError( + raise VLLMValidationError( f"n must be at most {max_n}, got {self.n}. " "To increase this limit, set the VLLM_MAX_N_SEQUENCES " "environment variable." ) if not -2.0 <= self.presence_penalty <= 2.0: - raise ValueError( + raise VLLMValidationError( f"presence_penalty must be in [-2, 2], got {self.presence_penalty}." ) if not -2.0 <= self.frequency_penalty <= 2.0: - raise ValueError( + raise VLLMValidationError( f"frequency_penalty must be in [-2, 2], got {self.frequency_penalty}." ) if not math.isfinite(self.repetition_penalty): - raise ValueError( + raise VLLMValidationError( "repetition_penalty must be a finite number, " f"got {self.repetition_penalty}." ) if self.repetition_penalty <= 0.0: - raise ValueError( + raise VLLMValidationError( "repetition_penalty must be greater than zero, got " f"{self.repetition_penalty}." ) @@ -561,15 +570,15 @@ def _verify_args(self) -> None: ) # quietly accept -1 as disabled, but prefer 0 if self.top_k < -1: - raise ValueError( + raise VLLMValidationError( f"top_k must be 0 (disable), or at least 1, got {self.top_k}." ) if not isinstance(self.top_k, int): - raise TypeError( + raise VLLMValidationError( f"top_k must be an integer, got {type(self.top_k).__name__}" ) if not 0.0 <= self.min_p <= 1.0: - raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.") + raise VLLMValidationError(f"min_p must be in [0, 1], got {self.min_p}.") if self.max_tokens is not None and self.max_tokens < 1: raise VLLMValidationError( f"max_tokens must be at least 1, got {self.max_tokens}.", @@ -577,14 +586,20 @@ def _verify_args(self) -> None: value=self.max_tokens, ) if self.min_tokens < 0: - raise ValueError( + raise VLLMValidationError( f"min_tokens must be greater than or equal to 0, got {self.min_tokens}." ) if self.max_tokens is not None and self.min_tokens > self.max_tokens: - raise ValueError( + raise VLLMValidationError( f"min_tokens must be less than or equal to " f"max_tokens={self.max_tokens}, got {self.min_tokens}." ) + if self.stream_interval is not None and self.stream_interval < 1: + raise VLLMValidationError( + f"stream_interval must be at least 1, got {self.stream_interval}.", + parameter="stream_interval", + value=self.stream_interval, + ) if self.logprobs is not None and self.logprobs != -1 and self.logprobs < 0: raise VLLMValidationError( f"logprobs must be non-negative or -1, got {self.logprobs}.", @@ -604,27 +619,29 @@ def _verify_args(self) -> None: ) assert isinstance(self.stop_token_ids, list) if not all(isinstance(st_id, int) for st_id in self.stop_token_ids): - raise ValueError( + raise VLLMValidationError( f"stop_token_ids must contain only integers, got {self.stop_token_ids}." ) assert isinstance(self.stop, list) if any(not stop_str for stop_str in self.stop): - raise ValueError("stop cannot contain an empty string.") + raise VLLMValidationError("stop cannot contain an empty string.") if self.stop and not self.detokenize: - raise ValueError( + raise VLLMValidationError( "stop strings are only supported when detokenize is True. " "Set detokenize=True to use stop." ) assert isinstance(self.bad_words, list) if any(not bad_word for bad_word in self.bad_words): - raise ValueError( + raise VLLMValidationError( f"bad_words cannot contain an empty string. " f"Got bad_words={self.bad_words}" ) def _verify_greedy_sampling(self) -> None: if self.n > 1: - raise ValueError(f"n must be 1 when using greedy sampling, got {self.n}.") + raise VLLMValidationError( + f"n must be 1 when using greedy sampling, got {self.n}." + ) def update_from_generation_config( self, @@ -876,7 +893,7 @@ def _validate_spec_decode( # Some sampling parameters are not yet compatible with spec decoding. if self.min_p > _SAMPLING_EPS or self.logit_bias: - raise ValueError( + raise VLLMValidationError( "The min_p and logit_bias sampling parameters " "are not yet supported with speculative decoding." ) @@ -897,7 +914,7 @@ def _validate_diffusion(self, model_config: ModelConfig) -> None: or self.bad_words or self.allowed_token_ids ): - raise ValueError( + raise VLLMValidationError( "The temperature, min_p, seed, min_tokens, logit_bias, " "bad_words, and allowed_token_ids sampling parameters " "are not yet supported with diffusion models." @@ -917,7 +934,7 @@ def _validate_structured_outputs( # rather than sampling left-to-right, which the grammar FSM # requires. Without this check, requests fail mid-generation # with an FSM rejection (HTTP 500). See issue #45436. - raise ValueError( + raise VLLMValidationError( "Structured outputs are not yet supported for diffusion " "language models. Remove the structured output constraint " "(e.g. `response_format`, `structured_outputs`) from the " @@ -925,7 +942,7 @@ def _validate_structured_outputs( ) if tokenizer is None: - raise ValueError( + raise VLLMValidationError( "Structured outputs requires a tokenizer so it can't be used with 'skip_tokenizer_init'" # noqa: E501 ) @@ -939,7 +956,7 @@ def _validate_structured_outputs( if backend != _backend and not ( backend == "auto" and self.structured_outputs._backend_was_auto ): - raise ValueError( + raise VLLMValidationError( "Request-level structured output backend selection is not " f"supported. The request specified '{_backend}', but vLLM " f"was initialised with '{backend}'. This error can be " @@ -954,7 +971,7 @@ def _validate_structured_outputs( and not self.structured_outputs.choice ): # It is invalid for choice to be an empty list - raise ValueError( + raise VLLMValidationError( f"Choice '{self.structured_outputs.choice}' cannot be an empty list" # noqa: E501 ) # Reject empty string grammar early to avoid engine-side crashes @@ -962,16 +979,20 @@ def _validate_structured_outputs( isinstance(self.structured_outputs.grammar, str) and self.structured_outputs.grammar.strip() == "" ): - raise ValueError("structured_outputs.grammar cannot be an empty string") + raise VLLMValidationError( + "structured_outputs.grammar cannot be an empty string" + ) # Reject empty string json schema early to avoid engine-side crashes if ( isinstance(self.structured_outputs.json, str) and self.structured_outputs.json.strip() == "" ): - raise ValueError("structured_outputs.json cannot be an empty string") + raise VLLMValidationError( + "structured_outputs.json cannot be an empty string" + ) # Reject json_object=False early to avoid engine-side crashes if self.structured_outputs.json_object is False: - raise ValueError( + raise VLLMValidationError( "structured_outputs.json_object must be True if set; omit " "structured_outputs to disable structured outputs" ) @@ -993,7 +1014,7 @@ def _validate_structured_outputs( validate_xgrammar_grammar(self) elif backend.startswith("guidance"): if _is_non_tekken_mistral(tokenizer=tokenizer): - raise ValueError( + raise VLLMValidationError( "Non-tekken Mistral tokenizers are not supported for the 'guidance'" " structured output backend. Please either use a more recent " "Mistral model, the ['xgrammar', 'outlines'] " @@ -1013,7 +1034,7 @@ def _validate_structured_outputs( elif backend == "lm-format-enforcer": # lm format enforcer backend if is_mistral_tokenizer(tokenizer): - raise ValueError( + raise VLLMValidationError( "Mistral tokenizer is not supported for the 'lm-format-enforcer' " "structured output backend. Please use ['xgrammar', 'outlines'] " "backends or tokenizer_mode='hf' instead." diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py new file mode 100644 index 000000000000..60fc53350702 --- /dev/null +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch +import torch.nn as nn + +from vllm.model_executor.custom_op import CustomOp +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n, + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr( + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr( + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = ( + (b_x - b_mean[:, None]) * b_rstd[:, None] + if not IS_RMS_NORM + else b_x * b_rstd[:, None] + ) + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + + +def layer_norm_gated_fwd( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, +): + if residual is not None: + residual_dtype = residual.dtype + T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 + if residual is not None: + assert residual.shape == (T, D) + if weight is not None: + assert weight.shape == (D,) + if bias is not None: + assert bias.shape == (D,) + # allocate output + y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype) + if residual is not None or ( + residual_dtype is not None and residual_dtype != x.dtype + ): + residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = ( + torch.empty((T,), dtype=torch.float, device=x.device) + if not is_rms_norm + else None + ) + rstd = torch.empty((T,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + if D <= 512: + BT = 16 + layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + T=T, + H=H, + g_stride_n=g_stride_n, + D=D, + BD=BD, + BT=BT, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + else: + layer_norm_gated_fwd_kernel1[(T,)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + D=D, + BD=BD, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +def rms_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.contiguous().reshape(-1, x.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.contiguous().reshape(-1, residual.shape[-1]) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float if residual_in_fp32 else None) + ) + y, _, _, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=weight, + bias=bias, + activation=activation, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, + ) + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + +@CustomOp.register("fused_rms_norm_gated") +class FusedRMSNormGated(CustomOp): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + activation: str = "swish", + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ["swish", "silu", "sigmoid"]: + raise ValueError(f"Unsupported activation: {self.activation}") + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward_native( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + """Decomposed PyTorch ops for torch.compile/inductor fusion.""" + # TODO(https://github.com/vllm-project/vllm/issues/36175): implement + # native residual/prenorm path and unify with RMSNormGated. + # For now, fall back to the triton kernel. + if residual is not None or prenorm: + return self.forward_cuda(x, g, residual, prenorm, residual_in_fp32) + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_normed = x_float * torch.rsqrt(variance + self.eps) + if self.weight is not None: + x_normed = x_normed * self.weight.float() + g_float = g.float() + if self.activation in ("swish", "silu"): + out = x_normed * g_float * torch.sigmoid(g_float) + else: # sigmoid + out = x_normed * torch.sigmoid(g_float) + return out.to(x.dtype) + + def forward_cuda( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index 920efa444178..c8004cb57173 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -12,7 +12,7 @@ from vllm.triton_utils import tl, triton -from .op import exp +from .op import exp, log @triton.heuristics( diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 10acc214f4d6..1701cd47aca8 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -166,6 +166,8 @@ def layer_norm_gated_fwd_kernel( rstd, # pointer to the 1/std eps, # epsilon to avoid division by zero T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n: tl.constexpr, D: tl.constexpr, # number of columns in x BT: tl.constexpr, BD: tl.constexpr, @@ -221,8 +223,13 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) if ACTIVATION == "swish" or ACTIVATION == "silu": b_y = b_y * b_g * tl.sigmoid(b_g) elif ACTIVATION == "sigmoid": @@ -320,10 +327,15 @@ def layer_norm_gated_fwd( out_dtype: torch.dtype = None, residual_dtype: torch.dtype = None, is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, ): if residual is not None: residual_dtype = residual.dtype T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 if residual is not None: assert residual.shape == (T, D) if weight is not None: @@ -349,10 +361,8 @@ def layer_norm_gated_fwd( BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) if D > BD: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - if D <= 512: - BT = 32 + BT = 16 layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( x=x, g=g, @@ -365,12 +375,14 @@ def layer_norm_gated_fwd( rstd=rstd, eps=eps, T=T, + H=H, + g_stride_n=g_stride_n, D=D, BD=BD, BT=BT, ACTIVATION=activation, IS_RMS_NORM=is_rms_norm, - num_warps=4, + num_warps=8, ) else: layer_norm_gated_fwd_kernel1[(T,)]( @@ -408,7 +420,16 @@ def rms_norm_gated( x_shape_og = x.shape # reshape input data into 2D tensor x = x.contiguous().reshape(-1, x.shape[-1]) - g = g.contiguous().reshape(-1, g.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D if residual is not None: assert residual.shape == x_shape_og residual = residual.contiguous().reshape(-1, residual.shape[-1]) @@ -427,6 +448,8 @@ def rms_norm_gated( residual=residual, residual_dtype=residual_dtype, is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, ) y = y.reshape(x_shape_og) return y if not prenorm else (y, residual_out.reshape(x_shape_og)) @@ -1187,6 +1210,7 @@ def kda_gate_cumsum_fwd_kernel( cu_seqlens, chunk_indices, cumsum_scale, + lower_bound, beta, threshold, T, @@ -1196,6 +1220,7 @@ def kda_gate_cumsum_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H @@ -1235,14 +1260,17 @@ def kda_gate_cumsum_fwd_kernel( b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - b_a = -tl.exp(tl.load(A + i_h).to(tl.float32)) - b_g_scaled = b_g * beta - b_softplus = tl.where( - b_g_scaled > threshold, - b_g, - (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), - ) - b_gate = b_a * b_softplus + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to @@ -1260,6 +1288,7 @@ def fused_kda_gate_chunk_cumsum( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, @@ -1293,16 +1322,17 @@ def grid(meta): # exp2-based kernels reproduce exp(g). Keep this in sync with the # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, beta=beta, threshold=threshold, T=T, H=H, D=D, BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, ) return y - def _chunk_kda_fwd_with_cumulative_g( q: torch.Tensor, k: torch.Tensor, @@ -1424,6 +1454,7 @@ def chunk_kda_with_fused_gate_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE @@ -1439,6 +1470,7 @@ def chunk_kda_with_fused_gate_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, + lower_bound=lower_bound, ) return _chunk_kda_fwd_with_cumulative_g( q=q, @@ -1500,6 +1532,7 @@ def chunk_kda_with_fused_gate( scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, + lower_bound: float | None = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, **kwargs, @@ -1523,6 +1556,7 @@ def chunk_kda_with_fused_gate( scale=scale, initial_state=initial_state.contiguous() if initial_state is not None else None, output_final_state=output_final_state, + lower_bound=lower_bound, cu_seqlens=cu_seqlens, ) return o, final_state @@ -1543,6 +1577,7 @@ def kda_gate_fwd_kernel( A, y, g_bias, + lower_bound, beta: tl.constexpr, threshold: tl.constexpr, T, @@ -1551,12 +1586,12 @@ def kda_gate_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_t, i_h = tl.program_id(0), tl.program_id(1) n_t = i_t * BT - b_a = tl.load(A + i_h).to(tl.float32) - b_a = -tl.exp(b_a) + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) stride_row = H * D stride_col = 1 @@ -1589,13 +1624,13 @@ def kda_gate_fwd_kernel( ) b_g = b_g + b_bias[None, :] - # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) - # When beta * x > threshold, use linear approximation x - # Use threshold to switch to linear when beta*x > threshold - g_scaled = b_g * beta - use_linear = g_scaled > threshold - sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) - b_y = b_a * sp + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) @@ -1607,6 +1642,7 @@ def fused_kda_gate( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, ) -> torch.Tensor: """ Forward pass for KDA gate: @@ -1634,6 +1670,7 @@ def grid(meta): A, y, g_bias, + lower_bound or 0.0, beta, threshold, T, @@ -1641,6 +1678,7 @@ def grid(meta): head_k_dim, BD=next_power_of_2(head_k_dim), HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, ) y = y.view(*orig_shape, H, head_k_dim) diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 1164f7c41a76..fe2cb3fd573b 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -21,7 +21,6 @@ ) from mistral_common.tokens.tokenizers.instruct import ( InstructTokenizerBase, - InstructTokenizerV13, ) from mistral_common.tokens.tokenizers.mistral import ( MistralTokenizer as MistralCommonTokenizer, @@ -45,20 +44,6 @@ logger = init_logger(__name__) -def _pop_unallowed_keys_and_warn( - dictionary: dict[str, Any], allowed_keys: set[str], err_dict_name: str -): - keys = list(dictionary.keys()) - for key in keys: - if key not in allowed_keys: - dictionary.pop(key) - logger.warning_once( - f"'{key=}' is not supported by mistral-common " - f"for {err_dict_name}. It has been popped from the " - "object." - ) - - def maybe_serialize_tool_calls(request: "MistralChatCompletionRequest"): # SEE: https://github.com/vllm-project/vllm/pull/9951 # Credits go to: @gcalmettes @@ -475,6 +460,7 @@ def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: def convert_tokens_to_string(self, tokens: list[str]) -> str: to_decode_special_tokens = { SpecialTokens.tool_calls, + SpecialTokens.args, SpecialTokens.begin_think, SpecialTokens.end_think, } @@ -531,11 +517,20 @@ def convert_ids_to_tokens( non_skip_special_tokens_ids = { self.tokenizer.get_special_token(SpecialTokens.tool_calls), } - if isinstance(self.instruct, InstructTokenizerV13): - if self.instruct.BEGIN_THINK: - non_skip_special_tokens_ids.add(self.instruct.BEGIN_THINK) - if self.instruct.END_THINK: - non_skip_special_tokens_ids.add(self.instruct.END_THINK) + # [ARGS] only exists in v11+ tool-call tokenizers; older tokenizers + # raise (Tekken) or return unk (SPM) for it. + if self.tokenizer.is_special(SpecialTokens.args): + non_skip_special_tokens_ids.add( + self.tokenizer.get_special_token(SpecialTokens.args) + ) + # [THINK]/[/THINK] only exist in v13+ reasoning tokenizers; use the + # same is_special gate as [ARGS] above so newer versions are covered + # without an isinstance check on the instruct tokenizer. + for think_token in (SpecialTokens.begin_think, SpecialTokens.end_think): + if self.tokenizer.is_special(think_token): + non_skip_special_tokens_ids.add( + self.tokenizer.get_special_token(think_token) + ) ids_kept = [ i diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index e6c12ccc3bc7..713c6228f919 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -18,6 +18,7 @@ ) from vllm.utils.import_utils import resolve_obj_by_qualname +from .hf import CachedHfTokenizer from .protocol import TokenizerLike if TYPE_CHECKING: @@ -39,10 +40,14 @@ } _VLLM_TOKENIZERS = { + # ``cohere`` mode uses the standard cached HF tokenizer; only the + # renderer (template stage) is replaced with a melody-based one. + "cohere": ("hf", "CachedHfTokenizer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), "hf": ("hf", "CachedHfTokenizer"), "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"), + "kimi_k3": ("hf", "CachedHfTokenizer"), "mistral": ("mistral", "MistralTokenizer"), # Inkling uses the plain HF tokenizer for token operations; the "inkling" # mode exists to select the InklingRenderer, which renders chat to @@ -204,17 +209,27 @@ def get_tokenizer( **kwargs, ) + if tokenizer_cls == TokenizerLike: + tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) + else: + tokenizer_cls_ = tokenizer_cls + # Ensure that, if the config were to come from vllm.transformers_utils.config, it is # registered with AutoConfig before the tokenizer is loaded. This is necessary since # tokenizer_cls_.from_pretrained will call AutoConfig.from_pretrained internally. # This may fail for paths that don't have a model config (e.g. LoRA adapters), # which is fine — those don't need custom config registration. + # HF-backed tokenizers must receive the HF config. In a dual-format Mistral + # repository, auto detection intentionally prefers params.json, but passing + # that generic config to AutoTokenizer can select the wrong tokenizer class. + config_format = "hf" if tokenizer_cls_ is CachedHfTokenizer else "auto" config = None with contextlib.suppress(ValueError, OSError): config = get_config( tokenizer_name, trust_remote_code=trust_remote_code, revision=revision, + config_format=config_format, ) # Some models have an incorrect tokenizer_class on the hub. @@ -228,10 +243,12 @@ def get_tokenizer( model_type, ) tokenizer_cls_ = TokenizersBackend - elif tokenizer_cls == TokenizerLike: - tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) - else: - tokenizer_cls_ = tokenizer_cls + + if config is not None and tokenizer_cls_ is CachedHfTokenizer: + # AutoTokenizer otherwise reloads config.json internally. Reuse the + # config that get_config just loaded successfully so a concurrent Hub + # cache refresh cannot invalidate the file between the two reads. + kwargs.setdefault("config", config) tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 64b1342fd253..4db6f675fb3b 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -102,6 +102,10 @@ "kimi_k2_tool_parser", "KimiK2ToolParser", ), + "kimi_k3": ( + "kimi_k3_tool_parser", + "KimiK3ToolParser", + ), "llama3_json": ( "llama_tool_parser", "Llama3JsonToolParser", diff --git a/vllm/tool_parsers/gemma4_utils.py b/vllm/tool_parsers/gemma4_utils.py index a72e16ea56f4..d9aad254c728 100644 --- a/vllm/tool_parsers/gemma4_utils.py +++ b/vllm/tool_parsers/gemma4_utils.py @@ -37,15 +37,8 @@ import regex as re -# Tool call delimiter tokens as they appear in decoded text. -# Standard format: <|tool_call>call:name{args} -_TOOL_CALL_START_TAG = "<|tool_call>" -_TOOL_CALL_END_TAG = "" _TOOL_RESPONSE_START_TAG = "<|tool_response>" -# Gemma4 escape token as it appears in decoded text. -_ESCAPE_TOKEN = '<|"|>' - def _parse_tool_arguments(args_str: str) -> dict[str, str]: """Parse tool call arguments from the Gemma4 compact format. diff --git a/vllm/tool_parsers/gptoss_tool_parser.py b/vllm/tool_parsers/gptoss_tool_parser.py index 6857e6bbe728..7321c2f049cb 100644 --- a/vllm/tool_parsers/gptoss_tool_parser.py +++ b/vllm/tool_parsers/gptoss_tool_parser.py @@ -22,6 +22,8 @@ class GptOssToolParser(ToolParser): capability declaration via HarmonyParser.tool_parser_cls. """ + structural_tag_model = "harmony" + def __init__(self, tokenizer: "TokenizerLike", tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index f4aaeef71a0f..7cbffcb1c0b2 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ Tool, ToolParser, ) -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json logger = init_logger(__name__) @@ -146,9 +146,17 @@ def extract_tool_calls_streaming( elif cur_arguments and not prev_arguments: cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(delta_text) + len(delta_text) - ] + match_start = cur_arguments_json.find(delta_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(delta_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None delta = DeltaMessage( tool_calls=[ DeltaToolCall( diff --git a/vllm/tool_parsers/jamba_tool_parser.py b/vllm/tool_parsers/jamba_tool_parser.py index dec3c88d934a..193a51faa3d9 100644 --- a/vllm/tool_parsers/jamba_tool_parser.py +++ b/vllm/tool_parsers/jamba_tool_parser.py @@ -24,7 +24,7 @@ from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json from vllm.utils.mistral import is_mistral_tokenizer logger = init_logger(__name__) @@ -266,9 +266,18 @@ def extract_tool_calls_streaming( cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) logger.debug("finding %s in %s", new_text, cur_arguments_json) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(new_text) + len(new_text) - ] + # `new_text` may not appear verbatim in the re-serialized JSON. + match_start = cur_arguments_json.find(new_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(new_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None logger.debug( "First tokens in arguments received: %s", arguments_delta ) diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py new file mode 100644 index 000000000000..9f688f2270e4 --- /dev/null +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tool-call parser for the Kimi K3 (XTML) chat format. + +This turns the generated XTML ``response`` and ``tools`` channels back into +OpenAI-compatible ``content`` and ``tool_calls``. + +K3 assistant tool calls live in a nested ``tools`` channel:: + + <|open|>tools<|sep|> + <|open|>call tool="python" index="1"<|sep|> + <|open|>argument key="code" type="string"<|sep|>print(1)<|close|>argument<|sep|> + <|open|>argument key="opts" type="object"<|sep|>{"a":1}<|close|>argument<|sep|> + <|close|>call<|sep|> + <|close|>tools<|sep|> + +where ``<|open|>``, ``<|close|>``, ``<|sep|>`` are dedicated special tokens. The plain +reply lives in a sibling ``response`` channel which is unwrapped into content. +``response`` always precedes ``tools`` (see the protocol channel grammar). + +Argument decoding mirrors the template's type tagging (inverse encoding): + * type="string" -> value is the RAW text, no unescaping (template emits it raw) + * other types -> value is JSON-decoded (number/boolean/null/object/array) + +Attribute values (``tool=``, ``index=``, ``key=``, ``type=``) ARE escaped on the +encode side (``&`` -> ``&``, ``"`` -> ``"``), so ``_attrs`` reverses +them on decode (``"`` before ``&`` -- reverse of encode order). + +Known limitation: because string argument and response bodies are emitted raw, +a value that literally contains ``<|close|>argument<|sep|>`` or +``<|close|>response<|sep|>`` is indistinguishable from a real closing marker. +""" + +import json +from collections.abc import Sequence + +import regex as re +from openai.types.responses import ToolChoiceFunction + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.exceptions import VLLMValidationError +from vllm.logger import init_logger +from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser + +logger = init_logger(__name__) + +_O, _C, _S = r"<\|open\|>", r"<\|close\|>", r"<\|sep\|>" +_TEXT_UNTIL_SEP = r"(?:(?!" + _S + r").)*?" + + +def _partial_tag_overlap(text: str, tag: str) -> int: + max_len = min(len(text), len(tag) - 1) + for n in range(max_len, 0, -1): + if text.endswith(tag[:n]): + return n + return 0 + + +class KimiK3ToolParser(ToolParser): + supports_required_and_named = False + # Enables the vLLM-side XTML structural tag builder + # (``get_kimi_k3_structural_tag`` in ``structural_tag_registry``). With + # ``VLLM_ENFORCE_STRICT_TOOL_CALLING`` on (default), ``_apply_structural_tag`` + # constrains generation to K3's ``<|open|>tools<|sep|>`` channel for + # ``required`` (and ``auto`` when a tool sets ``strict``) instead of the + # generic JSON guided decoding, which conflicts with the XTML format. + structural_tag_model = "kimi_k3" + + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + + self.tools_open = "<|open|>tools<|sep|>" + self.tools_close = "<|close|>tools<|sep|>" + self.response_open = "<|open|>response<|sep|>" + self.response_close = "<|close|>response<|sep|>" + + # Regexes operate on detokenized text. The XTML markers reach us as the + # literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them + # from being stripped and suppresses spacing between adjacent token + # pieces. As DEFENSE-IN-DEPTH we still tolerate optional whitespace + # WITHIN a marker (e.g. "<|open|> tools <|sep|>") in case some serving + # path leaves vLLM's added-token spacing on -- this `\s*` is a no-op on + # clean input, so the normal byte-exact path is unaffected. We do NOT + # strip body whitespace, so clean content stays byte-exact. + # All bodies use non-greedy (.*?) so each block stops at its own first + # closing marker -- see the module-level KNOWN LIMITATION about + # literal-marker values. + self._tools_open_re = re.compile(_O + r"\s*tools\s*" + _S) + self._tools_close_re = re.compile(_C + r"\s*tools\s*" + _S) + self._response_open_re = re.compile(_O + r"\s*response\s*" + _S) + self._response_close_re = re.compile(_C + r"\s*response\s*" + _S) + self._message_close_re = re.compile(_C + r"\s*message\s*" + _S) + self._call_re = re.compile( + _O + + r"\s*call\s+(?P" + + _TEXT_UNTIL_SEP + + r")" + + _S + + r"(?P.*?)" + + _C + + r"\s*call\s*" + + _S, + re.DOTALL, + ) + self._arg_re = re.compile( + _O + + r"\s*argument\s+(?P" + + _TEXT_UNTIL_SEP + + r")" + + _S + + r"(?P.*?)" + + _C + + r"\s*argument\s*" + + _S, + re.DOTALL, + ) + # attr segment: key="value" (value already escaped on the encode side) + self._attr_re = re.compile(r'(?P\w+)="(?P[^"]*)"') + self._response_re = re.compile( + _O + r"\s*response\s*" + _S + r"(?P.*?)" + _C + r"\s*response\s*" + _S, + re.DOTALL, + ) + + # streaming state + self._sent_content_idx = 0 + self._sent_tool_call_count = 0 + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ToolParser " + "constructor during construction." + ) + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + named = isinstance( + request.tool_choice, + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ) + structured_outputs = getattr(request, "structured_outputs", None) + has_structural_tag = ( + structured_outputs is not None + and structured_outputs.structural_tag is not None + ) + if named and not has_structural_tag: + # Without the XTML structural tag there is no way to force the + # named call (the generic JSON guided-decoding path conflicts + # with the XTML channel format). + raise VLLMValidationError( + "Named tool choice for Kimi K3 requires strict tool calling " + "(VLLM_ENFORCE_STRICT_TOOL_CALLING) so the XTML structural " + "tag can force the call. Otherwise use `tool_choice` set to " + '"auto", "required", or "none".', + parameter="tool_choice", + value=request.tool_choice, + ) + + if request.tools and (request.tool_choice == "required" or named): + # K3 emits tool calls in XTML. When strict tool calling is enabled, + # DelegatingParser._apply_structural_tag has already attached the K3 + # XTML structural tag (structural_tag_model = "kimi_k3"). We return + # early to skip the generic parent path, which would otherwise attach + # JSON guided decoding when strict calling is off -- that JSON + # constraint conflicts with the <|open|>tools<|sep|> channel. + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + request = super().adjust_request(request) + # The XTML markers (<|open|>/<|close|>/<|sep|>, + # <|open|>response<|sep|> ...) must + # reach this parser as CONTIGUOUS literal text. Two request flags govern + # that, and vLLM's detokenizer couples them: + # effective_spaces_between = + # skip_special_tokens OR spaces_between_special_tokens + # (vllm/v1/engine/detokenizer.py). + # The detokenizer treats these control tokens as separate sub-texts and, + # when the effective flag is True, joins them with spaces -> + # "<|open|> response <|sep|>", which the regexes below would NOT match. + # Forcing BOTH flags off is the only way to suppress that spacing. We + # set both unconditionally: the response channel is unwrapped from these + # markers even with no tools. + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + def _attrs(self, s: str) -> dict[str, str]: + return { + m["k"]: m["v"].replace(""", '"').replace("&", "&") + for m in self._attr_re.finditer(s) + } + + def _decode_call(self, attrs: str, body: str) -> ToolCall | None: + """Decode one ``call`` block into a :class:`ToolCall`. + + ``attrs`` is the text between ``<|open|>call `` and ``<|sep|>`` (carries + ``tool=`` / ``index=``); ``body`` is the sequence of ``argument`` blocks. + Each argument is re-typed per its ``type=`` tag: strings pass through + raw, everything else is JSON-decoded (falling back to raw text if the + JSON is malformed, so a partial stream never raises). Returns ``None`` + when no tool name is present (an empty/garbage block is dropped). + """ + call_attrs = self._attrs(attrs) + tool_name = call_attrs.get("tool", "") + arguments: dict = {} + for arg_match in self._arg_re.finditer(body): + arg_attrs = self._attrs(arg_match["attrs"]) + key = arg_attrs.get("key", "") + arg_type = arg_attrs.get("type", "string") + raw_value = arg_match["val"] + if arg_type == "string": + arguments[key] = raw_value + else: + try: + arguments[key] = json.loads(raw_value) + except json.JSONDecodeError: + arguments[key] = raw_value + if not tool_name: + return None + return ToolCall( + type="function", + function=FunctionCall( + name=tool_name, + arguments=json.dumps(arguments, ensure_ascii=False), + ), + ) + + def _strip_response_content(self, text: str) -> str | None: + """Strip XTML response/message markers from generated response text. + + In chat serving, ``<|open|>response<|sep|>`` is often part of the prompt + generation prefix, so the model output may only contain the body plus + ``<|close|>response<|sep|>``. Handle both that consumed-prefix shape and a + complete ``<|open|>response<|sep|>... <|close|>response<|sep|>`` wrapper. + """ + m_open = self._response_open_re.search(text) + if m_open is not None: + m_close = self._response_close_re.search(text, m_open.end()) + if m_close is not None: + text = text[m_open.end() : m_close.start()] + else: + text = text[m_open.end() :] + else: + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + return text or None + + def _content(self, model_output: str, before: str) -> str | None: + # prefer the unwrapped response channel; else the text before the tools + m = self._response_re.search(model_output) + if m is not None: + return m["c"] or None + return self._strip_response_content(before) + + def _extract_response_content(self, current_text: str) -> str | None: + # Streaming response text is computed from the accumulated text. This is + # what keeps split markers from leaking: + # <|open|> / response / <|sep|>Hi -> emit only "Hi" after open closes + # Hi<|open|> / tools / <|sep|>... -> emit "Hi", hold the tools marker + # Hi<|close|> / response / <|sep|>... -> emit "Hi", hold the close marker + # Tool calls are even simpler: they are not emitted until a full + # <|close|>call<|sep|> is present in current_text. + m_open = self._response_open_re.search(current_text) + # In the normal chat path, the response-open marker may be consumed as + # the generation prefix. Then generated text starts directly with the + # response body or with <|close|>response<|sep|> before a tools channel. + body_start = m_open.end() if m_open is not None else 0 + m_tools = self._tools_open_re.search(current_text, body_start) + m_rclose = self._response_close_re.search(current_text, body_start) + tools_start = m_tools.start() if m_tools else -1 + response_end = m_rclose.start() if m_rclose else -1 + + candidates = [i for i in (tools_start, response_end) if i != -1] + if candidates: + sendable_idx = min(candidates) + else: + overlap = max( + _partial_tag_overlap(current_text, self.response_open), + _partial_tag_overlap(current_text, self.response_close), + _partial_tag_overlap(current_text, self.tools_open), + ) + sendable_idx = len(current_text) - overlap + + if sendable_idx <= body_start: + return None + if self._sent_content_idx < body_start: + self._sent_content_idx = body_start + if sendable_idx <= self._sent_content_idx: + return None + + content = current_text[self._sent_content_idx : sendable_idx] + self._sent_content_idx = sendable_idx + return content or None + + def extract_tool_calls( + self, model_output: str, request: ChatCompletionRequest + ) -> ExtractedToolCallInformation: + m_open = self._tools_open_re.search(model_output) + if m_open is None: + # no tools channel -> content is the response channel (unwrapped) + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._content(model_output, model_output), + ) + try: + before = model_output[: m_open.start()] + start = m_open.end() + m_close = self._tools_close_re.search(model_output, start) + section = ( + model_output[start:] + if m_close is None + else model_output[start : m_close.start()] + ) + + tool_calls = [ + tc + for m in self._call_re.finditer(section) + if (tc := self._decode_call(m["attrs"], m["body"])) is not None + ] + if not tool_calls: + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._content(model_output, before), + ) + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=self._content(model_output, before), + ) + except Exception: + logger.exception("Error extracting K3 tool calls.") + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + # Conservative streaming: stream unwrapped response-channel text, then + # buffer tool calls and emit each call once its block closes. + content = self._extract_response_content(current_text) + + # tools channel is open: parse fully-closed calls we have not emitted yet + m_tools = self._tools_open_re.search(current_text) + if m_tools is None: + return DeltaMessage(content=content) if content else None + + section = current_text[m_tools.end() :] + calls = [ + tc + for m in self._call_re.finditer(section) + if (tc := self._decode_call(m["attrs"], m["body"])) is not None + ] + if len(calls) <= self._sent_tool_call_count: + return DeltaMessage(content=content) if content else None + new = calls[self._sent_tool_call_count :] + + deltas = [ + DeltaToolCall( + index=self._sent_tool_call_count + i, + id=tc.id, + type="function", + function=DeltaFunctionCall( + name=tc.function.name, arguments=tc.function.arguments + ).model_dump(exclude_none=True), + ) + for i, tc in enumerate(new) + ] + self._sent_tool_call_count = len(calls) + return DeltaMessage(content=content, tool_calls=deltas) diff --git a/vllm/tool_parsers/mistral_tool_parser.py b/vllm/tool_parsers/mistral_tool_parser.py index 026098a87358..01bddd4ed668 100644 --- a/vllm/tool_parsers/mistral_tool_parser.py +++ b/vllm/tool_parsers/mistral_tool_parser.py @@ -3,767 +3,47 @@ from __future__ import annotations -import json -from collections.abc import Sequence -from enum import Enum, auto -from random import choices -from string import ascii_letters, digits -from typing import Any +from typing import TYPE_CHECKING -import ijson -import regex as re -from mistral_common.protocol.instruct.tool_calls import ( - NamedToolChoice as MistralNamedToolChoice, -) -from mistral_common.protocol.instruct.tool_calls import ( - Tool as MistralTool, -) -from mistral_common.protocol.instruct.tool_calls import ( - ToolChoice as MistralToolChoice, -) -from mistral_common.protocol.instruct.tool_calls import ( - ToolChoiceEnum as MistralToolChoiceEnum, -) -from pydantic import Field +from vllm.parser.engine.registered_adapters import MistralParserToolAdapter -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.logger import init_logger -from vllm.sampling_params import StructuredOutputsParams -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) -from vllm.utils.mistral import is_mistral_tokenizer +if TYPE_CHECKING: + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) -ALPHANUMERIC = ascii_letters + digits +class MistralToolParser(MistralParserToolAdapter): # type: ignore[valid-type, misc] + # Marked so is_mistral_tool_parser() in vllm.utils.mistral recognises + # this adapter and lets adjust_request() fire for tool_choice="none" + # on grammar-capable Mistral tokenizers. + IS_MISTRAL_TOOL_PARSER = True -_DEFAULT_JSON_SCHEMA = {"anyOf": [{"type": "object"}, {"type": "array"}]} + # Mistral emits tool calls in the [TOOL_CALLS]name[ARGS]{...} format, not the + # plain JSON array / bare-args that the serving layer's generic required and + # named handlers expect. Route required/named through this parser's own + # extraction (treated like "auto") so the [TOOL_CALLS] format is parsed + # correctly instead of crashing or dumping the raw envelope into arguments. + supports_required_and_named = False + @property + def bot_token(self) -> str: + """Return the Mistral tool-call marker exposed by the legacy parser.""" + return self._parser_engine.bot_token -class StreamingState(Enum): - """Enum for tracking the current streaming parsing state.""" - - WAITING_FOR_TOOL_START = auto() - WAITING_FOR_TOOL_KEY = ( - auto() - ) # waiting for the "name" or "arguments" key to be complete - PARSING_NAME = auto() - PARSING_NAME_COMPLETED = auto() - WAITING_FOR_ARGUMENTS_START = auto() - PARSING_ARGUMENTS = auto() - PARSING_ARGUMENTS_COMPLETED = auto() - TOOL_COMPLETE = auto() - ALL_TOOLS_COMPLETE = auto() - - -class MistralToolCall(ToolCall): - id: str = Field(default_factory=lambda: MistralToolCall.generate_random_id()) - - @staticmethod - def generate_random_id(): - # Mistral Tool Call Ids must be alphanumeric with a length of 9. - # https://github.com/mistralai/mistral-common/blob/21ee9f6cee3441e9bb1e6ed2d10173f90bd9b94b/src/mistral_common/protocol/instruct/validator.py#L299 - return "".join(choices(ALPHANUMERIC, k=9)) - - @staticmethod - def is_valid_id(id: str) -> bool: - return id.isalnum() and len(id) == 9 - - -def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool: - if is_mistral_tokenizer(model_tokenizer): - return model_tokenizer.version < 11 - # For HF tokenizers, check if [ARGS] token exists in vocab - # which indicates a v11+ equivalent tokenizer - vocab: dict[str, int] = getattr(model_tokenizer, "get_vocab", lambda: {})() - return "[ARGS]" not in vocab - - -class MistralToolParser(ToolParser): - r"""Tool call parser for Mistral models, intended for use with either: - - - `mistral_common `_ - (recommended) - - the `examples/tool_chat_template_mistral.jinja` template. - - Used when `--enable-auto-tool-choice --tool-call-parser mistral` are all - set. - """ - - IS_MISTRAL_TOOL_PARSER = True # used by vllm.utils.mistral - - # Used to generate correct grammar in `adjust_request` - model_can_reason: bool = False - - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - - if not is_mistral_tokenizer(self.model_tokenizer): - logger.info("Non-Mistral tokenizer detected when using a Mistral model...") - - # initialize properties used for state when parsing tool calls in - # streaming mode - self.prev_tool_call_arr: list[dict[str, Any]] = [] - self.current_tool_id: int = -1 - self.streaming_state: StreamingState = StreamingState.WAITING_FOR_TOOL_START - - # For streaming pre v11 tokenizer tool calls - self.current_tool_name: str | None = None - self.current_tool_mistral_id: str | None = None - self.starting_new_tool = False - self._is_pre_v11 = _is_pre_v11_tokeniser(self.model_tokenizer) - if self._is_pre_v11: - self.parse_coro = ijson.parse_coro( - self.update_stream_state_pre_v11_tokenizer() - ) - - self.bot_token = "[TOOL_CALLS]" - self.bot_token_id = self.vocab.get(self.bot_token) - self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL) - - if self.bot_token_id is None: - raise RuntimeError( - "Mistral Tool Parser could not locate the tool call token in " - "the tokenizer!" - ) + @property + def bot_token_id(self) -> int | None: + """Return the token id for :attr:`bot_token`.""" + return self._parser_engine.bot_token_id def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - so_non_supported_attributes = [ - "regex", - "choice", - "grammar", - # whitespace_pattern is not a constraint type but an option; - # Mistral grammar factory does not support it. - "whitespace_pattern", - "structural_tag", - ] - any_so_non_supported_active = request.structured_outputs is not None and any( - getattr(request.structured_outputs, attribute) is not None - for attribute in so_non_supported_attributes - ) - response_format_non_supported_active = ( - isinstance(request, ResponsesRequest) - or request.response_format is not None - and request.response_format.type == "structural_tag" - ) - - if ( - not is_mistral_tokenizer(self.model_tokenizer) - or isinstance(request, ResponsesRequest) - or not self.model_tokenizer.supports_grammar - or any_so_non_supported_active - or response_format_non_supported_active - ): - request = super().adjust_request(request) - if request.tools and request.tool_choice != "none": - # Do not skip special tokens when using chat template - # with Mistral parser as TOOL_CALL token is needed - # for tool detection. - # Note: we don't want skip_special_tokens=False - # with MistralTokenizer as it is incompatible - request.skip_special_tokens = False - return request - - json_schema: dict[str, Any] | None = None - if request.structured_outputs is not None: - if request.structured_outputs.json_object is not None: - json_schema = _DEFAULT_JSON_SCHEMA - elif request.structured_outputs.json is not None: - if isinstance(request.structured_outputs.json, str): - json_schema = json.loads(request.structured_outputs.json) - else: - json_schema = request.structured_outputs.json - else: - raise ValueError( - "Unsupported request.structured_outputs for MistralToolParser. " - "Only `json` and `json_object` are supported." - ) - elif ( - request.response_format is not None - and request.response_format.type != "text" - ): - if request.response_format.type == "json_object": - json_schema = _DEFAULT_JSON_SCHEMA - elif request.response_format.type == "json_schema": - if request.response_format.json_schema is not None: - json_schema = request.response_format.json_schema.json_schema - else: - json_schema = _DEFAULT_JSON_SCHEMA - else: - raise ValueError( - "MistralToolParser only accepts `text`, `json_object` or " - f"`json_schema`, got {request.response_format=}" - ) - # Structured Outputs will be defined. - request.response_format = None - - grammar_factory = self.model_tokenizer.grammar_factory - - # TODO: Once unified parser, improve this. - # The issue is figuring out when a model is a reasoning one or not. - template = grammar_factory.select_jinja_template( - reasoning=self.model_can_reason - ) - - mistral_tools = ( - [MistralTool.from_openai(tool.model_dump()) for tool in request.tools] - if request.tools is not None - else None - ) - - tool_choice: MistralToolChoice - match request.tool_choice: - case "none" | "auto" | "required": - tool_choice = MistralToolChoiceEnum(request.tool_choice) - case None: - tool_choice = MistralToolChoiceEnum.auto - # _ == Named tool choice - case _: - tool_choice = MistralNamedToolChoice.model_validate( - { - "type": "function", - "function": {"name": request.tool_choice.function.name}, - } - ) - - # Rendering grammar is cached in mistral-common given tools, template and mode. - match tool_choice, json_schema is not None: - case MistralToolChoiceEnum.none, True: - lark_grammar = grammar_factory.get_lark_for_json_schema( - template=template, json_schema=json_schema - ) - case _, _: - lark_grammar = grammar_factory.get_lark_from_jinja( - template=template, - mode=tool_choice, - tools=mistral_tools, - json_schema=json_schema, - parallel_tool_calls=request.parallel_tool_calls, - json_only=False, - ) - - request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar) - request._grammar_from_tool_parser = True - return request - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. - - Content and tool calls formatting depends on the Mistral's tokenizer version - used to train the model: - - - < v11: `content[BOT] [{tool_call1},{tool_call2}]` - - >= v11: `content[BOT]tool_name1{args_call1}[BOT]tool_name2{args_call2}` - - with [BOT] the tool call token. - - Note: - For tokenizer versions >= v11, tool calls with arguments wrongly formatted - are still returned as tool calls. This is to allow the model to know it - tried to make a tool call. It reduces chance of another failure and - prevents that the context is filled with tool calls wrongly placed in - assistant message contents. - """ - - # If the tool call token is not present, return a text response - if self.bot_token not in model_output: - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - content_and_raw_tool_calls = model_output.split(self.bot_token) - content = content_and_raw_tool_calls[0] - raw_tool_calls = content_and_raw_tool_calls[1:] - - # >= v11: content[BOT]tool_name1{args_call1}[BOT]tool_name2{args_call2} - if not self._is_pre_v11: - tool_calls = [] - for raw_tool_call in raw_tool_calls: - if "{" not in raw_tool_call: - continue - - end_name = raw_tool_call.find("{") - tool_name, args = ( - raw_tool_call[:end_name], - raw_tool_call[end_name:], - ) - - # HF tokenizers may include [ARGS] in the text - tool_name = tool_name.replace("[ARGS]", "") - tool_calls.append({"name": tool_name, "arguments": args}) - - # < v11: content[BOT] [{tool_call1},{tool_call2}] - else: - if len(raw_tool_calls) != 1: - raise ValueError( - "Only one BOT token should have been outputted, " - f"but got {model_output}." - ) - stringified_tool_calls = raw_tool_calls[0].strip() - try: - # Use raw_decode to parse the first valid JSON value, - # ignoring trailing tokens the model may emit after - # the tool call array. - tool_calls, _ = json.JSONDecoder().raw_decode(stringified_tool_calls) - except json.JSONDecodeError: - try: - raw_tool_call = self.tool_call_regex.findall( - stringified_tool_calls - )[0] - tool_calls = json.loads(raw_tool_call) - tool_calls = [ - { - "name": tool_call["name"], - "arguments": json.dumps( - tool_call.get("arguments", {}), - ensure_ascii=False, - ), - } - for tool_call in tool_calls - ] - except (IndexError, json.JSONDecodeError): - logger.exception("Error in extracting tool call from response.") - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=stringified_tool_calls, - ) - else: - tool_calls = [ - { - "name": tool_call["name"], - "arguments": json.dumps( - tool_call.get("arguments", {}), - ensure_ascii=False, - ), - } - for tool_call in tool_calls - ] - - mistral_tool_calls: list[MistralToolCall] = [ - MistralToolCall( - type="function", - function=FunctionCall( - name=tool_call["name"], - arguments=tool_call.get("arguments", "{}"), - ), - ) - for tool_call in tool_calls - ] - - return ExtractedToolCallInformation( - tools_called=True, - tool_calls=mistral_tool_calls, - content=content if content.strip() else None, - ) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - has_bot_token = ( - self.bot_token_id in current_token_ids or self.bot_token in current_text - ) - if not has_bot_token: - # if the tool call token is not in the tokens generated so far, - # append output to contents since it's not a tool - return DeltaMessage(content=delta_text) - - # if the tool call token IS in the tokens generated so far, that - # means we're parsing as tool calls now - try: - if self._is_pre_v11: - return self._extract_tool_calls_streaming_pre_v11_tokenizer( - delta_text=delta_text, - delta_token_ids=delta_token_ids, - ) - else: - return self._extract_tool_calls_streaming( - delta_text=delta_text, delta_token_ids=delta_token_ids - ) - except Exception: - logger.exception("Error trying to handle streaming tool call.") - return None - - def _extract_tool_calls_streaming( self, - delta_text: str, - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - """ - Extracts tool calls for Mistral models - doing tool calls of the following format: - `[TOOL_CALLS]add{"a": 3.5, "b": 4}` - """ - additional_content: str = "" - if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START: - # this is the first tool call - if self.bot_token not in delta_text: - return DeltaMessage(content=delta_text) - if not delta_text.startswith(self.bot_token): - additional_content += delta_text.split(self.bot_token)[0] - delta_text = self.bot_token + "".join( - delta_text.split(self.bot_token)[1:] - ) - - delta_tool_calls = self._generate_delta_tool_call(delta_text) - if not additional_content and len(delta_tool_calls) == 0: - if self.streaming_state in [ - StreamingState.PARSING_ARGUMENTS, - StreamingState.PARSING_ARGUMENTS_COMPLETED, - StreamingState.TOOL_COMPLETE, - StreamingState.ALL_TOOLS_COMPLETE, - ]: - # Return an empty DeltaMessage once the tool calls are all done - # so that finish_reason gets set. - return DeltaMessage() - else: - # return None when the tool is not likely to be finished - # This can occur when the name is being parsed for example - # and we wait for the name to be complete - # before sending the function name - return None - - delta = DeltaMessage() - if additional_content: - delta.content = additional_content - if len(delta_tool_calls) > 0: - delta.tool_calls = delta_tool_calls - - return delta - - def _generate_delta_tool_call(self, delta_text: str) -> list[DeltaToolCall]: - if delta_text == "" or delta_text is None: - return [] - delta_function_name = None - tool_id = None - if self.streaming_state not in [ - StreamingState.PARSING_NAME, - StreamingState.PARSING_ARGUMENTS, - ] and delta_text.startswith(self.bot_token): - self.current_tool_id += 1 - self.streamed_args_for_tool.append("") - self.prev_tool_call_arr.append({}) - self.streaming_state = StreamingState.PARSING_NAME - delta_text = delta_text.replace(self.bot_token, "", 1) - if self.streaming_state == StreamingState.PARSING_NAME: - if self.current_tool_name is None: - self.current_tool_name = "" - # The name stops where the arguments start - # And the arguments start with the `{` char - if "{" in delta_text: - tool_id = MistralToolCall.generate_random_id() - delta_function_name = delta_text.split("{")[0] - self.current_tool_name += delta_function_name - # HF tokenizers may include [ARGS] in the text - self.current_tool_name = self.current_tool_name.replace("[ARGS]", "") - self.prev_tool_call_arr[self.current_tool_id]["name"] = ( - self.current_tool_name - ) - delta_text = delta_text[len(delta_function_name) :] - self.streaming_state = StreamingState.PARSING_ARGUMENTS - else: - # we want to send the tool name once it's complete - self.current_tool_name += delta_text - return [] - if self.streaming_state == StreamingState.PARSING_ARGUMENTS: - next_function_text = None - if self.bot_token in delta_text: - # current tool call is over - delta_arguments = "" - delta_arguments += delta_text.split(self.bot_token)[0] - next_function_text = delta_text[len(delta_arguments) :] - self.streaming_state = StreamingState.TOOL_COMPLETE - else: - delta_arguments = delta_text - self.streamed_args_for_tool[self.current_tool_id] += delta_arguments - self.prev_tool_call_arr[self.current_tool_id]["arguments"] = ( - self.streamed_args_for_tool[self.current_tool_id] - ) - ret = [] - if self.current_tool_name or delta_arguments: - ret += [ - DeltaToolCall( - index=self.current_tool_id, - type="function", - id=tool_id, - function=DeltaFunctionCall( - name=self.current_tool_name, arguments=delta_arguments - ).model_dump(exclude_none=True), - ) - ] - self.current_tool_name = None - if next_function_text: - ret += self._generate_delta_tool_call(next_function_text) - return ret - # Should not happen - return [] - - @ijson.coroutine - def update_stream_state_pre_v11_tokenizer(self): - while True: - (prefix, event, value) = yield - - 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": - self.current_tool_name = value - self.streaming_state = StreamingState.PARSING_NAME_COMPLETED - if prefix == "item" and event == "map_key" and value == "arguments": - self.streaming_state = StreamingState.WAITING_FOR_ARGUMENTS_START - if prefix == "item.arguments" and event == "start_map": - self.streaming_state = StreamingState.PARSING_ARGUMENTS - if prefix == "item.arguments" and event == "end_map": - self.streaming_state = StreamingState.PARSING_ARGUMENTS_COMPLETED - if prefix == "item" and event == "end_map": - self.streaming_state = StreamingState.TOOL_COMPLETE - if prefix == "" and event == "end_array": - self.streaming_state = StreamingState.ALL_TOOLS_COMPLETE - - def _extract_tool_calls_streaming_pre_v11_tokenizer( - self, - delta_text: str, - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - """ - Extracts tool calls for Mistral models - doing tool calls of the following format: - `[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5, "b": 4}}` - """ - assert self.parse_coro is not None - content = None - delta_tool_calls: list[DeltaToolCall] = [] - current_tool_call: DeltaToolCall = DeltaToolCall( - index=self.current_tool_id, type="function" - ) - current_tool_call_modified = False - if self.bot_token_id in delta_token_ids or self.bot_token in delta_text: - # this is the first tool call - if not delta_text.startswith(self.bot_token): - content = delta_text.split(self.bot_token)[0] - delta_text = "".join(delta_text.split(self.bot_token)[1:]) - - # Cut smartly the delta text to catch the ijson events - # as ijson does not give us the index in the text at each event. - # We need to cut so that we know - # where in the text the events are emitted from. - while len(delta_text) > 0: - streaming_state_before_parse = self.streaming_state - - if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_opening_curly_braces=1, - ) - elif self.streaming_state == StreamingState.WAITING_FOR_TOOL_KEY: - # Wait until another key is sent - # or the current tool is completed - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_colon=1, - stop_after_opening_curly_braces=1, - # if the tool ends, we want to separate - # at the start of the next tool - ) - elif self.streaming_state == StreamingState.PARSING_NAME: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_comma=1, - stop_after_closing_brackets=1, - ) - elif self.streaming_state == StreamingState.WAITING_FOR_ARGUMENTS_START: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_opening_curly_braces=1, - ) - elif self.streaming_state == StreamingState.PARSING_ARGUMENTS: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_closing_curly_braces=1, - # we could be more clever - # by listening to item.arguments.* start_map events - # and know how many curly braces we can allow - ) - elif self.streaming_state in [ - StreamingState.PARSING_ARGUMENTS_COMPLETED, - StreamingState.PARSING_NAME_COMPLETED, - ]: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_closing_curly_braces=1, - stop_after_closing_brackets=1, - ) - elif self.streaming_state == StreamingState.TOOL_COMPLETE: - delta_to_be_parsed, delta_text = self._split_delta( - delta_text=delta_text, - stop_after_opening_curly_braces=1, - stop_after_closing_brackets=1, - ) - elif self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE: - content = delta_text - delta_text = "" - else: - delta_to_be_parsed = delta_text - delta_text = "" - - if self.streaming_state != StreamingState.ALL_TOOLS_COMPLETE: - self.parse_coro.send(delta_to_be_parsed.encode("utf-8")) - - # Given the parsed text and the possible streaming state change, - # let's add to the tool delta - # 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 - self.current_tool_mistral_id = None - self._track_streamed_args_pre_v11(current_tool_call) - delta_tool_calls.append(current_tool_call) - current_tool_call_modified = False - self.current_tool_id += 1 - self.streamed_args_for_tool.append("") - self.prev_tool_call_arr.append({}) - self.current_tool_mistral_id = MistralToolCall.generate_random_id() - current_tool_call = DeltaToolCall( - index=self.current_tool_id, - type="function", - ) - if current_tool_call.function is None: - current_tool_call.function = DeltaFunctionCall() - - if self.current_tool_name is not None: - # we have the complete tool name - current_tool_call_modified = True - current_tool_call.function.name = self.current_tool_name - self.prev_tool_call_arr[self.current_tool_id]["name"] = ( - self.current_tool_name - ) - self.current_tool_name = None - if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED: - self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY - if self.streaming_state in [ - StreamingState.PARSING_ARGUMENTS, - StreamingState.PARSING_ARGUMENTS_COMPLETED, - ]: - if self.streaming_state == StreamingState.PARSING_ARGUMENTS_COMPLETED: - self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY - # the delta_to_be_parsed is part of arguments. - current_tool_call_modified = True - if current_tool_call.function.arguments is None: - current_tool_call.function.arguments = delta_to_be_parsed - else: - current_tool_call.function.arguments += delta_to_be_parsed - if streaming_state_before_parse != StreamingState.PARSING_ARGUMENTS: - # It's the first chunk of arg. let's lstrip it - current_tool_call.function.arguments = ( - current_tool_call.function.arguments.lstrip() - ) - - if current_tool_call_modified: - if self.current_tool_mistral_id is not None: - current_tool_call.id = self.current_tool_mistral_id - self.current_tool_mistral_id = None - self._track_streamed_args_pre_v11(current_tool_call) - delta_tool_calls.append(current_tool_call) - - if content or len(delta_tool_calls) > 0: - delta_message = DeltaMessage() - if content: - delta_message.content = content - if len(delta_tool_calls) > 0: - delta_message.tool_calls = delta_tool_calls - return delta_message - else: - if self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE: - return DeltaMessage() - else: - return None - - def _track_streamed_args_pre_v11(self, tool_call: DeltaToolCall) -> None: - r"""Accumulate `tool_call` arguments into the streaming state.""" - if tool_call.function is not None and tool_call.function.arguments is not None: - self.streamed_args_for_tool[self.current_tool_id] += ( - tool_call.function.arguments - ) - self.prev_tool_call_arr[self.current_tool_id]["arguments"] = ( - self.streamed_args_for_tool[self.current_tool_id] - ) - - def _split_delta( - self, - delta_text: str, - stop_after_quotes: int = -1, - stop_after_opening_curly_braces: int = -1, - stop_after_closing_curly_braces: int = -1, - stop_after_closing_brackets: int = -1, - stop_after_colon: int = -1, - stop_after_comma=-1, - ) -> tuple[str, str]: - delta_to_be_parsed = "" - for i, c in enumerate(delta_text): - if c in ['"', "'"]: - delta_to_be_parsed += c - stop_after_quotes -= 1 - if stop_after_quotes == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - elif c == "{": - delta_to_be_parsed += c - stop_after_opening_curly_braces -= 1 - if stop_after_opening_curly_braces == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - elif c == "}": - delta_to_be_parsed += c - stop_after_closing_curly_braces -= 1 - if stop_after_closing_curly_braces == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - elif c == "]": - delta_to_be_parsed += c - stop_after_closing_brackets -= 1 - if stop_after_closing_brackets == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - elif c == ":": - delta_to_be_parsed += c - stop_after_colon -= 1 - if stop_after_colon == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - elif c == ",": - delta_to_be_parsed += c - stop_after_comma -= 1 - if stop_after_comma == 0: - return (delta_to_be_parsed, delta_text[i + 1 :]) - else: - delta_to_be_parsed += c - - return (delta_to_be_parsed, "") + request: ChatCompletionRequest | ResponsesRequest, + ) -> ChatCompletionRequest | ResponsesRequest: + # Skip the base ToolParser tool_choice -> structured_outputs + # conversion. Mistral enforces tool_choice via the grammar ``mode`` + # (auto/none/required/named); forwarding a tool-derived json_schema + # would be unioned into the grammar by mistral-common, allowing raw + # JSON in place of a real [TOOL_CALLS] call. Genuine user-provided + # structured output on the request is still honored by the engine's + # adjust_request. + return self._parser_engine.adjust_request(request) diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 5ee7c6f6c291..d318d5b96376 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -47,12 +47,6 @@ def _bracket_level_state( return level, in_string, escaped -def _bracket_level(s: str, opening: str = "{", closing: str = "}") -> int: - """Calculate the current level of nested brackets in a string.""" - level, _, _ = _bracket_level_state(s, opening, closing) - return level - - def filter_delta_text( delta_text: str, previous_text: str, @@ -105,7 +99,7 @@ def extract_named_tool_call_streaming( else: if is_mistral_tokenizer(tokenizer): # Import mistral_common only if we need it. - from vllm.tool_parsers.mistral_tool_parser import MistralToolCall + from vllm.parser.mistral import MistralToolCall tool_call_id = MistralToolCall.generate_random_id() else: diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 99c92f8f0a2e..1f9a651ee46d 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -19,7 +19,12 @@ AnyTextFormat, ConstStringFormat, JSONSchemaFormat, + OptionalFormat, + OrFormat, + PlusFormat, + RegexFormat, SequenceFormat, + StarFormat, TagFormat, TagsWithSeparatorFormat, TriggeredTagsFormat, @@ -60,13 +65,12 @@ "qwen_3_5", "qwen_3_coder", "qwen_3", - "harmony", "deepseek_v3_2", "glm_4_7", "deepseek_v4", } ) -VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes"}) +VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "kimi_k3"}) SUPPORTED_STRUCTURAL_TAG_MODELS = ( XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS ) @@ -204,7 +208,7 @@ def _dump_allowed_tool_ref_for_xgrammar(tool_ref: AllowedToolRef) -> AllowedTool return tool_ref -def _get_function_parameters(function) -> dict[str, Any] | bool: +def get_function_parameters(function) -> dict[str, Any] | bool: if getattr(function, "strict", None) is False: return True return function.parameters if function.parameters is not None else True @@ -225,7 +229,7 @@ def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: TagFormat( begin=begin + tool.function.name + arguments_field_prefix, content=JSONSchemaFormat( - json_schema=_get_function_parameters(tool.function) + json_schema=get_function_parameters(tool.function) ), end=end, ) @@ -274,7 +278,7 @@ def _minimax_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: TagFormat( begin=f'\n', content=JSONSchemaFormat( - json_schema=_get_function_parameters(tool.function), + json_schema=get_function_parameters(tool.function), style="minimax_xml", ), end="\n", @@ -345,3 +349,259 @@ def get_minimax_structural_tag( ) return StructuralTag(format=suffix_tag) + + +# --------------------------------------------------------------------------- +# Kimi K3 (XTML channel format) +# --------------------------------------------------------------------------- +# K3 assistant output after the reasoning gate (``<|close|>think<|sep|>``): +# <|open|>response<|sep|> <|close|>response<|sep|> +# [ <|open|>tools<|sep|> +# <|open|>call tool="NAME" index="1"<|sep|> +# <|open|>argument key="K" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|> +# <|close|>call<|sep|> ... +# <|close|>tools<|sep|> ] +# See ``encoding_k3.py`` (_render_assistant_segments / _open_tag / _attr) and +# ``kimi_k3_tool_parser.py`` for the exact byte-level encoding this mirrors. +_K3_OPEN = "<|open|>" +_K3_CLOSE = "<|close|>" +_K3_SEP = "<|sep|>" +_K3_RESPONSE_OPEN = f"{_K3_OPEN}response{_K3_SEP}" +_K3_RESPONSE_CLOSE = f"{_K3_CLOSE}response{_K3_SEP}" +_K3_TOOLS_OPEN = f"{_K3_OPEN}tools{_K3_SEP}" +_K3_TOOLS_CLOSE = f"{_K3_CLOSE}tools{_K3_SEP}" +_K3_CALL_CLOSE = f"{_K3_CLOSE}call{_K3_SEP}" +_K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}" +# The model closes the assistant turn with <|close|>message<|sep|> right before +# the end-of-message token. It is generated (not part of the prompt prefix), so +# the tag must permit it or the FSM would mask the model's natural terminator. +_K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}" + +# JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors +# ``encoding_k3._xtml_type`` (integer collapses onto number). +_K3_JSON_TO_XTML_TYPE = { + "string": "string", + "integer": "number", + "number": "number", + "boolean": "boolean", + "null": "null", + "object": "object", + "array": "array", +} + + +def _k3_escape_attr(value: str) -> str: + """Mirror ``encoding_k3._escape_attr_value`` (``&`` then ``"``).""" + return str(value).replace("&", "&").replace('"', """) + + +_K3_STRING_ATOM = r"(?:[^<]|<[^|])" +"""One raw-string character: anything but the ambiguous "<|" marker prefix. +Allows '<' inside values (e.g. HTML snippets); a value *ending* in '<' or +containing a literal "<|" is not expressible and falls back to AnyText via +the pattern checks below never matching those cases at build time (schemas +cannot know values, so the only build-time effect is the length bound).""" + + +def _k3_bounded_string_regex(prop: dict[str, Any]) -> str | None: + """Length/pattern constraint for the raw string channel, if expressible. + + The XTML string channel emits values raw (not JSON-quoted), so + JSONSchemaFormat cannot enforce string constraints there; unconstrained + AnyText lets maxLength/pattern violations through (observed on the walle + verifier: over-long junk strings pass the grammar and fail validation). + xgrammar's regex engine has no lookahead, so the close marker is kept + unambiguous by excluding the "<|" prefix from value characters. + + Returns a regex for the value, or None to keep permissive AnyText. + """ + max_len = prop.get("maxLength") + min_len = prop.get("minLength", 0) + if not isinstance(max_len, int) or max_len < 0 or max_len > 4096: + return None + if not isinstance(min_len, int) or min_len < 0 or min_len > max_len: + min_len = 0 + return _K3_STRING_ATOM + f"{{{min_len},{max_len}}}" + + +def _k3_argument_tag( + key: str, + schema: dict[str, Any], + root_defs: dict[str, Any] | None = None, +) -> TagFormat | None: + """Build one ``argument`` XTML tag for property ``key``. + + ``string`` values are emitted raw (bounded by the close marker); every other + JSON type is emitted as JSON and validated against the property schema. A + property whose type is a union / missing is left permissive (any XTML type, + raw value) so a valid call is never rejected. + + ``root_defs`` carries the tool parameters' root-level ``$defs`` / + ``definitions``: slicing a property out of the parameters document orphans + its ``#/$defs/...`` references, so those tables must be re-attached to keep + the embedded schema self-contained. + """ + prop = schema if isinstance(schema, dict) else {} + json_type = prop.get("type") + xtml_type = ( + _K3_JSON_TO_XTML_TYPE.get(json_type) if isinstance(json_type, str) else None + ) + if xtml_type is None: + # Unknown / union type: constrain the key but keep the value permissive. + return None + begin = ( + f'{_K3_OPEN}argument key="{_k3_escape_attr(key)}" type="{xtml_type}"{_K3_SEP}' + ) + if xtml_type == "string": + # Raw string channel: JSONSchemaFormat can't apply (values are not + # JSON-quoted), but an enum/const of strings is a finite set that can + # be enforced exactly with const-string alternation. Enum semantics + # are exclusive, so this never over-rejects. Fall back to permissive + # AnyText for open-ended strings or non-representable enums. + enum_values = prop.get("enum") + if enum_values is None and isinstance(prop.get("const"), str): + enum_values = [prop["const"]] + if ( + isinstance(enum_values, list) + and enum_values + and len(enum_values) <= 256 + and all(isinstance(v, str) for v in enum_values) + and not any("<|" in v for v in enum_values) + ): + branches = [ConstStringFormat(value=v) for v in enum_values] + content: Any = ( + branches[0] if len(branches) == 1 else OrFormat(elements=branches) + ) + elif (bounded := _k3_bounded_string_regex(prop)) is not None: + content = RegexFormat(pattern=bounded) + else: + content = AnyTextFormat(excludes=[_K3_CLOSE]) + else: + embedded = prop + if root_defs: + embedded = dict(prop) + for defs_key, defs_value in root_defs.items(): + embedded.setdefault(defs_key, defs_value) + content = JSONSchemaFormat(json_schema=embedded) + return TagFormat(begin=begin, content=content, end=_K3_ARG_CLOSE) + + +def _k3_permissive_argument_tag() -> TagFormat: + """A key/type-agnostic ``argument`` tag: any attributes, raw value. + + Used as a fallback so tools with union/loose schemas still get the XTML + skeleton constrained without over-rejecting the value. + """ + return TagFormat( + begin=_K3_OPEN + "argument ", + content=SequenceFormat( + elements=[ + RegexFormat(pattern=r"[^<]*" + _K3_SEP.replace("|", r"\|")), + AnyTextFormat(excludes=[_K3_CLOSE]), + ] + ), + end=_K3_ARG_CLOSE, + ) + + +def _k3_arguments_block(parameters: dict[str, Any] | bool) -> Any: + """Build ``argument`` tags for a tool's parameter schema. + + Require at least one tag when the root schema declares required properties. + Otherwise, keep accepting zero-or-more tags. Arguments remain order-agnostic + and non-unique. + """ + if not isinstance(parameters, dict): + return StarFormat(content=_k3_permissive_argument_tag()) + props = parameters.get("properties") + if not isinstance(props, dict) or not props: + # No declared properties: allow any argument blocks (or none). + return StarFormat(content=_k3_permissive_argument_tag()) + root_defs = { + defs_key: parameters[defs_key] + for defs_key in ("$defs", "definitions") + if isinstance(parameters.get(defs_key), dict) + } + tags: list[TagFormat] = [] + for key, prop in props.items(): + tag = _k3_argument_tag(key, prop, root_defs) + tags.append(tag if tag is not None else _k3_permissive_argument_tag()) + inner = tags[0] if len(tags) == 1 else OrFormat(elements=list(tags)) + required = parameters.get("required") + if isinstance(required, list) and required: + return PlusFormat(content=inner) + return StarFormat(content=inner) + + +def _k3_call_tag(tool: FunctionToolParam) -> TagFormat: + """One ``call`` tag: ``<|open|>call tool="N" index=""<|sep|> args``.""" + function = tool.function + parameters = get_function_parameters(function) + begin = f'{_K3_OPEN}call tool="{_k3_escape_attr(function.name)}" index="' + return TagFormat( + begin=begin, + content=SequenceFormat( + elements=[ + RegexFormat(pattern=r"[0-9]+"), + ConstStringFormat(value=f'"{_K3_SEP}'), + _k3_arguments_block(parameters), + ] + ), + end=_K3_CALL_CLOSE, + ) + + +def _k3_response_prefix() -> list[Any]: + """The response channel that always precedes the tools channel. + + ``response`` is generated in thinking mode (prefix ends at + ``<|open|>think<|sep|>``) but is part of the generation prefix in + non-thinking mode, so its open marker is optional. The body is bounded by + the response close marker. + """ + return [ + OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)), + TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE), + ] + + +def _k3_tools_channel(tools: list[FunctionToolParam]) -> TagFormat: + return TagFormat( + begin=_K3_TOOLS_OPEN, + content=TagsWithSeparatorFormat( + tags=[_k3_call_tag(tool) for tool in tools], + separator="", + at_least_one=True, + ), + end=_K3_TOOLS_CLOSE, + ) + + +@register_vllm_structural_tag("kimi_k3") +def get_kimi_k3_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + del builtin_tools, reasoning + + trailer = OptionalFormat(content=ConstStringFormat(value=_K3_MESSAGE_CLOSE)) + + if not tools: + return StructuralTag( + format=SequenceFormat(elements=[*_k3_response_prefix(), trailer]) + ) + + if tool_choice == "auto": + tools_part: Any = OptionalFormat(content=_k3_tools_channel(tools)) + elif tool_choice == "forced": + # K3 rejects named tool choice upstream; treat defensively as a single + # mandatory call of the first tool. + tools_part = _k3_tools_channel(tools[:1]) + else: # required + tools_part = _k3_tools_channel(tools) + + return StructuralTag( + format=SequenceFormat(elements=[*_k3_response_prefix(), tools_part, trailer]) + ) diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 95769bafd7f3..a11d4a9eec7a 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -215,7 +215,7 @@ def iter_response_function_tool_dicts( namespace, namespaced_tool.name ) function_tools.append(tool_dict) - else: + elif isinstance(tool, FunctionTool): function_tools.append(tool.model_dump()) return function_tools diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 03103fbf5382..89da81ea918f 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -87,6 +87,7 @@ def __getitem__(self, key): deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", + k3_dspark="K3DSparkConfig", flex_olmo="FlexOlmoConfig", fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", @@ -100,6 +101,7 @@ def __getitem__(self, key): kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", kimi_k25="KimiK25Config", + kimi_k3="KimiK3Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) mlp_speculator="MLPSpeculatorConfig", @@ -124,7 +126,9 @@ def __getitem__(self, key): qwen3_asr="Qwen3ASRConfig", qwen3_next="Qwen3NextConfig", qwen3_5="Qwen3_5Config", + qwen3_5_text="Qwen3_5TextConfig", qwen3_5_moe="Qwen3_5MoeConfig", + qwen3_5_moe_text="Qwen3_5MoeTextConfig", laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, @@ -600,16 +604,6 @@ def _is_encoder_decoder(config: PretrainedConfig) -> bool: return _is_encoder_decoder(config) or _is_encoder_decoder(config.get_text_config()) -def is_interleaved(config: PretrainedConfig) -> bool: - """ - Detect if the model with this config is used with interleaved attention. - """ - text_config = config.get_text_config() - if layer_types := getattr(text_config, "layer_types", None): - return len(set(layer_types)) > 1 - return False - - def _maybe_update_auto_config_kwargs(kwargs: dict[str, Any], model_type: str): """ Update kwargs for AutoConfig initialization based on model_type @@ -732,13 +726,17 @@ def get_config( raise ValueError(error_message) from e config_parser = get_config_parser(config_format) - config_dict, config = config_parser.parse( - model, - trust_remote_code=trust_remote_code, - revision=revision, - code_revision=code_revision, - hf_overrides=hf_overrides_kw or hf_overrides_fn, - **kwargs, + # Retry to tolerate a concurrent HF cache refresh briefly hiding config.json. + config_dict, config = with_retry( + lambda: config_parser.parse( + model, + trust_remote_code=trust_remote_code, + revision=revision, + code_revision=code_revision, + hf_overrides=hf_overrides_kw or hf_overrides_fn, + **kwargs, + ), + f"Error parsing config for {model}", ) # Architecture mapping for models without explicit architectures field diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 4bb7674ddbb4..6685ee4f7883 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -35,6 +35,7 @@ "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", + "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark", "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr", "EAGLEConfig": "vllm.transformers_utils.configs.eagle", "FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid", @@ -72,6 +73,8 @@ "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", + "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", + "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", "NemotronConfig": "vllm.transformers_utils.configs.nemotron", "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h", "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid", @@ -124,6 +127,7 @@ "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", + "K3DSparkConfig", "DotsOCRConfig", "EAGLEConfig", "FlexOlmoConfig", @@ -156,6 +160,8 @@ "KimiLinearConfig", "KimiVLConfig", "KimiK25Config", + "KimiK3Config", + "KimiK3VisionConfig", "NemotronConfig", "NemotronHConfig", "OlmoHybridConfig", diff --git a/vllm/transformers_utils/configs/hunyuan_vl.py b/vllm/transformers_utils/configs/hunyuan_vl.py index a826ed9b5155..548dcfefcaf6 100644 --- a/vllm/transformers_utils/configs/hunyuan_vl.py +++ b/vllm/transformers_utils/configs/hunyuan_vl.py @@ -194,7 +194,6 @@ def __init__( self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling - # self._rope_scaling_validation() # TODO: Need validation? self.attention_bias = attention_bias self.attention_dropout = attention_dropout @@ -206,46 +205,6 @@ def __init__( **kwargs, ) - def _rope_scaling_validation(self): - """ - Validate the `rope_scaling` configuration. - """ - if self.rope_scaling is None: - return - - if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: - raise ValueError( - "`rope_scaling` must be a dictionary with with two fields, `type` and " - f"`factor` or `type` and `alpha`, got {self.rope_scaling}" - ) - rope_scaling_type = self.rope_scaling.get("type", None) - rope_scaling_factor = self.rope_scaling.get("factor", None) - rope_scaling_alpha = self.rope_scaling.get("alpha", None) - if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: - raise ValueError( - "`rope_scaling`'s type field must be one of ['linear', 'dynamic'], " - f"got {rope_scaling_type}" - ) - if rope_scaling_factor is None and rope_scaling_alpha is None: - raise ValueError( - "`rope_scaling`'s factor or alpha field must be have one, " - "got both of none" - ) - if rope_scaling_factor is not None and ( - not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0 - ): - raise ValueError( - "`rope_scaling`'s factor field must be a float > 1.0, " - f"got {rope_scaling_factor}" - ) - if rope_scaling_alpha is not None and ( - not isinstance(rope_scaling_alpha, float) or rope_scaling_alpha <= 1.0 - ): - raise ValueError( - "`rope_scaling`'s alpha field must be a float > 1.0, " - f"got {rope_scaling_alpha}" - ) - class HunYuanVLConfig(PretrainedConfig): model_type = "hunyuan_vl" diff --git a/vllm/transformers_utils/configs/k3_dspark.py b/vllm/transformers_utils/configs/k3_dspark.py new file mode 100644 index 000000000000..f1b21ed076ea --- /dev/null +++ b/vllm/transformers_utils/configs/k3_dspark.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers import DeepseekV2Config + + +class K3DSparkConfig(DeepseekV2Config): + """Configuration for a dense MLA DSpark draft model.""" + + model_type = "k3_dspark" + has_no_defaults_at_init = True + + def __init__( + self, + mla_use_nope: bool = False, + mla_use_output_gate: bool = False, + mla_use_qk_norm: bool = False, + rope_theta: float = 50000.0, + **kwargs, + ) -> None: + # DeepseekV2Config defaults to a MoE topology. Zero these fields so + # generic vLLM config logic also recognizes this draft as dense. + kwargs.setdefault("n_routed_experts", 0) + kwargs.setdefault("n_shared_experts", 0) + kwargs.setdefault("num_experts_per_tok", 0) + + rope_parameters = kwargs.get("rope_parameters") + if rope_parameters is None: + kwargs["rope_parameters"] = { + "rope_type": "default", + "rope_theta": rope_theta, + } + else: + rope_parameters = dict(rope_parameters) + rope_parameters.setdefault("rope_type", "default") + rope_parameters.setdefault("rope_theta", rope_theta) + kwargs["rope_parameters"] = rope_parameters + + super().__init__(**kwargs) + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + self.mla_use_qk_norm = mla_use_qk_norm + + unsupported = [ + name + for name in ( + "mla_use_nope", + "mla_use_output_gate", + "mla_use_qk_norm", + "dspark_bonus_anchor", + ) + if getattr(self, name, False) + ] + if self.q_lora_rank is None: + unsupported.append("q_lora_rank=None") + if unsupported: + raise ValueError( + "MLA DSpark does not support " + ", ".join(unsupported) + "." + ) + + self.draft_vocab_size = ( + getattr(self, "draft_vocab_size", None) or self.vocab_size + ) + if self.draft_vocab_size != self.vocab_size: + raise ValueError( + "MLA DSpark requires draft_vocab_size to equal vocab_size when " + "sharing the target embedding and LM head." + ) + + target_layer_ids = getattr(self, "target_layer_ids", None) + if not target_layer_ids or getattr(self, "num_target_layers", None) != len( + target_layer_ids + ): + raise ValueError( + "MLA DSpark requires non-empty target_layer_ids and a matching " + "num_target_layers." + ) diff --git a/vllm/transformers_utils/configs/kimi_k3.py b/vllm/transformers_utils/configs/kimi_k3.py new file mode 100644 index 000000000000..62de6dc03b40 --- /dev/null +++ b/vllm/transformers_utils/configs/kimi_k3.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from vllm.logger import init_logger +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +logger = init_logger(__name__) + + +class KimiK3VisionConfig(PretrainedConfig): + model_type = "kimi_k3_vision" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple[int, int] = (2, 2), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + mm_projector_type: str = "patchmergerv2", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + qkv_hidden_size: int = 1536, + norm_type: str = "rmsnorm", + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = "mlp2", + linear_bias: bool = False, + activation_func: str = "gelu_pytorch_tanh", + pos_emb_interpolation_mode: str = "bilinear", + text_hidden_size: int = 2304, + **kwargs, + ): + super().__init__(**kwargs) + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = tuple(merge_kernel_size) + self.video_attn_type = video_attn_type + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = ( + mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + ) + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + # Aliases consumed by the Kimi-K2.5 vision implementation. + self.num_attention_heads = vt_num_attention_heads + self.num_hidden_layers = vt_num_hidden_layers + self.hidden_size = vt_hidden_size + self.intermediate_size = vt_intermediate_size + + +class KimiK3Config(PretrainedConfig): + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig | None = None, + vision_config: dict | KimiK3VisionConfig | None = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + image_placeholder: str = "<|kimi_image_placeholder|>", + **kwargs, + ): + if text_config is None: + self.text_config = KimiLinearConfig() + elif isinstance(text_config, dict): + self.text_config = KimiLinearConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = KimiK3VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = KimiK3VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + # K3's vision projector output must match the text model's hidden size. + # Override any value provided in vision_config for safety. + if self.vision_config.text_hidden_size != self.text_config.hidden_size: + logger.info( + "Overriding vision_config.text_hidden_size from %s to %s " + "to match text_config.hidden_size", + self.vision_config.text_hidden_size, + self.text_config.hidden_size, + ) + self.vision_config.text_hidden_size = self.text_config.hidden_size + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.image_placeholder = image_placeholder + + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + @property + def vocab_size(self) -> int: + return self.text_config.vocab_size diff --git a/vllm/transformers_utils/configs/kimi_linear.py b/vllm/transformers_utils/configs/kimi_linear.py index 14894816801d..6c2fd335ead8 100644 --- a/vllm/transformers_utils/configs/kimi_linear.py +++ b/vllm/transformers_utils/configs/kimi_linear.py @@ -3,10 +3,6 @@ from transformers.configuration_utils import PretrainedConfig -from vllm.logger import init_logger - -logger = init_logger(__name__) - class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" @@ -49,8 +45,16 @@ def __init__( qk_rope_head_dim: int | None = None, v_head_dim: int | None = None, mla_use_nope: bool | None = False, + mla_use_output_gate: bool | None = False, num_nextn_predict_layers: int = 0, linear_attn_config: dict | None = None, + attn_res_block_size: int | None = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: int | None = None, + topk_method: str = "noaux_tc", **kwargs, ): self.model_type = model_type @@ -86,6 +90,7 @@ def __init__( self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate # moe config self.num_experts = num_experts self.num_experts_per_token = num_experts_per_token @@ -102,6 +107,14 @@ def __init__( self.topk_group = topk_group self.num_nextn_predict_layers = num_nextn_predict_layers + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + if linear_attn_config is not None: assert linear_attn_config["kda_layers"] is not None assert linear_attn_config["full_attn_layers"] is not None diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index e034cec9745a..fbd2c4518196 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -106,6 +106,9 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids (for gpu_model_runner) and dflash_config.target_layer_ids (for the DFlash model). + - sample_from_anchor: Whether to sample from the anchor position. Default + False (anchor is a bonus token, only mask tokens predict, yielding + block_size - 1 speculative tokens). """ pre_trained_config["architectures"] = ["DFlashDraftModel"] pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") @@ -119,6 +122,7 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: pre_trained_config["dflash_config"] = { "mask_token_id": config_dict["mask_token_id"], "target_layer_ids": [i - 1 for i in aux_layer_ids], + "sample_from_anchor": config_dict.get("sample_from_anchor", False), } # Enable causal masking in SWA for vllm-project/speculators models pre_trained_config["dflash_config"]["causal"] = not config_dict.get( @@ -146,10 +150,14 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None: - aux_hidden_state_layer_ids (required): target layer indices feeding the drafter. Mapped to both eagle_aux_hidden_state_layer_ids and target_layer_ids (DSpark's i-1 layer semantics). + - sample_from_anchor: Whether to sample from the anchor position. Default + False (anchor is a bonus token, only mask tokens predict, yielding + block_size - 1 speculative tokens). """ pre_trained_config["architectures"] = ["Qwen3DSparkModel"] - # Speculators DSpark uses the 1+N fill-in block (anchor is a bonus token). - pre_trained_config["dspark_bonus_anchor"] = True + pre_trained_config["sample_from_anchor"] = config_dict.get( + "sample_from_anchor", False + ) aux_layer_ids = config_dict["aux_hidden_state_layer_ids"] pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 70bb4caa5354..90c026cf9b85 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -267,6 +267,7 @@ def is_deepseek_mla(self) -> bool: "deepseek_v32", "deepseek_v4", "deepseek_mtp", + "k3_dspark", "glm_moe_dsa", "glm4_moe_lite", "glm4_moe_lite_mtp", @@ -300,8 +301,16 @@ def is_deepseek_mla(self) -> bool: ) return False - def is_mm_prefix_lm(self) -> bool: - """Whether to use bidirectional attention for mm positions.""" + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + """Whether to use bidirectional attention for mm positions. + + ``supports_multimodal`` is False when the deployment is configuration- + disabled for multimodal inputs (text-only serving). In that case + mm_prefix is unnecessary and must stay off so attention backends + without ``supports_mm_prefix()`` remain eligible. + """ + if not supports_multimodal: + return False if hasattr(self.hf_config, "is_mm_prefix_lm"): return bool(self.hf_config.is_mm_prefix_lm) # fallback to list of known models @@ -358,7 +367,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: derived_max_model_len = tmp_max_len return derived_max_model_len, max_len_key - def convert(self) -> ModelArchitectureConfig: + def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig: model_arch_config = ModelArchitectureConfig( architectures=self.get_architectures(), model_type=self.hf_config.model_type, @@ -372,7 +381,7 @@ def convert(self) -> ModelArchitectureConfig: num_experts=self.get_num_experts(), quantization_config=self.get_quantization_config(), is_deepseek_mla=self.is_deepseek_mla(), - is_mm_prefix_lm=self.is_mm_prefix_lm(), + is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal), rswa_window=self.rswa_window(), derived_max_model_len_and_key=self.derive_max_model_len_and_key(), ) @@ -401,7 +410,7 @@ def get_total_num_kv_heads(self) -> int: ) return enc_num_kv_heads - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: return False @@ -584,7 +593,9 @@ def get_num_hidden_layers(self) -> int: class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase): - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + if not supports_multimodal: + return False return ( getattr(self.hf_text_config, "use_bidirectional_attention", None) == "vision" diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index 60b56b75a00c..4fa7d2fee84f 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -27,6 +27,7 @@ "IsaacProcessor", "KimiAudioProcessor", "KimiK25Processor", + "KimiK3Processor", "MiMoOmniProcessor", "MiniCPMOProcessor", "MiniCPMVProcessor", @@ -65,6 +66,7 @@ "IsaacProcessor": "vllm.transformers_utils.processors.isaac", "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio", "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25", + "KimiK3Processor": "vllm.transformers_utils.processors.kimi_k3", "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni", "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo", "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv", diff --git a/vllm/transformers_utils/processors/kimi_k3.py b/vllm/transformers_utils/processors/kimi_k3.py new file mode 100644 index 000000000000..73af85eb090b --- /dev/null +++ b/vllm/transformers_utils/processors/kimi_k3.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import BaseImageProcessor, BatchFeature, TensorType +from transformers.processing_utils import ProcessorMixin + +from vllm.tokenizers.hf import HfTokenizer + + +class KimiK3Processor(ProcessorMixin): + """HF-style processor wrapper for the image-only Kimi-K3 model. + + K3 exposes the standard ``image`` modality, so vLLM calls this processor + with ``images=[PIL, ...]``. The underlying checkpoint image processor + (``KimiK3VisionProcessor``) works on ``{"type": "image", "image": PIL}`` + media dicts, so this wrapper adapts bare PIL images into that shape before + delegating to ``preprocess``. + + Text is only tokenized here; the single ``<|kimi_image_placeholder|>`` + token per image is expanded into the resolution-aware media block by the + model's ``_get_prompt_updates`` on the vLLM side. + """ + + attributes = ["image_processor", "tokenizer"] + + def __init__( + self, + image_processor: BaseImageProcessor, + tokenizer: HfTokenizer, + ) -> None: + self.image_processor = image_processor + self.tokenizer = tokenizer + + def __call__( + self, + text: str | list[str] | None = None, + images: object | list[object] | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + if images is not None: + if not isinstance(images, list): + images = [images] + medias = [{"type": "image", "image": image} for image in images] + mm_inputs = self.image_processor.preprocess( + medias, + return_tensors=return_tensors, + ) + else: + mm_inputs = {} + + if text is not None: + if not isinstance(text, list): + text = [text] + text_inputs = self.tokenizer(text) + else: + text_inputs = {} + + return BatchFeature( + data={**text_inputs, **mm_inputs}, + tensor_type=return_tensors, + ) diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py index 899e0402ba52..d9ac841cd86b 100644 --- a/vllm/transformers_utils/processors/minicpmo.py +++ b/vllm/transformers_utils/processors/minicpmo.py @@ -72,19 +72,6 @@ def __init__( self.version = getattr(image_processor, "version", None) self.pool_step = pool_step - def _safe_get_token_id(self, attr_name, default_token_str): - """Get token ID safely, with fallback to default.""" - val = getattr(self.tokenizer, attr_name, None) - if val is None: - val = self.tokenizer.convert_tokens_to_ids(default_token_str) - if val is None: - return -1 - return val - - def _safe_get_token_str(self, attr_name, default_token_str): - """Get token string safely, with fallback to default.""" - return getattr(self.tokenizer, attr_name, default_token_str) - def __call__( self, text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py index 13dbce5368f8..60fa7ed536b7 100644 --- a/vllm/transformers_utils/processors/minimax_m3.py +++ b/vllm/transformers_utils/processors/minimax_m3.py @@ -15,7 +15,6 @@ import math -import regex as re import torch from torchvision.transforms import InterpolationMode from transformers import AutoTokenizer, BatchFeature @@ -554,85 +553,6 @@ def __init__( self.VISION_END_TOKEN ) - def _prune_video_tokens( - self, - input_text: str, - video_segments: list[int], - video_token: str, - ) -> str: - """Prune video tokens by temporal_patch_size (e.g., 2:1). - - Expects the prompt to carry exactly sum(video_segments) video tokens - — i.e. one token per *sampled* frame — then drops tokens. - """ - # If no videos or temporal_patch_size <= 1, no pruning needed - if not video_segments or self.video_processor.temporal_patch_size <= 1: - return input_text - - # Split while keeping delimiters - special_tokens = [video_token] - pattern = "|".join(map(re.escape, special_tokens)) - parts = re.split(f"({pattern})", input_text) - - def is_timestamp(text: str) -> bool: - """Check if text ends with timestamp format like ']<]0.0 seconds[>['""" - return ( - text.endswith("seconds[>[") - or text.endswith("seconds[>[ ") - or text.endswith("seconds [>[") - or text.endswith("seconds [>[ ") - ) - - def extract_timestamp(text: str) -> str: - """Extract timestamp text from the end, starting from ']<]'""" - start_index = text.rfind("]<]") - if start_index == -1: - raise ValueError(f"Failed to extract timestamp: {text}") - return text[start_index:] - - # Build new text with pruned video tokens - final_parts = [] - current_seg_idx = 0 # Which video segment we're in - frame_in_seg = 0 # Frame index within current segment - last_timestamp_len = 0 # Length of timestamp to potentially remove - - for part in parts: - if part == video_token: - if current_seg_idx < len(video_segments): - if frame_in_seg % self.video_processor.temporal_patch_size == 0: - # Keep this video token - final_parts.append(part) - frame_in_seg += 1 - if frame_in_seg >= video_segments[current_seg_idx]: - current_seg_idx += 1 - frame_in_seg = 0 - last_timestamp_len = 0 - else: - # Skip this video token - frame_in_seg += 1 - if frame_in_seg >= video_segments[current_seg_idx]: - current_seg_idx += 1 - frame_in_seg = 0 - # Remove the timestamp that was already appended - if last_timestamp_len > 0: - assert len(final_parts) > 0 - final_parts[-1] = final_parts[-1][:-last_timestamp_len] - last_timestamp_len = 0 - else: - # No more video segments, keep as is - final_parts.append(part) - last_timestamp_len = 0 - else: - # Text part - final_parts.append(part) - # Check if this text ends with a timestamp - if is_timestamp(part): - last_timestamp_len = len(extract_timestamp(part)) - else: - last_timestamp_len = 0 - - return "".join(final_parts) - def __call__( self, images=None, diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index d48a29d6b43a..028d207a25ac 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -23,9 +23,9 @@ from transformers import BatchFeature, PretrainedConfig, TensorType from vllm.model_executor.models.parakeet import ParakeetExtractor -from vllm.multimodal.evs import compute_retained_tokens_count from vllm.multimodal.inputs import AudioItem from vllm.multimodal.processing.processor import PromptUpdateDetails +from vllm.multimodal.video_prune.evs import compute_retained_tokens_count from vllm.tokenizers.hf import HfTokenizer from .internvl import calculate_internvl_targets, get_internvl_target_ratios diff --git a/vllm/transformers_utils/processors/ovis2_5.py b/vllm/transformers_utils/processors/ovis2_5.py index 11ac0360e757..b9c2d112409c 100644 --- a/vllm/transformers_utils/processors/ovis2_5.py +++ b/vllm/transformers_utils/processors/ovis2_5.py @@ -78,7 +78,6 @@ def __init__( @cached_property def extra_special_tokens(self): - vocab = self.tokenizer.get_vocab() required_tokens = { "image_token": "", "video_token": "