From 7dc44a00161ace6f8ba7be2cc008f926e5a35fe2 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 24 Oct 2026 08:01:24 +0800 Subject: [PATCH 001/631] [Bugfix][LoRA][MoE] Process reloaded weights of LoRA-wrapped MoE layers (#63319) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../model_loader/test_reload.py | 62 +++++++++++++++++++ .../fused_moe/fused_moe_modular_method.py | 19 ++++++ 2 files changed, 81 insertions(+) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index fd27329f1d65..eac5245bd988 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -6,6 +6,7 @@ import sys import types from importlib import import_module +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock from weakref import WeakKeyDictionary, ref @@ -17,6 +18,7 @@ import vllm.model_executor.model_loader.reload.layerwise as reload_layerwise import vllm.model_executor.model_loader.reload.meta as reload_meta from vllm.config import ModelConfig +from vllm.lora.request import LoRARequest from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import QKVParallelLinear @@ -41,8 +43,10 @@ composed_weight_loader, default_weight_loader, load_attention_sink, + safetensors_weights_iterator, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api @pytest.mark.parametrize("tp_rank", [0, 1]) @@ -1367,6 +1371,64 @@ def test_reload_weights( assert add_perp < mul_perp +def _reload_with_lora_expert_names(worker) -> None: + """Reload the checkpoint under the names that LoRA-wrapped experts load.""" + + def lora_name(name: str) -> str: + parts = name.split(".") + if "experts" in parts[:-2] and parts[parts.index("experts") + 1].isdigit(): + parts.insert(parts.index("experts") + 3, "base_layer") + return ".".join(parts) + + path = Path(hf_api().snapshot_download(worker.vllm_config.model_config.model)) + files = sorted(str(f) for f in path.glob("*.safetensors")) + weights = safetensors_weights_iterator(files, use_tqdm_on_load=False) + worker.reload_weights( + weights_iterator=((lora_name(name), weight) for name, weight in weights) + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="LoRA on quantized MoE requires CUDA" +) +def test_reload_weights_with_lora_moe(vllm_runner, monkeypatch): + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + lora_path = hf_api().snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider") + prompts = [ + "<|im_start|>user\nCount the rows of table candidate.<|im_end|>\n" + "<|im_start|>assistant\n\n\n\n\nSELECT count(*) FROM candidate" + ] + + def prompt_logprobs(llm, lora_id): + lora_request = LoRARequest("sql", lora_id, lora_path) if lora_id else None + output = llm.generate_greedy_logprobs( + prompts, + max_tokens=1, + num_logprobs=None, + num_prompt_logprobs=0, + lora_request=lora_request, + )[0] + return [next(iter(token.values())).logprob for token in output[3][1:]] + + with vllm_runner( + model_name="RedHatAI/Qwen3-30B-A3B-quantized.w4a16", + enable_lora=True, + max_lora_rank=8, + max_loras=1, + enforce_eager=True, + enable_prefix_caching=False, + max_model_len=64, + max_num_seqs=1, + ) as llm: + before = [prompt_logprobs(llm, 0), prompt_logprobs(llm, 1)] + assert before[0] != before[1] + llm.collective_rpc(_reload_with_lora_expert_names) + after = [prompt_logprobs(llm, 0), prompt_logprobs(llm, 2)] + assert after == before + + def test_kv_scale_reload(vllm_runner): """Test reloading a checkpoint that contains k_scale/v_scale weights.""" if _fp8_reload_unsupported(): diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 5bf694218db6..10c8cf3aef48 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -78,6 +78,25 @@ def get_fused_moe_quant_config( ) -> FusedMoEQuantConfig | None: return self.moe_quant_config + def process_weights_after_loading(self, layer: "RoutedExperts") -> None: + old_kernel = self.old_quant_method.moe_kernel + self.old_quant_method.process_weights_after_loading(layer) + new_kernel = self.old_quant_method.moe_kernel + if new_kernel is old_kernel: + return + assert self.moe_kernel is not None and new_kernel is not None + for old, new in ( + (self.moe_kernel.fused_experts, new_kernel.fused_experts), + (self.moe_kernel.prepare_finalize, new_kernel.prepare_finalize), + ): + context = getattr(old, "_lora_context", None) + if context is not None: + assert hasattr(new, "set_lora_context") + new.set_lora_context(context) + self.moe_kernel.impl = new_kernel.impl + self.old_quant_method.moe_kernel = self.moe_kernel + self.moe_quant_config = self.old_quant_method.moe_quant_config + def apply( self, layer: "RoutedExperts", From debcf228431f0b04ee8f12cfc410b0fe3c34a4f7 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 23 Oct 2026 20:05:23 -0500 Subject: [PATCH 002/631] [CI][ROCm] Export kernel source maps for CI selection (#59100) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: Codex Co-authored-by: Claude Co-authored-by: Andreas Karatzas --- .buildkite/scripts/ci-bake-rocm.sh | 68 +- .buildkite/test_areas/docker.yaml | 8 + .buildkite/test_areas/misc.yaml | 1 + CMakeLists.txt | 12 + cmake/hipify.py | 46 +- cmake/utils.cmake | 7 +- docker/Dockerfile.rocm | 19 + docker/ci-rocm.hcl | 13 + tests/test_cmake_utils.py | 33 + .../tools/test_docker_build_metadata_args.py | 168 +++++ tests/tools/test_kernel_symbol_map.py | 562 ++++++++++++++ tools/ci/kernel_symbol_map.py | 24 + tools/ci/rocm_kernel_symbol_map.py | 693 ++++++++++++++++++ 13 files changed, 1645 insertions(+), 9 deletions(-) create mode 100644 tests/tools/test_kernel_symbol_map.py create mode 100644 tools/ci/rocm_kernel_symbol_map.py diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 5935938616bc..db3d0a4acfe4 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -864,6 +864,50 @@ should_export_rocm_smoke() { || "${TARGET}" == "smoke-test-rocm-ci" ]] } +should_export_kernel_symbol_map() { + [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" == "1" ]] || return 1 + case "${TARGET}" in + test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|\ + export-wheel-rocm|csrc-rocm-ci|smoke-test-rocm-ci|kernel-symbol-map-rocm) + return 0 ;; + *) return 1 ;; + esac +} + +export_kernel_symbol_map() { + should_export_kernel_symbol_map || return 0 + echo "--- :world_map: Exporting ROCm kernel symbol map" + rm -rf ./kernel-symbol-map-rocm || return 1 + docker buildx bake \ + "${BAKE_ALLOW_ARGS[@]}" "${BAKE_FILES[@]}" \ + --progress "${BUILDKIT_PROGRESS:-plain}" \ + kernel-symbol-map-rocm || return 1 + upload_kernel_symbol_map_if_present +} + +export_kernel_symbol_map_from_cache() { + should_export_kernel_symbol_map || return 0 + # A separate shell keeps setup failures fatal inside the optional export. + FORCE_BUILD=1 VLLM_KERNEL_SYMBOL_MAP=1 UPLOAD_ROCM_WHEEL_ARTIFACTS=0 \ + CI_HCL_SOURCE="${CI_HCL_PATH}" VLLM_BAKE_FILE="${VLLM_BAKE_FILE}" \ + BUILDER_NAME="${BUILDER_NAME}" \ + bash "${BASH_SOURCE[0]}" kernel-symbol-map-rocm \ + || echo "ROCm kernel symbol map export from cache failed; continuing without it" >&2 +} + +upload_kernel_symbol_map_if_present() { + should_export_kernel_symbol_map || return 0 + local map_dir="./kernel-symbol-map-rocm" + local map_name="kernel_symbol_map.rocm.json.gz" + if [[ ! -s "${map_dir}/${map_name}" ]]; then + echo "ROCm kernel symbol map export is missing: ${map_dir}/${map_name}" >&2 + return 1 + fi + if [[ "${BUILDKITE:-false}" == "true" ]]; then + (cd "${map_dir}" && buildkite-agent artifact upload "${map_name}") + fi +} + verify_rocm_smoke_export() { local marker="./build/rocm-smoke-export/vllm-smoke-ok" local expected_smoke_id="${BUILDKITE_BUILD_ID:-local}" @@ -1453,7 +1497,8 @@ maybe_skip_existing_image() { return 0 fi if ! is_ci_base_target \ - && { should_upload_wheel_artifacts || should_export_rocm_smoke; }; then + && { should_upload_wheel_artifacts || should_export_rocm_smoke \ + || [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; }; then echo "Local-output targets always run for the current build" return 0 fi @@ -1501,11 +1546,13 @@ maybe_skip_existing_image() { echo "Commit image already exists: ${IMAGE_TAG}" echo "Skipping build" + export_kernel_symbol_map_from_cache exit 0 fi echo "Image already exists: ${IMAGE_TAG}" echo "Skipping build" + export_kernel_symbol_map_from_cache exit 0 } @@ -1806,6 +1853,7 @@ EOF uses_rocm_csrc_cache() { case "${TARGET}" in csrc-rocm-ci \ + | kernel-symbol-map-rocm \ | test-rocm-ci \ | test-rocm-ci-with-wheel \ | test-rocm-ci-with-artifacts \ @@ -2252,6 +2300,15 @@ EOF cat <&2 + fi } if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 4c80f26ca52e..632c3aa1322d 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -9,12 +9,17 @@ steps: source_file_dependencies: - .buildkite/release-pipeline.yaml - .buildkite/scripts/docker-build-metadata-args.sh + - .buildkite/scripts/rocm/build-config.sh - docker/Dockerfile - docker/Dockerfile.cpu - docker/docker-bake.hcl - tests/tools/test_docker_build_metadata_args.py + - tests/tools/test_kernel_symbol_map.py + - tools/ci/kernel_symbol_map.py + - tools/ci/rocm_kernel_symbol_map.py commands: - pytest -v -s tools/test_docker_build_metadata_args.py + - pytest -v -s tools/test_kernel_symbol_map.py mirror: amd: label: ":amd: (MI250) Docker Build Metadata" @@ -32,3 +37,6 @@ steps: - docker/ci-rocm.hcl - docker/docker-bake-rocm.hcl - tests/tools/test_docker_build_metadata_args.py + - tests/tools/test_kernel_symbol_map.py + - tools/ci/kernel_symbol_map.py + - tools/ci/rocm_kernel_symbol_map.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index a81db91830ef..d3f47bdc76c8 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -644,6 +644,7 @@ steps: - vllm/v1/ - vllm/version.py - cmake/utils.cmake + - cmake/hipify.py - tests/test_uvicorn_logging.py - tests/test_attention_backend_registry.py - tests/test_audio_media_size_precheck.py diff --git a/CMakeLists.txt b/CMakeLists.txt index e1c6169d92be..92a2b37ecf68 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,9 @@ endif() # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py) set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM") +if (VLLM_TARGET_DEVICE STREQUAL "rocm" OR "$ENV{VLLM_KERNEL_SYMBOL_MAP}" STREQUAL "1") + set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +endif() message(STATUS "Build type: ${CMAKE_BUILD_TYPE}") message(STATUS "Target device: ${VLLM_TARGET_DEVICE}") @@ -406,6 +409,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") WORKING_DIRECTORY "${TORCH_SHADOW_INCLUDE_DIR}" COMMAND_ERROR_IS_FATAL ANY) + execute_process( + COMMAND ${Python_EXECUTABLE} -c + "import json,sys; from pathlib import Path; Path(sys.argv[1]).write_text(json.dumps({sys.argv[2]: sys.argv[3:]}))" + "${CMAKE_CURRENT_BINARY_DIR}/generated-source-deps.json" + "${TORCH_SHADOW_INCLUDE_DIR}/${_stable_hdr_rel}" + "${TORCH_INCLUDE_DIR}/${_stable_hdr_rel}" + "${CMAKE_CURRENT_LIST_DIR}/cmake/patches/pytorch_stable_string.patch" + COMMAND_ERROR_IS_FATAL ANY) + include_directories(BEFORE "${TORCH_SHADOW_INCLUDE_DIR}") endif() endif() diff --git a/cmake/hipify.py b/cmake/hipify.py index 2cf58a7b7d33..2135a625dfae 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -11,8 +11,10 @@ # import argparse +import json import os import shutil +from pathlib import Path from torch.utils.hipify.hipify_python import get_hip_file_path, hipify @@ -48,14 +50,29 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: parser.add_argument( "sources", help="Source files to hipify.", nargs="*", default=[] ) + parser.add_argument( + "--manifest", help="Write generated-to-original source/header provenance." + ) args = parser.parse_args() # Limit include scope to project_dir only includes = [os.path.join(args.project_dir, "*")] - # Get absolute path for all source files. - extra_files = [os.path.abspath(s) for s in args.sources] + project_dir = Path(args.project_dir).resolve() + output_dir = Path(args.output_dir).resolve() + provenance = { + str(output_dir / source.relative_to(project_dir)): str(source) + for source in project_dir.rglob("*") + if source.is_file() + } + source_copies = { + os.path.abspath(source): str( + output_dir / Path(source).resolve().relative_to(project_dir) + ) + for source in args.sources + } + extra_files = list(source_copies.values()) # Copy sources from project directory to output directory. # The directory might already exist to hold object files so we ignore that. @@ -75,16 +92,24 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: is_pytorch_extension=True, hipify_extra_files_only=True, ) + copied_sources = dict(provenance) + for source, result in hipify_result.items(): + if result.hipified_path is not None: + original = copied_sources.get( + os.path.abspath(source), os.path.abspath(source) + ) + provenance[os.path.abspath(result.hipified_path)] = original hipified_sources = [] for source in args.sources: s_abs = os.path.abspath(source) - if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: - path = hipify_result[s_abs].hipified_path + copied = source_copies[s_abs] + if copied in hipify_result and hipify_result[copied].hipified_path is not None: + path = hipify_result[copied].hipified_path # PyTorch skips writing when is_pytorch_extension and text unchanged; # hipified_path then stays *.cu. CMake expects *.hip under output_dir. if s_abs.endswith(".cu") and path.endswith(".cu"): - dest = _expected_hip_build_path(s_abs, args.output_dir) + dest = _expected_hip_build_path(copied, args.output_dir) if os.path.normpath(path) != os.path.normpath(dest): os.makedirs(os.path.dirname(dest), exist_ok=True) shutil.copy2(path, dest) @@ -92,10 +117,19 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: else: hipified_s_abs = path else: - hipified_s_abs = s_abs + hipified_s_abs = copied hipified_sources.append(hipified_s_abs) + provenance[os.path.abspath(hipified_s_abs)] = s_abs assert len(hipified_sources) == len(args.sources) + manifest = ( + Path(args.manifest) if args.manifest else output_dir / "hipify-source-map.json" + ) + manifest.parent.mkdir(parents=True, exist_ok=True) + temporary = manifest.with_suffix(manifest.suffix + ".tmp") + temporary.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n") + temporary.replace(manifest) + # Print hipified source files. print("\n".join(hipified_sources)) diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 83eed340eb7d..7d6987122f79 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -113,11 +113,14 @@ function (vllm_finalize_hipify_target) list(REMOVE_DUPLICATES ALL_BYPRODUCTS) set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc) + set(HIPIFY_MANIFEST ${CMAKE_CURRENT_BINARY_DIR}/hipify-source-map.json) add_custom_target( hipify_all - COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS} + COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} --manifest ${HIPIFY_MANIFEST} ${ALL_SRCS} DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS} - BYPRODUCTS ${ALL_BYPRODUCTS} + BYPRODUCTS ${ALL_BYPRODUCTS} ${HIPIFY_MANIFEST} + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} + VERBATIM COMMENT "Running hipify on all extension source files.") endfunction() diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0d3686f58c11..d1a9c99803ac 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -338,6 +338,23 @@ RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproj && echo "=== ccache stats after ROCm native build ===" \ && (ccache --show-stats || true) +FROM csrc-build AS kernel-symbol-map-build +ARG VLLM_KERNEL_SYMBOL_MAP=0 +ARG buildkite_commit="" +COPY tools/ci/kernel_symbol_map.py tools/ci/rocm_kernel_symbol_map.py tools/ci/ +RUN if [ "${VLLM_KERNEL_SYMBOL_MAP}" = "1" ]; then \ + uv venv --system-site-packages /tmp/kernel-map/.venv \ + && /tmp/kernel-map/.venv/bin/python tools/ci/kernel_symbol_map.py \ + --backend rocm --build-root build --source-root . \ + --commit "${buildkite_commit}" \ + --out dist/kernel_symbol_map.rocm.json.gz \ + && rm -rf /tmp/kernel-map; \ + fi + +FROM scratch AS kernel-symbol-map-rocm +ARG COMMON_WORKDIR +COPY --from=kernel-symbol-map-build ${COMMON_WORKDIR}/vllm/dist/kernel_symbol_map.rocm.json.gz / + # 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 dependency installation and HIP/C++ compilation when their @@ -370,6 +387,8 @@ 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/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +# tests/tools/test_kernel_symbol_map.py loads the symbol-map producers from here. +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/ci /tools/ci COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index c251641452f6..31c8ad6405c5 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -25,6 +25,10 @@ variable "BUILDKITE_BUILD_ID" { default = "" } +variable "VLLM_KERNEL_SYMBOL_MAP" { + default = "0" +} + variable "PARENT_COMMIT" { default = "" } @@ -267,6 +271,8 @@ target "_ci-rocm" { CI_BASE_IMAGE = CI_BASE_IMAGE ROCM_SMOKE_ID = BUILDKITE_BUILD_ID max_jobs = CI_MAX_JOBS + VLLM_KERNEL_SYMBOL_MAP = VLLM_KERNEL_SYMBOL_MAP + buildkite_commit = BUILDKITE_COMMIT } } @@ -302,6 +308,13 @@ target "csrc-rocm-ci" { output = ["type=cacheonly"] } +target "kernel-symbol-map-rocm" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "kernel-symbol-map-rocm" + cache-from = get_cache_from_rocm_csrc() + output = ["type=local,dest=./kernel-symbol-map-rocm"] +} + # Cache-only target for the Rust frontend build stage. Final-image cache # exports use mode=min and do not reliably persist intermediate cargo layers, # so Rust gets its own source-scoped cache target. diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index 5d217433db71..ad634f08d61b 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import shutil import subprocess import sys @@ -80,3 +81,35 @@ def test_clear_cuda_gencode_flags(tmp_path: Path): ) subprocess.run([_get_cmake_bin(), "-P", script], check=True) + + +def test_hipify_records_original_sources_and_headers_without_modifying_them(tmp_path): + project = tmp_path / "source" + output = tmp_path / "build" / "csrc" + project.mkdir() + source = project / "probe.cu" + header = project / "cuda_shared.h" + source.write_text('#include "cuda_shared.h"\n__global__ void probe() {}\n') + header.write_text("#include \n") + original = {path.name: path.read_bytes() for path in project.iterdir()} + script = Path(__file__).parents[1] / "cmake" / "hipify.py" + subprocess.run( + [ + sys.executable, + str(script), + "-p", + str(project), + "-o", + str(output), + str(source), + ], + check=True, + capture_output=True, + text=True, + ) + assert {path.name: path.read_bytes() for path in project.iterdir()} == original + provenance = json.loads((output / "hipify-source-map.json").read_text()) + assert provenance[str(output / "probe.hip")] == str(source) + assert provenance[str(output / "cuda_shared.h")] == str(header) + assert provenance[str(output / "hip_shared.h")] == str(header) + assert '#include "hip_shared.h"' in (output / "probe.hip").read_text() diff --git a/tests/tools/test_docker_build_metadata_args.py b/tests/tools/test_docker_build_metadata_args.py index d7b65790888b..1821c7498ec7 100644 --- a/tests/tools/test_docker_build_metadata_args.py +++ b/tests/tools/test_docker_build_metadata_args.py @@ -6,6 +6,8 @@ import subprocess from pathlib import Path +import pytest + REPO_ROOT = Path(__file__).resolve().parents[2] HELPER = REPO_ROOT / ".buildkite" / "scripts" / "docker-build-metadata-args.sh" ROCM_CI_BAKE = REPO_ROOT / ".buildkite" / "scripts" / "ci-bake-rocm.sh" @@ -362,3 +364,169 @@ def test_rocm_git_fetch_disables_automatic_maintenance(tmp_path: Path) -> None: "origin", "HEAD", ] + + +def run_rocm_map_build( + tmp_path: Path, *, target: str = "test-rocm-ci", **settings: str +) -> tuple[subprocess.CompletedProcess[str], list[list[str]]]: + # Exercise real orchestration/errexit; replace expensive build preparation + # and external services. Cached exports launch another copy of this harness. + harness = tmp_path / "bake.sh" + harness.write_text( + r""" +source "$1" +for fn in configure_ci_base_write_scope print_header validate_inputs load_ci_hcl \ + init_bake_files compute_ci_base_hash_if_needed configure_ci_base_image_refs \ + prepare_git_cache_metadata prepare_ci_build_context configure_custom_rocm_stages \ + extract_dependency_pins write_rocm_build_arg_override \ + compute_dependency_cache_keys \ + write_ci_base_label_override compute_rocm_csrc_content_hash_if_needed \ + compute_rocm_rust_content_hash_if_needed write_rocm_cache_override \ + write_build_context_override resolve_ci_base_dependency_targets print_bake_config \ + seed_dependency_caches_if_needed promote_stable_ci_base_tag \ + publish_ci_base_handoff_ref +do + eval "$fn() { :; }" +done +init_config() { + TARGET="$1" + BAKE_TARGETS=("$TARGET") + BAKE_ALLOW_ARGS=(--allow 'fs.read=/owned context') + BAKE_FILES=(-f 'resolved args and caches.hcl' -f 'owned context.hcl') + CI_HCL_PATH=ci.hcl VLLM_BAKE_FILE=bake.hcl BUILDER_NAME=builder +} +setup_builder() { + if [[ "$TARGET" == kernel-symbol-map-rocm && "$FAILURE" == setup ]]; then + false + fi + echo setup >> events +} +remote_image_exists() { [[ "$CACHED" == 1 ]]; } +get_remote_image_label() { echo "$BUILDKITE_COMMIT"; } +bash() { command bash "$0" "$@"; } +docker() { + printf '%q ' "$@" >> calls + printf '\n' >> calls + if [[ "${*: -1}" == kernel-symbol-map-rocm ]]; then + [[ ! -e kernel-symbol-map-rocm ]] || return 90 + [[ "$FAILURE" != export ]] || return 41 + [[ "$FAILURE" != missing ]] || return 0 + mkdir kernel-symbol-map-rocm + if [[ "$FAILURE" == empty ]]; then + touch kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz + else + echo current > kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz + fi + [[ "$FAILURE" != partial ]] || return 42 + else + [[ "$FAILURE" != primary ]] || return 43 + fi +} +buildkite-agent() { + printf '%q ' "$@" >> "$TEST_ROOT/calls" + printf '\n' >> "$TEST_ROOT/calls" + [[ "$(cat kernel_symbol_map.rocm.json.gz)" == current ]] || return 91 + [[ "$FAILURE" != upload ]] || return 44 +} +main "$2" +""" + ) + export = tmp_path / "kernel-symbol-map-rocm" + export.mkdir() + (export / "kernel_symbol_map.rocm.json.gz").write_text("stale") + (tmp_path / "calls").touch() + env = { + "PATH": os.environ["PATH"], + "TEST_ROOT": str(tmp_path), + "VLLM_KERNEL_SYMBOL_MAP": "1", + "BUILDKITE": "true", + "BUILDKITE_COMMIT": "current-commit", + "IMAGE_TAG": "test:current", + "CACHED": "0", + "FAILURE": "", + **settings, + } + result = subprocess.run( + ["bash", str(harness), str(ROCM_CI_BAKE), target], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + ) + return result, [ + shlex.split(line) for line in (tmp_path / "calls").read_text().splitlines() + ] + + +@pytest.mark.parametrize("cached", ["0", "1"]) +@pytest.mark.parametrize( + "failure", ["", "export", "partial", "missing", "empty", "upload"] +) +def test_rocm_optional_map_failures_preserve_image_success( + tmp_path: Path, cached: str, failure: str +) -> None: + result, calls = run_rocm_map_build(tmp_path, CACHED=cached, FAILURE=failure) + assert result.returncode == 0, result.stderr + targets = [call[-1] for call in calls if call[:2] == ["buildx", "bake"]] + assert targets == ( + ["test-rocm-ci", "kernel-symbol-map-rocm"] + if cached == "0" + else ["kernel-symbol-map-rocm"] + ) + if cached == "0": + assert calls[0][:-1] == calls[1][:-1] + uploads = [call for call in calls if call[:2] == ["artifact", "upload"]] + assert uploads == ( + [["artifact", "upload", "kernel_symbol_map.rocm.json.gz"]] + if failure in ("", "upload") + else [] + ) + assert ("continuing without it" in result.stderr) == bool(failure) + + +def test_rocm_cached_map_setup_stops_on_failure_without_failing_image( + tmp_path: Path, +) -> None: + result, calls = run_rocm_map_build(tmp_path, CACHED="1", FAILURE="setup") + assert result.returncode == 0, result.stderr + assert "continuing without it" in result.stderr + assert not calls + assert not (tmp_path / "events").exists() + + +@pytest.mark.parametrize( + ("target", "failure", "expected_status"), + [ + ("test-rocm-ci", "primary", 43), + ("kernel-symbol-map-rocm", "export", 41), + ("kernel-symbol-map-rocm", "upload", 44), + ("smoke-test-rocm-ci", "", 1), + ("export-wheel-rocm", "", 1), + ], +) +def test_rocm_required_build_failures_remain_fatal( + tmp_path: Path, target: str, failure: str, expected_status: int +) -> None: + result, calls = run_rocm_map_build(tmp_path, target=target, FAILURE=failure) + assert result.returncode == expected_status, result.stderr + assert [call[-1] for call in calls if call[:2] == ["buildx", "bake"]] == [target] + assert "continuing without it" not in result.stderr + + +@pytest.mark.parametrize("cached", ["0", "1"]) +@pytest.mark.parametrize("setting", ["VLLM_KERNEL_SYMBOL_MAP", "BAKE_PRINT_ONLY"]) +def test_rocm_disabled_or_print_only_build_skips_map_export( + tmp_path: Path, cached: str, setting: str +) -> None: + settings = {setting: "0" if setting == "VLLM_KERNEL_SYMBOL_MAP" else "1"} + result, calls = run_rocm_map_build(tmp_path, CACHED=cached, **settings) + assert result.returncode == 0, result.stderr + expected = ( + ["test-rocm-ci"] + if cached == "0" and setting == "VLLM_KERNEL_SYMBOL_MAP" + else [] + ) + assert [call[-1] for call in calls] == expected + assert ( + tmp_path / "kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz" + ).read_text() == "stale" diff --git a/tests/tools/test_kernel_symbol_map.py b/tests/tools/test_kernel_symbol_map.py new file mode 100644 index 000000000000..3bec8495251d --- /dev/null +++ b/tests/tools/test_kernel_symbol_map.py @@ -0,0 +1,562 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Conservative native ROCm source attribution, without requiring a GPU.""" + +import importlib.util +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +TOOLS = Path(__file__).resolve().parents[2] / "tools" / "ci" +spec = importlib.util.spec_from_file_location( + "rocm_kernel_symbol_map", TOOLS / "rocm_kernel_symbol_map.py" +) +assert spec is not None and spec.loader is not None +producer = importlib.util.module_from_spec(spec) +spec.loader.exec_module(producer) + +SYMBOL = "_Z14selector_probeIiEvPT_" +ELF_HEADER = "ELF Header:\n Type: REL (Relocatable file)\nSection Headers:\n" +ELF = f""" +Symbol table '.dynsym' contains 3 entries: + Num: Value Size Type Bind Vis Ndx Name + 1: 00001900 32 FUNC GLOBAL PROTECTED 7 {SYMBOL} + 2: 000007c0 64 OBJECT GLOBAL PROTECTED 6 {SYMBOL}.kd +AMDGPU Metadata: +amdhsa.kernels: + - .name: {SYMBOL} + .symbol: {SYMBOL}.kd +""" + + +def test_descriptor_suffix_matches_mangled_dispatch_name(): + assert producer.parse_kernel_symbols(ELF) == [SYMBOL] + + +@pytest.mark.parametrize( + "text", + [ + ELF.replace(f"7 {SYMBOL}", f"UND {SYMBOL}"), + ELF.replace(f".symbol: {SYMBOL}.kd", ".symbol: missing.kd"), + "Symbol table '.dynsym' contains 0 entries:", + "AMDGPU Metadata:\namdhsa.kernels:\n - .unknown: value", + ], +) +def test_unreadable_device_symbols_are_unknown_not_empty(text): + with pytest.raises(ValueError): + producer.parse_kernel_symbols(text) + + +@pytest.mark.parametrize( + "text", + [ + "probe.o: #deps 1, deps mtime 12 (STALE)\n source.hip\n", + "probe.o: #deps 2, deps mtime 12 (VALID)\n source.hip\n", + "unrecognized output\n", + ], +) +def test_incomplete_or_stale_dependencies_cannot_publish(text): + with pytest.raises(ValueError): + producer.parse_deps(text) + + +@pytest.fixture +def compilation(tmp_path, monkeypatch): + source = tmp_path / "source" + build = tmp_path / "build" + (source / "csrc").mkdir(parents=True) + (build / "csrc").mkdir(parents=True) + original = source / "csrc" / "cuda_kernel.cu" + original.write_text("__global__ void probe() {}\n") + header = source / "csrc" / "shared.h" + header.write_text("constexpr int kValue = 1;\n") + host = source / "csrc" / "host.cpp" + host.write_text('#include "shared.h"\n') + generated = build / "csrc" / "hip_kernel.hip" + generated.write_text(original.read_text()) + copied_header = build / "csrc" / "shared.h" + copied_header.write_text(header.read_text()) + commands = [] + dependency_records = [] + objects = [] + for name, src, deps in ( + ("device", generated, [generated, copied_header]), + ("host", host, [host, header]), + ): + obj = build / "CMakeFiles" / f"{name}.dir" / f"{name}.o" + obj.parent.mkdir(parents=True) + obj.write_bytes(b"object") + objects.append(obj) + commands.append({"directory": str(build), "file": str(src), "output": str(obj)}) + dependency_records.append( + f"{obj.relative_to(build)}: #deps {len(deps)}, deps mtime 1 (VALID)\n" + + "".join(f" {dep}\n" for dep in deps) + + "\n" + ) + (build / "compile_commands.json").write_text(json.dumps(commands)) + (build / ".ninja_deps").touch() + manifest = tmp_path / "hipify.json" + manifest.write_text( + json.dumps({str(generated): str(original), str(copied_header): str(header)}) + ) + deps_output = "".join(dependency_records) + monkeypatch.setattr( + producer, + "run", + lambda args: "tool version" if "--version" in args else deps_output, + ) + monkeypatch.setattr( + producer, + "device_symbols", + lambda obj, tools: ( + ([SYMBOL], ["hipv4-amdgcn-amd-amdhsa--gfx942"]) + if obj == objects[0] + else ([], []) + ), + ) + monkeypatch.setattr(producer, "device_dependencies", lambda *args: set()) + return { + "build_root": build, + "source_root": source, + "tools": {"ninja": "ninja"}, + "commit": "same-build", + "hipify_map": manifest, + }, objects + + +def test_hipified_source_and_header_keep_original_paths_and_host_use(compilation): + options, _ = compilation + result = producer.build_map(**options) + assert result["version"] == 2 and result["backend"] == "rocm" + assert not result["incomplete"], result.get("reason") + assert result["stats"]["objects"] == 2 + objects = {obj["source"]: obj for obj in result["objects"]} + native = objects["csrc/cuda_kernel.cu"] + assert native["device"] and native["symbols"] == [SYMBOL] + assert native["deps"] == ["csrc/cuda_kernel.cu", "csrc/shared.h"] + host = objects["csrc/host.cpp"] + assert not host["device"] and host["symbols"] == [] + assert host["deps"] == ["csrc/host.cpp", "csrc/shared.h"] + + +def test_missing_hipify_header_provenance_invalidates_every_object(compilation): + options, _ = compilation + manifest = options["hipify_map"] + manifest.write_text( + json.dumps( + { + k: v + for k, v in json.loads(manifest.read_text()).items() + if not k.endswith(".h") + } + ) + ) + result = producer.build_map(**options) + assert result["objects"] == [] + assert result["incomplete"] + assert "has no provenance" in result["reason"] + + +def test_missing_host_object_cannot_hide_its_header_dependencies(compilation): + options, objects = compilation + objects[1].unlink() + result = producer.build_map(**options) + assert result["objects"] == [] + assert result["incomplete"] + assert "missing compiled object" in result["reason"] + + +def test_host_object_omitted_by_compilation_database_invalidates_map(compilation): + options, _ = compilation + database = options["build_root"] / "compile_commands.json" + database.write_text(json.dumps(json.loads(database.read_text())[:1])) + result = producer.build_map(**options) + assert result["objects"] == [] + assert "absent from compilation databases" in result["reason"] + + +@pytest.mark.parametrize("conflicting", [False, True]) +def test_duplicate_object_records_must_agree_on_provenance(compilation, conflicting): + options, _ = compilation + database = options["build_root"] / "compile_commands.json" + commands = json.loads(database.read_text()) + duplicate = dict(commands[0]) + if conflicting: + duplicate["file"] = commands[1]["file"] + database.write_text(json.dumps([*commands, duplicate])) + result = producer.build_map(**options) + if conflicting: + assert result["objects"] == [] + assert "conflicting compilation records" in result["reason"] + else: + assert not result["incomplete"], result.get("reason") + assert result["stats"]["objects"] == 2 + + +def test_postbuild_source_edit_invalidates_old_provenance(compilation): + options, objects = compilation + source = options["source_root"] / "csrc" / "cuda_kernel.cu" + later = max(obj.stat().st_mtime_ns for obj in objects) + 1_000_000_000 + os.utime(source, ns=(later, later)) + result = producer.build_map(**options) + assert result["objects"] == [] + assert "newer than object" in result["reason"] + + +def test_extraction_failure_invalidates_the_whole_map(compilation, monkeypatch): + options, _ = compilation + + def unreadable(obj, tools): + raise ValueError("unrecognized fatbin") + + monkeypatch.setattr(producer, "device_symbols", unreadable) + result = producer.build_map(**options) + assert result["objects"] == [] + assert result["incomplete"] + assert result["reason"] == "unrecognized fatbin" + + +def test_objcopy_never_rewrites_the_compiled_object(tmp_path, monkeypatch): + commands = [] + + def fake_run(args): + commands.append(args) + if "--sections" in args: + return ELF_HEADER + ".hip_fatbin" + if "--list" in args: + return "hipv4-amdgcn-amd-amdhsa--gfx942\nhost-x86_64-unknown-linux-gnu-" + if "--symbols" in args: + return ELF + return "" + + monkeypatch.setattr(producer, "run", fake_run) + tools = { + name: name for name in ("llvm-readelf", "llvm-objcopy", "clang-offload-bundler") + } + obj = tmp_path / "probe.o" + obj.write_bytes(b"\x7fELF") + assert producer.device_symbols(obj, tools)[0] == [SYMBOL] + objcopy = next(args for args in commands if args[0] == "llvm-objcopy") + assert objcopy[-1] == os.devnull, "objcopy defaults to modifying its input in place" + + +@pytest.mark.parametrize("magic", [b"BC\xc0\xde", b"\xde\xc0\x17\x0b"]) +def test_bitcode_cannot_be_mistaken_for_host_only_elf(tmp_path, magic): + obj = tmp_path / "lto.hip.o" + obj.write_bytes(magic + bytes(64)) + with pytest.raises(ValueError, match="not an ELF object"): + producer.device_symbols(obj, {}) + + +@pytest.mark.parametrize( + "sections", + ["", ELF_HEADER + "__CLANG_OFFLOAD_BUNDLE__hip-amdgcn-amd-amdhsa--gfx942"], +) +def test_unsupported_native_formats_cannot_be_mistaken_for_host_only( + tmp_path, monkeypatch, sections +): + obj = tmp_path / "rdc.hip.o" + obj.write_bytes(b"\x7fELF") + monkeypatch.setattr(producer, "run", lambda args: sections) + with pytest.raises(ValueError, match="ELF header/sections|linked-code provenance"): + producer.device_symbols(obj, {"llvm-readelf": "llvm-readelf"}) + + +def test_automatic_build_manifest_is_used_without_manual_flag(compilation): + options, _ = compilation + options["hipify_map"].rename(options["build_root"] / "hipify-source-map.json") + del options["hipify_map"] + result = producer.build_map(**options) + assert not result["incomplete"], result.get("reason") + assert result["objects"][0]["source"] == "csrc/cuda_kernel.cu" + + +def test_installed_torch_hipified_headers_remain_external_dependencies(tmp_path): + header = tmp_path / "torch" / "HIPGuardImpl.h" + header.parent.mkdir() + header.write_text("// This file was automatically generated by hipify\n") + assert producer.source_path( + header, tmp_path / "vllm", [tmp_path / "build"], {} + ) == str(header) + + +def test_patched_vendor_header_keeps_both_original_and_patch_dependencies(compilation): + options, objects = compilation + build = options["build_root"] + source = options["source_root"] + patched = build / "csrc/shared.h" + manifest = options["hipify_map"] + provenance = json.loads(manifest.read_text()) + del provenance[str(patched)] + manifest.write_text(json.dumps(provenance)) + vendor = build.parent / "torch-header.h" + patch = source / "torch-fix.patch" + for path in (vendor, patch): + path.write_text("input to patched header\n") + earlier = objects[0].stat().st_mtime_ns - 1000 + os.utime(path, ns=(earlier, earlier)) + (build / "generated-source-deps.json").write_text( + json.dumps({str(patched): [str(vendor), str(patch)]}) + ) + result = producer.build_map(**options) + assert not result["incomplete"], result.get("reason") + native = next(obj for obj in result["objects"] if obj["device"]) + assert str(vendor) in native["deps"] + assert "torch-fix.patch" in native["deps"] + + +def test_cyclic_generated_provenance_is_rejected(tmp_path): + first, second = tmp_path / "first.h", tmp_path / "second.h" + with pytest.raises(ValueError, match="cyclic"): + producer.expand_dependencies(first, {first: {second}, second: {first}}) + + +def test_device_depfile_preserves_escaped_file_names(): + assert producer.parse_depfile( + "kernrec_device_deps: source.hip \\\n" + " path\\ with\\ spaces.h hash\\#dollar$$.h\n" + ) == ["source.hip", "path with spaces.h", "hash#dollar$.h"] + + +@pytest.mark.parametrize( + "text", + [ + "", + "other: file.h", + "kernrec_device_deps:", + "kernrec_device_deps: file.h\nother: more.h", + ], +) +def test_missing_or_ambiguous_device_depfile_is_rejected(text): + with pytest.raises(ValueError): + producer.parse_depfile(text) + + +def test_compiler_response_files_expand_relative_to_working_directory(tmp_path): + (tmp_path / "flags.rsp").write_text('@nested.rsp -o "object file.o"') + (tmp_path / "nested.rsp").write_text('-I"./include dir" -DDEVICE_HEADER=1') + command = {"directory": str(tmp_path), "arguments": ["clang++", "@flags.rsp"]} + assert producer.compilation_output(command) == tmp_path / "object file.o" + assert producer.compilation_arguments(command) == [ + "clang++", + "-I./include dir", + "-DDEVICE_HEADER=1", + "-o", + "object file.o", + ] + (tmp_path / "nested.rsp").write_text("@flags.rsp") + with pytest.raises(ValueError, match="cyclic"): + producer.compilation_arguments(command) + + +def test_response_quotes_follow_compiler_backslash_rules(): + assert producer.response_arguments(r'''"-DNAME=foo\q" '-DOTHER=bar\x' ""''') == [ + "-DNAME=fooq", + "-DOTHER=barx", + "", + ] + + +def test_compiler_frontend_diagnostics_unescape_shell_metacharacters(): + assert producer.response_arguments( + r'''"clang" "-D" "HEADER=\"dollar\$name.h\""''' + ) == ["clang", "-D", 'HEADER="dollar$name.h"'] + + +@pytest.mark.parametrize( + "error", + [ValueError("missing device target"), subprocess.TimeoutExpired("clang++", 180)], +) +def test_device_dependency_failure_invalidates_every_object( + compilation, monkeypatch, error +): + options, _ = compilation + + def fail(*args): + raise error + + monkeypatch.setattr(producer, "device_dependencies", fail) + result = producer.build_map(**options) + assert result["incomplete"] and result["objects"] == [] + + +@pytest.mark.parametrize("response_files", [False, True]) +def test_real_hip_device_headers_reach_all_dependent_kernels(tmp_path, response_files): + """Host depfiles omit device-only includes even when another TU includes them.""" + rocm = Path(os.environ.get("ROCM_PATH", "/opt/rocm")) + compiler = rocm / "llvm/bin/clang++" + tools = { + name: str(rocm / "llvm/bin" / name) + for name in ("llvm-readelf", "llvm-objcopy", "clang-offload-bundler") + } + ninja, cmake = shutil.which("ninja"), shutil.which("cmake") + if ( + not compiler.is_file() + or not all(Path(p).is_file() for p in tools.values()) + or not ninja + or not cmake + ): + pytest.skip("requires ROCm's HIP compiler and CMake/Ninja") + assert ninja is not None and cmake is not None + tools["ninja"] = ninja + source, build = tmp_path / "source with spaces", tmp_path / "build" + source.mkdir() + (source / "device only.h").write_text("constexpr int device_bias = 2;\n") + for arch in ("gfx942", "gfx950"): + (source / f"{arch}.h").write_text(f"constexpr int arch_bias = {arch[3:]};\n") + (source / "first.hip").write_text("""#include +#if defined(__HIP_DEVICE_COMPILE__) && ENABLE_DEVICE_HEADER +#include "device only.h" +#if defined(__gfx942__) +#include "gfx942.h" +#elif defined(__gfx950__) +#include "gfx950.h" +#endif +#endif +__global__ void first(int* output) { +#if defined(__HIP_DEVICE_COMPILE__) + output[0] = device_bias + arch_bias; +#endif +} +""") + (source / "second.hip").write_text("""#include +#include "device only.h" +__global__ void second(int* output) { output[0] = device_bias; } +""") + (source / "host.cpp").write_text("int host_only() { return 0; }\n") + (source / "CMakeLists.txt").write_text("""cmake_minimum_required(VERSION 3.26) +project(device_dependencies LANGUAGES CXX HIP) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +add_library(probe OBJECT first.hip second.hip host.cpp) +target_compile_definitions(probe PRIVATE ENABLE_DEVICE_HEADER=1) +""") + if response_files: + cmake_file = source / "CMakeLists.txt" + cmake_file.write_text( + cmake_file.read_text().replace( + "target_compile_definitions(probe PRIVATE ENABLE_DEVICE_HEADER=1)", + """file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/nested.rsp + "-DENABLE_DEVICE_HEADER=1") +file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/flags.rsp "@nested.rsp") +target_compile_options(probe PRIVATE @flags.rsp)""", + ) + ) + subprocess.run( + [ + cmake, + "-S", + str(source), + "-B", + str(build), + "-G", + "Ninja", + f"-DCMAKE_HIP_COMPILER={compiler}", + "-DCMAKE_HIP_ARCHITECTURES=gfx942:sramecc+:xnack-;gfx950:sramecc+:xnack-", + *(["-DCMAKE_NINJA_FORCE_RESPONSE_FILE=ON"] if response_files else []), + ], + check=True, + capture_output=True, + text=True, + ) + subprocess.run( + [cmake, "--build", str(build)], check=True, capture_output=True, text=True + ) + objects = {p: p.read_bytes() for p in build.rglob("*.o")} + result = producer.build_map(build, source, tools) + assert not result["incomplete"], result.get("reason") + by_source = {obj["source"]: obj for obj in result["objects"]} + assert {"device only.h", "gfx942.h", "gfx950.h"} <= set( + by_source["first.hip"]["deps"] + ) + assert "device only.h" in by_source["second.hip"]["deps"] + assert by_source["first.hip"]["symbols"] == ["_Z5firstPi"] + assert by_source["second.hip"]["symbols"] == ["_Z6secondPi"] + assert not by_source["host.cpp"]["device"] + assert all(p.read_bytes() == before for p, before in objects.items()) + + later = max(p.stat().st_mtime_ns for p in objects) + 1_000_000_000 + os.utime(source / "gfx942.h", ns=(later, later)) + stale = producer.build_map(build, source, tools) + assert stale["incomplete"] and stale["objects"] == [] + assert "newer than object" in stale["reason"] + + +@pytest.mark.parametrize("header", [r"a\b.h", "dollar$name.h"]) +def test_real_device_scan_preserves_outputs_and_response_header_choice( + tmp_path, header +): + compiler = Path(os.environ.get("ROCM_PATH", "/opt/rocm")) / "llvm/bin/clang++" + if not compiler.is_file(): + pytest.skip("requires ROCm's HIP compiler") + source = tmp_path / "probe.hip" + source.write_text("""#include +#ifdef __HIP_DEVICE_COMPILE__ +#include HEADER +#endif +__global__ void probe(int* output) { +#ifdef __HIP_DEVICE_COMPILE__ + output[0] = selected_value; +#endif +} +""") + (tmp_path / "ab.h").write_text("constexpr int selected_value = 7;\n") + (tmp_path / "dollar$name.h").write_text("constexpr int selected_value = 13;\n") + (tmp_path / "a").mkdir() + (tmp_path / "a/b.h").write_text("constexpr int selected_value = 11;\n") + (tmp_path / "flags.rsp").write_text(f'"-DHEADER=\\"{header}\\""') + args = [ + str(compiler), + "--offload-arch=gfx950", + "-c", + str(source), + "@flags.rsp", + "-o", + "probe.o", + "-MD", + "-MF", + "probe.d", + "-MT", + "probe.o", + "-MJ", + "fragment.json", + "-serialize-diagnostics", + "diagnostics.dia", + "-Xclang", + "-dependency-dot", + "-Xclang", + "dependencies.dot", + "-Xclang", + "-header-include-file", + "-Xclang", + "includes.txt", + "-Xclang", + "-diagnostic-log-file", + "-Xclang", + "diagnostics.log", + "-Xclang", + "-stats-file=stats.json", + ] + subprocess.run(args, cwd=tmp_path, check=True, capture_output=True, text=True) + before = { + p: (p.read_bytes(), p.stat().st_mtime_ns) + for p in tmp_path.rglob("*") + if p.is_file() + } + expanded = producer.compilation_arguments( + {"directory": str(tmp_path), "arguments": args} + ) + dependencies = producer.device_dependencies( + expanded, tmp_path, source, ["hipv4-amdgcn-amd-amdhsa--gfx950"] + ) + assert tmp_path / header.replace("\\", "") in dependencies + assert tmp_path / "a/b.h" not in dependencies + after = { + p: (p.read_bytes(), p.stat().st_mtime_ns) + for p in tmp_path.rglob("*") + if p.is_file() + } + assert after == before diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index 3177a10bba9b..6de6967b2d02 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -236,8 +236,32 @@ def main() -> int: ap.add_argument("--out", type=Path, required=True) ap.add_argument("--jobs", type=int, default=min(8, os.cpu_count() or 2)) ap.add_argument("--cuobjdump", default=None) + ap.add_argument("--backend", choices=("cuda", "rocm"), default="cuda") + ap.add_argument("--rocm-path", type=Path, default=Path("/opt/rocm")) + ap.add_argument("--hipify-map", type=Path) ap.add_argument("--commit", default="", help="commit the build is of") a = ap.parse_args() + if a.backend == "rocm": + from rocm_kernel_symbol_map import build_map + + tools = {} + for name in ("llvm-readelf", "llvm-objcopy", "clang-offload-bundler", "ninja"): + candidate = a.rocm_path / "llvm" / "bin" / name + tools[name] = ( + str(candidate) if candidate.is_file() else shutil.which(name) or name + ) + result = build_map( + a.build_root, + a.source_root, + tools, + commit=a.commit or os.environ.get("BUILDKITE_COMMIT", ""), + hipify_map=a.hipify_map, + ) + a.out.parent.mkdir(parents=True, exist_ok=True) + with gzip.open(a.out, "wt", encoding="utf-8") as stream: + json.dump(result, stream, separators=(",", ":")) + log(f"wrote {a.out}: {result['stats']}; {result.get('reason', '')}") + return 0 t0 = time.time() source_root = a.source_root.resolve() diff --git a/tools/ci/rocm_kernel_symbol_map.py b/tools/ci/rocm_kernel_symbol_map.py new file mode 100644 index 000000000000..04fa98ddc989 --- /dev/null +++ b/tools/ci/rocm_kernel_symbol_map.py @@ -0,0 +1,693 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build a ROCm kernel/source map from a completed CMake/Ninja build. + +Requires CMAKE_EXPORT_COMPILE_COMMANDS=ON, retained objects and .ninja_deps, +and ROCm's llvm-readelf, llvm-objcopy and clang-offload-bundler. HIPified +sources and copied headers use the generated-to-original manifests written +by the build, or an explicit --hipify-map. Relative paths use --source-root. +Provenance is recorded during generation, never guessed from path names. + +Device dependencies are collected by preprocessing every GPU target with +the original compiler's frontend arguments; Ninja only records host includes. +All compilation units are inspected, including host-only units. Missing +objects, dependency information, provenance, or unreadable device code +invalidate the entire map. This is an ahead-of-time map; an AITER JIT cache +without complete source and dependency provenance is not an input. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import gzip +import json +import os +import shlex +import shutil +import subprocess +import tempfile +import time +from pathlib import Path + +import regex as re + +DEPS_HEADER = re.compile(r"^(.+): #deps (\d+), deps mtime \d+ \((VALID|STALE)\)$") +TARGET = re.compile(r"(?:^|/)CMakeFiles/([^/]+)\.dir/") + + +def run(command: list[str], *, cwd: Path | None = None, diagnostics=False) -> str: + result = subprocess.run( + command, cwd=cwd, capture_output=True, text=True, timeout=180 + ) + if result.returncode: + detail = (result.stderr or result.stdout).strip()[-2000:] + raise ValueError( + f"{Path(command[0]).name} exited {result.returncode}: {detail}" + ) + return result.stdout + result.stderr if diagnostics else result.stdout + + +def parse_deps(text: str) -> dict[str, list[str]]: + """Read only complete, valid Ninja dependency records.""" + records: dict[str, list[str]] = {} + current: str | None = None + expected = 0 + + def finish(): + if current is not None and len(records[current]) != expected: + raise ValueError(f"incomplete dependency record: {current}") + + for line in text.splitlines(): + match = DEPS_HEADER.match(line) + if match: + finish() + current, count, status = match.groups() + if status != "VALID": + raise ValueError(f"stale dependency record: {current}") + if current in records: + raise ValueError(f"duplicate dependency record: {current}") + expected = int(count) + records[current] = [] + elif line.startswith(" ") and current is not None: + records[current].append(line.strip()) + elif not line.strip(): + finish() + current = None + else: + raise ValueError(f"unrecognized Ninja dependency output: {line}") + finish() + return records + + +def parse_kernel_symbols(text: str) -> list[str]: + """Join AMDGPU metadata descriptors to their defined kernel entry symbols.""" + if "AMDGPU Metadata:" not in text or "amdhsa.kernels:" not in text: + raise ValueError("no supported AMDGPU kernel metadata") + functions: set[str] = set() + descriptors: set[str] = set() + metadata: set[str] = set() + for line in text.splitlines(): + parts = line.split() + if len(parts) >= 8 and parts[0].rstrip(":").isdigit(): + _, _, _, kind, _, _, section, name = parts[:8] + if section == "UND": + continue + if kind == "FUNC": + functions.add(name) + elif kind == "OBJECT" and name.endswith(".kd"): + descriptors.add(name.removesuffix(".kd")) + match = re.match(r"^\s+\.symbol:\s+['\"]?([^\s'\"]+)['\"]?\s*$", line) + if match: + symbol = match.group(1) + if not symbol.endswith(".kd"): + raise ValueError(f"unsupported kernel descriptor: {symbol}") + metadata.add(symbol.removesuffix(".kd")) + if descriptors != metadata or not descriptors <= functions: + raise ValueError("kernel metadata and ELF entry symbols disagree") + if not metadata and not re.search(r"amdhsa\.kernels:\s*\[\s*\]", text): + raise ValueError("unrecognized AMDGPU kernel metadata") + return sorted(metadata) + + +def device_symbols(obj: Path, tools: dict[str, str]) -> tuple[list[str], list[str]]: + """Read HIP fatbins in a host object; retain names across every architecture.""" + with obj.open("rb") as stream: + magic = stream.read(4) + if magic != b"\x7fELF": + raise ValueError(f"not an ELF object (bitcode/LTO is unsupported): {obj}") + sections = run([tools["llvm-readelf"], "--sections", "--file-header", str(obj)]) + if ( + "ELF Header:" not in sections + or "Section Headers:" not in sections + or not re.search(r"^\s+Type:\s+REL\b", sections, re.MULTILINE) + ): + raise ValueError(f"cannot identify relocatable ELF header/sections: {obj}") + if any( + marker in sections + for marker in (".llvm.offloading", ".llvmbc", "__CLANG_OFFLOAD_BUNDLE__") + ): + raise ValueError( + "LLVM offloading/RDC/LTO objects require linked-code provenance" + ) + if ".hip_fatbin" not in sections: + if "AMDGPU" in sections: + raise ValueError("device-only/RDC objects require linked-code provenance") + return [], [] + with tempfile.TemporaryDirectory(prefix="rocm-symbol-map-") as temporary: + fatbin = Path(temporary) / "fatbin" + run( + [ + tools["llvm-objcopy"], + f"--dump-section=.hip_fatbin={fatbin}", + str(obj), + os.devnull, + ] + ) + targets = run( + [ + tools["clang-offload-bundler"], + "--type=o", + f"--input={fatbin}", + "--list", + ] + ).splitlines() + device_targets = [t for t in targets if t.startswith("hipv4-amdgcn-")] + unexpected = [ + t for t in targets if not t.startswith(("host-", "hipv4-amdgcn-")) + ] + if not device_targets or unexpected: + raise ValueError(f"unsupported offload bundles: {targets}") + symbols: set[str] = set() + for index, target in enumerate(device_targets): + code = Path(temporary) / f"device-{index}.hsaco" + run( + [ + tools["clang-offload-bundler"], + "--type=o", + f"--input={fatbin}", + "--unbundle", + f"--targets={target}", + f"--output={code}", + ] + ) + symbols.update( + parse_kernel_symbols( + run([tools["llvm-readelf"], "--symbols", "--notes", str(code)]) + ) + ) + return sorted(symbols), sorted(device_targets) + + +def absolute(path: str | Path, base: Path) -> Path: + return (base / path).resolve() + + +def relative(path: Path, source_root: Path) -> str: + try: + return path.relative_to(source_root).as_posix() + except ValueError: + return path.as_posix() + + +def load_provenance(path: Path | None, source_root: Path) -> dict[Path, Path]: + if path is None: + return {} + data = json.loads(path.read_text()) + if not isinstance(data, dict) or not all( + isinstance(k, str) and isinstance(v, str) for k, v in data.items() + ): + raise ValueError("HIPify manifest must map generated paths to original paths") + result: dict[Path, Path] = {} + for generated, original in data.items(): + key, value = absolute(generated, source_root), absolute(original, source_root) + if key in result and result[key] != value: + raise ValueError(f"ambiguous HIPify provenance: {key}") + if not key.is_file() or not value.is_file(): + raise ValueError( + f"missing file in HIPify provenance: {generated}: {original}" + ) + if not value.is_relative_to(source_root): + raise ValueError(f"HIPify original is outside source root: {original}") + result[key] = value + return result + + +def source_path( + path: Path, source_root: Path, build_dirs: list[Path], provenance: dict[Path, Path] +) -> str: + if path in provenance: + return relative(provenance[path], source_root) + if any(path.is_relative_to(build) for build in build_dirs): + raise ValueError(f"generated source/header has no provenance: {path}") + if not path.is_relative_to(source_root): + return path.as_posix() + # HIPify can also write alongside its input, outside the build tree. + if path.suffix in (".hip", ".cu", ".cpp", ".cc", ".c", ".h", ".hpp", ".cuh"): + with path.open("rb") as stream: + if b"automatically generated by hipify" in stream.read(256): + raise ValueError(f"HIPified source/header has no provenance: {path}") + return relative(path, source_root) + + +def generated_dependencies( + build_root: Path, source_root: Path +) -> dict[Path, set[Path]]: + """Read explicit inputs for build-generated headers, including vendor patches.""" + result: dict[Path, set[Path]] = {} + for manifest in build_root.rglob("generated-source-deps.json"): + data = json.loads(manifest.read_text()) + if not isinstance(data, dict): + raise ValueError(f"invalid generated dependency manifest: {manifest}") + for output, inputs in data.items(): + if ( + not isinstance(output, str) + or not isinstance(inputs, list) + or not inputs + ): + raise ValueError(f"invalid generated dependencies: {manifest}") + if not all(isinstance(path, str) for path in inputs): + raise ValueError(f"invalid generated dependency paths: {manifest}") + output = absolute(output, source_root) + paths = {absolute(path, source_root) for path in inputs} + if output in result and result[output] != paths: + raise ValueError(f"conflicting generated dependencies: {output}") + result[output] = paths + return result + + +def expand_dependencies( + path: Path, generated: dict[Path, set[Path]], active=() +) -> set[Path]: + if path not in generated: + return {path} + if path in active: + raise ValueError(f"cyclic generated dependencies: {path}") + return { + original + for dependency in generated[path] + for original in expand_dependencies(dependency, generated, (*active, path)) + } + + +def response_arguments(text: str) -> list[str]: + """Clang's GNU response syntax escapes the next character in either quote.""" + arguments, word = [], [] + quote = None + started = False + chars = iter(text) + for char in chars: + if char == "\\": + escaped = next(chars, None) + if escaped is None: + raise ValueError("unfinished compiler response escape") + word.append(escaped) + started = True + elif quote: + if char == quote: + quote = None + else: + word.append(char) + elif char in ("'", '"'): + quote = char + started = True + elif char in " \t\r\n\v\f": + if started: + arguments.append("".join(word)) + word, started = [], False + else: + word.append(char) + started = True + if quote: + raise ValueError("unfinished compiler response quote") + if started: + arguments.append("".join(word)) + return arguments + + +def compilation_arguments( + command: dict, response_files: set[Path] | None = None +) -> list[str]: + args = command.get("arguments") or shlex.split(command.get("command", "")) + if not isinstance(args, list) or not all(isinstance(a, str) for a in args): + raise ValueError("invalid compiler arguments") + directory = Path(command["directory"]).resolve() + + def expand(arguments, active=()): + expanded = [] + for arg in arguments: + if not arg.startswith("@"): + expanded.append(arg) + continue + response = absolute(arg[1:], directory) + if response in active or len(active) >= 16: + raise ValueError(f"cyclic or deeply nested response file: {response}") + if response_files is not None: + response_files.add(response) + expanded.extend( + expand(response_arguments(response.read_text()), (*active, response)) + ) + return expanded + + expanded = expand(args) + if any(arg.startswith("--rsp-quoting") for arg in expanded): + raise ValueError("explicit compiler response quoting is unsupported") + return expanded + + +def compilation_output(command: dict) -> Path: + directory = Path(command["directory"]).resolve() + output = command.get("output") + args = compilation_arguments(command) + if "-fgpu-rdc" in args or any(a.startswith("-flto") for a in args): + raise ValueError("RDC/LTO needs linked-code provenance") + if not output: + for index, arg in enumerate(args[:-1]): + if arg == "-o": + output = args[index + 1] + break + if not output: + raise ValueError(f"compilation has no object output: {command.get('file')}") + return absolute(output, directory) + + +def parse_depfile(text: str) -> list[str]: + """Read Clang's single-rule Make depfile, including escaped file names.""" + text = text.replace("\\\r\n", "").replace("\\\n", "") + prefix = "kernrec_device_deps:" + if not text.startswith(prefix): + raise ValueError("unexpected device dependency target") + paths, word = [], [] + chars = iter(text[len(prefix) :]) + for char in chars: + if char == "\\": + escaped = next(chars, "") + if escaped not in (" ", "\t", "#", ":", "\\"): + raise ValueError("unsupported device dependency escape") + word.append(escaped) + elif char == "$": + if next(chars, "") != "$": + raise ValueError("unsupported device dependency variable") + word.append("$") + elif char in "#:;\0": + raise ValueError("unexpected device dependency rule") + elif char.isspace(): + if word: + paths.append("".join(word)) + word = [] + else: + word.append(char) + if word: + paths.append("".join(word)) + if not paths: + raise ValueError("empty device dependencies") + return paths + + +def device_dependencies( + args: list[str], directory: Path, source: Path, targets: list[str] +) -> set[Path]: + """Preprocess the actual AMDGPU frontends without writing build outputs.""" + if not targets: + return set() + if not args: + raise ValueError(f"missing compiler command for device source: {source}") + expected = set() + for target in targets: + match = re.fullmatch( + r"hipv4-amdgcn-amd-amdhsa--(gfx[0-9a-z]+)((?::(?:sramecc|xnack)[+-])*)", + target, + ) + if not match: + raise ValueError(f"unsupported device dependency target: {target}") + expected.add((match[1], tuple(sorted(filter(None, match[2].split(":")))))) + + def option(frontend, name): + if name not in frontend or frontend.index(name) + 1 == len(frontend): + raise ValueError(f"missing device compiler argument: {name}") + return frontend[frontend.index(name) + 1] + + # Clang creates/removes driver output files even with -###. + driver = [] + arguments = iter(args) + for arg in arguments: + if arg in ("-MJ", "-gen-cdb-fragment-path", "-serialize-diagnostics"): + if next(arguments, None) is None: + raise ValueError(f"missing compiler output argument: {arg}") + elif arg.startswith( + ("-MJ", "-gen-cdb-fragment-path=", "-serialize-diagnostics=") + ): + continue + else: + driver.append(arg) + frontends = {} + for line in run([*driver, "-###"], cwd=directory, diagnostics=True).splitlines(): + if '"-cc1"' not in line: + continue + frontend = response_arguments(line) + if "-fcuda-is-device" not in frontend: + continue + + if ( + option(frontend, "-triple") != "amdgcn-amd-amdhsa" + or "-emit-obj" not in frontend + ): + raise ValueError("unsupported device compiler frontend") + features = [] + for index, arg in enumerate(frontend[:-1]): + if arg == "-target-feature": + feature = frontend[index + 1] + if feature.startswith(("+", "-")) and feature[1:] in ( + "sramecc", + "xnack", + ): + features.append(feature[1:] + feature[0]) + arch = ( + option(frontend, "-target-cpu"), + tuple(sorted(features)), + ) + if arch in frontends: + raise ValueError(f"duplicate device compiler frontend: {arch}") + frontends[arch] = frontend + if set(frontends) != expected: + raise ValueError("compiler device targets do not match object offload targets") + + dependencies = set() + with tempfile.TemporaryDirectory(prefix="rocm-device-deps-") as temporary: + for index, frontend in enumerate(frontends.values()): + command = [] + arguments = iter(frontend) + for arg in arguments: + if arg in ( + "-o", + "-dependency-file", + "-MT", + "-MQ", + "-serialize-diagnostic-file", + "-coverage-notes-file", + "-coverage-data-file", + "-split-dwarf-output", + "-opt-record-file", + "-dependency-dot", + "-diagnostic-log-file", + "-header-include-file", + ): + if next(arguments, None) is None: + raise ValueError(f"missing compiler output argument: {arg}") + elif arg in ("-emit-obj", "-MP") or arg.startswith( + ("-ftime-trace", "-stats-file=") + ): + continue + elif arg.startswith(("-fmodule", "-include-pch")): + raise ValueError("device PCH/modules need expanded dependencies") + else: + command.append(arg) + depfile = Path(temporary) / f"device-{index}.d" + run( + [ + *command, + "-Eonly", + "-dependency-file", + str(depfile), + "-MT", + "kernrec_device_deps", + "-sys-header-deps", + ], + cwd=directory, + ) + paths = {absolute(p, directory) for p in parse_depfile(depfile.read_text())} + if source not in paths: + raise ValueError(f"device dependencies omit compiled source: {source}") + dependencies.update(paths) + return dependencies + + +def build_map( + build_root: Path, + source_root: Path, + tools: dict[str, str], + *, + commit: str = "", + hipify_map: Path | None = None, +) -> dict: + start = time.monotonic() + source_root, build_root = source_root.resolve(), build_root.resolve() + result: dict = { + "version": 2, + "backend": "rocm", + "commit": commit, + "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"), + "incomplete": False, + "objects": [], + "errors": [], + "build_dirs": [], + "tools": {}, + } + try: + for name, executable in tools.items(): + version = run([executable, "--version"]).splitlines() + if not version: + raise ValueError(f"no version information from {name}") + result["tools"][name] = version[0] + databases = sorted(build_root.rglob("compile_commands.json")) + if not databases: + raise ValueError( + "no compile_commands.json; enable CMAKE_EXPORT_COMPILE_COMMANDS" + ) + build_dirs = sorted({build_root, *(p.parent for p in databases)}) + result["build_dirs"] = [relative(p, source_root) for p in build_dirs] + manifests = ( + [hipify_map] + if hipify_map + else sorted(build_root.rglob("hipify-source-map.json")) + ) + provenance: dict[Path, Path] = {} + for manifest in manifests: + for generated, original in load_provenance(manifest, source_root).items(): + if generated in provenance and provenance[generated] != original: + raise ValueError(f"conflicting HIPify provenance: {generated}") + provenance[generated] = original + result["hipify_manifests"] = [relative(p, source_root) for p in manifests] + generated = generated_dependencies(build_root, source_root) + seen: dict[Path, tuple] = {} + for database in databases: + build_dir = database.parent + if not (build_dir / ".ninja_deps").is_file(): + raise ValueError(f"missing Ninja dependency log: {build_dir}") + dependencies = { + absolute(obj, build_dir): paths + for obj, paths in parse_deps( + run([tools["ninja"], "-C", str(build_dir), "-t", "deps"]) + ).items() + } + commands = json.loads(database.read_text()) + if not isinstance(commands, list) or not commands: + raise ValueError(f"empty or invalid compilation database: {database}") + for command in commands: + obj = compilation_output(command) + directory = Path(command["directory"]).resolve() + source = absolute(command["file"], directory) + response_files: set[Path] = set() + args = compilation_arguments(command, response_files) + identity = (source, directory, tuple(args)) + if obj in seen: + if seen[obj] != identity: + raise ValueError(f"conflicting compilation records: {obj}") + continue + seen[obj] = identity + if not obj.is_file(): + raise ValueError(f"missing compiled object: {obj}") + paths = dependencies.get(obj) + if not paths: + raise ValueError(f"missing object dependencies: {obj}") + deps = {absolute(p, build_dir) for p in paths} + if source not in deps: + raise ValueError(f"dependency log omits compiled source: {obj}") + symbols, targets = device_symbols(obj, tools) + deps.update(device_dependencies(args, directory, source, targets)) + original_deps = { + original + for path in deps + for original in expand_dependencies(path, generated) + } + for path in ( + deps + | original_deps + | response_files + | {provenance[p] for p in original_deps if p in provenance} + ): + if not path.is_file(): + raise ValueError(f"missing source/header: {path}") + if path.stat().st_mtime_ns > obj.stat().st_mtime_ns: + raise ValueError(f"source/header is newer than object: {path}") + mapped_deps = sorted( + { + source_path(p, source_root, build_dirs, provenance) + for p in original_deps + } + ) + mapped_source = source_path(source, source_root, build_dirs, provenance) + target = TARGET.search(obj.as_posix()) + result["objects"].append( + { + "source": mapped_source, + "compiled_source": relative(source, source_root), + "target": target.group(1) if target else "", + "object": relative(obj, build_dir), + "device": bool(targets), + "symbols": symbols, + "deps": mapped_deps, + "offload_targets": targets, + } + ) + for obj in build_root.rglob("*.o"): + if any( + "CompilerId" in part or part == "CMakeScratch" for part in obj.parts + ): + continue + if obj.resolve() not in seen: + raise ValueError(f"object is absent from compilation databases: {obj}") + result["objects"].sort(key=lambda entry: (entry["source"], entry["target"])) + except ( + OSError, + ValueError, + KeyError, + TypeError, + subprocess.SubprocessError, + ) as error: + result["objects"] = [] + result["incomplete"] = True + result["reason"] = str(error) + result["errors"].append({"error": str(error)}) + result["stats"] = { + "objects": len(result["objects"]), + "device_objects": sum(bool(o["device"]) for o in result["objects"]), + "symbols": sum(len(o["symbols"]) for o in result["objects"]), + "seconds": round(time.monotonic() - start, 2), + } + result["offload_targets"] = sorted( + {target for obj in result["objects"] for target in obj["offload_targets"]} + ) + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--build-root", type=Path, required=True) + parser.add_argument("--source-root", type=Path, required=True) + parser.add_argument("--hipify-map", type=Path) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--commit", default=os.environ.get("BUILDKITE_COMMIT", "")) + parser.add_argument("--rocm-path", type=Path, default=Path("/opt/rocm")) + args = parser.parse_args() + tools = {} + for name in ("llvm-readelf", "llvm-objcopy", "clang-offload-bundler", "ninja"): + candidate = args.rocm_path / "llvm" / "bin" / name + tools[name] = ( + str(candidate) if candidate.is_file() else shutil.which(name) or name + ) + result = build_map( + args.build_root, + args.source_root, + tools, + commit=args.commit, + hipify_map=args.hipify_map, + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + with gzip.open(args.out, "wt", encoding="utf-8") as stream: + json.dump(result, stream, separators=(",", ":")) + print( + json.dumps( + { + "out": str(args.out), + "stats": result["stats"], + "reason": result.get("reason", ""), + } + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From ae065918640a198a56f6f17eb45e251c3229fe41 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:32:26 -0400 Subject: [PATCH 003/631] [Kimi K3] KDA RecoverSSM: CuTe verify and commit kernels, records outside the KV cache (#63045) Signed-off-by: wzhao18 Signed-off-by: Wei Zhao Co-authored-by: Claude Opus 5.5 --- tests/models/kimi_k3/test_kda.py | 40 +- tests/models/kimi_k3/test_kda_metadata.py | 51 +- .../models/kimi_k3/test_recoverssm_kernels.py | 722 ++++++++++++++++++ vllm/config/vllm.py | 4 + .../layers/mamba/mamba_utils.py | 30 - vllm/models/kimi_k3/nvidia/kda.py | 259 +++++-- vllm/models/kimi_k3/nvidia/kda_metadata.py | 1 + vllm/models/kimi_k3/nvidia/model.py | 18 +- vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 126 ++- .../kimi_k3/nvidia/recoverssm/__init__.py | 2 + .../nvidia/recoverssm/cutedsl/__init__.py | 2 + .../nvidia/recoverssm/cutedsl/commit.py | 292 +++++++ .../nvidia/recoverssm/cutedsl/verify.py | 526 +++++++++++++ .../nvidia/recoverssm/cutedsl/verify_ws.py | 722 ++++++++++++++++++ .../kimi_k3/nvidia/recoverssm/kernels.py | 374 +++++++++ 15 files changed, 2957 insertions(+), 212 deletions(-) create mode 100644 tests/models/kimi_k3/test_recoverssm_kernels.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/commit.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify_ws.py create mode 100644 vllm/models/kimi_k3/nvidia/recoverssm/kernels.py diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 0b250654450e..2dfc08eb0497 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.mamba.kda_checkpoint import ( FlashKDAPrefillCheckpointExporter, ) +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator 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, @@ -45,6 +46,7 @@ from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( KDARecoverSSMCommitContext, + KDARecoverSSMRecords, kda_recoverssm_verify, ) from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( @@ -134,16 +136,15 @@ def test_kda_recoverssm_config_state_layout(): speculative_config=SimpleNamespace(num_speculative_tokens=2), ) + # The verify records are kept outside the KV cache: the Mamba page holds the + # stock conv and SSM state. assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == ( torch.bfloat16, torch.float32, - torch.float32, - torch.bfloat16, - ) - assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == ( - (4, 3, 32), - (4, 3, 64), ) + assert KimiLinearForCausalLM.get_mamba_state_shape_from_config( + vllm_config + ) == MambaStateShapeCalculator.kda_state_shape(1, 4, 32, num_spec=2) @torch.inference_mode() @@ -799,29 +800,17 @@ def test_kda_recoverssm_verify_and_group_commit( else (num_blocks, history_len + query_len - 1, conv_dim) ) conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE) - correction_cache = torch.empty( - num_blocks, + records = KDARecoverSSMRecords.allocate( + global_num_accepted.shape[0], num_heads, query_len, dim, - dtype=torch.float32, - device=DEVICE, - ) - kg_cache = torch.empty( - num_blocks, - num_heads, - query_len, - 2 * dim, - dtype=torch.bfloat16, - device=DEVICE, + torch.bfloat16, + DEVICE, ) layer = SimpleNamespace( - kv_cache=( - conv_state, - checkpoint, - correction_cache, - kg_cache, - ), + kv_cache=(conv_state, checkpoint), + recoverssm_records=records, A_log=A_log, dt_bias=dt_bias, local_num_heads=num_heads, @@ -842,8 +831,7 @@ def test_kda_recoverssm_verify_and_group_commit( dt_bias=dt_bias, lower_bound=lower_bound, checkpoint_state=checkpoint, - correction_cache=correction_cache, - kg_cache=kg_cache, + records=records, query_start_loc=query_start_loc, state_indices=state_indices, spec_query_len=query_len, diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index f71633438309..d46aab2f6acd 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -15,6 +15,7 @@ ) from vllm.config import SpeculativeConfig from vllm.config.compilation import CUDAGraphMode +from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.kda_metadata import ( KimiK3KDAAttentionBackend, KimiK3KDAMetadata, @@ -23,6 +24,8 @@ stage_spec_decode_metadata, ) from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, @@ -185,7 +188,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): ) # This is the same compatibility check performed while initializing the - # model runner. RecoverSSM adds two workspace states that are not copied. + # model runner. RecoverSSM keeps its verify records outside the KV cache. validate_mamba_state_copy_funcs( {kv_cache_spec: [0]}, { @@ -194,7 +197,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): ) }, ) - assert len(kv_cache_spec.shapes) == 4 + assert len(kv_cache_spec.shapes) == 2 builder_config = SimpleNamespace( model_config=SimpleNamespace( @@ -258,7 +261,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): # Creating the commit context is the first metadata path that consumes the # builder's layer_names. Mock only the GPU-cache-dependent constructor. builder.recoverssm_context = None - forward_layer = Mock() + forward_layer = Mock(recoverssm_backend="triton") builder.vllm_config.compilation_config.static_forward_context["layer.0"] = ( forward_layer ) @@ -293,6 +296,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): [forward_layer], spec_query_len=3, max_num_reqs=builder.vllm_config.scheduler_config.max_num_seqs, + backend="triton", ) @@ -1087,3 +1091,44 @@ def test_cudagraph_capture_batch_stays_decode_only(): assert staged is not None assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr() torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0]) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Kimi-K3 NVIDIA KDA layer") +def test_recoverssm_layer_keeps_records_outside_kv_cache( + monkeypatch: pytest.MonkeyPatch, default_vllm_config +): + """A RecoverSSM KDA layer stores the stock conv and SSM state in its Mamba + pages and its verify records in a buffer with one row per schedulable + request.""" + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", + lambda: SimpleNamespace(world_size=1, rank_in_group=0), + ) + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + default_vllm_config.speculative_config = SimpleNamespace(num_speculative_tokens=2) + default_vllm_config.scheduler_config.max_num_seqs = 5 + default_vllm_config.cache_config.use_replayssm = True + default_vllm_config.cache_config.use_kda_recoverssm = True + config = KimiLinearConfig( + hidden_size=16, + linear_attn_config={ + "kda_layers": [0], + "full_attn_layers": [], + "head_dim": 32, + "num_heads": 4, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + }, + ) + layer = KimiK3DeltaAttention( + config, default_vllm_config, prefix="model.layers.0.self_attn" + ) + # The CuTe kernels need 128-wide heads: this layer runs the Triton ones. + assert layer.recoverssm_backend == "triton" + assert len(layer.get_state_shape()) == len(layer.get_state_dtype()) == 2 + records = layer.recoverssm_records + assert records is not None + assert records.correction.shape == (5, 4, 3, 32) + assert records.correction.dtype == torch.float32 + assert records.key_gate.shape == (5, 4, 3, 64) + assert records.key_gate.dtype == torch.bfloat16 diff --git a/tests/models/kimi_k3/test_recoverssm_kernels.py b/tests/models/kimi_k3/test_recoverssm_kernels.py new file mode 100644 index 000000000000..adad82174ad3 --- /dev/null +++ b/tests/models/kimi_k3/test_recoverssm_kernels.py @@ -0,0 +1,722 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The CuTe KDA RecoverSSM verify and state commit against the Triton reference +(``causal_conv1d_update`` + ``kda_recoverssm_verify``, and the Triton commit), on +Mamba page views as vLLM binds them (one padded page per slot holding the conv +and SSM state) and per-layer verify records indexed by batch row.""" + +from math import prod +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) +): + pytest.skip("The CuTe RecoverSSM kernels need SM10x", allow_module_level=True) +pytest.importorskip("cutlass") + +from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( # noqa: E402 + causal_conv1d_update, +) +from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( # noqa: E402 + KDARecoverSSMCommitContext, + KDARecoverSSMRecords, + kda_recoverssm_verify, +) +from vllm.models.kimi_k3.nvidia.recoverssm.kernels import ( # noqa: E402 + CuteRecoverSSMVerify, + compile_state_commit, +) + +D = 128 +LOWER_BOUND = -5.0 +DEVICE = torch.device("cuda") +# Extra bytes per page: vLLM pads Mamba pages to the attention page size. +PAGE_PADDING = 4096 + + +def _pages(slots: int, heads: int, window: int, dim_first: bool): + """Page views of (conv state, SSM state) with a padded slot stride, filled + with random values.""" + conv_dim, conv_len = 3 * heads * D, window + 2 + conv_shape = ( + (slots, conv_dim, conv_len) if dim_first else (slots, conv_len, conv_dim) + ) + specs = [ + (torch.bfloat16, conv_shape), + (torch.bfloat16, (slots, heads, D, D)), + ] + sizes = [dtype.itemsize * prod(shape[1:]) for dtype, shape in specs] + page = sum(sizes) + PAGE_PADDING + raw = torch.randn(slots, page // 4, device=DEVICE).view(torch.uint8) + views = [] + for i, (dtype, shape) in enumerate(specs): + inner = torch.empty(shape[1:], device="meta").stride() + views.append( + raw.view(dtype) + .view(-1) + .as_strided( + shape, + (page // dtype.itemsize, *inner), + sum(sizes[:i]) // dtype.itemsize, + ) + ) + conv, state = views + # Bounded values: random bits are not finite floats. + conv.copy_(torch.randn_like(conv)) + state.copy_(torch.randn_like(state) * 0.1) + return (raw, *views) + + +def _records(rows: int, heads: int, window: int) -> KDARecoverSSMRecords: + records = KDARecoverSSMRecords.allocate( + rows, heads, window, D, torch.bfloat16, DEVICE + ) + records.correction.normal_() + records.key_gate.normal_() + return records + + +def _clone_records(records: KDARecoverSSMRecords) -> KDARecoverSSMRecords: + return KDARecoverSSMRecords(records.correction.clone(), records.key_gate.clone()) + + +def _clone_pages(pages): + raw = pages[0].clone() + base = pages[0].data_ptr() + return ( + raw, + *( + raw.view(t.dtype) + .view(-1) + .as_strided(t.shape, t.stride(), (t.data_ptr() - base) // t.element_size()) + for t in pages[1:] + ), + ) + + +def _inputs( + batch: int, heads: int, window: int, varlen: bool, row_pad: int, row_offset: int +): + """Raw projections laid out as the KDA in-projection output: q/k/v, then the + output gate, f_a, beta and the alignment padding in one row; the spec-decode + tokens start ``row_offset`` rows in (after a batch's non-spec tokens).""" + lens = torch.full((batch,), window, dtype=torch.int32) + if varlen and batch > 2: + lens[0] = window - 2 + lens[1] = 0 + query_start = torch.zeros(batch + 1, dtype=torch.int32) + query_start[1:] = torch.cumsum(lens, 0) + tokens = int(query_start[-1]) + width = 4 * heads * D + D + heads + row_pad + projected = torch.randn( + row_offset + tokens, width, device=DEVICE, dtype=torch.bfloat16 + )[row_offset:] + mixed_qkv = projected[:, : 3 * heads * D] + raw_beta = projected[:, 4 * heads * D + D : 4 * heads * D + D + heads] + raw_gate = torch.randn(tokens, heads * D, device=DEVICE, dtype=torch.bfloat16) + return lens, query_start.to(DEVICE), mixed_qkv, raw_gate, raw_beta + + +def _compiled_verify(heads, window, dim_first, conv_weight_dtype): + verify = CuteRecoverSSMVerify( + num_heads=heads, + window=window, + conv_dim_first=dim_first, + conv_weight_dtype=conv_weight_dtype, + device=DEVICE, + ) + verify.compile() + return verify + + +def _check_verify( + heads=96, + dim_first=True, + conv_weight_dtype=torch.float32, + batch=8, + varlen=False, + row_pad=0, + recoverssm_accepted=True, + window=5, + row_offset=0, + dead_slot=0, +): + """Run the CuTe verify and the Triton reference on the same pages; returns + whether the CuTe verify launched (it declines unaligned inputs).""" + torch.manual_seed(batch * 7 + heads + row_pad + window) + slots = 3 * batch + 1 + lens, query_start, mixed_qkv, raw_gate, raw_beta = _inputs( + batch, heads, window, varlen, row_pad, row_offset + ) + tokens = mixed_qkv.shape[0] + A_log = torch.randn(heads, device=DEVICE) * 0.5 + dt_bias = torch.randn(heads * D, device=DEVICE) * 0.5 + conv_weight = (torch.randn(3 * heads * D, 4, device=DEVICE) * 0.5).to( + conv_weight_dtype + ) + accepted = ( + torch.ones(batch, dtype=torch.int32, device=DEVICE) + if recoverssm_accepted + else torch.randint(1, window + 1, (batch,), dtype=torch.int32, device=DEVICE) + ) + state_indices = torch.randperm(slots - 1, device=DEVICE)[:batch].to(torch.int32) + 1 + if varlen and batch > 4: + # Padded requests: CUDA graph batches pad with the null block. + state_indices[3:5] = dead_slot + + pages = _pages(slots, heads, window, dim_first) + ref = _clone_pages(pages) + initial = _clone_pages(pages) + _, conv, state = pages + _, ref_conv, ref_state = ref + # More rows than requests: rows past the batch are left alone. + records = _records(batch + 2, heads, window) + ref_records = _clone_records(records) + initial_records = _clone_records(records) + + # Reference: the stock convolution, then the Triton verify. (The stock + # convolution writes out of bounds for slot -1, so it gets the null block.) + ref_indices = state_indices.clamp(min=0) + ref_out = torch.full( + (1, tokens, heads, D), 7.0, device=DEVICE, dtype=torch.bfloat16 + ) + conv_out = torch.empty_like(mixed_qkv) + causal_conv1d_update( + mixed_qkv, + ref_conv if dim_first else ref_conv.transpose(-1, -2), + conv_weight, + None, + activation="silu", + conv_state_indices=ref_indices, + num_accepted_tokens=accepted, + query_start_loc=query_start, + max_query_len=window, + validate_data=False, + out=conv_out, + ) + q, k, v = ( + x.unflatten(-1, (heads, D)).unsqueeze(0) for x in conv_out.split(heads * D, -1) + ) + kda_recoverssm_verify( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, tokens, heads, D), + raw_beta=raw_beta.unsqueeze(0), + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + checkpoint_state=ref_state, + records=ref_records, + query_start_loc=query_start, + state_indices=ref_indices, + spec_query_len=window, + out=ref_out, + ) + + verify = _compiled_verify(heads, window, dim_first, conv_weight_dtype) + out = torch.full((tokens, heads * D), 7.0, device=DEVICE, dtype=torch.bfloat16) + launched = verify( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + state=state, + records=records, + query_start_loc=query_start, + state_indices=state_indices, + num_accepted_tokens=accepted, + conv_state=conv, + conv_weight=conv_weight, + out=out, + ) + torch.accelerator.synchronize() + if not launched: + assert (out == 7.0).all() + for page, before in zip(pages[1:], initial[1:]): + assert torch.equal(page, before) + return False + assert int(verify.counter) == 0 # left zeroed for the next launch + + live = [b for b in range(batch) if int(state_indices[b]) > 0 and int(lens[b]) > 0] + bounds = query_start.tolist() + live_tokens = [t for b in live for t in range(bounds[b], bounds[b + 1])] + dead_tokens = [ + t + for b in range(batch) + if int(state_indices[b]) <= 0 + for t in range(bounds[b], bounds[b + 1]) + ] + ref_out = ref_out.view(tokens, heads * D) + torch.testing.assert_close( + out[live_tokens].float(), ref_out[live_tokens].float(), atol=3e-2, rtol=3e-2 + ) + if dead_tokens: + assert (out[dead_tokens] == 0).all() + for b in live: + slot, n = int(state_indices[b]), int(lens[b]) + torch.testing.assert_close( + records.correction[b, :, :n], + ref_records.correction[b, :, :n], + atol=3e-2, + rtol=3e-2, + ) + torch.testing.assert_close( + records.key_gate[b, :, :n], ref_records.key_gate[b, :, :n] + ) + # The rolled conv state: two history rows, then the raw tokens. + if dim_first: + torch.testing.assert_close( + conv[slot, :, : n + 2], ref_conv[slot, :, : n + 2] + ) + else: + torch.testing.assert_close(conv[slot, : n + 2], ref_conv[slot, : n + 2]) + # The checkpoint is never written; other slots' pages and other rows' + # records are untouched. + torch.testing.assert_close(state, ref_state, atol=0, rtol=0) + untouched = sorted(set(range(slots)) - {int(state_indices[b]) for b in live}) + for page, before in zip(pages[1:], initial[1:]): + assert torch.equal(page[untouched], before[untouched]) + for record, before in zip( + (records.correction, records.key_gate), + (initial_records.correction, initial_records.key_gate), + ): + assert torch.equal(record[batch:], before[batch:]) + return True + + +@pytest.mark.parametrize("heads", [96, 12]) +@pytest.mark.parametrize("dim_first", [True, False]) +@pytest.mark.parametrize("conv_weight_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("batch,varlen", [(1, False), (8, False), (13, True)]) +@pytest.mark.parametrize("row_pad", [0, 3]) +@pytest.mark.parametrize("recoverssm_accepted", [True, False]) +@torch.inference_mode() +def test_cute_verify_matches_triton( + heads, dim_first, conv_weight_dtype, batch, varlen, row_pad, recoverssm_accepted +): + """``row_pad`` 3 misaligns the projection rows, so the dim-first layout falls + back from the warp-specialized verify to the plain tcgen05 one. + ``recoverssm_accepted`` feeds RecoverSSM's accepted counts (always 1: the + commit compacts the conv state); otherwise the stock spec-decode counts.""" + assert _check_verify( + heads, dim_first, conv_weight_dtype, batch, varlen, row_pad, recoverssm_accepted + ) + + +@pytest.mark.parametrize("window", [2, 8]) +@pytest.mark.parametrize("dim_first", [True, False]) +@torch.inference_mode() +def test_cute_verify_window_sizes(window, dim_first): + """1 and 7 speculative tokens (the 1-8 token windows the kernels support).""" + assert _check_verify(window=window, dim_first=dim_first, batch=13, varlen=True) + + +@pytest.mark.parametrize("row_pad,launches", [(0, True), (3, False)]) +@torch.inference_mode() +def test_cute_verify_token_offset(row_pad, launches): + """Spec-decode tokens after a mixed batch's other tokens start one row into + the projection: with 16-byte aligned rows the CuTe verify runs, otherwise it + declines (writing nothing) and the caller runs the Triton verify.""" + assert _check_verify(row_pad=row_pad, row_offset=1) == launches + + +@torch.inference_mode() +def test_cute_verify_pad_slot_id(): + """Requests padded with PAD_SLOT_ID (-1) are dead like null-block ones.""" + assert _check_verify(batch=13, varlen=True, dead_slot=-1) + + +@torch.inference_mode() +def test_cute_verify_rejects_unsupported_state_pages(): + """The CuTe verify rejects state pages the TMA cannot stream: padded value + rows, or a page off a 16-byte boundary.""" + heads, window, batch = 16, 5, 2 + _, query_start, mixed_qkv, raw_gate, raw_beta = _inputs( + batch, heads, window, False, 0, 0 + ) + tokens = mixed_qkv.shape[0] + _, conv, state = _pages(4, heads, window, True) + verify = _compiled_verify(heads, window, True, torch.float32) + args = dict( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=torch.zeros(heads, device=DEVICE), + dt_bias=torch.zeros(heads * D, device=DEVICE), + lower_bound=LOWER_BOUND, + state=state, + records=_records(batch, heads, window), + query_start_loc=query_start, + state_indices=torch.arange(1, batch + 1, dtype=torch.int32, device=DEVICE), + num_accepted_tokens=torch.ones(batch, dtype=torch.int32, device=DEVICE), + conv_state=conv, + conv_weight=torch.zeros(3 * heads * D, 4, device=DEVICE), + out=torch.empty(tokens, heads * D, device=DEVICE, dtype=torch.bfloat16), + ) + # Padded value rows: the state TMA reads K-element rows back to back. + padded = torch.zeros(*state.shape[:-1], D + 8, device=DEVICE, dtype=torch.bfloat16)[ + ..., :D + ] + with pytest.raises(ValueError, match="state strides"): + verify(**{**args, "state": padded}) + # A state 8 bytes off a 16-byte boundary. + shifted = torch.zeros(state.numel() + 8, device=DEVICE, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="16-byte boundaries"): + verify(**{**args, "state": shifted[4 : 4 + state.numel()].view(state.shape)}) + + +def _layer(pages, records: KDARecoverSSMRecords, heads: int) -> SimpleNamespace: + _, conv, state = pages + return SimpleNamespace( + pages=pages, + kv_cache=(conv, state), + recoverssm_records=records, + A_log=torch.randn(heads, device=DEVICE) * 0.5, + dt_bias=torch.randn(heads * D, device=DEVICE) * 0.5, + local_num_heads=heads, + head_dim=D, + gate_lower_bound=LOWER_BOUND, + ) + + +@pytest.mark.parametrize("heads", [96, 12, 3]) +@pytest.mark.parametrize("batch", [1, 8, 13]) +@pytest.mark.parametrize("window", [2, 5, 8]) +@torch.inference_mode() +def test_cute_state_commit_matches_triton(heads, batch, window, monkeypatch): + """The CuTe fold of every layer's accepted tokens equals the Triton fold, + with varlen windows, rejected-everything (0 accepted) and dead requests.""" + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + lambda: "DS", + ) + torch.manual_seed(heads + batch + window) + num_layers = 3 + slots = 2 * batch + 1 + layers = [] + for _ in range(num_layers): + layers.append( + _layer( + _pages(slots, heads, window, dim_first=True), + _records(batch, heads, window), + heads, + ) + ) + lens = torch.full((batch,), window, dtype=torch.int32) + if batch > 2: + lens[0] = min(2, window) + query_start = torch.zeros(batch + 1, dtype=torch.int32) + query_start[1:] = torch.cumsum(lens, 0) + query_start = query_start.to(DEVICE) + accepted = torch.randint(0, window + 1, (batch,), dtype=torch.int32, device=DEVICE) + state_indices = torch.randperm(slots - 1, device=DEVICE)[:batch].to(torch.int32) + 1 + if batch > 3: + state_indices[3] = 0 + + compile_state_commit(heads, window, DEVICE, align=False) + initial = [_clone_pages(layer.pages) for layer in layers] + results = {} + for cute in (False, True): + for layer, before in zip(layers, initial): + layer.pages[0].copy_(before[0]) + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=window, + max_num_reqs=batch, + backend="cute" if cute else "triton", + ) + context.commit(accepted, state_indices, query_start) + torch.accelerator.synchronize() + results[cute] = [layer.kv_cache[1].clone() for layer in layers] + states_before = [before[2] for before in initial] + + for triton_state, cute_state, before in zip( + results[False], results[True], states_before + ): + torch.testing.assert_close( + cute_state.float(), triton_state.float(), atol=2e-2, rtol=2e-2 + ) + # Requests with no accepted tokens keep their state, as do unused slots. + for b in range(batch): + slot = int(state_indices[b]) + if slot > 0 and int(accepted[b]) == 0: + assert torch.equal(cute_state[slot], before[slot]) + unused = sorted(set(range(slots)) - set(state_indices.tolist())) + assert torch.equal(cute_state[unused], before[unused]) + + +@pytest.mark.parametrize("heads", [96, 16]) +@torch.inference_mode() +def test_recoverssm_steps_match_stock_spec_decode(heads, monkeypatch): + """Several verify -> sample -> commit steps of RecoverSSM (CuTe verify and + commit, one state slot per request) produce the outputs of stock + speculative decoding (``causal_conv1d_update`` + ``fused_recurrent_kda`` + over 1 + k state slots per request).""" + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import fused_recurrent_kda + + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + lambda: "DS", + ) + torch.manual_seed(heads) + window, batch, steps = 5, 6, 4 + tokens = batch * window + query_start = torch.arange(0, tokens + 1, window, dtype=torch.int32, device=DEVICE) + conv_weight = torch.randn(3 * heads * D, 4, device=DEVICE) * 0.5 + + # RecoverSSM: one slot per request. + layer = _layer( + _pages(batch + 1, heads, window, dim_first=True), + _records(batch, heads, window), + heads, + ) + rss_slots = torch.arange(1, batch + 1, dtype=torch.int32, device=DEVICE) + # Stock: 1 + k slots per request; the request's state starts in column 0. + _, ref_conv, ref_state = _pages(batch * window + 1, heads, window, True) + ref_slots = torch.arange( + 1, batch * window + 1, dtype=torch.int32, device=DEVICE + ).view(batch, window) + ref_conv[ref_slots[:, 0]] = layer.kv_cache[0][rss_slots] + ref_state[ref_slots[:, 0]] = layer.kv_cache[1][rss_slots] + + verify = _compiled_verify(heads, window, True, torch.float32) + compile_state_commit(heads, window, DEVICE, align=False) + context = KDARecoverSSMCommitContext.create( + [layer], spec_query_len=window, max_num_reqs=batch, backend="cute" + ) + ones = torch.ones(batch, dtype=torch.int32, device=DEVICE) + accepted = ones.clone() # the stock path's counts from the previous step + for _ in range(steps): + _, _, mixed_qkv, raw_gate, raw_beta = _inputs(batch, heads, window, False, 0, 0) + out = torch.empty(tokens, heads * D, device=DEVICE, dtype=torch.bfloat16) + assert verify( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + lower_bound=LOWER_BOUND, + state=layer.kv_cache[1], + records=layer.recoverssm_records, + query_start_loc=query_start, + state_indices=rss_slots, + num_accepted_tokens=ones, + conv_state=layer.kv_cache[0], + conv_weight=conv_weight, + out=out, + ) + conv_out = torch.empty_like(mixed_qkv) + causal_conv1d_update( + mixed_qkv, + ref_conv, + conv_weight, + None, + activation="silu", + conv_state_indices=ref_slots[:, 0], + num_accepted_tokens=accepted, + query_start_loc=query_start, + max_query_len=window, + validate_data=False, + out=conv_out, + ) + q, k, v = ( + x.unflatten(-1, (heads, D)).unsqueeze(0) + for x in conv_out.split(heads * D, -1) + ) + ref_out, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, tokens, heads, D), + raw_beta=raw_beta.unsqueeze(0), + A_log=layer.A_log, + dt_bias=layer.dt_bias, + lower_bound=LOWER_BOUND, + initial_state=ref_state, + cu_seqlens=query_start, + ssm_state_indices=ref_slots, + num_accepted_tokens=accepted, + ) + torch.testing.assert_close( + out.float(), ref_out.view(tokens, -1).float(), atol=3e-2, rtol=3e-2 + ) + # Sample: accept 1..k+1 tokens per request, then commit them. + accepted = torch.randint( + 1, window + 1, (batch,), dtype=torch.int32, device=DEVICE + ) + context.commit(accepted, rss_slots, query_start) + + +@pytest.mark.parametrize("window", [2, 8]) +@torch.inference_mode() +def test_cute_verify_cuda_graph_replay(window): + """A captured verify replays with new inputs (every launch, including the + captured one, leaves the work counter zeroed).""" + heads, batch = 96, 8 + torch.manual_seed(window) + _, query_start, mixed_qkv, raw_gate, raw_beta = _inputs( + batch, heads, window, False, 0, 0 + ) + tokens = mixed_qkv.shape[0] + A_log = torch.randn(heads, device=DEVICE) * 0.5 + dt_bias = torch.randn(heads * D, device=DEVICE) * 0.5 + conv_weight = torch.randn(3 * heads * D, 4, device=DEVICE) * 0.5 + slots = torch.arange(1, batch + 1, dtype=torch.int32, device=DEVICE) + ones = torch.ones(batch, dtype=torch.int32, device=DEVICE) + pages = _pages(batch + 1, heads, window, True) + initial = _clone_pages(pages) + _, conv, state = pages + records = _records(batch, heads, window) + verify = _compiled_verify(heads, window, True, torch.float32) + out = torch.empty(tokens, heads * D, device=DEVICE, dtype=torch.bfloat16) + + def run(): + assert verify( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + state=state, + records=records, + query_start_loc=query_start, + state_indices=slots, + num_accepted_tokens=ones, + conv_state=conv, + conv_weight=conv_weight, + out=out, + ) + + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + run() + stream.synchronize() + with torch.cuda.graph(graph, stream=stream): + run() + for replay in range(2): + pages[0].copy_(initial[0]) + mixed_qkv.normal_() + graph.replay() + torch.accelerator.synchronize() + ref = _clone_pages(initial) + ref_records = _records(batch, heads, window) + conv_out = torch.empty_like(mixed_qkv) + causal_conv1d_update( + mixed_qkv, + ref[1], + conv_weight, + None, + activation="silu", + conv_state_indices=slots, + num_accepted_tokens=ones, + query_start_loc=query_start, + max_query_len=window, + validate_data=False, + out=conv_out, + ) + q, k, v = ( + x.unflatten(-1, (heads, D)).unsqueeze(0) + for x in conv_out.split(heads * D, -1) + ) + ref_out = kda_recoverssm_verify( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, tokens, heads, D), + raw_beta=raw_beta.unsqueeze(0), + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + checkpoint_state=ref[2], + records=ref_records, + query_start_loc=query_start, + state_indices=slots, + spec_query_len=window, + ) + torch.testing.assert_close( + out.float(), ref_out.view(tokens, -1).float(), atol=3e-2, rtol=3e-2 + ) + torch.testing.assert_close( + records.correction, ref_records.correction, atol=3e-2, rtol=3e-2 + ) + assert int(verify.counter) == 0, replay + + +@pytest.mark.parametrize("window", [2, 5, 8]) +@pytest.mark.parametrize("batch", [1, 13]) +@torch.inference_mode() +def test_cute_state_commit_align_matches_triton(window, batch, monkeypatch): + """With prefix caching the commit stores the window's final state to the + block holding its last token and, when the window crosses a block boundary, + the state at the boundary to the boundary's block.""" + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + lambda: "DS", + ) + torch.manual_seed(window * 31 + batch) + heads, block_size, width = 12, 8, 4 + slots = 1 + batch * width + layers = [ + _layer( + _pages(slots, heads, window, dim_first=True), + _records(batch, heads, window), + heads, + ) + for _ in range(2) + ] + query_start = torch.arange( + 0, (batch + 1) * window, window, dtype=torch.int32, device=DEVICE + ) + accepted = torch.randint(0, window + 1, (batch,), dtype=torch.int32, device=DEVICE) + # Each request owns blocks 1 + b * width ...; its state is in the block of + # its last computed token. + block_table = ( + 1 + torch.arange(batch * width, dtype=torch.int32, device=DEVICE) + ).view(batch, width) + num_computed = torch.randint( + 1, 2 * block_size, (batch,), dtype=torch.int32, device=DEVICE + ) + # Request 0 crosses a block boundary with every accepted token. + num_computed[0] = block_size - 1 + accepted[0] = window + state_indices = block_table.gather( + 1, ((num_computed - 1) // block_size).long().unsqueeze(1) + ).squeeze(1) + if batch > 3: + state_indices[3] = 0 + compile_state_commit(heads, window, DEVICE, align=True) + initial = [_clone_pages(layer.pages) for layer in layers] + results = {} + for cute in (False, True): + for layer, before in zip(layers, initial): + layer.pages[0].copy_(before[0]) + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=window, + max_num_reqs=batch, + backend="cute" if cute else "triton", + ) + context.commit( + accepted, + state_indices, + query_start, + block_table=block_table, + num_computed_tokens=num_computed, + mamba_block_size=block_size, + ) + torch.accelerator.synchronize() + results[cute] = [layer.kv_cache[1].clone() for layer in layers] + for triton_state, cute_state in zip(results[False], results[True]): + torch.testing.assert_close( + cute_state.float(), triton_state.float(), atol=2e-2, rtol=2e-2 + ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 991ecc386d6b..248fd9b359df 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3536,6 +3536,10 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": if ( self.kv_transfer_config is not None and self.kv_transfer_config.is_kv_transfer_instance + # A KDA RecoverSSM request's state is its conv and SSM state in the + # stock Mamba page (the verify records are kept outside the KV + # cache), so KV connectors transfer it as usual. + and not self.cache_config.use_kda_recoverssm ): raise ValueError( "--use-replayssm is incompatible with KV connectors " diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 5c423739b7a0..0ede48005af3 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -148,15 +148,6 @@ def kda_state_dtype( ) return (conv_state_dtype, recurrent_state_dtype) - @classmethod - def append_kda_recoverssm_record( - cls, - base_dtypes: tuple[torch.dtype, ...], - model_dtype: ModelDType | torch.dtype, - ) -> tuple[torch.dtype, ...]: - activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) - return (*base_dtypes, torch.float32, activation_dtype) - class MambaStateShapeCalculator: @classmethod @@ -327,27 +318,6 @@ def kda_state_shape( recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) return (conv_state_shape, recurrent_state_shape) - @classmethod - def append_kda_recoverssm_record( - cls, - base_shapes: tuple[tuple[int, int], tuple[int, int, int]], - num_heads: int, - head_dim: int, - tp_world_size: int, - spec_query_len: int, - ) -> tuple[ - tuple[int, int], - tuple[int, int, int], - tuple[int, int, int], - tuple[int, int, int], - ]: - local_num_heads = divide(num_heads, tp_world_size) - return ( - *base_shapes, - (local_num_heads, spec_query_len, head_dim), - (local_num_heads, spec_query_len, 2 * head_dim), - ) - @dataclass class MambaCopySpec: diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index b8be78c75e1e..a921d1bd2879 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import replace +from typing import Literal import torch from einops import rearrange @@ -50,6 +51,15 @@ KimiK3KDAAttentionBackend, KimiK3KDAMetadata, ) +from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMRecords, + kda_recoverssm_verify, +) +from vllm.models.kimi_k3.nvidia.recoverssm.kernels import ( + CuteRecoverSSMVerify, + can_use_cute_recoverssm, + compile_state_commit, +) 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 @@ -563,36 +573,22 @@ def get_state_dtype( ) -> tuple[torch.dtype, ...]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") - base_dtypes = MambaStateDtypeCalculator.kda_state_dtype( + return MambaStateDtypeCalculator.kda_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) - if self.cache_config.use_kda_recoverssm: - return MambaStateDtypeCalculator.append_kda_recoverssm_record( - base_dtypes, self.model_config.dtype - ) - return base_dtypes def get_state_shape( self, ) -> tuple[tuple[int, ...], ...]: - base_shapes = MambaStateShapeCalculator.kda_state_shape( + 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, ) - if self.cache_config.use_kda_recoverssm: - return MambaStateShapeCalculator.append_kda_recoverssm_record( - base_shapes, - self.num_heads, - self.head_dim, - tp_world_size=self.tp_size, - spec_query_len=1 + self.num_spec, - ) - return base_shapes def __init__( self, @@ -723,6 +719,70 @@ def __init__( self.gate_lower_bound, self.use_recoverssm, ) + # RecoverSSM: this layer's verify records of a speculative step, and its + # verify and commit kernels: "cute" (SM10x, the short convolution fused + # into the verify) where they can run, else "triton". + self.recoverssm_backend: Literal["cute", "triton"] | None = None + self.recoverssm_records: KDARecoverSSMRecords | None = None + self.recoverssm_verify: CuteRecoverSSMVerify | None = None + if self.use_recoverssm: + self.recoverssm_records = KDARecoverSSMRecords.allocate( + vllm_config.scheduler_config.max_num_seqs, + self.local_num_heads, + self.spec_query_len, + self.head_dim, + vllm_config.model_config.dtype, + self.conv1d.weight.device, + ) + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + records_bytes = len(kda_config["kda_layers"]) * sum( + t.numel() * t.element_size() + for t in ( + self.recoverssm_records.correction, + self.recoverssm_records.key_gate, + ) + ) + logger.info_once( + "Allocated the KDA RecoverSSM verify records: %.2f GiB in total " + "for %d requests (%.2f MiB per request).", + records_bytes / 2**30, + max_num_seqs, + records_bytes / max_num_seqs / 2**20, + ) + supported, reason = can_use_cute_recoverssm( + head_dim=self.head_dim, + window=self.spec_query_len, + conv_width=self.conv_size, + has_conv_bias=self.conv1d.bias is not None, + activation_dtype=vllm_config.model_config.dtype, + conv_state_dtype=conv_state_dtype, + state_dtype=recurrent_state_dtype, + conv_weight_dtype=self.conv1d.weight.dtype, + lower_bound=self.gate_lower_bound, + ) + self.recoverssm_backend = "cute" if supported else "triton" + if self.recoverssm_backend == "triton": + logger.info_once( + "Using the Triton KDA RecoverSSM verify and commit: the CuTe " + "kernels cannot run (%s).", + reason, + ) + else: + logger.info_once("Using the CuTe KDA RecoverSSM verify and commit.") + self.recoverssm_verify = CuteRecoverSSMVerify( + num_heads=self.local_num_heads, + window=self.spec_query_len, + conv_dim_first=is_conv_state_dim_first(), + conv_weight_dtype=self.conv1d.weight.dtype, + device=self.conv1d.weight.device, + ) + self.recoverssm_verify.compile() + compile_state_commit( + self.local_num_heads, + self.spec_query_len, + self.conv1d.weight.device, + align=self.cache_config.mamba_cache_mode == "align", + ) decode_conv1d_weight = None if ( self.kda_decode_backend != "triton" @@ -912,6 +972,94 @@ def forward( return self.gemm_rs_ar.apply(core_attn_out, self.o_proj) return self.o_proj(core_attn_out)[0] + def _recoverssm_verify( + self, + m: KimiK3KDAMetadata, + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + conv_weight: torch.Tensor, + query_start_loc: torch.Tensor, + out: torch.Tensor | None, + ) -> torch.Tensor: + """RecoverSSM verify of the speculative requests: the short convolution + and the KDA recurrence from each request's checkpoint, which is left + unchanged. The verify records it writes let the commit after sampling + fold the accepted tokens in (KimiK3KDAMetadata.commit_recoverssm_state). + """ + num_tokens = mixed_qkv.shape[0] + if out is None: + out = mixed_qkv.new_empty( + (1, num_tokens, self.local_num_heads, self.head_dim) + ) + assert m.spec_state_indices_tensor is not None + assert m.num_accepted_tokens is not None + state_indices = m.spec_state_indices_tensor[: m.num_spec_decodes, 0] + records = self.recoverssm_records + assert records is not None + conv_state, state = self.kv_cache[:2] + if self.recoverssm_backend == "cute": + assert self.recoverssm_verify is not None + assert self.gate_lower_bound is not None + if self.recoverssm_verify( + mixed_qkv=mixed_qkv, + raw_gate=raw_g[0].flatten(-2), + raw_beta=raw_beta[0], + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + state=state, + records=records, + query_start_loc=query_start_loc, + state_indices=state_indices, + num_accepted_tokens=m.num_accepted_tokens, + conv_state=conv_state, + conv_weight=conv_weight, + out=out[0].flatten(-2), + ): + return out + # The CuTe verify declined misaligned activations: a mixed batch's + # speculative tokens can start at an odd row of a projection whose + # rows are not 16-byte aligned (e.g. the beta of the projection + # overlap path under graph capture). The Triton verify runs instead, + # leaving the same conv state, records and SSM state. + + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_state, + conv_weight, + self.conv1d.bias, + activation="silu", + conv_state_indices=state_indices, + num_accepted_tokens=m.num_accepted_tokens, + query_start_loc=query_start_loc, + max_query_len=self.spec_query_len, + validate_data=False, + out=torch.empty_like(mixed_qkv), + ) + q, k, v = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv.split(self.local_projection_size, dim=-1) + ) + return kda_recoverssm_verify( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + checkpoint_state=state, + records=records, + query_start_loc=query_start_loc, + state_indices=state_indices, + spec_query_len=self.spec_query_len, + out=out, + ) + @eager_break_during_capture def _forward( self, @@ -955,7 +1103,7 @@ def _forward( g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] - conv_state, recurrent_state, *recoverssm_records = self.kv_cache + conv_state, recurrent_state = self.kv_cache[:2] # The convolution kernels consume (..., dim, width - 1). if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) @@ -1086,65 +1234,44 @@ def _forward( 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 = ( - self.spec_query_len - if self.use_recoverssm - else spec_state_indices_tensor.size(-1) - ) - spec_conv_out = torch.empty_like(mixed_qkv_spec) - mixed_qkv_spec = causal_conv1d_update( - mixed_qkv_spec, - 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]] + core_attn_out[:, : mixed_qkv_spec.shape[0]] if m.num_prefills == 0 and m.num_decodes == 0 else None ) if spec_token_start is not None: spec_out = core_attn_out[:, spec_slice] if self.use_recoverssm: - from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( - kda_recoverssm_verify, - ) - - if len(recoverssm_records) != 2: - raise ValueError( - "KDA RecoverSSM requires correction and key/gate buffers" - ) - core_attn_out_spec = kda_recoverssm_verify( - q=q_spec, - k=k_spec, - v=v_spec, - raw_g=g1_spec, - raw_beta=beta_spec, - A_log=self.A_log, - dt_bias=self.dt_bias, - lower_bound=self.gate_lower_bound, - checkpoint_state=recurrent_state, - correction_cache=recoverssm_records[0], - kg_cache=recoverssm_records[1], - query_start_loc=spec_cu_seqlens, - state_indices=spec_state_indices_tensor[: m.num_spec_decodes, 0], - spec_query_len=self.spec_query_len, - out=spec_out, + core_attn_out_spec = self._recoverssm_verify( + m, + mixed_qkv_spec, + g1_spec, + beta_spec, + conv_weights, + spec_cu_seqlens, + spec_out, ) else: + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=spec_state_indices_tensor[:, 0][ + : m.num_spec_decodes + ], + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_state_indices_tensor.size(-1), + validate_data=False, + out=torch.empty_like(mixed_qkv_spec), + ) + 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) + ) core_attn_out_spec, _ = fused_recurrent_kda( q=q_spec, k=k_spec, diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 58c74e625818..a6fbd89e505b 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -351,6 +351,7 @@ def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": layers, spec_query_len=1 + self.vllm_config.num_speculative_tokens, max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, + backend=layers[0].recoverssm_backend, ) self.recoverssm_context = context return context diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 6368cbaa0f77..2073c0fbe71a 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1496,16 +1496,11 @@ def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", ) -> tuple[torch.dtype, ...]: - dtypes = MambaStateDtypeCalculator.kda_state_dtype( + return MambaStateDtypeCalculator.kda_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, ) - if vllm_config.cache_config.use_kda_recoverssm: - dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( - dtypes, vllm_config.model_config.dtype - ) - return dtypes @classmethod def get_mamba_state_shape_from_config( @@ -1519,22 +1514,13 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) - shapes = MambaStateShapeCalculator.kda_state_shape( + 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, ) - if vllm_config.cache_config.use_kda_recoverssm: - return MambaStateShapeCalculator.append_kda_recoverssm_record( - shapes, - hf_config.linear_attn_config["num_heads"], - hf_config.linear_attn_config["head_dim"], - tp_world_size=tp_size, - spec_query_len=1 + num_spec, - ) - return shapes @classmethod def get_mamba_state_copy_func( diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index 7eda21f1a5ab..38cba0b81c19 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -4,11 +4,12 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import Any +from typing import Any, Literal import torch from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.kimi_k3.nvidia.recoverssm.kernels import cute_recoverssm_commit from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -85,11 +86,11 @@ def _kda_recoverssm_verify_kernel( stride_state_head, stride_state_v, stride_state_k, - stride_correction_block, + stride_correction_row, stride_correction_head, stride_correction_pos, stride_correction_dim, - stride_kg_block, + stride_kg_row, stride_kg_head, stride_kg_pos, stride_kg_dim, @@ -195,7 +196,7 @@ def _kda_recoverssm_verify_kernel( correction_ptr = ( correction_cache_ptr - + state_idx * stride_correction_block + + pid_b * stride_correction_row + pid_h * stride_correction_head + token_offset * stride_correction_pos ) @@ -207,7 +208,7 @@ def _kda_recoverssm_verify_kernel( if pid_v == 0: kg_ptr = ( kg_cache_ptr - + state_idx * stride_kg_block + + pid_b * stride_kg_row + pid_h * stride_kg_head + token_offset * stride_kg_pos ) @@ -472,18 +473,16 @@ def _commit_kda_state_kernel( correction_cache_ptr = correction_cache_base_addr.to( tl.pointer_type(correction_cache_ref_ptr.dtype.element_ty) ) + # The records are those of request pid_b of the step. correction_cache_ptr += ( - source_state_idx * correction_cache_block_stride - + pid_h * stride_correction_cache_head + pid_b * correction_cache_block_stride + pid_h * stride_correction_cache_head ) kg_cache_base_addr = tl.load(kg_cache_base_addrs_ptr + pid_l) kg_cache_block_stride = tl.load(kg_cache_block_strides_ptr + pid_l) kg_cache_ptr = kg_cache_base_addr.to( tl.pointer_type(kg_cache_ref_ptr.dtype.element_ty) ) - kg_cache_ptr += ( - source_state_idx * kg_cache_block_stride + pid_h * stride_kg_cache_head - ) + kg_cache_ptr += pid_b * kg_cache_block_stride + pid_h * stride_kg_cache_head offs_k = tl.arange(0, BK) offs_v = pid_v * BV + tl.arange(0, BV) @@ -583,6 +582,36 @@ def _commit_kda_state_kernel( tl.store(final_ptrs, state, mask=mask_state) +@dataclass(frozen=True) +class KDARecoverSSMRecords: + """One KDA layer's RecoverSSM verify records. + + Per request of a speculative step (its row in the step's speculative batch), + head and verified token: the FP32 correction ``u_t`` [V] and the activation + dtype key and raw gate [2K]. A verify writes them and the commit after sampling + reads them; they are not part of a request's state and so are not stored in + the KV cache.""" + + correction: torch.Tensor # [max_num_reqs, H, Q, V] float32 + key_gate: torch.Tensor # [max_num_reqs, H, Q, 2K] activation dtype + + @classmethod + def allocate( + cls, + max_num_reqs: int, + num_heads: int, + spec_query_len: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, + ) -> "KDARecoverSSMRecords": + shape = (max_num_reqs, num_heads, spec_query_len) + return cls( + torch.empty((*shape, head_dim), dtype=torch.float32, device=device), + torch.empty((*shape, 2 * head_dim), dtype=dtype, device=device), + ) + + def kda_recoverssm_verify( q: torch.Tensor, k: torch.Tensor, @@ -593,14 +622,14 @@ def kda_recoverssm_verify( dt_bias: torch.Tensor, lower_bound: float | None, checkpoint_state: torch.Tensor, - correction_cache: torch.Tensor, - kg_cache: torch.Tensor, + records: "KDARecoverSSMRecords", query_start_loc: torch.Tensor, state_indices: torch.Tensor, spec_query_len: int, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Verify a KDA speculative window without modifying its checkpoint.""" + """Verify a KDA speculative window without modifying its checkpoint. Request + ``b`` writes its verify records to row ``b`` of ``records``.""" if q.ndim != 4 or q.shape[0] != 1: raise ValueError("KDA RecoverSSM q must have shape [1, tokens, heads, dim]") _, total_tokens, num_heads, key_dim = q.shape @@ -613,32 +642,13 @@ def kda_recoverssm_verify( raise ValueError("KDA RecoverSSM q, k, and gate heads must be contiguous") if v.stride()[2:] != (value_dim, 1) or raw_beta.stride(2) != 1: raise ValueError("KDA RecoverSSM v and beta heads must be contiguous") - num_blocks = checkpoint_state.shape[0] if checkpoint_state.shape[1:] != ( num_heads, value_dim, key_dim, ): raise ValueError("KDA RecoverSSM checkpoint shape is incompatible") - expected_correction_shape = ( - num_blocks, - num_heads, - spec_query_len, - value_dim, - ) - if correction_cache.shape != expected_correction_shape: - raise ValueError( - f"KDA RecoverSSM correction buffer needs shape {expected_correction_shape}" - ) - expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) - if kg_cache.shape != expected_kg_shape: - raise ValueError( - f"KDA RecoverSSM key/gate buffer needs shape {expected_kg_shape}" - ) - if correction_cache.dtype != torch.float32: - raise ValueError("KDA RecoverSSM correction buffer must use float32") - if kg_cache.dtype != k.dtype: - raise ValueError("KDA RecoverSSM key/gate buffer must match activation dtype") + correction_cache, kg_cache = records.correction, records.key_gate if A_log.shape != (num_heads,) or dt_bias.numel() != num_heads * key_dim: raise ValueError("KDA RecoverSSM gate parameters are incompatible") if not A_log.is_contiguous() or not dt_bias.is_contiguous(): @@ -754,6 +764,9 @@ class KDARecoverSSMCommitContext: dt_bias: torch.Tensor lower_bound: float | None spec_query_len: int + # The kernel folding the accepted tokens into the state (the commit plan + # and the conv state compaction are Triton for both). + backend: Literal["cute", "triton"] = "triton" @classmethod def create( @@ -762,21 +775,16 @@ def create( *, spec_query_len: int, max_num_reqs: int, + backend: Literal["cute", "triton"] = "triton", ) -> "KDARecoverSSMCommitContext": if not layers: raise ValueError("KDA RecoverSSM commit requires at least one layer") - if any(len(layer.kv_cache) != 4 for layer in layers): - raise ValueError( - "KDA RecoverSSM pages must contain conv, state, correction, " - "and key/gate" - ) - conv_states = [layer.kv_cache[0] for layer in layers] if not is_conv_state_dim_first(): conv_states = [state.transpose(-1, -2) for state in conv_states] checkpoints = [layer.kv_cache[1] for layer in layers] - correction_caches = [layer.kv_cache[2] for layer in layers] - kg_caches = [layer.kv_cache[3] for layer in layers] + correction_caches = [layer.recoverssm_records.correction for layer in layers] + kg_caches = [layer.recoverssm_records.key_gate for layer in layers] A_log = [layer.A_log for layer in layers] dt_bias = [ layer.dt_bias.view(layer.local_num_heads, layer.head_dim) @@ -800,36 +808,6 @@ def create( raise ValueError( "KDA RecoverSSM layers need matching checkpoint layout" ) - expected_correction_shape = ( - num_blocks, - num_heads, - spec_query_len, - value_dim, - ) - correction_ref = correction_caches[0] - for correction_cache in correction_caches: - if ( - correction_cache.shape != expected_correction_shape - or correction_cache.dtype != torch.float32 - or correction_cache.device != state_ref.device - or correction_cache.stride()[1:] != correction_ref.stride()[1:] - ): - raise ValueError( - "KDA RecoverSSM correction buffers need float32 shape " - f"{expected_correction_shape}" - ) - expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) - kg_ref = kg_caches[0] - for kg_cache in kg_caches: - if ( - kg_cache.shape != expected_kg_shape - or kg_cache.dtype != kg_ref.dtype - or kg_cache.device != state_ref.device - or kg_cache.stride()[1:] != kg_ref.stride()[1:] - ): - raise ValueError( - f"KDA RecoverSSM key/gate buffers need shape {expected_kg_shape}" - ) if any(param.shape != (num_heads,) for param in A_log): raise ValueError("KDA RecoverSSM A_log shape is incompatible") if any(param.shape != (num_heads, key_dim) for param in dt_bias): @@ -905,6 +883,7 @@ def _block_strides(tensors: Sequence[torch.Tensor]) -> torch.Tensor: dt_bias=torch.stack(tuple(dt_bias)).contiguous(), lower_bound=lower_bounds.pop(), spec_query_len=spec_query_len, + backend=backend, ) def commit( @@ -1010,6 +989,11 @@ def commit( num_warps=4, ) + if self.backend == "cute": + cute_recoverssm_commit( + self, state_indices, batch, align=block_table is not None + ) + return state_ref = self.checkpoints[0] _, num_heads, value_dim, key_dim = state_ref.shape block_k = triton.next_power_of_2(key_dim) diff --git a/vllm/models/kimi_k3/nvidia/recoverssm/__init__.py b/vllm/models/kimi_k3/nvidia/recoverssm/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/__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/recoverssm/cutedsl/__init__.py b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/__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/recoverssm/cutedsl/commit.py b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/commit.py new file mode 100644 index 000000000000..ed84f947b5a3 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/commit.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KDA RecoverSSM state commit: one CTA per (head chunk, request, layer). + +Folds the n accepted tokens of a verify window into the checkpoint: + + S_n = S_0 * e^{L_n} + sum_{t 0) & (source > Int64(null_slot)) & (destination > Int64(null_slot)) + m = Int32(0) + boundary = Int64(null_slot) + if cutlass.const_expr(ALIGN): + m = boundary_lens[request] + boundary = Int64(boundaries[request]) + boundary_live = (m > 0) & (boundary > Int64(null_slot)) + if live: + V: cutlass.Constexpr = K + first_head = chunk_id * HEADS_PER_CTA + state_base = Int64(state_addrs[layer]) + stride = Int64(state_strides[layer]) * 2 # bytes per slot + src_base = state_base + source * stride + dst_base = state_base + destination * stride + boundary_base = state_base + boundary * stride + corr_ptr = cute.make_ptr( + Float32, corr_addrs[layer], cute.AddressSpace.gmem, assumed_align=16 + ) + kg_ptr = cute.make_ptr( + cutlass.BFloat16, + kg_addrs[layer], + cute.AddressSpace.gmem, + assumed_align=16, + ) + corr_slot = corr_ptr + Int64(request) * Int64(corr_strides[layer]) + kg_slot = kg_ptr + Int64(request) * Int64(kg_strides[layer]) + + if tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator, PAGE_BYTES) + page0 = cute.make_ptr( + cutlass.BFloat16, + src_base + Int64(first_head) * PAGE_BYTES, + cute.AddressSpace.gmem, + assumed_align=16, + ) + copy_g2s(tiles.iterator, page0, size=PAGE_BYTES, barrier=bars.iterator) + + for h in cutlass.range(HEADS_PER_CTA): + head = first_head + h + buf = h % 2 + tile = tiles.iterator + buf * PAGE + # ---- The next head's page streams in while this head folds. ---- + if tid == 0: # noqa: SIM102 (traced: `and` does not combine traced values) + if h + 1 < HEADS_PER_CTA: + # Its buffer's last bulk store must have read it out. + bulk_read_done() + nxt = (h + 1) % 2 + cute.arch.mbarrier_arrive_and_expect_tx( + bars.iterator + nxt, PAGE_BYTES + ) + page = cute.make_ptr( + cutlass.BFloat16, + src_base + Int64(head + 1) * PAGE_BYTES, + cute.AddressSpace.gmem, + assumed_align=16, + ) + copy_g2s( + tiles.iterator + nxt * PAGE, + page, + size=PAGE_BYTES, + barrier=bars.iterator + nxt, + ) + + corr = cute.make_tensor( + corr_slot + head * Q * V, cute.make_layout((Q, V), stride=(V, 1)) + ) + kg = cute.make_tensor( + kg_slot + head * Q * 2 * K, + cute.make_layout((Q, 2 * K), stride=(2 * K, 1)), + ) + for t in cutlass.range_constexpr(Q): + value = Float32(0) + if t < n: + value = corr[t, tid] + corrections[t, tid] = value + + # ---- Thread c: key channel c's decays and scaled keys. ---- + bias = dt_bias[layer, head, tid].to(Float32) + a = cute.math.exp(a_log[layer, head].to(Float32)) + cum = Float32(0) + for t in cutlass.range_constexpr(Q): + kv[t] = Float32(0) + if t < n: + kv[t] = kg[t, tid].to(Float32) + cum = cum + lower_bound * sigmoid( + a * (kg[t, K + tid].to(Float32) + bias), True + ) + logs[t] = cum + for t in cutlass.range_constexpr(Q): + s = warp_sum(kv[t] * kv[t]) + if lane == 0: + norms[warp, t] = s + cute.arch.barrier() + for t in cutlass.range_constexpr(Q): + total = norms[0, t] + norms[1, t] + norms[2, t] + norms[3, t] + kn[t] = kv[t] * cute.math.rsqrt(total + Float32(1e-6)) + # Token t's update decays by the tokens after it: e^{L_n - L_t}. + scaled[t, tid] = kn[t] * cute.math.exp(cum - logs[t]) + scaled[Q, tid] = cute.math.exp(cum) + if cutlass.const_expr(ALIGN): + # S_m: the first m tokens (m <= n), decayed to token m - 1. + cum_b = Float32(0) + for t in cutlass.range_constexpr(Q): + if t < m: + cum_b = logs[t] + for t in cutlass.range_constexpr(Q): + value = Float32(0) + if t < m: + value = kn[t] * cute.math.exp(cum_b - logs[t]) + scaled_b[t, tid] = value + scaled_b[Q, tid] = cute.math.exp(cum_b) + cute.arch.barrier() + for t in cutlass.range_constexpr(Q + 1): + for j in cutlass.range_constexpr(KEYS_PER_LANE): + coef[t, j] = scaled[t, col + j] + if cutlass.const_expr(ALIGN): + coef_b[t, j] = scaled_b[t, col + j] + + # ---- Rows: 16 lanes per row, 8 keys per lane, folded in place. ---- + wait(bars.iterator + buf, (h // 2) % 2, scope="cta") + for it in cutlass.range_constexpr(PASSES): + r = it * ROWS_PER_PASS + group + for t in cutlass.range_constexpr(Q): + u[t] = corrections[t, r] + pointer = tile + r * K + col + words = load_words(pointer, count=4, shared=True, mutable=True) + out = [] + out_b = [] + for w in cutlass.range_constexpr(4): + pair = [] + pair_b = [] + for half in cutlass.range_constexpr(2): + j = 2 * w + half + initial = bf16(words[w], half == 1) + value = initial * coef[Q, j] + for t in cutlass.range_constexpr(Q): + value = value + u[t] * coef[t, j] + pair.append(value) + if cutlass.const_expr(ALIGN): + value_b = initial * coef_b[Q, j] + for t in cutlass.range_constexpr(Q): + value_b = value_b + u[t] * coef_b[t, j] + pair_b.append(value_b) + out.append(bf16x2(pair[0], pair[1])) + if cutlass.const_expr(ALIGN): + out_b.append(bf16x2(pair_b[0], pair_b[1])) + store_words(pointer, tuple(out), shared=True) + if cutlass.const_expr(ALIGN): # noqa: SIM102 (a compile-time branch) + if boundary_live: + row = Int64(head) * PAGE_BYTES + Int64(r * K + col) * 2 + store_words(boundary_base + row, tuple(out_b)) + proxy_fence() + cute.arch.barrier() + if tid == 0: + store_bulk(dst_base + Int64(head) * PAGE_BYTES, tile, PAGE_BYTES) + if tid == 0: + bulk_done() + cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify.py b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify.py new file mode 100644 index 000000000000..b746dfec40a9 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify.py @@ -0,0 +1,526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KDA RecoverSSM verify on the fifth-generation tensor cores: one CTA (4 warps) per +(head, request). + +With per-channel log decays ``L_t = g_1 + ... + g_t`` and normalized ``k_t``, ``q_t`` +(``q_t`` carries the 1/sqrt(K) scale), the verify window satisfies + + u_t = beta_t * (v_t - S_0 p_t - sum_{i, d_it = and w_i = e^{-L_i} k_i. + +CTAs are persistent and walk the (head, request) items. Per item: + +1. The state S_0 [V, K] arrives by TMA in the no-swizzle K-major UMMA layout (core + matrices of 8 rows x 16 bytes, K-adjacent ones 2 KB apart). The next item's + state is loaded into the same buffer as soon as this item's MMAs complete, under + this item's solve and the next item's probes (the first item's is issued before + the PDL wait: the state is written by the previous step's commit). +2. Thread c owns channel c: the optional short convolution of q, k, v (the stock + ``causal_conv1d_update`` speculative step), then p_t, r_t (BF16 high and low + halves, the B operands) and w_t; two threads per (probe, i) pair reduce the + coupling scalars. +3. One thread issues the M128 x N16 x K16 BF16 MMAs (high and low halves accumulate + into the same FP32 accumulator): D[row, column] = S_0 row . probe column. +4. Thread r reads state row r's 16 projections from tensor memory, solves the + row's triangular system and stores its outputs and corrections (coalesced). + +Outputs and records (``u_t``, key and raw gate) match ``kda_recoverssm_verify``. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32, Uint64 +from cutlass.cute.nvgpu import cpasync + +from vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.primitives import ( + alloc, + asm, + at, + bf16, + fma, + load_words, + shared_address, + sigmoid, + warp_sum, +) + +K = 128 +THREADS = 128 # one per key channel / value row +CTAS_PER_SM = 5 # resident CTAs per SM (shared memory bound) +PROBES = 16 # MMA N: 2 Q <= 16 probe columns, zero padded +TMEM_COLS = 32 # the smallest allocation; the accumulator uses PROBES columns +A_KGROUP_BYTES = K * 16 # one 8-wide K group of all 128 rows +# B K groups 272 bytes apart (not 256): the 4 channel groups of a warp's stores land +# in different banks. +B_KGROUP_BYTES = PROBES * 16 + 16 + + +def state_smem_layout(): + """State tile in shared memory: element (row, k) at (k // 8) * 1024 + row * 8 + + k % 8 (BF16 elements).""" + return cute.make_layout((K, (8, K // 8)), stride=(8, (1, K * 8))) + + +# Instruction descriptor, kind::f16: FP32 accumulator, BF16 A and B, both K-major, +# N = 16, M = 128. +IDESC = (1 << 4) | (1 << 7) | (1 << 10) | ((PROBES >> 3) << 17) | ((K >> 4) << 24) + + +@asm() +def tmem_alloc(dst): + code = f"tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$0], {TMEM_COLS};" + return code, "r,~{memory}", (dst,) + + +@asm() +def tmem_relinquish(): + return "tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;", "", () + + +@asm() +def tmem_dealloc(address): + code = f"tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0, {TMEM_COLS};" + return code, "r", (Uint32(address),) + + +@asm() +def tc_fence_before(): + return "tcgen05.fence::before_thread_sync;", "~{memory}", () + + +@asm() +def tc_fence_after(): + return "tcgen05.fence::after_thread_sync;", "~{memory}", () + + +@asm() +def proxy_fence(): + """Generic-proxy shared memory writes (``st.shared``, ``cp.async``) visible to + the tensor cores' async proxy.""" + return "fence.proxy.async.shared::cta;", "~{memory}", () + + +@asm(Uint32) +def take_ticket(counter): + """Relaxed GPU-scope fetch-and-increment of the work counter.""" + return "atom.relaxed.gpu.global.add.u32 $0, [$1], 1;", "=r,l,~{memory}", (counter,) + + +@asm() +def reset_counter(counter): + return "st.relaxed.gpu.global.u32 [$0], 0;", "l,~{memory}", (counter,) + + +@asm() +def cp_async_wait_all(): + return "cp.async.wait_all;", "~{memory}", () + + +@asm() +def umma(d_address, a_desc, b_desc, accumulate): + code = "{ .reg .pred p; setp.ne.b32 p, $3, 0; " + code += f"tcgen05.mma.cta_group::1.kind::f16 [$0], $1, $2, {IDESC}, p; }}" + return code, "r,l,l,r", (Uint32(d_address), a_desc, b_desc, Int32(accumulate)) + + +@asm() +def umma_commit(barrier): + code = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$0];" + return code, "r,~{memory}", (barrier,) + + +@asm((Float32,) * PROBES) +def tmem_load_row(address): + """Lane i of warp w: TMEM lane 32 w + i, columns [0, 16) of ``address``.""" + regs = ",".join(f"${i}" for i in range(PROBES)) + code = f"tcgen05.ld.sync.aligned.32x32b.x16.b32 {{{regs}}}, [${PROBES}]; " + code += "tcgen05.wait::ld.sync.aligned;" + return code, ",".join(["=f"] * PROBES + ["r"]), (Uint32(address),) + + +@cute.jit +def smem_desc(address, lead_bytes: cutlass.Constexpr, stride_bytes: cutlass.Constexpr): + """No-swizzle K-major UMMA shared memory descriptor (SM100 version bits).""" + static = ((lead_bytes >> 4) << 16) | ((stride_bytes >> 4) << 32) | (1 << 46) + return Uint64((address >> 4) & 0x3FFF) | Uint64(static) + + +@cute.jit +def bf16_split(value: Float32): + """BF16 (high, low) with high + low ~= value to ~2^-17.""" + high = value.to(cutlass.BFloat16) + low = (value - high.to(Float32)).to(cutlass.BFloat16) + return high, low + + +class KdaTcVerify: + """Args: qkv [T, 3HK], raw_g [T, HK], raw_beta [T, H], a_log [H], dt_bias [HK], + state [slots, H, V, K], corr [rows, H, Q, V], kg [rows, H, Q, 2K] (request b's + records in row b), out [T, HV], + query_start [B + 1], request slots [B], (conv state [slots, Q + 2, 3HK] token + major, or [slots, 3HK, Q + 2] with ``conv_dim_first``, conv weight [3HK, 4] + (BF16, or FP32 with ``conv_weight_f32``: the stock conv1d parameter), + accepted tokens [B] when ``conv``; qkv is then the raw projection), a zeroed + int32 work counter [1] (left zeroed), + lower_bound, null slot, stream. ``ctas``: the persistent grid bound (CTAs + resident per SM x SMs).""" + + def __init__( + self, + heads: int, + window: int, + conv: bool = False, + ctas: int = 740, + conv_dim_first: bool = False, + conv_weight_f32: bool = False, + ): + assert 1 <= window <= 8 + self.heads, self.window, self.conv, self.ctas = heads, window, conv, ctas + self.conv_dim_first = conv_dim_first + self.conv_weight_f32 = conv_weight_f32 + + @cute.jit + def __call__(self, *args): + # The state as ((V, H, slots), K) rows with its head and slot strides (vLLM + # pads the slot stride to the cache page; V = K rows of a head are + # contiguous); one TMA box per (slot, head): all V rows + # into the no-swizzle K-major UMMA layout (row 16 bytes after row, K groups + # of 8 a whole 128-row group apart). + state = args[5] + state_rows = cute.make_tensor( + state.iterator, + cute.make_layout( + ((K, state.shape[1], state.shape[0]), K), + stride=((K, state.stride[1], state.stride[0]), 1), + ), + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), state_rows, state_smem_layout(), (K, K) + ) + items = args[10].shape[0] * self.heads + self.kernel(args[:-1], tma_atom, tma_tensor).launch( + grid=(cutlass.min(items, self.ctas), 1, 1), + block=(THREADS, 1, 1), + stream=args[-1], + use_pdl=True, + ) + + @cute.jit + def load_state(self, tma_atom, tma_tensor, a_tile, barrier, slot, head, lane): + """Warp-wide: TMA state (slot, head) into ``a_tile`` on ``barrier``.""" + H: cutlass.Constexpr = self.heads + if lane == 0: + cute.arch.mbarrier_arrive_and_expect_tx(barrier, K * K * 2) + cute.arch.sync_warp() + s_state = cute.make_tensor(a_tile.iterator, state_smem_layout()) + g_state = cute.local_tile(tma_tensor, (K, K), (Int32(slot) * H + head, 0)) + t_smem, t_gmem = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(s_state, 0, 2), + cute.group_modes(g_state, 0, 2), + ) + cute.copy(tma_atom, t_gmem, t_smem, tma_bar_ptr=barrier) + + @cute.kernel + def kernel(self, args, tma_atom: cute.CopyAtom, tma_tensor: cute.Tensor): + qkv, raw_g, raw_beta, a_log, dt_bias, state, corr, kg, out = args[:9] + query_start, slots = args[9:11] + counter, lower_bound, null_slot = args[-3], args[-2], args[-1] + Q: cutlass.Constexpr = self.window + H: cutlass.Constexpr = self.heads + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + ctas, _, _ = cute.arch.grid_dim() + warp, lane = tid // 32, tid % 32 + c = tid # thread c owns channel c of q, k, v and g, and state row c + items = slots.shape[0] * H + smem = cutlass.utils.SmemAllocator() + a_tile = alloc(smem, cutlass.BFloat16, (K * K,), align=1024) + b_hi = alloc(smem, cutlass.BFloat16, (K // 8 * B_KGROUP_BYTES // 2,), align=128) + b_lo = alloc(smem, cutlass.BFloat16, (K // 8 * B_KGROUP_BYTES // 2,), align=128) + coef = alloc(smem, Float32, (2 * Q * Q,)) # c[t * Q + i] then d[...] + # Coupling pairs: c_it (i < t), then d_it (i <= t), in rounds of 32. + pairs = [(0, t, i) for t in range(Q) for i in range(t)] + pairs += [(1, t, i) for t in range(Q) for i in range(t + 1)] + slot_of = [k * Q * Q + t * Q + i for k, t, i in pairs] + ROUNDS: cutlass.Constexpr = (len(pairs) + 31) // 32 + partial = alloc(smem, Float32, (ROUNDS, 4, 32)) + betas = alloc(smem, Float32, (8,)) + norms = alloc(smem, Float32, (4, 2 * Q)) + mma_done = alloc(smem, cutlass.Int64, (1,), align=8) + state_full = alloc(smem, cutlass.Int64, (1,), align=8) + tmem_base = alloc(smem, Uint32, (1,), align=4) + next_item = alloc(smem, Int32, (1,)) + + if warp == 0: + cute.arch.sync_warp() + tmem_alloc(tmem_base.iterator) + tmem_relinquish() + if tid == 0: + cute.arch.mbarrier_init(mma_done.iterator, 1) + cute.arch.mbarrier_init(state_full.iterator, 1) + cute.arch.mbarrier_init_fence() + tc_fence_before() + cute.arch.barrier() + tc_fence_after() + tmem = tmem_base[0] + if warp == 0: + first_slot = Int64(slots[cta // H]) + if first_slot > Int64(null_slot): + self.load_state( + tma_atom, tma_tensor, a_tile, state_full.iterator, first_slot, + cta % H, lane, + ) # fmt: skip + cute.arch.griddepcontrol_wait() + + state_phase = Int32(0) + mma_phase = Int32(0) + # The first item is the CTA's own; the rest come from the work counter, so + # CTAs that start late (beyond the resident ones) or run slow items take + # fewer. Exactly ``items`` tickets are taken per launch (one per item run, + # the last of each CTA past the end); the taker of the last one resets the + # counter for the next launch. + counter_address = Int64(counter.iterator.toint()) + item = cta + while item < items: + if tid == 0: + ticket = Int32(take_ticket(counter_address)) + if ticket == items - 1: + reset_counter(counter_address) + next_item[0] = ticket + ctas + request, head = item // H, item % H + slot = Int64(slots[request]) + valid = slot > Int64(null_slot) + bos = query_start[request] + qlen = query_start[request + 1] - bos + live = valid & (qlen > 0) + qv = cute.make_rmem_tensor((Q,), Float32) + kv = cute.make_rmem_tensor((Q,), Float32) + vv = cute.make_rmem_tensor((Q,), Float32) + logs = cute.make_rmem_tensor((Q,), Float32) + if live: + # ---- 2. Activations (thread c owns channel c of q, k, v and g). ---- + bias = dt_bias[head * K + c].to(Float32) + a = cute.math.exp(a_log[head].to(Float32)) + if cutlass.const_expr(self.conv): + conv_state, conv_w, accepted = args[11:14] + hist = cute.make_rmem_tensor((3, 3), cutlass.BFloat16) + taps = cute.make_rmem_tensor((3, 2), Uint32) + taps32 = cute.make_rmem_tensor((3, 4), Float32) + first = accepted[request] - 1 + for x in cutlass.range_constexpr(3): + channel = x * H * K + head * K + c + for i in cutlass.range_constexpr(3): + if cutlass.const_expr(self.conv_dim_first): + hist[x, i] = conv_state[slot, channel, first + i] + else: + hist[x, i] = conv_state[slot, first + i, channel] + if cutlass.const_expr(self.conv_weight_f32): + for j in cutlass.range_constexpr(4): + taps32[x, j] = conv_w[channel, j] + else: + w2 = load_words(at(conv_w, channel, 0), count=2) + taps[x, 0], taps[x, 1] = w2[0], w2[1] + raw = cute.make_rmem_tensor((3, Q), cutlass.BFloat16) + g_raw = cute.make_rmem_tensor((Q,), cutlass.BFloat16) + for t in cutlass.range_constexpr(Q): + row = bos + cutlass.min(t, qlen - 1) + for x in cutlass.range_constexpr(3): + raw[x, t] = qkv[row, x * H * K + head * K + c] + g_raw[t] = raw_g[row, head * K + c] + if cutlass.const_expr(self.conv): + # The rolled state: the last 2 history rows, then the raw tokens. + for x in cutlass.range_constexpr(3): + channel = x * H * K + head * K + c + for i in cutlass.range_constexpr(2): + if cutlass.const_expr(self.conv_dim_first): + conv_state[slot, channel, i] = hist[x, i + 1] + else: + conv_state[slot, i, channel] = hist[x, i + 1] + for t in cutlass.range_constexpr(Q): + if t < qlen: + if cutlass.const_expr(self.conv_dim_first): + conv_state[slot, channel, 2 + t] = raw[x, t] + else: + conv_state[slot, 2 + t, channel] = raw[x, t] + # Width-4 convolution and SiLU in the stock kernel's order: with + # BF16 taps the products round to BF16, with FP32 taps they stay + # FP32; the output rounds to BF16. + for x in cutlass.range_constexpr(3): + if cutlass.const_expr(self.conv_weight_f32): + w = [taps32[x, j] for j in range(4)] + else: + w = [bf16(taps[x, j // 2], j % 2) for j in range(4)] + win = [hist[x, i].to(Float32) for i in range(3)] + for t in cutlass.range_constexpr(Q): + win.append(raw[x, t].to(Float32)) + acc = Float32(0) + for j in cutlass.range_constexpr(4): + if cutlass.const_expr(self.conv_weight_f32): + acc = fma(win[t + j], w[j], acc) + else: + prod = (win[t + j] * w[j]).to(cutlass.BFloat16) + acc = acc + prod.to(Float32) + raw[x, t] = (acc / (Float32(1) + cute.math.exp(-acc))).to( + cutlass.BFloat16 + ) + cum = Float32(0) + for t in cutlass.range_constexpr(Q): + vv[t] = raw[2, t].to(Float32) + qv[t], kv[t] = Float32(0), Float32(0) + if t < qlen: + kg[request, head, t, c] = raw[1, t] + kg[request, head, t, K + c] = g_raw[t] + qv[t], kv[t] = raw[0, t].to(Float32), raw[1, t].to(Float32) + gate = sigmoid(a * (g_raw[t].to(Float32) + bias), True) + cum = cum + lower_bound * gate + logs[t] = cum + for t in cutlass.range_constexpr(Q): + sq = warp_sum(qv[t] * qv[t]) + sk = warp_sum(kv[t] * kv[t]) + if lane == 0: + norms[warp, 2 * t] = sq + norms[warp, 2 * t + 1] = sk + if tid < Q: + beta = Float32(0) + if tid < qlen: + beta = sigmoid(raw_beta[bos + tid, head].to(Float32), True) + betas[tid] = beta + cute.arch.barrier() + nxt = next_item[0] # rewritten only after this item's end barrier + if live: + # ---- Probes: B operands (K-major, no swizzle) and the coupling inputs. + b_offset = (c // 8) * (B_KGROUP_BYTES // 2) + c % 8 + for t in cutlass.range_constexpr(Q): + sq = (norms[0, 2 * t] + norms[1, 2 * t]) + ( + norms[2, 2 * t] + norms[3, 2 * t] + ) + sk = (norms[0, 2 * t + 1] + norms[1, 2 * t + 1]) + ( + norms[2, 2 * t + 1] + norms[3, 2 * t + 1] + ) + q_n = qv[t] * cute.math.rsqrt(sq + Float32(1e-6)) * Float32(K**-0.5) + k_n = kv[t] * cute.math.rsqrt(sk + Float32(1e-6)) + decay = cute.math.exp(logs[t]) + p_t, r_t = decay * k_n, decay * q_n + for n, value in ((t, p_t), (Q + t, r_t)): + b_hi[b_offset + n * 8], b_lo[b_offset + n * 8] = bf16_split( + value + ) + # From here on qv holds r_t, kv holds p_t and logs holds w_t. + qv[t], kv[t] = r_t, p_t + logs[t] = cute.math.exp(-logs[t]) * k_n + for n in cutlass.range_constexpr(2 * Q, PROBES): + b_hi[b_offset + n * 8] = cutlass.BFloat16(0) + b_lo[b_offset + n * 8] = cutlass.BFloat16(0) + # The probes visible to the MMA issuer. + proxy_fence() + cute.arch.barrier() + if live: + # ---- 3. The projections: 8 K steps x (high, low) MMAs. ---- + if tid == 0: + cute.arch.mbarrier_wait(state_full.iterator, state_phase) + tc_fence_after() + a_base = shared_address(a_tile.iterator) + hi_base = shared_address(b_hi.iterator) + lo_base = shared_address(b_lo.iterator) + for kt in cutlass.range_constexpr(K // 16): + a_desc = smem_desc( + a_base + kt * 2 * A_KGROUP_BYTES, A_KGROUP_BYTES, 128 + ) + for h, base in enumerate((hi_base, lo_base)): + b_desc = smem_desc( + base + kt * 2 * B_KGROUP_BYTES, B_KGROUP_BYTES, 128 + ) + umma(tmem, a_desc, b_desc, 0 if kt == 0 and h == 0 else 1) + umma_commit(mma_done.iterator) + # The coupling scalars while the tensor cores run: this channel's + # products, reduce-scattered over the warp (lane l ends with pair + # 32 round + l), then summed over warps. + for rnd in cutlass.range_constexpr(ROUNDS): + vals = [] + for j in cutlass.range_constexpr(32): + value = Float32(0) + if cutlass.const_expr(32 * rnd + j < len(pairs)): + kind, t, i = pairs[32 * rnd + j] + value = logs[i] * (kv[t] if kind == 0 else qv[t]) + vals.append(value) + for shift in (16, 8, 4, 2, 1): + upper = (lane & shift) != 0 + kept = [] + for j in cutlass.range_constexpr(shift): + lo_v, hi_v = vals[j], vals[j + shift] + send, keep = hi_v, lo_v + if upper: + send, keep = lo_v, hi_v + kept.append(keep + cute.arch.shuffle_sync_bfly(send, shift)) + vals = kept + partial[rnd, warp, lane] = vals[0] + cute.arch.barrier() + if live: + for rnd in cutlass.range_constexpr(ROUNDS): + if tid < len(pairs) - 32 * rnd: + total = (partial[rnd, 0, tid] + partial[rnd, 1, tid]) + ( + partial[rnd, 2, tid] + partial[rnd, 3, tid] + ) + target = Int32(0) # pair index -> coefficient slot + for j in cutlass.range_constexpr( + min(32, len(pairs) - 32 * rnd) + ): + if tid == j: + target = Int32(slot_of[32 * rnd + j]) + coef[target] = total + cute.arch.barrier() # the coupling scalars + proj = [Float32(0)] * PROBES + if live: + cute.arch.mbarrier_wait(mma_done.iterator, mma_phase) + tc_fence_after() + cute.arch.sync_warp() # tcgen05.ld is warp aligned + proj = list(tmem_load_row(tmem + ((warp * 32) << 16))) + mma_phase = mma_phase ^ 1 + elif valid: + if tid == 0: + cute.arch.mbarrier_wait(state_full.iterator, state_phase) # qlen 0 + if valid: + state_phase = state_phase ^ 1 + # The state buffer is free: the next item's state streams under the solve. + if warp == 0: # noqa: SIM102 (traced: `and` does not combine traced values) + if nxt < items: + next_slot = Int64(slots[nxt // H]) + if next_slot > Int64(null_slot): + self.load_state( + tma_atom, tma_tensor, a_tile, state_full.iterator, + next_slot, nxt % H, lane, + ) # fmt: skip + if live: + # ---- 4. Thread r solves state row r. ---- + r = tid + u = cute.make_rmem_tensor((Q,), Float32) + for t in cutlass.range_constexpr(Q): + acc = vv[t] - proj[t] + for i in cutlass.range_constexpr(t): + acc = acc - coef[t * Q + i] * u[i] + u[t] = betas[t] * acc + o = proj[Q + t] + for i in cutlass.range_constexpr(t + 1): + o = fma(coef[Q * Q + t * Q + i], u[i], o) + if t < qlen: + out[bos + t, head * K + r] = o.to(out.element_type) + corr[request, head, t, r] = u[t] + if slot <= Int64(null_slot): + for t in cutlass.range_constexpr(Q): + if t < qlen: + out[bos + t, head * K + tid] = out.element_type(0) + tc_fence_before() + cute.arch.barrier() # shared scalars and TMEM reused by the next item + tc_fence_after() + item = nxt + + cute.arch.griddepcontrol_launch_dependents() + if warp == 0: + cute.arch.sync_warp() + tmem_dealloc(tmem) diff --git a/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify_ws.py b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify_ws.py new file mode 100644 index 000000000000..8e550232cfe5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/cutedsl/verify_ws.py @@ -0,0 +1,722 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KDA RecoverSSM verify with the fused short convolution on the fifth-generation +tensor cores, warp specialized (the dim-first conv state layout). + +The math and the per-item MMA / solve are those of ``verify.py`` (read its +docstring first). What differs is how an item's inputs reach the CTA: + +* A producer warp (warp 4) walks the CTA's items (the first is the CTA's own, the + rest come from the work counter) one item ahead of the 4 compute warps. Per item + it bulk-copies everything the activation phase reads into shared memory: the + conv state rows of the item's 3 x 128 channels (7 BF16 each, contiguous in the + dim-first layout), the Q raw q / k / v and gate rows (256 bytes each), the conv + taps and dt_bias, and loads the item's scalars (slot, tokens, accepted, beta, + A_log) into a small record. The state tile follows by TMA once the previous + item's MMAs have read the buffer. The first item's state, conv state and taps + are requested before the PDL wait. +* The compute warps therefore never wait on a global load: thread c reads channel + c from shared memory, writes the rolled conv state back in place, and one thread + bulk-stores it (3 x 1792 bytes) instead of 21 strided 2-byte stores per thread. +* The probes are unnormalized (p_t = e^{L_t} k_t, r_t = e^{L_t} q_t): the q / k + squared norms join the coupling dot products in one warp reduce-scatter, and the + norms scale the projections and couplings afterwards (no separate norm pass). +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 + +from vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.primitives import ( + alloc, + asm, + bf16, + fma, + load_words, + prefetch_bytes, + shared_address, + sigmoid, +) + +from .commit import bulk_done, bulk_read_done, store_bulk +from .verify import ( + A_KGROUP_BYTES, + B_KGROUP_BYTES, + PROBES, + K, + bf16_split, + proxy_fence, + reset_counter, + smem_desc, + state_smem_layout, + take_ticket, + tc_fence_after, + tc_fence_before, + tmem_alloc, + tmem_dealloc, + tmem_load_row, + tmem_relinquish, + umma, + umma_commit, +) + +CONSUMERS = 128 # 4 compute warps: thread c owns channel c / state row c +PRODUCER = 4 # the producer warp +THREADS = CONSUMERS + 32 +# Compute threads that load and activate the betas: past the Q * Q + 2 Q (at +# most 80) threads that reduce the couplings and norms. +BETA_TID = 96 +CTAS_PER_SM = 3 # shared memory bound (~60 KB per CTA) +CONV_W = 4 # conv taps + + +EPS = 1e-6 + + +def state_sw128_layout(): + """State tile in shared memory in the 128-byte swizzled K-major UMMA layout: two + 16 KB K halves (64 keys each) of 128 rows x 128 bytes, rows 128 bytes apart, + 16-byte chunks XOR-swizzled by row % 8. TMA writes it in 2 boxes of 64 x 128 + (the no-swizzle layout takes 16 boxes of 8 x 128, i.e. 16-byte TMA rows).""" + from cutlass.cute.nvgpu import tcgen05 + + atom = tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + return cute.tile_to_shape(atom, (K, K), (0, 1)) + + +@asm() +def bulk_load(dst, src, size, barrier): + """``size`` bytes global -> shared on ``barrier`` (complete_tx).""" + code = "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + code += " [$0], [$1], $2, [$3];" + return code, "r,l,r,r,~{memory}", (dst, Int64(src), Int32(size), barrier) + + +@asm() +def arrive_expect(barrier, size): + code = "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], $1;" + return code, "r,r,~{memory}", (barrier, Int32(size)) + + +@asm() +def arrive(barrier): + return "mbarrier.arrive.shared::cta.b64 _, [$0];", "r,~{memory}", (barrier,) + + +@asm(Int32) +def try_wait_cta(barrier, phase): + code = "{ .reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p,[$1],$2; " + code += "selp.u32 $0,1,0,p; }" + return code, "=r,r,r,~{memory}", (barrier, Int32(phase)) + + +@cute.jit +def wait_bar(barrier, phase): + done = try_wait_cta(barrier, phase) + while done == 0: + done = try_wait_cta(barrier, phase) + + +@asm(Float32, pure=True) +def div_full(a, b): + return "div.full.f32 $0, $1, $2;", "=f,f,f", (Float32(a), Float32(b)) + + +@asm() +def consumer_sync(): + """Named barrier 1 over the 4 compute warps.""" + return f"bar.sync 1, {CONSUMERS};", "~{memory}", () + + +class KdaTcVerify2: + """Same arguments as ``KdaTcVerify`` with ``conv=True`` and the dim-first conv + state [slots, 3HK, Q + 2] (channel rows contiguous). Requires 16-byte aligned + qkv / raw_g rows, conv state channel blocks, conv weight and dt_bias (checked by + the host wrapper).""" + + def __init__( + self, heads: int, window: int, ctas: int, conv_weight_f32: bool, swizzle=True + ): + assert 1 <= window <= 8 + self.heads, self.window, self.ctas = heads, window, ctas + self.conv_weight_f32 = conv_weight_f32 + self.swizzle = swizzle + + def smem_layout(self): + return state_sw128_layout() if self.swizzle else state_smem_layout() + + @cute.jit + def __call__(self, *args): + state = args[5] + state_rows = cute.make_tensor( + state.iterator, + cute.make_layout( + ((K, state.shape[1], state.shape[0]), K), + stride=((K, state.stride[1], state.stride[0]), 1), + ), + ) + from cutlass.cute.nvgpu import cpasync + + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), state_rows, self.smem_layout(), (K, K) + ) + items = args[10].shape[0] * self.heads + self.kernel(args[:-1], tma_atom, tma_tensor).launch( + grid=(cutlass.min(items, self.ctas), 1, 1), + block=(THREADS, 1, 1), + stream=args[-1], + use_pdl=True, + ) + + @cute.jit + def load_state(self, tma_atom, tma_tensor, a_tile, barrier, slot, head, lane): + """Warp-wide: TMA state (slot, head) into ``a_tile`` on ``barrier``.""" + from cutlass.cute.nvgpu import cpasync + + H: cutlass.Constexpr = self.heads + if lane == 0: + cute.arch.mbarrier_arrive_and_expect_tx(barrier, K * K * 2) + cute.arch.sync_warp() + s_state = cute.make_tensor(a_tile.iterator, self.smem_layout()) + g_state = cute.local_tile(tma_tensor, (K, K), (Int32(slot) * H + head, 0)) + t_smem, t_gmem = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(s_state, 0, 2), + cute.group_modes(g_state, 0, 2), + ) + cute.copy(tma_atom, t_gmem, t_smem, tma_bar_ptr=barrier) + + @cute.kernel + def kernel(self, args, tma_atom: cute.CopyAtom, tma_tensor: cute.Tensor): + qkv, raw_g, raw_beta, a_log, dt_bias, state, corr, kg, out = args[:9] + query_start, slots, conv_state, conv_w, accepted = args[9:14] + counter, lower_bound, null_slot = args[-3], args[-2], args[-1] + Q: cutlass.Constexpr = self.window + H: cutlass.Constexpr = self.heads + S: cutlass.Constexpr = Q + 2 # conv state length + W32: cutlass.Constexpr = self.conv_weight_f32 + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + ctas, _, _ = cute.arch.grid_dim() + warp, lane = tid // 32, tid % 32 + items = slots.shape[0] * H + + smem = cutlass.utils.SmemAllocator() + a_tile = alloc(smem, cutlass.BFloat16, (K * K,), align=1024) + b_hi = alloc(smem, cutlass.BFloat16, (K // 8 * B_KGROUP_BYTES // 2,), align=128) + b_lo = alloc(smem, cutlass.BFloat16, (K // 8 * B_KGROUP_BYTES // 2,), align=128) + # The activation stage (one item): raw q / k / v rows, gate rows, conv + # state rows (dim first: channel c's S values at c * S), taps, dt_bias. + qkv_s = alloc(smem, cutlass.BFloat16, (3, Q, K), align=128) + g_s = alloc(smem, cutlass.BFloat16, (Q, K), align=128) + conv_s = alloc(smem, cutlass.BFloat16, (3, K * S), align=128) + w_dtype = Float32 if W32 else cutlass.BFloat16 + w_s = alloc(smem, w_dtype, (3, K, CONV_W), align=128) + bias_s = alloc(smem, Float32, (K,), align=128) + # Item records (double buffered): item, slot, bos, qlen, first, loaded, + # valid; and A, raw beta[Q]. + meta = alloc(smem, Int32, (2, 8)) + meta_f = alloc(smem, Float32, (2, 8)) + # Reduced values: coupling c_it (i < t), d_it (i <= t), then |k_t|^2, |q_t|^2. + specs = [(0, t, i) for t in range(Q) for i in range(t)] + specs += [(1, t, i) for t in range(Q) for i in range(t + 1)] + NC = len(specs) + specs += [(2, t, 0) for t in range(Q)] + [(3, t, 0) for t in range(Q)] + NV = len(specs) + ROUNDS: cutlass.Constexpr = (NV + 31) // 32 + partial = alloc(smem, Float32, (ROUNDS, 4, 32)) + coef = alloc(smem, Float32, (2 * Q * Q,)) # c[t * Q + i] then d[...] + scales = alloc(smem, Float32, (2 * Q,)) # 1/|k_t|, K^-1/2 / |q_t| + betas = alloc(smem, Float32, (8,)) + table = alloc(smem, Int32, (NV,)) # per coupling: slot | na << 8 | nb << 16 + bars = alloc(smem, cutlass.Int64, (4,), align=8) + tmem_base = alloc(smem, Uint32, (1,), align=4) + act_full = bars.iterator + act_empty = bars.iterator + 1 + state_full = bars.iterator + 2 + mma_done = bars.iterator + 3 + + if warp == 0: + cute.arch.sync_warp() + tmem_alloc(tmem_base.iterator) + tmem_relinquish() + if tid == 0: + cute.arch.mbarrier_init(act_full, 2) + cute.arch.mbarrier_init(act_empty, 1) + cute.arch.mbarrier_init(state_full, 1) + cute.arch.mbarrier_init(mma_done, 1) + cute.arch.mbarrier_init_fence() + if tid < CONSUMERS: + # The zero probe columns, once. + b_offset = (tid // 8) * (B_KGROUP_BYTES // 2) + tid % 8 + for n in cutlass.range_constexpr(2 * Q, PROBES): + b_hi[b_offset + n * 8] = cutlass.BFloat16(0) + b_lo[b_offset + n * 8] = cutlass.BFloat16(0) + if tid < NC: + code = Int32(0) + for j in cutlass.range_constexpr(NC): + kind, t, i = specs[j] + na = NC + i # |k_i|^2 + nb = NC + t if kind == 0 else NC + Q + t # |k_t|^2 or |q_t|^2 + word = (kind * Q * Q + t * Q + i) | (na << 8) | (nb << 16) + if tid == j: + code = Int32(word) + table[tid] = code + proxy_fence() + tc_fence_before() + cute.arch.barrier() + tc_fence_after() + tmem = tmem_base[0] + + if warp == PRODUCER: + # ================= producer ================= + # Item scalars are read one item ahead (they are step metadata, written + # before this layer: the first item's are read before the PDL wait like + # the slots), so after the stage frees only bulk copies are issued. + T = qkv.shape[0] + item = cta + request = item // H + c_slot = Int64(slots[request]) + c_bos = query_start[request] + c_end = query_start[request + 1] + c_first = accepted[request] - 1 + c_a = a_log[item % H].to(Float32) + issued = Int32(0) + loads = Int32(0) # state loads so far + p_store = Int32(0) # the previous item's conv rows to store + p_slot = Int64(0) + p_head = Int32(0) + while item < items: + ticket = Int32(0) + if lane == 0: + ticket = Int32(take_ticket(Int64(counter.iterator.toint()))) + if ticket == items - 1: + reset_counter(Int64(counter.iterator.toint())) + head = item % H + slot, bos, qlen = c_slot, c_bos, c_end - c_bos + loaded = (slot > Int64(null_slot)) & (T > 0) + mi = issued % 2 + if issued > 0: + wait_bar(act_empty, (issued - 1) % 2) + # The previous item's rolled conv rows (written in place). + if p_store != 0: + if lane < 3: + channel = lane * H * K + p_head * K + store_bulk( + at_address(conv_state, p_slot, channel, 0), + conv_s.iterator + lane * (K * S), + K * S * 2, + ) + bulk_read_done() + cute.arch.sync_warp() + # Pre-dependency copies: conv state rows, taps, dt_bias. + if lane == 0: + if loaded: + w_bytes = 3 * K * CONV_W * (4 if W32 else 2) + arrive_expect(act_full, 3 * K * S * 2 + w_bytes + K * 4) + else: + arrive(act_full) + cute.arch.sync_warp() + if loaded: + if lane < 3: + channel = lane * H * K + head * K + bulk_load( + shared_address(conv_s.iterator) + lane * (K * S * 2), + at_address(conv_state, slot, channel, 0), + K * S * 2, + act_full, + ) + elif lane < 6: + x = lane - 3 + wb: cutlass.Constexpr = K * CONV_W * (4 if W32 else 2) + bulk_load( + shared_address(w_s.iterator) + x * wb, + at_address(conv_w, (x * H + head) * K, 0), + wb, + act_full, + ) + elif lane == 6: + bulk_load( + shared_address(bias_s.iterator), + at_address(dt_bias, head * K), + K * 4, + act_full, + ) + cute.arch.sync_warp() + if issued == 0: + if loaded: + self.load_state( + tma_atom, tma_tensor, a_tile, state_full, slot, head, lane + ) + cute.arch.griddepcontrol_wait() + if lane == 0: + meta[mi, 0] = item + meta[mi, 1] = Int32(slot) + meta[mi, 2] = bos + meta[mi, 3] = qlen + meta[mi, 4] = c_first + meta[mi, 5] = Int32(loaded) + meta[mi, 6] = Int32(slot > Int64(null_slot)) + meta_f[mi, 0] = cute.math.exp(c_a) + cute.arch.sync_warp() + if lane == 0: + if loaded: + arrive_expect(act_full, 4 * Q * K * 2) + else: + arrive(act_full) + cute.arch.sync_warp() + if loaded: # noqa: SIM102 (traced: `and` does not combine traced values) + if lane < 4 * Q: + x, t = lane // Q, lane % Q + src = Int64(0) + dst = Uint32(0) + row = bos + cutlass.min(t, cutlass.max(qlen - 1, 0)) + row = cutlass.min(row, T - 1) + if x < 3: + src = at_address(qkv, row, x * H * K + head * K) + dst = shared_address(qkv_s.iterator) + (x * Q + t) * (K * 2) + else: + src = at_address(raw_g, row, head * K) + dst = shared_address(g_s.iterator) + t * (K * 2) + bulk_load(dst, src, K * 2, act_full) + # The state, once the previous item's MMAs have read the buffer. + if loaded: + if issued > 0: + if loads > 0: + wait_bar(mma_done, (loads - 1) % 2) + self.load_state( + tma_atom, tma_tensor, a_tile, state_full, slot, head, lane + ) + loads += 1 + p_store = Int32(loaded & (qlen > 0)) + p_slot, p_head = slot, head + # The next item: its scalars, and its rows and state head to L2 + # while this one computes. + nxt = cute.arch.shuffle_sync(ticket, 0) + ctas + if nxt < items: + n_request, n_head = nxt // H, nxt % H + c_slot = Int64(slots[n_request]) + c_bos = query_start[n_request] + c_end = query_start[n_request + 1] + c_first = accepted[n_request] - 1 + c_a = a_log[n_head].to(Float32) + if (c_slot > Int64(null_slot)) & (T > 0): + n_qlen = c_end - c_bos + if lane < 4 * Q: + x, t = lane // Q, lane % Q + row = c_bos + cutlass.min(t, cutlass.max(n_qlen - 1, 0)) + row = cutlass.min(row, T - 1) + src = Int64(0) + if x < 3: + src = at_address(qkv, row, x * H * K + n_head * K) + else: + src = at_address(raw_g, row, n_head * K) + prefetch_bytes(src, K * 2) + elif lane < 4 * Q + 3: + channel = (lane - 4 * Q) * H * K + n_head * K + prefetch_bytes( + at_address(conv_state, c_slot, channel, 0), K * S * 2 + ) + elif lane == 4 * Q + 3: + prefetch_bytes( + at_address(state, c_slot, n_head, 0, 0), K * K * 2 + ) + issued += 1 + item = nxt + # The last conv rows, then the end record. + if issued > 0: + wait_bar(act_empty, (issued - 1) % 2) + if p_store != 0: + if lane < 3: + channel = lane * H * K + p_head * K + store_bulk( + at_address(conv_state, p_slot, channel, 0), + conv_s.iterator + lane * (K * S), + K * S * 2, + ) + cute.arch.sync_warp() + if lane == 0: + meta[issued % 2, 0] = items + arrive(act_full) + arrive(act_full) + if lane < 3: + bulk_done() + else: + # ================= compute warps ================= + cute.arch.griddepcontrol_wait() + c = tid + k_step = Int32(0) + state_phase = Int32(0) + mma_phase = Int32(0) + wait_bar(act_full, 0) + cur = meta[0, 0] + b_offset = (c // 8) * (B_KGROUP_BYTES // 2) + c % 8 + while cur < items: + mi = k_step % 2 + bos = meta[mi, 2] + qlen = meta[mi, 3] + loaded = meta[mi, 5] != 0 + valid = meta[mi, 6] + request, head = cur // H, cur % H + live = loaded & (qlen > 0) + qv = cute.make_rmem_tensor((Q,), Float32) + kv = cute.make_rmem_tensor((Q,), Float32) + vv = cute.make_rmem_tensor((Q,), Float32) + ev = cute.make_rmem_tensor((Q,), Float32) # e^{L_t} + wv = cute.make_rmem_tensor((Q,), Float32) # e^{-L_t} k_t + beta_raw = Float32(0) + if live: + if tid >= BETA_TID and tid < BETA_TID + Q: # noqa: SIM102 (traced: `and` does not combine traced values) + if tid - BETA_TID < qlen: + beta_raw = raw_beta[bos + tid - BETA_TID, head].to(Float32) + first = meta[mi, 4] + a = meta_f[mi, 0] + bias = bias_s[c] + hist = cute.make_rmem_tensor((3, 3), Float32) + raw = cute.make_rmem_tensor((3, Q), cutlass.BFloat16) + g_raw = cute.make_rmem_tensor((Q,), cutlass.BFloat16) + for x in cutlass.range_constexpr(3): + for i in cutlass.range_constexpr(3): + hist[x, i] = conv_s[x, c * S + first + i].to(Float32) + for t in cutlass.range_constexpr(Q): + raw[x, t] = qkv_s[x, t, c] + for t in cutlass.range_constexpr(Q): + g_raw[t] = g_s[t, c] + # The rolled conv state, in place (thread c owns channel c's + # S values); bulk-stored after the barrier. + for x in cutlass.range_constexpr(3): + for i in cutlass.range_constexpr(2): + conv_s[x, c * S + i] = hist[x, i + 1].to(cutlass.BFloat16) + for t in cutlass.range_constexpr(Q): + if t < qlen: + conv_s[x, c * S + 2 + t] = raw[x, t] + # Width-4 convolution and SiLU (the stock kernel's order). + for x in cutlass.range_constexpr(3): + if cutlass.const_expr(W32): + w = load_words( + at_ptr(w_s, x, c, 0), + count=4, + shared=True, + mutable=True, + result_type=Float32, + word_type="f32", + ) + else: + ww = load_words( + at_ptr(w_s, x, c, 0), count=2, shared=True, mutable=True + ) + w = [bf16_word(ww[j // 2], j % 2) for j in range(4)] + win = [hist[x, i] for i in range(3)] + for t in cutlass.range_constexpr(Q): + win.append(raw[x, t].to(Float32)) + acc = Float32(0) + for j in cutlass.range_constexpr(CONV_W): + if cutlass.const_expr(W32): + acc = fma(win[t + j], w[j], acc) + else: + prod = (win[t + j] * w[j]).to(cutlass.BFloat16) + acc = acc + prod.to(Float32) + # The stock kernel's SiLU: ex2.approx and div.full. + raw[x, t] = div_full( + acc, Float32(1) + cute.math.exp(-acc, fastmath=True) + ).to(cutlass.BFloat16) + cum = Float32(0) + for t in cutlass.range_constexpr(Q): + vv[t] = raw[2, t].to(Float32) + qv[t], kv[t] = Float32(0), Float32(0) + if t < qlen: + kg[request, head, t, c] = raw[1, t] + kg[request, head, t, K + c] = g_raw[t] + qv[t], kv[t] = raw[0, t].to(Float32), raw[1, t].to(Float32) + gate = sigmoid(a * (g_raw[t].to(Float32) + bias), True) + cum = cum + lower_bound * gate + ev[t] = cute.math.exp(cum, fastmath=True) + wv[t] = cute.math.exp(-cum, fastmath=True) * kv[t] + # Unnormalized probes p_t = e^{L_t} k_t, r_t = e^{L_t} q_t. + for t in cutlass.range_constexpr(Q): + for n, value in ((t, ev[t] * kv[t]), (Q + t, ev[t] * qv[t])): + hi, lo = bf16_split(value) + b_hi[b_offset + n * 8] = hi + b_lo[b_offset + n * 8] = lo + # This channel's products, reduce-scattered over the warp. + for rnd in cutlass.range_constexpr(ROUNDS): + count = min(32, NV - 32 * rnd) + LOG: cutlass.Constexpr = (count - 1).bit_length() + width: cutlass.Constexpr = 1 << LOG + vals = [] + for j in cutlass.range_constexpr(width): + value = Float32(0) + if cutlass.const_expr(j < count): + kind, t, i = specs[32 * rnd + j] + if cutlass.const_expr(kind == 0): + value = wv[i] * (ev[t] * kv[t]) + elif cutlass.const_expr(kind == 1): + value = wv[i] * (ev[t] * qv[t]) + elif cutlass.const_expr(kind == 2): + value = kv[t] * kv[t] + else: + value = qv[t] * qv[t] + vals.append(value) + for st in cutlass.range_constexpr(LOG): + shift = 16 >> st + half = width >> (st + 1) + upper = (lane & shift) != 0 + kept = [] + for j in cutlass.range_constexpr(half): + lo_v, hi_v = vals[j], vals[j + half] + send, keep = hi_v, lo_v + if upper: + send, keep = lo_v, hi_v + kept.append( + keep + cute.arch.shuffle_sync_bfly(send, shift) + ) + vals = kept + total = vals[0] + for st in cutlass.range_constexpr(LOG, 5): + total = total + cute.arch.shuffle_sync_bfly(total, 16 >> st) + stride: cutlass.Constexpr = 32 // width + if lane % stride == 0: + partial[rnd, warp, lane // stride] = total + # Probes and conv rows visible to the async proxy; TMEM reads of the + # previous item ordered before this item's MMAs. + proxy_fence() + tc_fence_before() + consumer_sync() + if tid == 0: + arrive(act_empty) # the stage is read (conv rows: the producer) + if loaded: + wait_bar(state_full, state_phase) + if live: + tc_fence_after() + a_base = shared_address(a_tile.iterator) + hi_base = shared_address(b_hi.iterator) + lo_base = shared_address(b_lo.iterator) + for kt in cutlass.range_constexpr(K // 16): + if cutlass.const_expr(self.swizzle): + # K step kt: 32 bytes into its 128-byte atom row. + a_desc = sw128_desc( + a_base + (kt // 4) * (K * 128) + (kt % 4) * 32 + ) + else: + a_desc = smem_desc( + a_base + kt * 2 * A_KGROUP_BYTES, + A_KGROUP_BYTES, + 128, + ) + for h, base in enumerate((hi_base, lo_base)): + b_desc = smem_desc( + base + kt * 2 * B_KGROUP_BYTES, + B_KGROUP_BYTES, + 128, + ) + umma( + tmem, + a_desc, + b_desc, + 0 if kt == 0 and h == 0 else 1, + ) + umma_commit(mma_done) + else: + arrive(mma_done) # the buffer is free (no tokens) + if loaded: + state_phase = state_phase ^ 1 + if live: + # Coupling scalars and norm scales from the warp partials. + if tid < NV: + r0, j0 = tid // 32, tid % 32 + total = (partial[r0, 0, j0] + partial[r0, 1, j0]) + ( + partial[r0, 2, j0] + partial[r0, 3, j0] + ) + if tid < NC: + word = table[tid] + na, nb = (word >> 8) & 0xFF, (word >> 16) & 0xFF + sa = ( + partial[na // 32, 0, na % 32] + + partial[na // 32, 1, na % 32] + ) + ( + partial[na // 32, 2, na % 32] + + partial[na // 32, 3, na % 32] + ) + sb = ( + partial[nb // 32, 0, nb % 32] + + partial[nb // 32, 1, nb % 32] + ) + ( + partial[nb // 32, 2, nb % 32] + + partial[nb // 32, 3, nb % 32] + ) + scale = cute.math.rsqrt( + sa + Float32(EPS) + ) * cute.math.rsqrt(sb + Float32(EPS)) + if word & 0xFF >= Q * Q: + scale = scale * Float32(K**-0.5) + coef[word & 0xFF] = total * scale + else: + scale = cute.math.rsqrt(total + Float32(EPS)) + if tid >= NC + Q: + scale = scale * Float32(K**-0.5) + scales[tid - NC] = scale + elif tid >= BETA_TID and tid < BETA_TID + Q: + beta = Float32(0) + if tid - BETA_TID < qlen: + beta = sigmoid(beta_raw, True) + betas[tid - BETA_TID] = beta + consumer_sync() + if live: + wait_bar(mma_done, mma_phase) + tc_fence_after() + cute.arch.sync_warp() + proj = list(tmem_load_row(tmem + ((warp * 32) << 16))) + r = tid + u = cute.make_rmem_tensor((Q,), Float32) + for t in cutlass.range_constexpr(Q): + acc = vv[t] - proj[t] * scales[t] + for i in cutlass.range_constexpr(t): + acc = acc - coef[t * Q + i] * u[i] + u[t] = betas[t] * acc + o = proj[Q + t] * scales[Q + t] + for i in cutlass.range_constexpr(t + 1): + o = fma(coef[Q * Q + t * Q + i], u[i], o) + if t < qlen: + out[bos + t, head * K + r] = o.to(out.element_type) + corr[request, head, t, r] = u[t] + if loaded: + mma_phase = mma_phase ^ 1 + if valid == 0: + for t in cutlass.range_constexpr(Q): + if t < qlen: + out[bos + t, head * K + tid] = out.element_type(0) + k_step += 1 + wait_bar(act_full, k_step % 2) + cur = meta[k_step % 2, 0] + + cute.arch.griddepcontrol_launch_dependents() + tc_fence_before() + cute.arch.barrier() + if warp == 0: + tc_fence_after() + cute.arch.sync_warp() + tmem_dealloc(tmem) + + +@cute.jit +def sw128_desc(address): + """128-byte swizzled K-major UMMA descriptor: 8-row groups 1024 bytes apart.""" + static = (1 << 16) | (64 << 32) | (1 << 46) | (2 << 61) + return cutlass.Uint64((address >> 4) & 0x3FFF) | cutlass.Uint64(static) + + +def at_ptr(tensor, *coord): + return tensor.iterator + cute.crd2idx(coord, tensor.layout) + + +def at_address(tensor, *coord): + """64-bit global address of ``tensor[coord]``.""" + return Int64(at_ptr(tensor, *coord).toint()) + + +def bf16_word(word, high): + return bf16(word, high) diff --git a/vllm/models/kimi_k3/nvidia/recoverssm/kernels.py b/vllm/models/kimi_k3/nvidia/recoverssm/kernels.py new file mode 100644 index 000000000000..0144dfe8eb4f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/recoverssm/kernels.py @@ -0,0 +1,374 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host side of the CuTe KDA RecoverSSM kernels (SM10x tensor cores). + +* ``CuteRecoverSSMVerify``: the speculative verify of one KDA layer with the short + convolution fused in (``cutedsl/verify_ws.py`` for the dim-first conv state, + ``cutedsl/verify.py`` otherwise). It writes the outputs, the rolled conv state + and the verify records, like ``causal_conv1d_update`` followed by + ``kda_recoverssm_verify``. +* ``cute_recoverssm_commit``: the state fold of a prepared commit plan for all + KDA layers in one launch (``cutedsl/commit.py``), replacing the Triton fold of + ``KDARecoverSSMCommitContext.commit``. + +The kernels compile when the KV cache is bound, before CUDA graph capture; +nothing compiles under capture. +""" + +from functools import cache +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + KDARecoverSSMRecords, + ) + +K = 128 # KDA head dim (keys = values) +CONV_WIDTH = 4 +MAX_WINDOW = 8 + + +def can_use_cute_recoverssm( + *, + head_dim: int, + window: int, + conv_width: int, + has_conv_bias: bool, + activation_dtype: torch.dtype, + conv_state_dtype: torch.dtype, + state_dtype: torch.dtype, + conv_weight_dtype: torch.dtype, + lower_bound: float | None, +) -> tuple[bool, str | None]: + """Whether the CuTe verify and commit can serve a KDA layer, and if not, + why.""" + if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ): + return False, "it needs SM10x tensor cores (tcgen05)" + if head_dim != K or conv_width != CONV_WIDTH: + return False, f"it needs head_dim {K} and a width-{CONV_WIDTH} convolution" + if has_conv_bias: + return False, "it does not apply a convolution bias" + if not 1 <= window <= MAX_WINDOW: + return False, f"it verifies 1 to {MAX_WINDOW} tokens per step" + if lower_bound is None: + return False, "it needs a bounded KDA gate (gate_lower_bound)" + if (activation_dtype, conv_state_dtype, state_dtype) != (torch.bfloat16,) * 3: + return False, "it needs bfloat16 activations, conv state and SSM state" + if conv_weight_dtype not in (torch.bfloat16, torch.float32): + return False, "it needs a bfloat16 or float32 convolution weight" + return True, None + + +def _architecture(device: torch.device) -> str: + from vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.host import ( + architecture, + ) + + return architecture(device) + + +def _ctas(device: torch.device, per_sm: int) -> int: + return per_sm * torch.cuda.get_device_properties(device).multi_processor_count + + +@cache +def _compile_verify( + heads: int, + window: int, + arch: str, + ctas: int, + warp_specialized: bool, + conv_dim_first: bool, + conv_weight_dtype: torch.dtype, +): + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream + + from vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.host import ( + dynamic_fake, + ) + + bf16, f32, i32 = torch.bfloat16, torch.float32, torch.int32 + fakes = ( + dynamic_fake(2, bf16), # raw q / k / v projection [T, 3HK] + dynamic_fake(2, bf16), # raw gate [T, HK] + dynamic_fake(2, bf16), # raw beta [T, H] + dynamic_fake(1, f32), # A_log [H] + dynamic_fake(1, f32), # dt_bias [HK] + dynamic_fake(4, bf16), # SSM state [slots, H, V, K] + dynamic_fake(4, f32), # corrections [rows, H, Q, V] + dynamic_fake(4, bf16), # key / gate records [rows, H, Q, 2K] + dynamic_fake(2, bf16), # output [T, HV] + dynamic_fake(1, i32), # query start [B + 1] + dynamic_fake(1, i32), # state slots [B] + dynamic_fake(3, bf16), # conv state + dynamic_fake(2, conv_weight_dtype), # conv weight [3HK, 4] + dynamic_fake(1, i32), # accepted tokens [B] + dynamic_fake(1, i32), # work counter [1] + ) + if warp_specialized: + from .cutedsl.verify_ws import KdaTcVerify2 + + kernel: Any = KdaTcVerify2(heads, window, ctas, conv_weight_dtype == f32) + else: + from .cutedsl.verify import KdaTcVerify + + kernel = KdaTcVerify( + heads, + window, + conv=True, + ctas=ctas, + conv_dim_first=conv_dim_first, + conv_weight_f32=conv_weight_dtype == f32, + ) + return cute.compile( + kernel, + *fakes, + cutlass.Float32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options=f"--gpu-arch={arch} --enable-tvm-ffi", + ) + + +def _aligned(*tensors: torch.Tensor) -> bool: + return all(tensor.data_ptr() % 16 == 0 for tensor in tensors) + + +def _check_layouts(state: torch.Tensor, *tensors: torch.Tensor) -> None: + """The KV cache pages, records and parameters the CuTe kernels read: the + state streams in by TMA (K-element rows of contiguous heads, slots a multiple + of 16 bytes apart), and every tensor starts on a 16-byte boundary.""" + if state.stride()[1:] != (K * K, K, 1) or state.stride(0) * 2 % 16: + raise ValueError( + f"KDA RecoverSSM CuTe kernels: state strides {state.stride()} need " + "contiguous heads and 16-byte aligned slots" + ) + if not _aligned(state, *tensors): + raise ValueError( + "KDA RecoverSSM CuTe kernels: the KV cache pages, records and " + "parameters must start on 16-byte boundaries" + ) + + +class CuteRecoverSSMVerify: + """The CuTe verify for one KDA layer's shapes (compiled kernels are shared).""" + + def __init__( + self, + *, + num_heads: int, + window: int, + conv_dim_first: bool, + conv_weight_dtype: torch.dtype, + device: torch.device, + ) -> None: + self.num_heads = num_heads + self.window = window + self.conv_dim_first = conv_dim_first + self.conv_weight_dtype = conv_weight_dtype + self.device = device + # Each launch takes one ticket per work item and leaves the counter + # zeroed for the next one. + self.counter = torch.zeros(1, dtype=torch.int32, device=device) + self._compiled: dict[bool, Any] = {} + # Page and parameter layouts already checked (fixed per KV cache). + self._checked: set[tuple] = set() + + def compile(self) -> None: + """Compile the variants this layer launches (outside graph capture).""" + from .cutedsl import verify, verify_ws + + arch = _architecture(self.device) + variants = [False] + ([True] if self.conv_dim_first else []) + for warp_specialized in variants: + per_sm = verify_ws.CTAS_PER_SM if warp_specialized else verify.CTAS_PER_SM + self._compiled[warp_specialized] = _compile_verify( + self.num_heads, + self.window, + arch, + _ctas(self.device, per_sm), + warp_specialized, + self.conv_dim_first, + self.conv_weight_dtype, + ) + + def __call__( + self, + *, + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + state: torch.Tensor, + records: "KDARecoverSSMRecords", + query_start_loc: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + out: torch.Tensor, + ) -> bool: + """Verify ``state_indices.shape[0]`` requests. ``mixed_qkv`` [T, 3HK] is + the raw projection (before the convolution), ``raw_gate`` [T, HK], + ``raw_beta`` [T, H], ``out`` [T, HV]; ``conv_state`` is the stored + [slots, 3HK, Q + 2] (dim first) or [slots, Q + 2, 3HK] page view. + Request ``b`` writes its verify records to row ``b`` of ``records``. + + Returns False without launching when an activation does not start on a + 16-byte boundary (e.g. for an odd token offset into projection rows that + are not 16-byte aligned): the caller then runs the Triton verify.""" + correction, key_gate = records.correction, records.key_gate + static = (state, A_log, dt_bias, correction, key_gate, conv_state, conv_weight) + key = tuple((t.data_ptr(), t.stride()) for t in static) + if key not in self._checked: + _check_layouts(*static) + self._checked.add(key) + if not _aligned( + mixed_qkv, + raw_gate, + raw_beta, + out, + query_start_loc, + state_indices, + num_accepted_tokens, + ): + return False + warp_specialized = self.conv_dim_first and self._bulk_copy_aligned( + mixed_qkv, raw_gate, conv_state, conv_weight + ) + # TVM-FFI launch on the current stream, straight from the torch tensors. + self._compiled[warp_specialized]( + mixed_qkv, + raw_gate, + raw_beta, + A_log, + dt_bias, + state, + correction, + key_gate, + out, + query_start_loc, + state_indices, + conv_state, + conv_weight, + num_accepted_tokens, + self.counter, + lower_bound, + NULL_BLOCK_ID, + ) + return True + + def _bulk_copy_aligned( + self, + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + ) -> bool: + """The warp-specialized verify bulk-copies q / k / v and gate rows (16-byte + aligned rows), conv state channel blocks (contiguous per slot) and conv + taps; otherwise the plain tcgen05 verify runs.""" + conv_len = self.window + CONV_WIDTH - 2 + return ( + mixed_qkv.stride(0) % 8 == 0 + and raw_gate.stride(0) % 8 == 0 + and conv_weight.stride() == (CONV_WIDTH, 1) + and conv_state.stride()[1:] == (conv_len, 1) + and conv_state.stride(0) * conv_state.element_size() % 16 == 0 + ) + + +@cache +def _compile_commit( + heads: int, window: int, arch: str, param_dtype: torch.dtype, align: bool +): + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream + + from vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.host import ( + dynamic_fake, + ) + + from .cutedsl.commit import KdaCommit + + i32, i64 = torch.int32, torch.int64 + fakes = ( + dynamic_fake(1, i64), # state base addresses [L] + dynamic_fake(1, i64), # correction base addresses [L] + dynamic_fake(1, i64), # key / gate base addresses [L] + dynamic_fake(1, i64), # state slot strides [L] + dynamic_fake(1, i64), # correction row strides [L] + dynamic_fake(1, i64), # key / gate row strides [L] + dynamic_fake(2, param_dtype), # A_log [L, H] + dynamic_fake(3, param_dtype), # dt_bias [L, H, K] + dynamic_fake(1, i32), # source slots [B] + dynamic_fake(1, i32), # commit lengths [B] + dynamic_fake(1, i32), # destination slots [B] + dynamic_fake(1, i32), # boundary slots [B] + dynamic_fake(1, i32), # boundary lengths [B] + ) + return cute.compile( + KdaCommit(heads, window, align), + *fakes, + cutlass.Float32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options=f"--gpu-arch={arch} --enable-tvm-ffi", + ) + + +def compile_state_commit( + num_heads: int, window: int, device: torch.device, align: bool +) -> None: + """Compile the CuTe state commit (before serving: it runs eagerly).""" + _compile_commit(num_heads, window, _architecture(device), torch.float32, align) + + +def cute_recoverssm_commit( + context: "KDARecoverSSMCommitContext", + state_indices: torch.Tensor, + batch: int, + align: bool, +) -> None: + """Fold the accepted tokens of a prepared commit plan into the state of all + of ``context``'s layers in one launch. It runs after sampling, outside CUDA + graphs, on the current stream (TVM-FFI).""" + state = context.checkpoints[0] + assert context.lower_bound is not None + compiled = _compile_commit( + state.shape[1], + context.spec_query_len, + _architecture(state.device), + torch.float32, + align, + ) + compiled( + context.state_base_addrs, + context.correction_cache_base_addrs, + context.kg_cache_base_addrs, + context.state_block_strides, + context.correction_cache_block_strides, + context.kg_cache_block_strides, + context.A_log, + context.dt_bias, + state_indices[:batch], + context.commit_lens[:batch], + context.final_state_indices[:batch], + context.boundary_state_indices[:batch], + context.boundary_recovery_lens[:batch], + context.lower_bound, + NULL_BLOCK_ID, + ) From c50aca14416ef29f6ed0ed256302f30646726dcd Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 23 Oct 2026 20:43:48 -0700 Subject: [PATCH 004/631] [Doc][Recipe] MiniMax-M3.1-Flash: DEP2 + DEP4 PD on GB200/GB300, 7-token FLASHINFER_TS DSpark on NVIDIA (#63428) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- .../models/MiniMaxAI/MiniMax-M3.1-Flash.yaml | 150 +++++++++++++++--- 1 file changed, 126 insertions(+), 24 deletions(-) diff --git a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml index 6d95ca6305bc..6e4c5052a567 100644 --- a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml +++ b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml @@ -2,7 +2,7 @@ meta: title: "MiniMax-M3.1-Flash" slug: "minimax-m3.1-flash" provider: "MiniMax" - description: "MiniMax M3.1-Flash vision-language MoE (~428B total / 23B active) with MSA sparse attention throughout and a built-in DSpark draft; quantization-aware trained, MXFP8 dense layers and NVFP4 experts. Served TP4 on NVIDIA Blackwell or AMD MI355X (single node or prefill/decode disaggregated), and TP8 on NVIDIA Hopper." + description: "MiniMax M3.1-Flash vision-language MoE (~428B total / 23B active) with MSA sparse attention throughout and a built-in DSpark draft; quantization-aware trained, MXFP8 dense layers and NVFP4 experts. Served TP4 on NVIDIA Blackwell or AMD MI355X, prefill/decode disaggregated as TP4 + TP4 on MI355X or DEP2 + DEP4 on GB200/GB300, and TP8 on NVIDIA Hopper." date_added: 2026-10-23 date_updated: 2026-10-23 tasks: @@ -52,7 +52,7 @@ features: - "--reasoning-parser" - "minimax_m3" spec_decoding: - description: "DSpark speculative decoding with the draft shipped in the checkpoint (`dspark/`), three speculative tokens, FP8 draft KV cache. Draft attention backend is per-brand: FLASH_ATTN on NVIDIA, TRITON_ATTN on AMD." + description: "DSpark speculative decoding with the draft shipped in the checkpoint (`dspark/`), FP8 draft KV cache. Per brand: seven speculative tokens with FLASHINFER_TS draft attention on NVIDIA, three with TRITON_ATTN on AMD." # Per-brand draft attention backend (a feature's top-level `args` is the # fallback when no hardware_overrides. matches; both brands are # covered here). @@ -63,7 +63,7 @@ features: nvidia: args: - "--speculative-config" - - '{"method":"dspark","num_speculative_tokens":3,"attention_backend":"FLASH_ATTN","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' + - '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' amd: args: - "--speculative-config" @@ -278,9 +278,46 @@ strategy_overrides: # never decodes) and offloads KV to host memory in 1 GiB HugeTLB pages; both # nodes use the BLHNC KV layout, so NIXL reads one region per block instead # of one per layer. + # GB200/GB300 1P1D as in the guide's NVIDIA PD section: a DEP2 prefill on + # one tray and a DEP4 decode on a second tray. vllm-router takes + # one DP size for both pools, so it sends each request to the endpoint and + # each engine picks its own DP rank. pd_cluster: + router: + intra_node_data_parallel_size: 1 prefill: + nodes: + gb200: 1 + gb300: 1 + parallelism: + default: tp + gb200: dep + gb300: dep hardware_overrides: + gb200: &m31_gb_prefill + extra_args: + - "--data-parallel-size" + - "2" + - "--data-parallel-size-local" + - "2" + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--numa-bind" + - "--max-num-seqs" + - "64" + - "--max-num-batched-tokens" + - "32768" + - "--long-prefill-token-threshold" + - "8192" + - "--gpu-memory-utilization" + - "0.9" + - "--enable-prefix-caching" + - "--kv-transfer-config" + - '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"engine_ttl":0}}' + extra_env: + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_ENGINE_READY_TIMEOUT_S: "3600" + gb300: *m31_gb_prefill mi355x: extra_args: - "--tensor-parallel-size" @@ -295,7 +332,37 @@ strategy_overrides: CUDA_VISIBLE_DEVICES: "0,1,2,3" VLLM_KV_CACHE_LAYOUT: "BLHNC" decode: + nodes: + gb200: 1 + gb300: 1 + parallelism: + default: tp + gb200: dep + gb300: dep hardware_overrides: + # 28 sequences per rank; graphs up to 28 x (1 + 7 DSpark tokens). + gb200: &m31_gb_decode + extra_args: + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--numa-bind" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":224}' + - "--max-num-seqs" + - "28" + - "--max-num-batched-tokens" + - "8192" + - "--gpu-memory-utilization" + - "0.95" + - "--watermark" + - "0.05" + - "--no-enable-prefix-caching" + - "--kv-transfer-config" + - '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"engine_ttl":0}}' + extra_env: + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_ENGINE_READY_TIMEOUT_S: "3600" + gb300: *m31_gb_decode mi355x: extra_args: - "--tensor-parallel-size" @@ -407,7 +474,7 @@ guide: | --kv-cache-dtype nvfp4 \ --attention-config '{"indexer_kv_dtype": "nvfp4", "use_prefill_query_quantization": true}' \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ - --speculative-config '{"method":"dspark","num_speculative_tokens":3,"attention_backend":"FLASH_ATTN","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ + --speculative-config '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ --max-num-seqs 256 \ --disable-hybrid-kv-cache-manager ``` @@ -501,51 +568,86 @@ guide: | echo 2240 | sudo tee /sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages ``` - ### Prefill/decode disaggregation (1P1D, TP4 + TP4, NVIDIA GB200) + ### Prefill/decode disaggregation (1P1D, DEP2 + DEP4, NVIDIA GB200 / GB300) - The NVIDIA pair runs NVFP4 KV end to end over NIXL (kv_buffer_device - `cuda`); the prefill node skips CUDA graphs (`--enforce-eager`) and carries - a 32768-token batch budget, the decode node runs FULL_DECODE_ONLY graphs. + A DEP2 prefill on one 4-GPU tray feeds a DEP4 decode on a second tray over + NIXL. Both roles run attention data parallel with the NVFP4 experts + on the DeepGEMM MegaMoE kernel over expert parallelism, the V2 model runner, + NUMA binding, and a seven-token DSpark draft on the FlashInfer PrimiTS + (`FLASHINFER_TS`) backend. The settings were tuned on GB300; GB200 uses the + same commands. Prefill node: ```bash - VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_NIXL_SIDE_CHANNEL_HOST= \ + VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_USE_V2_MODEL_RUNNER=1 \ + VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_NIXL_SIDE_CHANNEL_HOST= \ vllm serve MiniMaxAI/MiniMax-M3.1-Flash \ - -tp 4 \ + --data-parallel-size 2 \ + --enable-expert-parallel \ + --moe-backend deep_gemm_mega_moe \ + --numa-bind \ --reasoning-parser minimax_m3 \ --block-size 128 \ --kv-cache-dtype nvfp4 \ - --attention-config '{"indexer_kv_dtype": "nvfp4"}' \ - --speculative-config '{"method":"dspark","num_speculative_tokens":3,"attention_backend":"FLASH_ATTN","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ - --max-num-seqs 16 \ + --attention-config '{"indexer_kv_dtype": "nvfp4", "use_prefill_query_quantization": true}' \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --speculative-config '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ + --max-num-seqs 64 \ --max-num-batched-tokens 32768 \ + --long-prefill-token-threshold 8192 \ --gpu-memory-utilization 0.9 \ - --enforce-eager \ + --enable-prefix-caching \ + --disable-hybrid-kv-cache-manager \ --port 8100 \ - --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"kv_lease_duration":120}}' + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"engine_ttl":0}}' ``` Decode node: ```bash - VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_NIXL_SIDE_CHANNEL_HOST= \ + VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_USE_V2_MODEL_RUNNER=1 \ + VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_NIXL_SIDE_CHANNEL_HOST= \ vllm serve MiniMaxAI/MiniMax-M3.1-Flash \ - -tp 4 \ + --data-parallel-size 4 \ + --enable-expert-parallel \ + --moe-backend deep_gemm_mega_moe \ + --numa-bind \ --reasoning-parser minimax_m3 \ --block-size 128 \ --kv-cache-dtype nvfp4 \ --attention-config '{"indexer_kv_dtype": "nvfp4", "use_prefill_query_quantization": true}' \ - --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}' \ - --speculative-config '{"method":"dspark","num_speculative_tokens":3,"attention_backend":"FLASH_ATTN","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ - --max-num-seqs 88 \ - --max-num-batched-tokens 512 \ - --gpu-memory-utilization 0.9 \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "max_cudagraph_capture_size": 224}' \ + --speculative-config '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ + --max-num-seqs 28 \ + --max-num-batched-tokens 8192 \ + --gpu-memory-utilization 0.95 \ + --watermark 0.05 \ + --no-enable-prefix-caching \ --disable-hybrid-kv-cache-manager \ --port 8200 \ - --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"kv_lease_duration":120}}' + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"engine_ttl":0}}' ``` + - `VLLM_ENGINE_READY_TIMEOUT_S=3600`: with data parallel, each API server + waits this long for its engines to start (default 600 s). A cold prefill + start (JIT builds, CUDA graph capture, FlashInfer autotune) takes longer + than 10 minutes. + - `--max-num-seqs` is per DP rank. The decode CUDA graphs cover 28 requests + of 1 + 7 tokens (224). + - `--watermark 0.05` keeps 5% of the decode KV cache free when admitting + requests, so extra requests wait in the queue instead of preempting + running ones. On GB300 (about 4.5M KV tokens per decode rank) at a + 200K-token mean context, expect about 19-21 running requests per rank. + - Keep the decode `--max-num-batched-tokens` at 8192: a preempted request + recomputes its context in chunks of that size, and every DP rank runs + eager steps until it finishes. + - `engine_ttl: 0` turns off NIXL's cleanup of idle remote engines; with the + default one-hour TTL, a decode crashed an hour after its prefill went away. + - Prefix caching stays on the prefill node only. For multi-turn traffic, the + prefill node can also offload KV to host memory, 180 GiB per rank: + `--kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_connector_extra_config":{"engine_ttl":0}},{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":193273528320,"lazy_offload":true}}]}}'`. + For a quick test, put the toy NIXL proxy from the vLLM repository in front of the two nodes and send requests to port 8000: @@ -595,7 +697,7 @@ guide: | cache block size must match (otherwise `No common block size for 16`). - **Parsers.** `--tool-call-parser` and `--reasoning-parser` both use `minimax_m3`. - - **HugeTLB pages (PD prefill node).** See the PD section above. + - **HugeTLB pages (PD prefill node).** See the AMD PD section above. ## References From 4920df120d8ef6deafc3954bb48cf4ae2e927b92 Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:21:37 +0900 Subject: [PATCH 005/631] [Bugfix][Model] Fix Gemma4 encoder CG with `max_soft_tokens` set (#60687) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> --- .../multimodal/processing/test_gemma4.py | 13 +++++++ vllm/model_executor/models/gemma4_mm.py | 38 ++++++++----------- 2 files changed, 29 insertions(+), 22 deletions(-) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index 5d5670e67e31..f9d935991e90 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -447,3 +447,16 @@ def test_encoder_cudagraph_mixed_width_images_pack_by_padded_width(): assert selected[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] == (70,) selected = model.select_encoder_cudagraph_items(mm_kwargs, [0, 1]) assert selected[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] == (280,) + + +@pytest.mark.parametrize( + "mm_processor_kwargs", + [{"max_soft_tokens": 560}, {"images_kwargs": {"max_soft_tokens": 560}}], +) +def test_encoder_cudagraph_config_respects_max_soft_tokens(mm_processor_kwargs): + model = _gemma4_for_encoder_cudagraph() + model.config.text_config = SimpleNamespace(hidden_size=8) + model.multimodal_config = SimpleNamespace(mm_processor_kwargs=mm_processor_kwargs) + + config = model.get_encoder_cudagraph_config() + assert config.capture_axes == ((70, 140, 280, 560),) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index f11217d8becd..85f271b83dd4 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -1303,6 +1303,19 @@ def _encoder_chunk( # Image processing # ------------------------------------------------------------------ # + def _get_max_soft_tokens(self) -> int: + max_soft_tokens = int(self.config.vision_config.default_output_length) + mm_processor_kwargs = getattr( + getattr(self, "multimodal_config", None), + "mm_processor_kwargs", + None, + ) + if isinstance(mm_processor_kwargs, Mapping): + value, _ = _get_max_soft_tokens(mm_processor_kwargs) + if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS: + max_soft_tokens = value + return max_soft_tokens + def _process_image_input( self, image_input: Gemma4ImageInputs, @@ -1333,17 +1346,7 @@ def _process_image_input( pool_position_ids = pixel_position_ids if self._enable_mm_lora: - max_soft_tokens = vision_cfg.default_output_length - mm_processor_kwargs = getattr( - getattr(self, "multimodal_config", None), - "mm_processor_kwargs", - None, - ) - if isinstance(mm_processor_kwargs, Mapping): - value, _ = _get_max_soft_tokens(mm_processor_kwargs) - if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS: - max_soft_tokens = value - + max_soft_tokens = self._get_max_soft_tokens() max_patches = max_soft_tokens * pooling_k2 padded_position_ids: list[torch.Tensor] = [] for idx in range(total_images): @@ -1663,7 +1666,7 @@ def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None: batch_size, num_patches = src.shape[0], src.shape[1] dst[:batch_size, :num_patches].copy_(src) - configured_max_tokens = self.config.vision_config.default_output_length + configured_max_tokens = self._get_max_soft_tokens() # One graph per row width (soft tokens of one image or video frame), # so a budget's rows are not all padded to the widest image. row_widths = tuple( @@ -2304,16 +2307,7 @@ def get_mm_lora_token_counts( if modality == "image": pixel_values_key = "pixel_values" - max_soft_tokens = int(vision_config.default_output_length) - mm_processor_kwargs = getattr( - getattr(self, "multimodal_config", None), - "mm_processor_kwargs", - None, - ) - if isinstance(mm_processor_kwargs, Mapping): - val, _ = _get_max_soft_tokens(mm_processor_kwargs) - if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: - max_soft_tokens = val + max_soft_tokens = self._get_max_soft_tokens() else: pixel_values_key = "pixel_values_videos" max_soft_tokens = _VIDEO_MAX_SOFT_TOKENS From a5b71c1eb3b0e4def2f0effd8cc4fd0c5e18cc3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Sat, 24 Oct 2026 07:46:56 +0200 Subject: [PATCH 006/631] [Security] Refresh vulnerable runtime and test dependencies (#62835) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Codex Co-authored-by: vllm-agent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- requirements/common.txt | 7 +++--- requirements/kv_connectors.txt | 2 +- requirements/kv_connectors_cu12.txt | 6 +++-- requirements/test/cpu.txt | 23 ++++++++++-------- requirements/test/cuda.in | 10 ++++---- requirements/test/cuda.txt | 24 +++++++++++-------- requirements/test/nightly-torch.txt | 4 ++-- requirements/test/rocm.in | 12 +++++----- requirements/test/rocm.txt | 37 ++++++++++++++++------------- requirements/test/xpu.in | 6 ++--- requirements/test/xpu.txt | 21 +++++++++------- rust/Cargo.lock | 4 ++-- 12 files changed, 86 insertions(+), 70 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 6b67eca537c5..d235d9e8bd81 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -16,12 +16,13 @@ protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33 protobuf == 7.36.1; platform_machine == "s390x" fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides. starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1 -aiohttp >= 3.14.3 +aiohttp >= 3.14.4 +multidict >= 6.9.1 openai >= 2.25.0 # For Responses API with namespace tools pydantic >= 2.12.0 orjson>=3.11.6 prometheus_client >= 0.18.0 -pillow >= 12.3.0 # Required for image processing +pillow >= 13.0.0 # Required for image processing prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0 tiktoken >= 0.6.0 # Required for DBRX tokenizer lm-format-enforcer == 0.11.3 @@ -54,7 +55,7 @@ setproctitle # Used to set process names for better debugging and monitoring oss-harmony >= 0.0.11 # Required for gpt-oss; includes embedded vocabularies anthropic >= 0.71.0 model-hosting-container-standards >= 0.1.14, < 1.0.0 -mcp >= 2.0.0, < 3.0.0 +mcp >= 2.2.0, < 3.0.0 opentelemetry-sdk >= 1.27.0 opentelemetry-api >= 1.27.0 opentelemetry-exporter-otlp >= 1.27.0 diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index f6cacd200b78..aff67fce2ea1 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,4 +1,4 @@ -lmcache >= 0.3.9 +lmcache >= 0.5.6 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. cupy-cuda13x != 14.1.0 diff --git a/requirements/kv_connectors_cu12.txt b/requirements/kv_connectors_cu12.txt index c1c064609671..4f5c26fc3bbc 100644 --- a/requirements/kv_connectors_cu12.txt +++ b/requirements/kv_connectors_cu12.txt @@ -1,5 +1,7 @@ -# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here. -lmcache == 0.4.4 +# The CUDA 12.9 wheels are published separately from the CUDA 13 PyPI build. +--find-links https://github.com/LMCache/LMCache/releases/expanded_assets/v0.5.6-cu129 +lmcache == 0.5.6 +cuda-python >= 12, < 13 cupy-cuda12x != 14.1.0 nixl == 1.5.0 mooncake-transfer-engine >= 0.3.12 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 599b98a43f9e..1fe62a75d014 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -8,7 +8,7 @@ accelerate==1.13.0 # via peft aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.14.3 +aiohttp==3.14.5 # via # -r requirements/test/../common.txt # aiohttp-cors @@ -160,7 +160,7 @@ cryptography==50.0.1 # pyjwt cycler==0.12.1 # via matplotlib -datamodel-code-generator==0.79.0 +datamodel-code-generator==0.86.0 # via -r requirements/test/cuda.in dataproperty==1.0.1 # via @@ -502,9 +502,9 @@ mbstrdecoder==1.1.3 # dataproperty # pytablewriter # typepy -mcp==2.1.1 +mcp==2.3.0 # via -r requirements/test/../common.txt -mcp-types==2.1.1 +mcp-types==2.3.0 # via mcp mdurl==0.1.2 # via markdown-it-py @@ -533,8 +533,9 @@ msgspec==0.21.1 # via -r requirements/test/../common.txt mteb==2.8.3 # via -r requirements/test/cuda.in -multidict==6.1.0 +multidict==7.1.0 # via + # -r requirements/test/../common.txt # aiohttp # yarl multiprocess==0.70.16 @@ -596,6 +597,7 @@ numpy==2.2.6 # scikit-learn # scipy # segmentation-models-pytorch + # sentence-transformers # soxr # statsmodels # tensorizer @@ -716,7 +718,7 @@ perceptron==0.1.4 # via -r requirements/test/cuda.in perf-analyzer==0.1.0 # via genai-perf -pillow==12.3.0 +pillow==13.0.0 # via # -r requirements/test/../common.txt # genai-perf @@ -758,7 +760,7 @@ prometheus-client==0.22.0 # ray prometheus-fastapi-instrumentator==8.0.2 # via -r requirements/test/../common.txt -propcache==0.2.0 +propcache==0.5.4 # via # aiohttp # yarl @@ -1054,7 +1056,7 @@ scipy==1.13.1 # vocos segmentation-models-pytorch==0.5.0 # via -r requirements/test/cuda.in -sentence-transformers==5.2.0 +sentence-transformers==5.6.0 # via # -r requirements/test/cuda.in # mteb @@ -1259,7 +1261,7 @@ typer==0.26.8 # transformers types-requests==2.33.0.20260712 # via cohere -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -r requirements/test/../common.txt # aiohttp @@ -1274,6 +1276,7 @@ typing-extensions==4.15.0 # azure-storage-blob # chz # cohere + # datamodel-code-generator # fastapi # grpcio # httpx2 @@ -1362,7 +1365,7 @@ xxhash==3.5.0 # via # datasets # evaluate -yarl==1.17.1 +yarl==1.25.1 # via aiohttp zipp==3.23.0 # via importlib-metadata diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index eb61e8967fa4..299fee0864b5 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -23,7 +23,7 @@ vocos # required for minicpmo_26 test peft>=0.20.0 # required for phi-4-mm test pqdm ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) -sentence-transformers>=5.2.0 # required for embedding tests +sentence-transformers>=5.6.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests tblib # for pickling test exceptions @@ -36,7 +36,7 @@ matplotlib # required for qwen-vl 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>=0.64.0 # required for minicpm3 test +datamodel_code_generator>=0.81.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test @@ -84,14 +84,14 @@ plotly # required for perf comparison html report # Dependabot GHSA floors for transitive test deps (excl. dependabot.yml ignores: # torch/torchvision/xformers/lm-format-enforcer/compressed-tensors/ray/lm-eval). # Skipped (no patched release yet): sqlitedict, accelerate, lmcache, nltk GHSA-8mgp-746c-j5xp. -urllib3>=2.7.0 +urllib3>=2.8.0 cryptography>=50.0.0 pyjwt>=2.14.0 pyasn1>=0.6.4 python-multipart>=0.0.31 anyio>=4.14.2 httpx2>=2.12.0 -werkzeug>=3.1.6 +werkzeug>=3.1.9 mako>=1.3.12 black>=26.3.1 msgpack>=1.2.1 @@ -101,6 +101,6 @@ orjson>=3.11.6 hpack>=4.2.0 h2>=4.4.1 pydantic-settings>=2.14.2 -virtualenv>=20.36.1 +virtualenv>=21.7.13 fonttools>=4.60.2 pygments>=2.20.0 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index dc268719a6c1..58706a39d07e 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -8,7 +8,7 @@ accelerate==1.13.0 # via peft aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.14.3 +aiohttp==3.14.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -174,7 +174,7 @@ cuda-toolkit==13.0.3.0 # via torch cycler==0.12.1 # via matplotlib -datamodel-code-generator==0.79.0 +datamodel-code-generator==0.86.0 # via -r requirements/test/cuda.in dataproperty==1.0.1 # via @@ -532,11 +532,11 @@ mbstrdecoder==1.1.3 # dataproperty # pytablewriter # typepy -mcp==2.1.1 +mcp==2.3.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -mcp-types==2.1.1 +mcp-types==2.3.0 # via mcp mdurl==0.1.2 # via markdown-it-py @@ -568,8 +568,10 @@ msgspec==0.21.1 # via -r requirements/test/../common.txt mteb==2.8.3 # via -r requirements/test/cuda.in -multidict==6.1.0 +multidict==7.1.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # aiohttp # yarl multiprocess==0.70.16 @@ -632,6 +634,7 @@ numpy==2.2.6 # scikit-learn # scipy # segmentation-models-pytorch + # sentence-transformers # soxr # statsmodels # tensorizer @@ -806,7 +809,7 @@ perceptron==0.1.4 # via -r requirements/test/cuda.in perf-analyzer==0.1.0 # via genai-perf -pillow==12.3.0 +pillow==13.0.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -852,7 +855,7 @@ prometheus-fastapi-instrumentator==8.0.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt -propcache==0.2.0 +propcache==0.5.4 # via # aiohttp # yarl @@ -1153,7 +1156,7 @@ scipy==1.13.1 # vocos segmentation-models-pytorch==0.5.0 # via -r requirements/test/cuda.in -sentence-transformers==5.2.0 +sentence-transformers==5.6.0 # via # -r requirements/test/cuda.in # mteb @@ -1371,7 +1374,7 @@ typer==0.26.8 # transformers types-requests==2.33.0.20260712 # via cohere -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1387,6 +1390,7 @@ typing-extensions==4.15.0 # azure-storage-blob # chz # cohere + # datamodel-code-generator # fastapi # grpcio # httpx2 @@ -1477,7 +1481,7 @@ xxhash==3.5.0 # via # datasets # evaluate -yarl==1.17.1 +yarl==1.25.1 # via aiohttp zipp==3.23.0 # via importlib-metadata diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index af452a8f83fb..ed4ad5f011a9 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -17,7 +17,7 @@ vocos # required for minicpmo_26 test peft pqdm ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) -sentence-transformers>=5.2.0 # required for embedding tests +sentence-transformers>=5.6.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests timm # required for internvl test @@ -26,7 +26,7 @@ matplotlib # required for qwen-vl 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>=0.64.0 # required for minicpm3 test +datamodel_code_generator>=0.81.0 # 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.20.2 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 9f52d0b77cb8..8d090b29b3c5 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -21,8 +21,8 @@ vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test peft>=0.20.0 # required for phi-4-mm test pqdm -ray[default]==2.55.1 # Includes worker startup getenv/setenv race fix -sentence-transformers>=5.2.0 # required for embedding tests +ray[default]>=2.56.0,<2.57 # Match the CUDA/CPU test series. +sentence-transformers>=5.6.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests tblib # for pickling test exceptions @@ -32,7 +32,7 @@ matplotlib # required for qwen-vl 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>=0.64.0 # required for minicpm3 test +datamodel_code_generator>=0.81.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test @@ -82,14 +82,14 @@ plotly # required for perf comparison html report # Dependabot GHSA floors for transitive test deps (excl. dependabot.yml ignores). # Skipped (no patched release yet): sqlitedict, accelerate, lmcache, nltk GHSA-8mgp-746c-j5xp. -urllib3>=2.7.0 +urllib3>=2.8.0 cryptography>=50.0.0 pyjwt>=2.14.0 pyasn1>=0.6.4 python-multipart>=0.0.31 anyio>=4.14.2 httpx2>=2.12.0 -werkzeug>=3.1.6 +werkzeug>=3.1.9 mako>=1.3.12 black>=26.3.1 msgpack>=1.2.1 @@ -99,7 +99,7 @@ orjson>=3.11.6 hpack>=4.2.0 h2>=4.4.1 pydantic-settings>=2.14.2 -virtualenv>=20.36.1 +virtualenv>=21.7.13 fonttools>=4.60.2 pygments>=2.20.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index fb7390335f48..ba87f1c3a336 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -6,7 +6,7 @@ accelerate==1.13.0 # via peft aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.14.3 +aiohttp==3.14.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -171,7 +171,7 @@ cryptography==50.0.1 # pyjwt cycler==0.12.1 # via matplotlib -datamodel-code-generator==0.79.0 +datamodel-code-generator==0.86.0 # via -r requirements/test/rocm.in dataproperty==1.1.0 # via @@ -200,7 +200,7 @@ dill==0.3.8 # evaluate # lm-eval # multiprocess -distlib==0.4.0 +distlib==0.4.4 # via virtualenv distro==1.9.0 # via @@ -530,11 +530,11 @@ mbstrdecoder==1.1.4 # dataproperty # pytablewriter # typepy -mcp==2.1.1 +mcp==2.3.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -mcp-types==2.1.1 +mcp-types==2.3.0 # via mcp mdurl==0.1.2 # via markdown-it-py @@ -570,8 +570,10 @@ msgspec==0.21.0 # via -r requirements/test/../common.txt mteb==2.11.5 # via -r requirements/test/rocm.in -multidict==6.7.1 +multidict==7.1.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # aiohttp # yarl multiprocess==0.70.16 @@ -735,7 +737,7 @@ outlines-core==0.2.14 # via # -c requirements/common.txt # -r requirements/test/../common.txt -packaging==26.2 +packaging==26.3 # via # -c requirements/rocm.txt # accelerate @@ -760,6 +762,7 @@ packaging==26.2 # statsmodels # transformers # typepy + # virtualenv pandas==3.0.1 # via # datasets @@ -782,7 +785,7 @@ perceptron==0.1.4 # via -r requirements/test/rocm.in perf-analyzer==0.1.0 # via genai-perf -pillow==12.3.0 +pillow==13.0.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -795,12 +798,11 @@ pillow==12.3.0 # scikit-image # segmentation-models-pytorch # torchvision -platformdirs==4.3.6 +platformdirs==4.13.5 # via # black # jupyter-core # pooch - # python-discovery # virtualenv plotly==6.6.0 # via @@ -968,7 +970,7 @@ python-dateutil==2.9.0.post0 # matplotlib # pandas # typepy -python-discovery==1.2.0 +python-discovery==1.6.4 # via virtualenv python-dotenv==1.2.2 # via @@ -1019,7 +1021,7 @@ rapidfuzz==3.12.1 # via # -r requirements/test/rocm.in # jiwer -ray==2.55.1 +ray==2.56.1 # via -r requirements/test/rocm.in redis==7.3.0 # via tensorizer @@ -1131,7 +1133,7 @@ scipy==1.17.1 # vocos segmentation-models-pytorch==0.5.0 # via -r requirements/test/rocm.in -sentence-transformers==5.3.0 +sentence-transformers==5.6.0 # via # -r requirements/test/rocm.in # mteb @@ -1320,7 +1322,7 @@ typer==0.24.1 # transformers types-requests==2.33.0.20260712 # via cohere -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1336,6 +1338,7 @@ typing-extensions==4.15.0 # azure-storage-blob # chz # cohere + # datamodel-code-generator # fastapi # grpcio # httpx2 @@ -1397,7 +1400,7 @@ uvloop==0.22.1 # via uvicorn vector-quantize-pytorch==1.28.0 # via -r requirements/test/rocm.in -virtualenv==21.2.0 +virtualenv==21.14.12 # via # -r requirements/test/rocm.in # ray @@ -1411,7 +1414,7 @@ wcwidth==0.6.0 # via ftfy websockets==16.0 # via uvicorn -werkzeug==3.1.6 +werkzeug==3.1.9 # via # -r requirements/test/rocm.in # schemathesis @@ -1427,7 +1430,7 @@ xxhash==3.6.0 # via # datasets # evaluate -yarl==1.23.0 +yarl==1.25.1 # via aiohttp z3-solver==4.15.4.0 # via tilelang diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index e92863789c9b..e1f7fa5e3e44 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -40,7 +40,7 @@ kaldi-native-fbank>=1.18.7 # required for fireredasr2 test bm25s pystemmer mteb[bm25s]>=2, <3 # required for mteb test -sentence-transformers==5.2.0 # required for embedding tests +sentence-transformers==5.6.0 # required for embedding tests num2words pqdm @@ -64,14 +64,14 @@ cohere_melody>=0.11.1 # required for cohere command reasoning/tool parser tests cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) # Dependabot GHSA floors for transitive test deps (excl. dependabot.yml ignores). -urllib3>=2.7.0 +urllib3>=2.8.0 cryptography>=50.0.0 pyjwt>=2.14.0 pyasn1>=0.6.4 python-multipart>=0.0.31 anyio>=4.14.2 httpx2>=2.12.0 -werkzeug>=3.1.6 +werkzeug>=3.1.9 mako>=1.3.12 msgpack>=1.2.1 lxml>=6.1.0 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index bcbada52456e..881402b5672f 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -8,7 +8,7 @@ accelerate==1.14.0 # via -r requirements/test/xpu.in aiohappyeyeballs==2.7.1 # via aiohttp -aiohttp==3.14.3 +aiohttp==3.14.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -478,11 +478,11 @@ mbstrdecoder==1.1.5 # dataproperty # pytablewriter # typepy -mcp==2.1.1 +mcp==2.3.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -mcp-types==2.1.1 +mcp-types==2.3.0 # via mcp mdurl==0.1.2 # via markdown-it-py @@ -523,8 +523,10 @@ msgspec==0.21.1 # via -r requirements/test/../common.txt mteb==2.19.4 # via -r requirements/test/xpu.in -multidict==6.7.1 +multidict==7.1.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # aiohttp # yarl multiprocess==0.70.19 @@ -573,6 +575,7 @@ numpy==2.2.6 # sacrebleu # scikit-learn # scipy + # sentence-transformers # soundfile # soxr # tensorizer @@ -699,7 +702,7 @@ partial-json-parser==0.2.1.1.post7 # via -r requirements/test/../common.txt pathvalidate==3.3.1 # via pytablewriter -pillow==12.3.0 +pillow==13.0.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -988,7 +991,7 @@ scipy==1.18.0 # pytrec-eval-terrier # scikit-learn # sentence-transformers -sentence-transformers==5.2.0 +sentence-transformers==5.6.0 # via # -r requirements/test/xpu.in # mteb @@ -1228,7 +1231,7 @@ umf==1.1.0 # via # intel-cmplr-lib-ur # torch -urllib3==2.7.0 +urllib3==2.8.0 # via # -r requirements/test/xpu.in # blobfile @@ -1255,7 +1258,7 @@ wcwidth==0.8.3 # via ftfy websockets==17.0.1 # via uvicorn -werkzeug==3.1.8 +werkzeug==3.1.9 # via # -r requirements/test/xpu.in # schemathesis @@ -1269,5 +1272,5 @@ xxhash==4.0.1 # via # datasets # evaluate -yarl==1.24.5 +yarl==1.25.1 # via aiohttp diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 62b82da1dd8b..d32bc06422d6 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3796,9 +3796,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.37" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "once_cell", "rustls-pki-types", From 49d634bd6a3dc7433d8fb464821c95c76db85350 Mon Sep 17 00:00:00 2001 From: Avishek Goswami <86944690+GOavi101@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:28:06 +0530 Subject: [PATCH 007/631] [Perf][Multimodal] Make encoder cache free linear in items per request (#62850) Signed-off-by: Avishek Goswami Co-authored-by: Avishek Goswami Co-authored-by: vllm-agent --- tests/v1/core/test_encoder_cache_manager.py | 25 +++++++++++++++ vllm/v1/core/encoder_cache_manager.py | 35 +++++++++++++++------ 2 files changed, 51 insertions(+), 9 deletions(-) diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index 3c61f3367f98..9d1fcdfd3ed6 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -445,6 +445,31 @@ def test_duplicate_mm_hash_is_not_evicted_before_last_use(): assert manager.get_freed_mm_hashes() == [] +def test_many_inputs_release_each_hash_at_its_last_occurrence(): + """A request with many items, one repeated, releases every entry exactly + once its last cached occurrence is freed.""" + num_items = 64 + hashes = [f"img{i}" for i in range(num_items)] + ["img0"] + manager = EncoderCacheManager(cache_size=num_items) + req = MockRequest("r1", hashes, [1] * len(hashes)) + + for i in range(num_items): + manager.allocate(req, i) + assert manager.check_and_update_cache(req, num_items) + + manager.free_encoder_input(req, 0) + assert "img0" not in manager.freeable + for i in range(1, num_items): + manager.free_encoder_input(req, i) + assert manager.freeable[f"img{i}"] == 1 + + manager.free_encoder_input(req, num_items) + assert manager.freeable["img0"] == 1 + assert manager.num_freeable_slots == num_items + assert "r1" not in manager.request_cached_ids + assert "r1" not in manager.request_hash_refs + + def test_encoder_decoder_cache_manager_reset(): manager = EncoderDecoderCacheManager(cache_size=20) diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index 24cf609d22b9..4bc995f2d83f 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -86,6 +86,9 @@ def __init__(self, cache_size: int): self.num_embeds_by_hash: dict[str, int] = {} # request_id => set of input_ids cached for that request self.request_cached_ids: dict[str, set[int]] = {} + # request_id => mm_hash => number of that request's cached input_ids + # with the mm_hash + self.request_hash_refs: dict[str, dict[str, int]] = {} # mm_hash of mm_data => num_encoder_embeds of the mm_data self.freeable: OrderedDict[str, int] = OrderedDict() @@ -100,6 +103,7 @@ def reset(self) -> None: self.cached.clear() self.num_embeds_by_hash.clear() self.request_cached_ids.clear() + self.request_hash_refs.clear() self.freeable.clear() self.freed.clear() self.num_free_slots = self.cache_size @@ -132,7 +136,7 @@ def check_and_update_cache(self, request: Request, input_id: int) -> bool: self.num_freeable_slots -= num_encoder_embeds self.cached[mm_hash].add(request.request_id) - self.request_cached_ids.setdefault(request.request_id, set()).add(input_id) + self._add_cached_input_id(request.request_id, input_id, mm_hash) return True def get_cached_num_encoder_embeds( @@ -235,10 +239,20 @@ def allocate(self, request: Request, input_id: int) -> None: assert self.num_freeable_slots >= num_encoder_embeds self.cached[mm_hash].add(request_id) - self.request_cached_ids.setdefault(request_id, set()).add(input_id) + self._add_cached_input_id(request_id, input_id, mm_hash) self.num_free_slots -= num_encoder_embeds self.num_freeable_slots -= num_encoder_embeds + def _add_cached_input_id( + self, request_id: str, input_id: int, mm_hash: str + ) -> None: + cached_ids = self.request_cached_ids.setdefault(request_id, set()) + if input_id in cached_ids: + return + cached_ids.add(input_id) + hash_refs = self.request_hash_refs.setdefault(request_id, {}) + hash_refs[mm_hash] = hash_refs.get(mm_hash, 0) + 1 + def get_cached_input_ids(self, request: Request) -> set[int]: """Get all cached multimodal input IDs for a request.""" return self.request_cached_ids.get(request.request_id, set()) @@ -257,10 +271,16 @@ def free_encoder_input(self, request: Request, input_id: int) -> None: mm_hash = request.mm_features[input_id].identifier # Always clean up request_cached_ids, even if the mm_hash was # already evicted from cache (e.g. by can_allocate). - if req_id in self.request_cached_ids: - self.request_cached_ids[req_id].discard(input_id) - if not self.request_cached_ids[req_id]: + cached_ids = self.request_cached_ids.get(req_id) + if cached_ids is not None and input_id in cached_ids: + cached_ids.remove(input_id) + hash_refs = self.request_hash_refs[req_id] + hash_refs[mm_hash] -= 1 + if not hash_refs[mm_hash]: + del hash_refs[mm_hash] + if not cached_ids: del self.request_cached_ids[req_id] + del self.request_hash_refs[req_id] # The mm_hash not in cache or the req_id set is empty if not self.cached.get(mm_hash, None): return @@ -270,10 +290,7 @@ def free_encoder_input(self, request: Request, input_id: int) -> None: # occurrence is freed: dropping it at the first makes the entry # evictable while the request still needs it, and the encoder then # recomputes an item it already has. - if any( - request.mm_features[other_id].identifier == mm_hash - for other_id in self.request_cached_ids.get(req_id, ()) - ): + if mm_hash in self.request_hash_refs.get(req_id, ()): return self.cached[mm_hash].discard(req_id) if not self.cached[mm_hash]: From f959137753c4734a6b47b7bed4d99e853e61639b Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Sat, 24 Oct 2026 07:36:02 +0100 Subject: [PATCH 008/631] [Bugfix][Frontend] Apply batch tool_choice rules to streamed tool calls (#63345) Signed-off-by: Martin Hickey --- .../chat_completion/test_stream_processor.py | 50 +++++++++++++++-- .../scale_out/derender/test_derender.py | 44 +++++++++------ .../derender/test_derender_stream.py | 17 +++++- .../chat_completion/message_assembly.py | 56 ++++++++++++++----- .../openai/chat_completion/serving.py | 10 +++- .../chat_completion/stream_processor.py | 26 +++++++-- vllm/renderers/online_derenderer.py | 2 + 7 files changed, 165 insertions(+), 40 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_stream_processor.py b/tests/entrypoints/openai/chat_completion/test_stream_processor.py index c1e059b64476..49ed2d30f9ae 100644 --- a/tests/entrypoints/openai/chat_completion/test_stream_processor.py +++ b/tests/entrypoints/openai/chat_completion/test_stream_processor.py @@ -3,8 +3,12 @@ """`ChatStreamProcessor` applies the per choice rules that chat streaming and streaming chat derender must share which means both emit the same deltas.""" +from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock +import pytest + from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -15,6 +19,14 @@ ChatStreamProcessor, ) +TOOLS = [ + { + "type": "function", + "function": {"name": "f", "parameters": {"type": "object"}}, + } +] +PARSER_CLS: Any = SimpleNamespace(tool_parser_cls=object) + def _request(**kwargs) -> ChatCompletionRequest: return ChatCompletionRequest( @@ -22,6 +34,12 @@ def _request(**kwargs) -> ChatCompletionRequest: ) +def _processor(request, parser) -> ChatStreamProcessor: + return ChatStreamProcessor( + request, parser, enable_auto_tools=True, parser_cls=PARSER_CLS + ) + + def _tool_call(index: int) -> DeltaToolCall: return DeltaToolCall( id=f"call_{index}", @@ -38,7 +56,9 @@ def test_tool_calls_set_finish_reason_and_respect_parallel_tool_calls(): DeltaMessage(tool_calls=[_tool_call(0), _tool_call(1)]), None, ] - processor = ChatStreamProcessor(_request(parallel_tool_calls=False), parser) + processor = _processor( + _request(tools=TOOLS, tool_choice="auto", parallel_tool_calls=False), parser + ) delta, finish_reason = processor.step("x", [1], None, [0]) assert delta is not None @@ -51,11 +71,33 @@ def test_tool_calls_set_finish_reason_and_respect_parallel_tool_calls(): assert processor.num_output_tokens == 2 +@pytest.mark.parametrize("tool_choice", ["none", None]) +def test_drops_tool_calls_like_non_streaming(tool_choice): + """Tool calls the parser extracts for `"none"` (Harmony ignores it) or + null are dropped like on the non-streaming path. A delta left empty is + not sent and `finish_reason` stays `stop`.""" + parser = MagicMock() + parser.parse_delta.side_effect = [ + DeltaMessage(content="text", tool_calls=[_tool_call(0)]), + DeltaMessage(tool_calls=[_tool_call(0)]), + ] + request = _request(tools=TOOLS) + request.tool_choice = tool_choice + processor = _processor(request, parser) + + delta, _ = processor.step("x", [1], None, [0]) + assert delta == DeltaMessage(content="text") + + delta, finish_reason = processor.step("y", [2], "stop", [0]) + assert delta is None + assert finish_reason == "stop" + + def test_finished_without_output_skips_parser(): """A choice that finishes before generating anything sends an empty delta without asking the parser.""" parser = MagicMock() - processor = ChatStreamProcessor(_request(), parser) + processor = _processor(_request(), parser) delta, finish_reason = processor.step("", [], "length", [0]) @@ -67,5 +109,5 @@ def test_finished_without_output_skips_parser(): def test_hidden_reasoning_hides_metadata_only_with_a_parser(): """Logprobs and token IDs are hidden only when there is reasoning to hide.""" request = _request(include_reasoning=False) - assert ChatStreamProcessor(request, MagicMock()).hide_metadata - assert not ChatStreamProcessor(request, None).hide_metadata + assert _processor(request, MagicMock()).hide_metadata + assert not _processor(request, None).hide_metadata diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 667635eb9fc7..dd312af88bdf 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -1414,13 +1414,14 @@ async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer): @pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("tool_choice", ["auto", "none"]) async def test_e2e_harmony_tool_choice_none_drops_tool_calls( - harmony_client, harmony_tokenizer, tool_choice + harmony_client, harmony_tokenizer, tool_choice, stream ): """HarmonyParser extracts tool calls whatever the tool_choice, so - /derender must drop them for "none", like coupled chat. "auto" checks - the tokens really parse as a tool call.""" + /derender must drop them for "none", like coupled chat, batch and + streaming alike. "auto" checks the tokens really parse as a tool call.""" messages = [{"role": "user", "content": "Weather in Paris?"}] gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages) @@ -1440,23 +1441,32 @@ async def test_e2e_harmony_tool_choice_none_drops_tool_calls( if "get_weather" not in harmony_tokenizer.decode(output_ids): pytest.skip("Harmony template did not render the tool call") - resp = await harmony_client.post( - "/v1/chat/completions/derender", - json={ + body: dict = { + "model": HARMONY_MODEL, + "prompt_tokens": len(gen_req["token_ids"]), + "prompt_token_ids": gen_req["token_ids"], + "chat_request": { "model": HARMONY_MODEL, - "generate_response": _e2e_generate_response(output_ids), - "prompt_tokens": len(gen_req["token_ids"]), - "chat_request": { - "model": HARMONY_MODEL, - "messages": messages, - "tools": _E2E_TOOLS, - "tool_choice": tool_choice, - }, + "messages": messages, + "tools": _E2E_TOOLS, + "tool_choice": tool_choice, }, - ) + } + if stream: + body["stream"] = True + body["generate_chunk"] = _e2e_generate_response(output_ids) + else: + body["generate_response"] = _e2e_generate_response(output_ids) + + resp = await harmony_client.post("/v1/chat/completions/derender", json=body) assert resp.status_code == 200, resp.text - choice = resp.json()["choices"][0] - assert bool(choice["message"].get("tool_calls")) == (tool_choice == "auto") + if stream: + choice = resp.json()["chunk"]["choices"][0] + tool_calls = choice["delta"].get("tool_calls") + else: + choice = resp.json()["choices"][0] + tool_calls = choice["message"].get("tool_calls") + assert bool(tool_calls) == (tool_choice == "auto") assert choice["finish_reason"] == ( "tool_calls" if tool_choice == "auto" else "stop" ) diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index c58f9cf0034f..c843b4e89f95 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -112,6 +112,15 @@ def _merge_fake_deltas( return accumulated +class _FakeToolParser: + """Tool parser slot filler so `tool_choice="auto"` keeps tool calls.""" + + engine_based_streaming = True + + def __init__(self, tokenizer, tools): + pass + + class _FakeParser(Parser): """Deterministic `Parser` stub for unit testing the replay/merge/pin mechanics of `OnlineDerenderer._derender_chat_stream_parsed` in @@ -132,6 +141,8 @@ class _FakeParser(Parser): - Anything else emits nothing. """ + tool_parser_cls = _FakeToolParser # type: ignore[assignment] + TOOL_START = 9 TOOL_ARG = 10 REASON = 11 @@ -306,6 +317,7 @@ def parsed_derenderer(tokenizer): request_logger=None, chat_template=None, chat_template_content_format="string", + enable_auto_tools=True, ) dr.parser = _FakeParser return dr @@ -1590,7 +1602,10 @@ async def test_logprobs_attached_to_each_parsed_delta( ids, logprobs=_generate_logprobs(ids) ), state=state, - chat_request=_chat_request(), + chat_request=_chat_request( + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tool_choice="auto", + ), prompt_token_ids=[1, 2], ) choice = chunk.choices[0] diff --git a/vllm/entrypoints/openai/chat_completion/message_assembly.py b/vllm/entrypoints/openai/chat_completion/message_assembly.py index f902fa6c600a..66bb6845378a 100644 --- a/vllm/entrypoints/openai/chat_completion/message_assembly.py +++ b/vllm/entrypoints/openai/chat_completion/message_assembly.py @@ -53,30 +53,60 @@ def assemble_chat_message_fields( if not request.include_reasoning: reasoning = None - tool_choice = request.tool_choice - if ( - type(tool_choice) is ChatCompletionNamedToolChoiceParam - or tool_choice == "required" + if not keeps_tool_calls( + request, enable_auto_tools=enable_auto_tools, parser_cls=parser_cls ): - return ChatMessageFields( - reasoning, content or "", _to_tool_calls(tool_calls, request) - ) + return ChatMessageFields(reasoning, content, []) + if _forces_tool_calls(request): + content = content or "" + return ChatMessageFields(reasoning, content, _to_tool_calls(tool_calls, request)) + + +def keeps_tool_calls( + request: ChatCompletionRequest, + *, + enable_auto_tools: bool, + parser_cls: type[Parser] | None, +) -> bool: + """Whether a chat choice returns the tool calls its parser extracts. + + Shared by the non-streaming and streaming chat paths so both apply + `tool_choice` the same way. Parsers don't all gate on it themselves: + `HarmonyParser` ignores `"none"` and `null` is parsed like `"auto"`. + + Args: + request: Supplies `tools` and `tool_choice`. + enable_auto_tools: The server's `--enable-auto-tool-choice`. + parser_cls: The server's parser class, if any. + + Returns: + True for `"required"`, a named tool, or `"auto"` with tools and auto + tool choice enabled. + """ + if _forces_tool_calls(request): + return True auto_tools = ( enable_auto_tools and parser_cls is not None and parser_cls.tool_parser_cls is not None ) - if auto_tools and tool_choice == "auto": - if request.tools: - return ChatMessageFields( - reasoning, content, _to_tool_calls(tool_calls, request) - ) + if not auto_tools or request.tool_choice != "auto": + return False + if not request.tools: logger.error( "Cannot determine if tools should be extracted. " "Returning a standard chat completion." ) - return ChatMessageFields(reasoning, content, []) + return False + return True + + +def _forces_tool_calls(request: ChatCompletionRequest) -> bool: + return ( + type(request.tool_choice) is ChatCompletionNamedToolChoiceParam + or request.tool_choice == "required" + ) def _to_tool_calls( diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 652235b0e3e2..3fbddeefa73c 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -544,12 +544,20 @@ async def chat_completion_stream_generator_raw( ChatStreamProcessor( request, self._make_parser(request, tokenizer, chat_template_kwargs), + enable_auto_tools=self.enable_auto_tools, + parser_cls=self.parser_cls, ) for _ in range(num_choices) ] else: processors = [ - ChatStreamProcessor(request, None) for _ in range(num_choices) + ChatStreamProcessor( + request, + None, + enable_auto_tools=self.enable_auto_tools, + parser_cls=None, + ) + for _ in range(num_choices) ] except Exception as e: logger.exception("Error in parser creation.") diff --git a/vllm/entrypoints/openai/chat_completion/stream_processor.py b/vllm/entrypoints/openai/chat_completion/stream_processor.py index 917c69c73331..4e90b273609f 100644 --- a/vllm/entrypoints/openai/chat_completion/stream_processor.py +++ b/vllm/entrypoints/openai/chat_completion/stream_processor.py @@ -3,6 +3,9 @@ from typing import NamedTuple from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.entrypoints.openai.chat_completion.message_assembly import ( + keeps_tool_calls, +) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.serve.utils.tool_calls_utils import ( maybe_filter_parallel_tool_calls, @@ -27,9 +30,19 @@ class ChatStreamProcessor: `step` once per chunk, in order and with the chunking the engine produced. """ - def __init__(self, request: ChatCompletionRequest, parser: Parser | None) -> None: + def __init__( + self, + request: ChatCompletionRequest, + parser: Parser | None, + *, + enable_auto_tools: bool, + parser_cls: type[Parser] | None, + ) -> None: self.request = request self.parser = parser + self.keep_tool_calls = keeps_tool_calls( + request, enable_auto_tools=enable_auto_tools, parser_cls=parser_cls + ) self.num_output_tokens = 0 self.tools_streamed = False # Logprobs and token IDs would leak hidden reasoning through decoded @@ -52,8 +65,8 @@ def step( prompt_token_ids: Prompt token IDs if known. Returns: - The delta to stream, filtered for `parallel_tool_calls` and the - chat `finish_reason`. + The delta to stream, filtered for `tool_choice` and + `parallel_tool_calls`, and the chat `finish_reason`. """ delta: DeltaMessage | None @@ -69,7 +82,12 @@ def step( finished=finish_reason is not None, ) if delta is not None and delta.tool_calls: - self.tools_streamed = True + if self.keep_tool_calls: + self.tools_streamed = True + else: + delta.tool_calls = [] + if not delta.content and not delta.reasoning: + delta = None else: delta = DeltaMessage(content=delta_text) self.num_output_tokens += len(delta_token_ids) diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index cad4877ec2e2..efbc4f50b898 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -745,6 +745,8 @@ def _derender_chat_stream_parsed( chat_template_kwargs=self._resolve_chat_template_kwargs(chat_request), model_config=self.model_config, ), + enable_auto_tools=self.enable_auto_tools, + parser_cls=parser_cls, ) # Ephemeral incremental detok window, local to this call. Threaded From e9acf48614a2c0d027d7aa2f0e38d963117617a6 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 24 Oct 2026 14:44:57 +0800 Subject: [PATCH 009/631] [Platform] Use pluggable layer for ViT mrope preparation (#63017) Signed-off-by: Isotr0py Co-authored-by: vllm-agent --- .../dummy_custom_ops.py | 11 ++ tests/plugins_tests/test_platform_plugins.py | 17 +++ .../rotary_embedding/mrope_vit_setup.py | 131 +++++++++++------- vllm/models/minimax_m3/common/vision_tower.py | 6 +- 4 files changed, 111 insertions(+), 54 deletions(-) diff --git a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py index b73028574526..13e55f6e9484 100644 --- a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py +++ b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py @@ -4,6 +4,9 @@ import torch from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.model_executor.layers.rotary_embedding.mrope_vit_setup import ( + VitMropeSetup, +) # Register CustomRotaryEmbedding to CustomOP. @@ -17,3 +20,11 @@ def __init__(self, *args, **kwargs): def forward_oot(self, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: return super().forward_oot(*args, **kwargs) + + +# Register a replacement for the ViT mrope setup pluggable layer. +@VitMropeSetup.register_oot +class DummyVitMropeSetup(VitMropeSetup): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.addition_config = True diff --git a/tests/plugins_tests/test_platform_plugins.py b/tests/plugins_tests/test_platform_plugins.py index a75e36a7d1c8..3e3883565e8d 100644 --- a/tests/plugins_tests/test_platform_plugins.py +++ b/tests/plugins_tests/test_platform_plugins.py @@ -62,3 +62,20 @@ def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch): "Expected DummyRotaryEmbedding to have an 'addition_config' attribute, " "which is set by the custom op." ) + + +def test_oot_pluggable_layer(default_vllm_config): + load_general_plugins() + from vllm.model_executor.layers.rotary_embedding.mrope_vit_setup import ( + VitMropeSetup, + ) + + layer = VitMropeSetup() + assert layer.__class__.__name__ == "DummyVitMropeSetup", ( + f"Expected DummyVitMropeSetup, got {layer.__class__.__name__}, " + "possibly because the pluggable layer is not registered correctly." + ) + assert hasattr(layer, "addition_config"), ( + "Expected DummyVitMropeSetup to have an 'addition_config' attribute, " + "which is set by the OOT pluggable layer." + ) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py b/vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py index 6453e8c83545..965ffe6e842e 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py @@ -3,6 +3,7 @@ import torch +from vllm.model_executor.custom_op import PluggableLayer from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import async_tensor_h2d @@ -91,63 +92,89 @@ def _vit_mrope_setup_kernel( tl.store(sin_out_ptr + 2 * plane + row, tl.sin(freq), mask=mask) -def vit_mrope_setup( - inv_freq_t: torch.Tensor, - inv_freq_h: torch.Tensor, - inv_freq_w: torch.Tensor, - grid_thw: list[list[int]], - spatial_merge_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope. +@PluggableLayer.register("vit_mrope_setup") +class VitMropeSetup(PluggableLayer): + """Builds triton_mrope's (3, N, half_rot) cos/sin t/h/w planes. One fused kernel: positions are computed arithmetically from the grids and the rotation comes from the per-axis inv_freq, so there is no cos/sin cache to size or grow. - Args: - inv_freq_t: t-axis inverse frequencies, (t_dim // 2,). - inv_freq_h: h-axis inverse frequencies, (h_dim // 2,). - inv_freq_w: w-axis inverse frequencies, (w_dim // 2,). - grid_thw: per-grid [t, h, w] patch counts (host metadata). - spatial_merge_size: spatial merge factor of the ViT. - + Stateless: the caller's inv_freq buffers are passed per call. OOT + platforms replace the whole layer via ``@VitMropeSetup.register_oot`` + with a subclass overriding :meth:`forward` (e.g. to retune the kernel + tiling for their hardware). """ - device = inv_freq_t.device - cu = [0] - for t, h, w in grid_thw: - cu.append(cu[-1] + t * h * w) - n = cu[-1] - g = len(grid_thw) - # grid_thw is per-request host metadata, so one H2D is unavoidable; - # pack cu_seqlens and grids into a single pinned async copy. - buf = async_tensor_h2d( - cu + [d for grid in grid_thw for d in grid], device=device, dtype=torch.int32 - ) - cu_seqlens = buf[: g + 1] - grids = buf[g + 1 :].view(g, 3) - half_t = inv_freq_t.numel() - half_h = inv_freq_h.numel() - half_w = inv_freq_w.numel() - half_rot = half_t + half_h + half_w - cos = torch.empty(3, n, half_rot, device=device, dtype=torch.float32) - sin = torch.empty(3, n, half_rot, device=device, dtype=torch.float32) - _vit_mrope_setup_kernel[(triton.cdiv(n, 32),)]( - inv_freq_t, - inv_freq_h, - inv_freq_w, - cos, - sin, - cu_seqlens, - grids, - n, - g, - half_t=half_t, - half_h=half_h, - half_w=half_w, - spatial_merge=spatial_merge_size, - BLOCK_N=32, - BLOCK_D=triton.next_power_of_2(half_rot), - num_warps=8, + def forward( + self, + inv_freq_t: torch.Tensor, + inv_freq_h: torch.Tensor, + inv_freq_w: torch.Tensor, + grid_thw: list[list[int]], + spatial_merge_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes. + + Args: + inv_freq_t: t-axis inverse frequencies, (t_dim // 2,). + inv_freq_h: h-axis inverse frequencies, (h_dim // 2,). + inv_freq_w: w-axis inverse frequencies, (w_dim // 2,). + grid_thw: per-grid [t, h, w] patch counts (host metadata). + spatial_merge_size: spatial merge factor of the ViT. + + """ + device = inv_freq_t.device + cu = [0] + for t, h, w in grid_thw: + cu.append(cu[-1] + t * h * w) + n = cu[-1] + g = len(grid_thw) + # grid_thw is per-request host metadata, so one H2D is unavoidable; + # pack cu_seqlens and grids into a single pinned async copy. + buf = async_tensor_h2d( + cu + [d for grid in grid_thw for d in grid], + device=device, + dtype=torch.int32, + ) + cu_seqlens = buf[: g + 1] + grids = buf[g + 1 :].view(g, 3) + + half_t = inv_freq_t.numel() + half_h = inv_freq_h.numel() + half_w = inv_freq_w.numel() + half_rot = half_t + half_h + half_w + cos = torch.empty(3, n, half_rot, device=device, dtype=torch.float32) + sin = torch.empty(3, n, half_rot, device=device, dtype=torch.float32) + _vit_mrope_setup_kernel[(triton.cdiv(n, 32),)]( + inv_freq_t, + inv_freq_h, + inv_freq_w, + cos, + sin, + cu_seqlens, + grids, + n, + g, + half_t=half_t, + half_h=half_h, + half_w=half_w, + spatial_merge=spatial_merge_size, + BLOCK_N=32, + BLOCK_D=triton.next_power_of_2(half_rot), + num_warps=8, + ) + return cos, sin + + +def vit_mrope_setup( + inv_freq_t: torch.Tensor, + inv_freq_h: torch.Tensor, + inv_freq_w: torch.Tensor, + grid_thw: list[list[int]], + spatial_merge_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compatibility wrapper around :class:`VitMropeSetup`.""" + return VitMropeSetup()( + inv_freq_t, inv_freq_h, inv_freq_w, grid_thw, spatial_merge_size ) - return cos, sin diff --git a/vllm/models/minimax_m3/common/vision_tower.py b/vllm/models/minimax_m3/common/vision_tower.py index fd9cff2c6ddd..c7aba484f2a3 100644 --- a/vllm/models/minimax_m3/common/vision_tower.py +++ b/vllm/models/minimax_m3/common/vision_tower.py @@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope from vllm.model_executor.layers.rotary_embedding.mrope_vit_setup import ( - vit_mrope_setup, + VitMropeSetup, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import maybe_prefix @@ -352,6 +352,8 @@ def __init__( self.register_buffer("inv_freq_t", inv_freq_t, persistent=False) self.register_buffer("inv_freq_h", inv_freq_h, persistent=False) self.register_buffer("inv_freq_w", inv_freq_w, persistent=False) + # Pluggable so OOT platforms can substitute the setup kernel. + self.mrope_setup = VitMropeSetup() self.embeddings = MiniMaxVLPatchEmbed(config) self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) @@ -456,7 +458,7 @@ def prepare_encoder_metadata( # 3D RoPE cos/sin: (3, total_N, half_rot_dim) fp32 t/h/w planes for # triton_mrope from one fused setup kernel. Kept fp32 so the kernel # computes the rotation in fp32, matching the reference. - rotary_cos, rotary_sin = vit_mrope_setup( + rotary_cos, rotary_sin = self.mrope_setup( self.inv_freq_t, self.inv_freq_h, self.inv_freq_w, From 7a64fb8fc9239cb65e0df9313fa759594dfb206a Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 24 Oct 2026 02:54:50 -0500 Subject: [PATCH 010/631] [CI] Fix the Domyn Edge (ROCm) and Qwen2 embed HF references (#63452) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- tests/models/language/generation/test_common.py | 10 +++++++++- tests/models/language/pooling/test_embedding.py | 12 ++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 493f8c425fad..5ca2531ed322 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -153,7 +153,15 @@ def test_models( trust_remote_code=model_info.trust_remote_code, ) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( - example_prompts, max_tokens, num_logprobs + example_prompts, + max_tokens, + num_logprobs, + # Domyn's remote code decodes with torch varlen_attn, whose ROCm + # causal mask is top-left aligned when max_q < max_k: + # https://github.com/pytorch/pytorch/pull/201930 + use_cache=not ( + current_platform.is_rocm() and model == "domyn/Domyn-Edge-Base" + ), ) prompt_embeds: list[torch.Tensor] | None = [] if use_prompt_embeds else None diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index 7d48880e9c06..47e1a0a7bad4 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/tests/models/language/pooling/test_embedding.py @@ -76,6 +76,18 @@ def test_models( example_prompts = [str(s).strip() for s in example_prompts] with hf_runner(model, is_sentence_transformer=True) as hf_model: + if model == "ssmits/Qwen2-7B-Instruct-embed-base": + # This checkpoint has no sentence-transformers config. Since + # sentence-transformers 5.6, a causal LM then gets its chat + # template applied and last-token pooling. Keep the plain text and + # the MEAN pooling that vLLM uses above. + from sentence_transformers.models import Pooling + + st_model = hf_model.model + getattr(st_model[0], "modality_config", {}).pop("message", None) + st_model[1] = Pooling( + st_model.get_sentence_embedding_dimension(), pooling_mode="mean" + ) hf_outputs = hf_model.encode(example_prompts) with vllm_runner( From 1bd56eaef38d0db35852fc38932d0c63003ef9f1 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sat, 24 Oct 2026 18:52:33 +0800 Subject: [PATCH 011/631] [CI][Kimi-K3] Expect no trailing-block drop in DSpark resend tests (#63449) Signed-off-by: Jared Wen Co-authored-by: Claude Opus 5.5 --- tests/models/kimi_k3/test_prefix_cache.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index e9940c8e383c..a0664ff9a1bc 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -378,7 +378,7 @@ def _format(turns: list[Turn]) -> str: def _check_resend_reuses_prompt_checkpoint( deployment: Deployment, prompt_len: int ) -> None: - """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" + """Resends reuse the last prompt checkpoint before the final token.""" mode = Mode(128, True) salt = str(uuid4()) rng = random.Random(0) @@ -391,8 +391,8 @@ def _check_resend_reuses_prompt_checkpoint( resend = _complete(url, prompt, salt, max_tokens=1) recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) - # EAGLE's dropped hash unit already leaves tokens to recompute on resend. - expected = (prompt_len // 128 - 1) * 128 + # DSpark keeps the trailing hash unit, unlike EAGLE. + expected = (prompt_len - 1) // 128 * 128 actual = _cached(resend) assert _cached(cold) == _cached(recompute) == 0 assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}" @@ -408,13 +408,11 @@ def _check_resend_reuses_prompt_checkpoint( "prompt_len", [ # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: - # ensure lookup reads the cached attention proof at 7040 despite - # the final hit limit of 7039, then applies EAGLE's 128-token drop. - pytest.param(7040, id="eagle-proof-scan"), + # the final hit limit of 7039 puts the checkpoint at 6912. + pytest.param(7040, id="pmu-aligned-tail"), # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: - # ensure the Mamba checkpoint is 7168, not 7040 from applying both - # the P - 1 cap and EAGLE's 128-token drop. - pytest.param(7296, id="eagle-double-cap"), + # ensure the Mamba checkpoint is 7168, not 7040 from capping twice. + pytest.param(7296, id="single-cap"), ], ) def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: @@ -474,9 +472,9 @@ def mooncake_store(tmp_path, monkeypatch) -> Iterator[None]: @pytest.mark.usefixtures("mooncake_store") def test_mooncake_resend_reuses_prompt_checkpoint() -> None: - """Blocks=6144, PMU=128. P=7449 needs Mamba state at 7296 and attention proof - at 7424, both beyond the normal Mooncake save boundary of 6144: ensure they - remain reusable from the store after a GPU cache reset.""" + """Blocks=6144, PMU=128. P=7449 needs Mamba state at 7424, beyond the normal + Mooncake save boundary of 6144: ensure it remains reusable from the store + after a GPU cache reset.""" _check_resend_reuses_prompt_checkpoint( Deployment(Instance(kv_config=MOONCAKE), offload=True), 7449 ) From 717cf843426ce85dbb7170aff7d7a833ad7d3c09 Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Sat, 24 Oct 2026 19:12:56 +0800 Subject: [PATCH 012/631] [CI] Label recipe PRs and link their recipes.vllm.ai preview (#63248) Signed-off-by: esmeetu Signed-off-by: Roy Wang Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .github/mergify.yml | 13 +++++++++++++ recipes/README.md | 4 ++++ 2 files changed, 17 insertions(+) diff --git a/.github/mergify.yml b/.github/mergify.yml index bd364fd4e4ac..0f545a4a0784 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -14,6 +14,19 @@ pull_request_rules: comment: message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/" +- name: label-recipe + description: Automatically apply recipe label + conditions: + - label != stale + - -closed + - files~=^recipes/ + actions: + label: + add: + - recipe + comment: + message: "Recipe preview: https://vllm-recipes-git-preview-vllm-pr-{{number}}-inferact-inc.vercel.app" + - name: comment-pre-commit-failure description: Comment on PR when pre-commit check fails conditions: diff --git a/recipes/README.md b/recipes/README.md index d9f9b0e19d07..d96c9096ce96 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -38,3 +38,7 @@ To preview the site with your changes, add a gitignored `sources.local.yaml` to that checkout with `overrides: { vllm: ~/vllm/recipes }` and run `pnpm dev`. This works once vllm-project/recipes reads this directory as its `vllm` source (vllm-project/recipes#1095). + +A pull request that changes this directory gets the `recipe` label and a +comment linking its own preview deployment, rebuilt a few minutes after each +push. From 2aa7269a837556bad5a16f1d6aa32916f7d14446 Mon Sep 17 00:00:00 2001 From: TJian Date: Sat, 24 Oct 2026 19:53:20 +0800 Subject: [PATCH 013/631] [Bugfix][ROCm][MoRIIO] Release every prefill rank a READ holds (#62473) Signed-off-by: tjtanaa Signed-off-by: TJian Co-authored-by: Claude Opus 5.5 (1M context) --- .buildkite/test-amd.yaml | 25 +- .buildkite/test_areas/misc.yaml | 11 +- .../test_write_dummy_e2e_rocm.py | 214 ++++++-- .../unit/test_moriio_hetero_tp_write.py | 1 + .../kv_connector/unit/test_moriio_read_tp.py | 498 ++++++++++++++++++ .../unit/test_moriio_routing_fairness.py | 3 + .../kv_connector/unit/test_moriio_tp_ack.py | 4 +- .../v1/moriio/moriio_connector.py | 174 ++++-- .../kv_connector/v1/moriio/moriio_layout.py | 57 +- 9 files changed, 870 insertions(+), 117 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_read_tp.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 6392e04f0540..9a35984acd55 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -5468,7 +5468,30 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +# The same P/D between different TP sizes (WRITE TP2 -> TP1, READ of MLA +# TP1 <-> TP2), which needs 3 GPUs: 8 cases, ~36 min on MI355. +- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - vllm/model_executor/model_loader/dummy_loader.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/moriio_integration/ + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp - label: ":amd: (MI355 DPX) V1 Metrics + LM Eval" timeout_in_minutes: 65 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index d3f47bdc76c8..3d1f1e5152c7 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -211,15 +211,15 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py -# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate -# GPUs over xGMI, compared exactly against a direct serve. +# Single-node MoRIIO P/D on dummy weights, WRITE and READ mode: prefill and +# decode on separate GPUs over xGMI, compared exactly against a direct serve. - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" key: amd-moriio-write-dummy-e2e-mi355 dind: false device: mi355_2 num_devices: 2 - # ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the - # largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports + # ~3.5-5 min per case; pytest-shard hashes test ids, so with 28 cases the + # largest of 6 shards holds 7 (~32 min on MI355). Every case picks free ports # for its servers and proxy and its own compile-cache dir. parallelism: 6 timeout_in_minutes: 45 @@ -239,7 +239,8 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -# The same P/D with prefill TP2 writing into decode TP1, which needs 3 GPUs. +# The same P/D between different TP sizes (WRITE TP2 -> TP1, READ of MLA +# TP1 <-> TP2), which needs 3 GPUs: 8 cases, ~36 min on MI355. - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP" key: amd-moriio-write-dummy-hetero-tp-mi355 dind: false diff --git a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py index 45b065488245..6f61d62d33d5 100644 --- a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py +++ b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Single-node MoRIIO WRITE P/D checks on dummy weights. +"""Single-node MoRIIO P/D checks on dummy weights, in WRITE and READ mode. Each case serves a real model architecture with ``--load-format dummy``, widened by ``moriio_wide_dummy`` (see ``DUMMY_WEIGHT_RANGE``): the @@ -13,7 +13,10 @@ ``--disable-hybrid-kv-cache-manager``, so all layers, sliding-window ones included, share a single cache group. -Prefill runs on one GPU, decode on another, and MoRI moves the KV over xGMI. +Prefill runs on one GPU, decode on another, and MoRI moves the KV over xGMI: +prefill pushes it in WRITE mode, and decode pulls it in READ mode (cases +marked ``READ``). The proxy runs a READ request's prefill to completion before +it sends the request to decode. Either way decode computes the last prompt token itself. For hybrid models prefill computes only the first ``n - 1`` tokens and transfers the KV and the recurrent state after them. For attention-only models prefill computes and @@ -47,9 +50,11 @@ import random import shutil import signal +import socket import subprocess import sys import tempfile +import threading import time from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -86,6 +91,8 @@ # MoRI transport between the single-node prefill and decode GPUs. BACKEND = os.getenv("MORIIO_DUMMY_E2E_BACKEND", "xgmi") TOP_K = 20 +# Well under the 60 s deferred-send reaper, which frees unreleased blocks. +PREFILL_FREE_TIMEOUT_S = 20 MAX_TOKENS = 4 MAX_MODEL_LEN = 4096 @@ -267,6 +274,8 @@ class Case: decode_tp: int = 1 # Extra arguments for every server of the case. server_args: tuple[str, ...] = () + # Decode pulls the KV (READ) instead of prefill pushing it (WRITE). + read_mode: bool = False @property def extra_args(self) -> tuple[str, ...]: @@ -288,6 +297,7 @@ def id(self) -> str: if self.heterogeneous_tp else "" ) + + ("-READ" if self.read_mode else "") ) @@ -340,6 +350,39 @@ def id(self) -> str: server_args=("--block-size", "2048"), ), ] +# READ mode, equal TP: one case per cache-group kind (KDA + MLA, Gated DeltaNet +# + full attention, full attention, and single-group MLA and sliding window). +READ_CASES = [ + Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT, read_mode=True), + Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, _DEFAULT, read_mode=True), + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT, read_mode=True), + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT, True, read_mode=True), + Case("deepseek_v3", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT, True, read_mode=True), + Case("gemma3", "TRITON_ATTN", _DEFAULT, _DEFAULT, True, read_mode=True), +] +# READ mode between different TP sizes: MLA latents are replicated, so decode +# TP2 reads its copy from prefill TP1 and decode TP1 reads one copy from +# prefill TP2, releasing both. +HETEROGENEOUS_TP_READ_CASES = [ + Case( + "deepseek_v3", + "ROCM_AITER_MLA", + _DEFAULT, + _DEFAULT, + True, + decode_tp=2, + read_mode=True, + ), + Case( + "deepseek_v3", + "ROCM_AITER_MLA", + _DEFAULT, + _DEFAULT, + True, + prefill_tp=2, + read_mode=True, + ), +] def _graph_args(mode: str) -> list[str]: @@ -405,7 +448,7 @@ def _server( extra_args: tuple[str, ...] = (), kv_transfer_config: dict[str, Any] | None = None, ) -> RemoteOpenAIServer: - port = port or get_open_port() + port = port or _open_port() args = [ *arch.args, *extra_args, @@ -444,7 +487,8 @@ def _server( "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"], **_device_env(selector, devices), "VLLM_HOST_IP": "127.0.0.1", - "VLLM_PORT": str(get_open_port()), + # vLLM scans upward from VLLM_PORT for the server internal sockets. + "VLLM_PORT": str(_open_port_block(32)), "VLLM_WORKER_MULTIPROC_METHOD": "spawn", "VLLM_SSM_CONV_STATE_LAYOUT": "DS", "VLLM_CACHE_ROOT": str(_CACHE_ROOT / role), @@ -460,19 +504,61 @@ def _server( ) -def _moriio_config(role: str, http_port: int, ping_port: int) -> dict[str, Any]: +# Ports handed to a server of this test session, so later picks avoid them. +_HANDED_OUT: set[int] = set() +# Prefill and decode start in parallel threads, so picking a port and recording +# it must be one step. +_HANDED_OUT_LOCK = threading.Lock() +# MoRIIO binds ``port + tp_rank`` (times the DP rank) on every rank, so a +# handshake or notify port needs a free block, not one free port. +_PORT_BLOCK = 8 + + +def _bindable(port: int) -> bool: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("", port)) + except OSError: + return False + return True + + +def _open_port() -> int: + with _HANDED_OUT_LOCK: + while True: + port = get_open_port() + if port not in _HANDED_OUT: + _HANDED_OUT.add(port) + return port + + +def _open_port_block(size: int = _PORT_BLOCK) -> int: + """First port of ``size`` consecutive free ports no other server got.""" + with _HANDED_OUT_LOCK: + while True: + base = get_open_port() + block = range(base, base + size) + if _HANDED_OUT.isdisjoint(block) and all(_bindable(p) for p in block): + _HANDED_OUT.update(block) + return base + + +def _moriio_config( + role: str, http_port: int, ping_port: int, read_mode: bool = False +) -> dict[str, Any]: return { "kv_connector": "MoRIIOConnector", "kv_role": role, "kv_connector_extra_config": { "backend": BACKEND, - "read_mode": False, + "read_mode": read_mode, "host_ip": "127.0.0.1", "proxy_ip": "127.0.0.1", "proxy_ping_port": str(ping_port), "http_port": str(http_port), - "handshake_port": str(get_open_port()), - "notify_port": str(get_open_port()), + # Every TP rank binds its port plus its rank offset. + "handshake_port": str(_open_port_block()), + "notify_port": str(_open_port_block()), }, } @@ -481,7 +567,7 @@ def _moriio_config(role: str, http_port: int, ping_port: int) -> dict[str, Any]: def _proxy(ping_port: int) -> Iterator[str]: # Free ports for both sockets, so cases from parallel shards on one host # do not collide. - port = get_open_port() + port = _open_port() process = subprocess.Popen( [ sys.executable, @@ -521,14 +607,14 @@ def _pd_pair( prefill_overrides: dict[str, Any], decode_overrides: dict[str, Any], case: Case, -) -> Iterator[tuple[str, RemoteOpenAIServer]]: +) -> Iterator[tuple[str, RemoteOpenAIServer, RemoteOpenAIServer]]: selector, devices = _visible_devices() prefill_devices = devices[: case.prefill_tp] decode_devices = devices[case.prefill_tp : case.prefill_tp + case.decode_tp] assert len(decode_devices) == case.decode_tp, ( f"needs {case.prefill_tp} GPUs for prefill and {case.decode_tp} for decode" ) - prefill_port, decode_port, ping_port = (get_open_port() for _ in range(3)) + prefill_port, decode_port, ping_port = (_open_port() for _ in range(3)) with _proxy(ping_port) as proxy_url, contextlib.ExitStack() as stack: servers: list[RemoteOpenAIServer] = [] stack.callback(RemoteOpenAIServer.shutdown_many, servers) @@ -545,7 +631,7 @@ def _pd_pair( port=prefill_port, extra_args=case.extra_args, kv_transfer_config=_moriio_config( - "kv_producer", prefill_port, ping_port + "kv_producer", prefill_port, ping_port, case.read_mode ), ), pool.submit( @@ -559,14 +645,14 @@ def _pd_pair( port=decode_port, extra_args=case.extra_args, kv_transfer_config=_moriio_config( - "kv_consumer", decode_port, ping_port + "kv_consumer", decode_port, ping_port, case.read_mode ), ), ] for future in futures: servers.append(future.result()) _wait_for_registration(proxy_url, arch.model) - yield proxy_url, servers[1] + yield proxy_url, servers[0], servers[1] def _wait_for_registration(proxy_url: str, model: str) -> None: @@ -670,6 +756,22 @@ def _distance(ref: dict[str, float], got: dict[str, float]) -> float: return max(abs((ref[t] - ref_mean) - (got[t] - got_mean)) for t in common) +def _wait_for_prefill_blocks_freed(prefill: RemoteOpenAIServer) -> None: + """Prefill frees a finished request's blocks once every decode rank that + holds a share of it has released it. A missed release leaves them in use + until the deferred-send reaper frees them after 60 s.""" + deadline = time.monotonic() + PREFILL_FREE_TIMEOUT_S + while True: + usage = sum(s.value for s in _metric(prefill, "vllm:kv_cache_usage_perc")) + if usage == 0: + return + assert time.monotonic() < deadline, ( + f"prefill still holds {usage:.2%} of its KV cache " + f"{PREFILL_FREE_TIMEOUT_S} s after the last request finished" + ) + time.sleep(0.5) + + def _external_hits(server: RemoteOpenAIServer) -> float: samples = _metric(server, "vllm:external_prefix_cache_hits_total") return sum(sample.value for sample in samples) @@ -680,12 +782,20 @@ class Reference: prompts: list[list[int]] logprobs: list[dict[str, float]] hybrid: bool + read_mode: bool = False @property def expected_external_hits(self) -> int: - # Hybrid prefill stops one token short; attention-only prefill hands - # decode the whole prompt. - return sum(len(p) - int(self.hybrid) for p in self.prompts) + return _expected_external_hits(self.prompts, self.hybrid, self.read_mode) + + +def _expected_external_hits( + prompts: list[list[int]], hybrid: bool, read_mode: bool +) -> int: + # Hybrid prefill stops one token short, and a READ decode always computes + # the last prompt token itself; an attention-only WRITE hands decode the + # whole prompt. + return sum(len(p) - int(hybrid or read_mode) for p in prompts) def _reference( @@ -724,10 +834,13 @@ def _reference( f"{moved:.4g} for len {len(prompt)}); a wrong prefix would go " "unnoticed" ) - # Determinism of the reference itself (served from its prefix cache). - again = _first_step(url, arch.model, prompts[-1]) - print(f"[noise] repeat distance={_distance(refs[-1], again)}") - return Reference(prompts, refs, arch.hybrid) + # Determinism of the reference itself (served from its prefix cache), + # for every prompt so a P/D mismatch at any length can be told apart + # from noise. + for prompt, ref in zip(prompts, refs): + again = _distance(ref, _first_step(url, arch.model, prompt)) + print(f"[noise] len={len(prompt)} repeat distance={again}") + return Reference(prompts, refs, arch.hybrid, case.read_mode) def _pd_distances( @@ -739,13 +852,18 @@ def _pd_distances( ) -> tuple[list[float], float]: with _pd_pair(arch, prefill_overrides, decode_overrides, case) as ( proxy_url, + _, decode, ): hits_before = _external_hits(decode) - distances = [ - _distance(expected, _first_step(proxy_url, arch.model, prompt)) - for prompt, expected in zip(ref.prompts, ref.logprobs) - ] + distances = [] + for prompt, expected in zip(ref.prompts, ref.logprobs): + got = _first_step(proxy_url, arch.model, prompt) + distances.append(_distance(expected, got)) + if not distances[-1] <= TOLERANCE: + # Keep the evidence: a wrong prefix and numerical noise look + # the same in the distance alone. + print(f"[pd] {case.id} len={len(prompt)} ref={expected} got={got}") hits = _external_hits(decode) - hits_before lengths = [len(p) for p in ref.prompts] print(f"[pd] {case.id} lens={lengths} distances={distances} external_hits={hits}") @@ -756,8 +874,7 @@ def _pd_distances( # Tests -@pytest.mark.parametrize("case", CASES, ids=lambda c: c.id) -def test_write_matches_direct_serve(case: Case): +def _check_matches_direct_serve(case: Case) -> None: arch = ARCHS[case.arch] overrides = _arch_overrides(arch) ref = _reference(arch, overrides, case) @@ -773,31 +890,45 @@ def test_write_matches_direct_serve(case: Case): assert not bad, f"P/D first-token log-probs differ from direct serve: {bad}" -@pytest.mark.parametrize("case", HETEROGENEOUS_TP_CASES, ids=lambda c: c.id) -def test_heterogeneous_tp_write_hands_decode_every_prompt_token(case: Case): - """Prefill and decode of different TP sizes complete WRITE transfers. +@pytest.mark.parametrize("case", CASES, ids=lambda c: c.id) +def test_write_matches_direct_serve(case: Case): + _check_matches_direct_serve(case) + + +@pytest.mark.parametrize("case", READ_CASES, ids=lambda c: c.id) +def test_read_matches_direct_serve(case: Case): + _check_matches_direct_serve(case) + + +@pytest.mark.parametrize( + "case", HETEROGENEOUS_TP_CASES + HETEROGENEOUS_TP_READ_CASES, ids=lambda c: c.id +) +def test_heterogeneous_tp_hands_decode_every_prompt_token(case: Case): + """Prefill and decode of different TP sizes complete WRITE and READ. Unlike the equal-TP cases this cannot compare against a direct serve: a different TP size sums in a different order, and on dummy weights that changes the top tokens entirely, so log-probs only show the transfer did not hang and decode consumed the transferred KV. That every head and state slice lands at the right offset is pinned byte for byte by - tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py. + tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py and + tests/v1/kv_connector/unit/test_moriio_read_tp.py. """ _, devices = _visible_devices() if len(devices) < case.prefill_tp + case.decode_tp: pytest.skip(f"needs {case.prefill_tp + case.decode_tp} GPUs") arch = ARCHS[case.arch] overrides = _arch_overrides(arch) - with _pd_pair(arch, overrides, overrides, case) as (proxy_url, decode): + with _pd_pair(arch, overrides, overrides, case) as (proxy_url, prefill, decode): lengths = _prompt_lengths(arch, decode) prompts = _prompts(lengths) hits_before = _external_hits(decode) tops = [_first_step(proxy_url, arch.model, prompt) for prompt in prompts] hits = _external_hits(decode) - hits_before + _wait_for_prefill_blocks_freed(prefill) print(f"[pd] {case.id} lens={lengths} external_hits={hits}") - expected = sum(len(p) - int(arch.hybrid) for p in prompts) + expected = _expected_external_hits(prompts, arch.hybrid, case.read_mode) assert hits == expected, ( f"decode reported {hits} external prefix-cache hits, expected {expected}" ) @@ -810,10 +941,25 @@ def test_heterogeneous_tp_write_hands_decode_every_prompt_token(case: Case): @pytest.mark.parametrize("arch_name", sorted(ARCHS)) def test_corrupted_prefill_state_is_detected(arch_name: str): """Negative control: KV computed with different values must fail.""" + _check_corrupted_prefill_state_is_detected(arch_name, read_mode=False) + + +@pytest.mark.parametrize("arch_name", ["kimi_linear", "qwen3_dense"]) +def test_corrupted_prefill_state_is_detected_read(arch_name: str): + """The same negative control with decode pulling the KV.""" + _check_corrupted_prefill_state_is_detected(arch_name, read_mode=True) + + +def _check_corrupted_prefill_state_is_detected(arch_name: str, read_mode: bool) -> None: arch = ARCHS[arch_name] default = next(c for c in CASES if c.arch == arch_name) case = Case( - arch_name, default.attention_backend, _DEFAULT, _DEFAULT, default.hma_off + arch_name, + default.attention_backend, + _DEFAULT, + _DEFAULT, + default.hma_off, + read_mode=read_mode, ) overrides = _arch_overrides(arch) ref = _reference(arch, overrides, case) diff --git a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py index 2c19a4ef0162..a1aeedf2fff5 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py @@ -220,6 +220,7 @@ def _build_rank(model, tp_rank, tp_size, layout, hma, state, is_producer) -> Ran worker._has_mamba = any(layer.kind == "mamba" for layer in layers) worker.model_config = SimpleNamespace(get_total_num_kv_heads=lambda: heads) worker.kv_caches = caches + worker.num_blocks = NUM_BLOCKS worker.layer_to_spec = specs worker.layer_to_group = { layer.name: groups.index(layer.group if hma else "full") for layer in layers diff --git a/tests/v1/kv_connector/unit/test_moriio_read_tp.py b/tests/v1/kv_connector/unit/test_moriio_read_tp.py new file mode 100644 index 000000000000..03891d1bbeb7 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_read_tp.py @@ -0,0 +1,498 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mock end-to-end MoRIIO READ: what lands on decode, and who is released. + +Every decode rank runs the real read path (source resolution, offsets, +completion polling) on real cache views; a fake MoRI transport copies the +planned bytes out of the prefill ranks' caches. Checks that every decode rank +ends up with exactly its shard, and that every prefill rank holding the +request is released as many times as it waits for. +""" + +import threading +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import msgspec +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1 import ssm_conv_transfer_utils +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + MoRIIOAgentMetadata, + MoRIIOConnectorMetadata, + MoRIIOError, + MoRIIOMode, + MoRIIOTransferAck, + ReqMeta, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + get_moriio_read_holder_tp_ranks, + resolve_moriio_transfer_ack, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import ( + HeadLayout, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +from .test_moriio_hetero_tp_write import ( + DECODE_BLOCKS, + MODELS, + NUM_BLOCKS, + PREFILL_BLOCKS, + Rank, + _build_rank, + _expected, + _global_state, + _groups, +) + +pytestmark = pytest.mark.cpu_test + +NOTIFY_PORT = 30000 +LAYOUTS = [ + (KVCacheLayout.LBNHC, KVCacheLayout.LBNHC), + (KVCacheLayout.LBHNC, KVCacheLayout.LBHNC), + (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC), +] +LAYOUT_IDS = ["token-major", "head-major", "mixed"] + + +class _Status: + def __init__(self, ok: bool = True): + self.ok = ok + + def Succeeded(self): + return self.ok + + def Failed(self): + return not self.ok + + def Message(self): + return "" if self.ok else "failed" + + def Code(self): + return 0 if self.ok else 1 + + +class _Session: + def __init__(self, local: torch.Tensor, remote: torch.Tensor, tp_rank: int): + self.local, self.remote, self.tp_rank = local, remote, tp_rank + + +def _metadata(rank: Rank) -> MoRIIOAgentMetadata: + worker = rank.worker + metadata = MoRIIOAgentMetadata( + engine_id="prefill", + agent_metadata=b"", + kv_caches_base_addr=[], + num_blocks=NUM_BLOCKS, + block_len=0, + attn_backend_name="test", + kv_head_layouts={ + name: layout.to_wire() for name, layout in worker.kv_head_layouts.items() + }, + mamba_state_layouts={ + name: layout.to_wire() + for name, layout in worker.mamba_state_layouts.items() + }, + attention_block_sizes=worker.attention_block_sizes, + ) + # The handshake carries the metadata as msgpack. + return msgspec.msgpack.decode( + msgspec.msgpack.encode(metadata), type=MoRIIOAgentMetadata + ) + + +def _make_reader( + consumer: Rank, + producers: list[Rank], + backend: str, + fail: bool = False, + prewarmed: dict[str, int] | None = None, +) -> list: + """Route the decode rank's reads to the prefill ranks; return its releases. + + With ``prewarmed`` (session key -> dialed prefill rank, filled by + ``_prewarm``), a read may only use a session the eager handshake opened. + """ + worker = consumer.worker + released: list = [] + read_from: list[int] = [] + + def get_built_session(engine_id: str): + if prewarmed is not None: + assert engine_id in prewarmed, (engine_id, sorted(prewarmed)) + tp_rank = prewarmed[engine_id] + elif "_tp" in engine_id: + tp_rank = int(engine_id.rsplit("_tp", 1)[1]) + else: + tp_rank = worker._remote_tp_rank(len(producers)) + remote = producers[tp_rank] + sessions = [ + _Session(local, remote_region, tp_rank) + for local, remote_region in zip( + consumer.regions, remote.regions, strict=True + ) + ] + return sessions, _metadata(remote) + + def read_remote_data(sizes, local_offsets, remote_offsets, session): + read_from.append(session.tp_rank) + for size, local, remote in zip(sizes, local_offsets, remote_offsets): + session.local[local : local + size] = session.remote[remote : remote + size] + return _Status(not fail) + + def send_notify(transfer_id, host, port, message_type=None, message_fields=None): + released.append((int(port) - NOTIFY_PORT, message_type, message_fields)) + + worker.mode = MoRIIOMode.READ + worker.is_producer = False + worker.dp_rank = 0 + worker._num_mamba_transfer_groups = 1 if worker._has_mamba else 0 + worker.moriio_config = SimpleNamespace( + backend=backend, transfer_timeout=30.0, recv_abort_timeout=600.0 + ) + worker._get_built_session = get_built_session + worker.moriio_wrapper = SimpleNamespace( + lock=threading.Lock(), + read_remote_data=read_remote_data, + send_notify=send_notify, + poll_transfer_batch=lambda statuses: MoRIIOWrapper.poll_transfer_batch( + None, statuses + ), + ) + worker._recving_transfers = defaultdict(dict) + worker._recving_transfers_start = {} + worker._recving_transfers_callback_addr = {} + worker._recving_local_blocks = {} + worker._reads_issued_this_step = [] + worker._mamba_reads_this_step = [] + worker._invalid_block_ids = set() + consumer.sent = read_from + return released + + +def _meta(groups: list[str], producer_tp: int) -> ReqMeta: + return ReqMeta( + transfer_id="t", + local_block_ids=[DECODE_BLOCKS[g] for g in groups], + remote_block_ids=[PREFILL_BLOCKS[g] for g in groups], + remote_host="prefill", + remote_port=0, + remote_handshake_port=0, + remote_notify_port=NOTIFY_PORT, + # start_load_kv keys the remote engine as host:handshake_port. + remote_engine_id="prefill:0", + tp_size=producer_tp, + remote_dp_size=1, + ) + + +def _prewarm(consumer: Rank, meta: ReqMeta, monkeypatch) -> dict[str, int]: + """Run the eager handshake; return session key -> dialed prefill rank.""" + worker = consumer.worker + prewarmed: dict[str, int] = {} + + def handshake(host, port, tp_size, engine_id, dp_rank, tp_rank): + prewarmed[engine_id] = ( + worker._remote_tp_rank(tp_size) if tp_rank is None else tp_rank + ) + return set() + + worker._moriio_handshake = handshake + worker._handshake_initiation_executor = ThreadPoolExecutor(1) + worker._handshake_lock = threading.Lock() + worker._eager_handshaked_engines = set() + worker._remote_agents = {} + worker.layer_name_to_remote_kv_cache_metadata = {} + worker.tp_group = SimpleNamespace(cpu_group=None) + monkeypatch.setattr(torch.distributed, "all_reduce", lambda *a, **k: None) + metadata = MoRIIOConnectorMetadata() + metadata.reqs_to_recv["req"] = meta + worker._eager_handshake_all_dp_ranks(metadata) + worker._handshake_initiation_executor.shutdown() + assert prewarmed + return prewarmed + + +def _ranks(model, producer_tp, consumer_tp, hma, layouts, state): + producers = [ + _build_rank(model, p, producer_tp, layouts[0], hma, state, True) + for p in range(producer_tp) + ] + consumers = [ + _build_rank(model, d, consumer_tp, layouts[1], hma, state, False) + for d in range(consumer_tp) + ] + return producers, consumers + + +def _read_all( + model, producer_tp, consumer_tp, hma, layouts, backend="xgmi", advertised_tp=None +): + """Run every decode rank's READ; return the releases per prefill rank. + + ``advertised_tp`` is the prefill TP size the request carries, which is 0 + when the proxy does not forward it. + """ + if advertised_tp is None: + advertised_tp = producer_tp + heads, layers = MODELS[model] + state = _global_state(layers, heads) + groups = _groups(model, hma) + producers, consumers = _ranks(model, producer_tp, consumer_tp, hma, layouts, state) + releases: dict[int, list] = defaultdict(list) + for consumer in consumers: + released = _make_reader(consumer, producers, backend) + consumer.worker._read_blocks_for_req("req", _meta(groups, advertised_tp)) + assert consumer.worker._pop_done_transfers() == {"t"} + for p, message_type, fields in released: + releases[p].append((message_type, fields)) + for layer in layers: + cache = consumer.caches[layer.name] + received = cache.view(NUM_BLOCKS, -1) if layer.kind == "mamba" else cache + expected = _expected( + layer, + layer.group if hma else "full", + state, + heads, + consumer.tp_rank, + consumer_tp, + received, + ) + assert torch.equal(received, expected), (layer.name, consumer.tp_rank) + return releases + + +def _assert_released(releases: dict[int, list], producer_tp: int) -> None: + """Every prefill rank frees the request exactly once, after its last ACK.""" + for p in range(producer_tp): + counts: dict = {} + completed: set = set() + resolved = [ + resolve_moriio_transfer_ack( + MoRIIOTransferAck("t", fields["consumer_tp_size"]), + producer_tp_size=producer_tp, + live_transfer_ids={"t"}, + notification_counts=counts, + completed_transfer_ids=completed, + ) + for message_type, fields in releases[p] + if message_type == "release" + ] + assert resolved and resolved[-1] == "t", (p, releases[p]) + assert resolved.count("t") == 1, (p, releases[p]) + + +@pytest.fixture(autouse=True) +def _conv_dim_first(monkeypatch): + monkeypatch.setattr( + ssm_conv_transfer_utils, "is_conv_state_dim_first", lambda: True + ) + + +@pytest.mark.parametrize("layouts", LAYOUTS, ids=LAYOUT_IDS) +@pytest.mark.parametrize("tp", [1, 2]) +@pytest.mark.parametrize( + ("model", "hma"), + [ + (model, hma) + for model in MODELS + for hma in (True, False) + # Recurrent state needs the hybrid KV cache manager. + if hma or "hybrid" not in model + ], +) +def test_equal_tp_read_lands_every_byte(model, hma, tp, layouts): + """Prefill and decode may lay heads and tokens out differently.""" + _assert_released(_read_all(model, tp, tp, hma, layouts), tp) + + +@pytest.mark.parametrize("layouts", LAYOUTS, ids=LAYOUT_IDS) +@pytest.mark.parametrize("tp", [1, 2]) +@pytest.mark.parametrize("model", list(MODELS)) +def test_read_without_advertised_prefill_tp_is_homogeneous(model, tp, layouts): + """A request without the prefill TP size (0) reads as equal TP.""" + _assert_released(_read_all(model, tp, tp, True, layouts, advertised_tp=0), tp) + + +@pytest.mark.parametrize("layouts", LAYOUTS, ids=LAYOUT_IDS) +@pytest.mark.parametrize( + ("model", "producer_tp", "consumer_tp"), + [ + *(("mla", p, d) for p, d in [(1, 2), (2, 4), (2, 1), (4, 2), (1, 4), (4, 1)]), + # Two KV heads are replicated once both sides have at least two ranks. + ("gqa_replicated", 2, 4), + ("gqa_replicated", 4, 2), + ], +) +def test_replicated_read_releases_every_holder( + model, producer_tp, consumer_tp, layouts +): + """In fan-in each decode rank reads one copy and releases every holder.""" + releases = _read_all(model, producer_tp, consumer_tp, True, layouts) + _assert_released(releases, producer_tp) + + +def test_flexible_mirror_spreads_reads_and_releases_every_holder(monkeypatch): + """MLA decode TP1 under prefill TP4 over RDMA: the round-robin picks a + different prefill rank per request, and all four are released each time.""" + heads, layers = MODELS["mla"] + state = _global_state(layers, heads) + producers, (consumer,) = _ranks("mla", 4, 1, True, LAYOUTS[0], state) + consumer.worker.moriio_config = SimpleNamespace(backend="rdma") + prewarmed = _prewarm(consumer, _meta(["full"], 4), monkeypatch) + assert sorted(prewarmed.values()) == [0, 1, 2, 3] + released = _make_reader(consumer, producers, "rdma", prewarmed=prewarmed) + picked = [] + for _ in range(4): + consumer.sent.clear() + released.clear() + consumer.worker._read_blocks_for_req("req", _meta(["full"], 4)) + assert consumer.worker._pop_done_transfers() == {"t"} + picked.append(set(consumer.sent)) + assert sorted(p for p, _, _ in released) == [0, 1, 2, 3] + assert picked == [{0}, {1}, {2}, {3}] + + +def test_xgmi_mirror_reads_the_fixed_source(monkeypatch): + """Over xGMI there is no NIC to balance, and MoRI cannot open sessions to + every peer GPU, so the replicated MLA copy is read from one fixed rank, + over the one session the eager handshake opened.""" + heads, layers = MODELS["mla"] + state = _global_state(layers, heads) + producers, (consumer,) = _ranks("mla", 4, 1, True, LAYOUTS[0], state) + consumer.worker.moriio_config = SimpleNamespace(backend="xgmi") + prewarmed = _prewarm(consumer, _meta(["full"], 4), monkeypatch) + assert len(prewarmed) == 1 + released = _make_reader(consumer, producers, "xgmi", prewarmed=prewarmed) + for _ in range(3): + consumer.sent.clear() + released.clear() + consumer.worker._read_blocks_for_req("req", _meta(["full"], 4)) + assert consumer.worker._pop_done_transfers() == {"t"} + assert set(consumer.sent) == set(prewarmed.values()) + assert sorted(p for p, _, _ in released) == [0, 1, 2, 3] + + +@pytest.mark.parametrize( + ("producer_tp", "consumer_tp", "advertised_tp"), + [(4, 2, 4), (2, 4, 2), (2, 2, 0)], +) +def test_aborted_read_releases_every_holder(producer_tp, consumer_tp, advertised_tp): + heads, layers = MODELS["mla"] + state = _global_state(layers, heads) + producers, consumers = _ranks( + "mla", producer_tp, consumer_tp, True, LAYOUTS[0], state + ) + for consumer in consumers: + released = _make_reader(consumer, producers, "xgmi") + consumer.worker._read_blocks( + local_block_ids=[], + remote_block_ids=[PREFILL_BLOCKS["full"]], + dst_engine_id="prefill", + request_id="req", + transfer_id="t", + remote_host="prefill", + remote_notify_port=NOTIFY_PORT, + remote_tp_size=advertised_tp, + ) + holders = get_moriio_read_holder_tp_ranks( + consumer.tp_rank, consumer_tp, producer_tp + ) + assert sorted(p for p, _, _ in released) == sorted(holders) + assert consumer.sent == [] + + +def test_failed_read_releases_every_holder(): + heads, layers = MODELS["mla"] + state = _global_state(layers, heads) + producers, (consumer,) = _ranks("mla", 2, 1, True, LAYOUTS[0], state) + released = _make_reader(consumer, producers, "xgmi", fail=True) + consumer.worker._read_blocks_for_req("req", _meta(["full"], 2)) + # A failed read is a load error, not a completion. + assert consumer.worker._pop_done_transfers() == set() + assert sorted(p for p, _, _ in released) == [0, 1] + assert consumer.worker._invalid_block_ids == set(DECODE_BLOCKS["full"]) + + +def _mixed_layout_pair(): + """Equal-TP GQA ranks whose layouts differ, so reads go head by head.""" + heads, layers = MODELS["gqa"] + state = _global_state(layers, heads) + (producer,), (consumer,) = _ranks("gqa", 1, 1, True, LAYOUTS[2], state) + return producer, consumer + + +def test_head_read_rejects_peer_with_other_head_count(): + producer, consumer = _mixed_layout_pair() + metadata = _metadata(producer) + layout = HeadLayout.from_wire(metadata.kv_head_layouts["attn0"]) + metadata.kv_head_layouts["attn0"] = layout._replace( + num_heads=layout.num_heads // 2 + ).to_wire() + with pytest.raises(MoRIIOError, match="KV heads per rank"): + consumer.worker._compute_read_attention_offsets( + "attn0", + DECODE_BLOCKS["full"], + PREFILL_BLOCKS["full"], + metadata, + remote_tp_size=1, + remote_engine_id="prefill", + ) + + +@pytest.mark.parametrize( + ("local", "remote", "match"), + [ + ([NUM_BLOCKS], [0], f"local block id {NUM_BLOCKS} is outside"), + ([0], [NUM_BLOCKS], f"remote block id {NUM_BLOCKS} is outside"), + ([0], [-1], "remote block id -1 is outside"), + ], +) +def test_head_read_rejects_block_ids_outside_the_cache(local, remote, match): + producer, consumer = _mixed_layout_pair() + with pytest.raises(ValueError, match=match): + consumer.worker._compute_read_attention_offsets( + "attn0", + local, + remote, + _metadata(producer), + remote_tp_size=1, + remote_engine_id="prefill", + ) + + +def test_head_read_rejects_more_local_than_remote_blocks(): + producer, consumer = _mixed_layout_pair() + remote = PREFILL_BLOCKS["full"] + with pytest.raises(ValueError, match="longer than remote_block_ids"): + consumer.worker._compute_read_attention_offsets( + "attn0", + DECODE_BLOCKS["full"][: len(remote)] + [0], + remote, + _metadata(producer), + remote_tp_size=1, + remote_engine_id="prefill", + ) + + +@pytest.mark.parametrize( + ("local_tp_rank", "local_tp_size", "remote_tp_size", "expected"), + [ + (1, 2, 2, (1,)), + (3, 4, 2, (1,)), + (1, 2, 4, (2, 3)), + (0, 1, 8, tuple(range(8))), + ], +) +def test_read_holder_tp_ranks(local_tp_rank, local_tp_size, remote_tp_size, expected): + assert ( + get_moriio_read_holder_tp_ranks(local_tp_rank, local_tp_size, remote_tp_size) + == expected + ) diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py index e38ec3f7e9be..d1e4b9b3070b 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -30,6 +30,7 @@ from collections import Counter from dataclasses import dataclass +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -88,6 +89,8 @@ def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int): w.tp_rank = tp_rank w.dp_rank = dp_rank w.use_mla = True + # The round-robin spreads reads over prefill NICs, so it runs over RDMA. + w.moriio_config = SimpleNamespace(backend="rdma") return w diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index c917732b6cf9..fe6b116476e8 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -340,7 +340,7 @@ def shutdown(self): dict, {"req": {"layer0": [DoneStatus(), DoneStatus()]}} ) worker._recving_transfers_callback_addr = { - "req": ("127.0.0.1", "7000", "tx-release") + "req": ("127.0.0.1", ("7000",), "tx-release") } # Transfer-timeout reaping state consulted by _pop_done_transfers. worker._recving_transfers_start = {} @@ -443,7 +443,7 @@ def shutdown(self): dict, {"req": {"kda_layer": [Status(done=True), Status(done=False)]}} ) worker._recving_transfers_callback_addr = { - "req": ("127.0.0.1", "7000", "tx-pending") + "req": ("127.0.0.1", ("7000",), "tx-pending") } worker._recving_transfers_start = {"req": float("inf")} worker._recving_local_blocks = {"req": [1, 2]} 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 d6039e09e382..8abbea1a2b98 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 @@ -82,6 +82,7 @@ is_mla_cache_layer, iter_layer_registration_regions, kda_conv_ssm, + validate_transfer_block_ids, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( MoRIIOWriteTPMapping, @@ -188,6 +189,19 @@ def get_moriio_remote_tp_rank( return local_tp_rank // (local_tp_size // remote_tp_size) +def get_moriio_read_holder_tp_ranks( + local_tp_rank: int, local_tp_size: int, remote_tp_size: int +) -> tuple[int, ...]: + """Remote TP ranks that hold a local rank's share of a READ request. + + Each of them frees its blocks only once this rank releases it. + """ + first = get_moriio_remote_tp_rank(local_tp_rank, local_tp_size, remote_tp_size) + if remote_tp_size <= local_tp_size: + return (first,) + return tuple(range(first, first + remote_tp_size // local_tp_size)) + + def validate_moriio_heterogeneous_tp_kv_heads( local_tp_size: int, remote_tp_size: int, @@ -1896,8 +1910,10 @@ def __init__( # In-progress READ transfers: req_id -> {layer_name: [statuses]}. # A KDA layer posts two reads (conv + ssm), so a layer maps to a list. self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict) - # Values are (remote_host, remote_notify_port, transfer_id). - self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {} + # (host, notify ports of the prefill ranks to release, transfer id). + self._recving_transfers_callback_addr: dict[ + ReqId, tuple[str, tuple[str, ...], str] + ] = {} # Monotonic-clock start times for each in-flight recv transfer. # Used by _pop_done_transfers to abort transfers whose RDMA # completion is lost, instead of hanging forever. @@ -2911,15 +2927,9 @@ def _pop_done_transfers(self) -> set[str]: ] state = self.moriio_wrapper.poll_transfer_batch(statuses) if statuses and state is TransferBatchState.DONE: - host, port, xfer_id = self._recving_transfers_callback_addr[req_id] + host, ports, xfer_id = self._recving_transfers_callback_addr[req_id] done_req_ids.add(xfer_id) - self.moriio_wrapper.send_notify( - xfer_id, - host, - port, - message_type="release", - message_fields={"consumer_tp_size": self.world_size}, - ) + self._send_read_release(xfer_id, host, ports) to_remove.append(req_id) elif state is TransferBatchState.FAILED: failed_status = next( @@ -2934,15 +2944,9 @@ def _pop_done_transfers(self) -> set[str]: failed_status.Code() if failed_status else "unknown", ) self._record_failed_recv(req_id) - host, port, xfer_id = self._recving_transfers_callback_addr[req_id] + host, ports, xfer_id = self._recving_transfers_callback_addr[req_id] try: - self.moriio_wrapper.send_notify( - xfer_id, - host, - port, - message_type="release", - message_fields={"consumer_tp_size": self.world_size}, - ) + self._send_read_release(xfer_id, host, ports) except Exception: logger.exception( "Failed to send error notification for request %s", @@ -3179,7 +3183,8 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No # 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 + self._spreads_replicated_reads() + and self.world_size == 1 and self.use_mla and remote_dp_size == 1 and tp_size > 1 @@ -3369,6 +3374,15 @@ def _next_flex_tp_rank(self, remote_tp_size: int) -> int: self._flex_tp_rr = rr + 1 return rr % remote_tp_size + def _spreads_replicated_reads(self) -> bool: + """Whether the flexible mirror round-robins over prefill TP ranks. + + The round-robin spreads RDMA reads over the prefill NICs. xGMI uses + none, and MoRI cannot open an xGMI session to a peer GPU whose device + index this process does not see, so xGMI reads the fixed source. + """ + return self.moriio_config.backend != "xgmi" + def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]: """Resolve (chosen_tp, flexible) for reading this request's KV. @@ -3384,7 +3398,8 @@ def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]: """ remote_tp_size = int(meta.tp_size) flexible = ( - self.world_size == 1 + self._spreads_replicated_reads() + and self.world_size == 1 and self.use_mla and int(meta.remote_dp_size) == 1 and remote_tp_size > 1 @@ -3627,6 +3642,64 @@ def _block_copy_matches_peer( return True return local.is_dense and HeadLayout.from_wire(remote_wire) == local + def _compute_read_attention_offsets( + self, + layer_name: str, + local_block_ids: list[int], + remote_block_ids: list[int], + remote_moriio_meta: MoRIIOAgentMetadata, + remote_tp_size: int, + remote_engine_id: EngineId, + ) -> tuple[list[int], list[int], list[int]]: + """(local, remote, sizes) reading this rank's whole shard of a layer. + + A raw block copy only lands right when both sides lay out heads and + tokens the same way; otherwise every head is moved one by one. + """ + if self._block_copy_matches_peer(layer_name, remote_moriio_meta): + return self._compute_block_transfer_offsets( + layer_name, + local_block_ids, + remote_block_ids, + remote_moriio_meta, + remote_tp_size=remote_tp_size, + remote_engine_id=remote_engine_id, + ) + local_layout = self.kv_head_layouts[layer_name] + remote_layout = HeadLayout.from_wire( + remote_moriio_meta.kv_head_layouts[layer_name] + ) + if remote_layout.num_heads != local_layout.num_heads: + raise MoRIIOError( + f"Layer {layer_name} holds {local_layout.num_heads} KV heads per " + f"rank here and {remote_layout.num_heads} on the peer" + ) + validate_transfer_block_ids( + local_block_ids, + remote_block_ids, + self.num_blocks, + remote_moriio_meta.num_blocks, + ) + # The prefill side is the source: compute its offsets as local and swap. + source_offsets, local_offsets, sizes = compute_head_slice_offsets( + remote_layout, + local_layout, + list(remote_block_ids[: len(local_block_ids)]), + list(local_block_ids), + 0, + 0, + local_layout.num_heads, + merge_fn=lambda src, dst, sizes: self.merge_contiguous_blocks( + src, dst, sizes, assume_sorted=False + ), + ) + return self._apply_layer_mr_offsets( + layer_name, + (local_offsets, source_offsets, sizes), + remote_moriio_meta, + remote_engine_id, + ) + def _compute_head_slice_transfer_offsets( self, layer_name: str, @@ -3877,6 +3950,8 @@ def _read_blocks( ) -> None: if self.mode == MoRIIOMode.WRITE: return + # 0 means the peer did not advertise its TP size: homogeneous TP. + remote_tp_size = int(remote_tp_size) or self.world_size # Read from the prefill rank that actually computed this request's KV # (forwarded by the proxy). Hardcoding DP0 reads from a different rank's @@ -3893,25 +3968,28 @@ def _read_blocks( if chosen_tp is not None else self._remote_tp_rank(remote_tp_size) ) + # Every prefill rank holding this rank's share frees its blocks only + # once released, including the ones whose identical copy is not read. + holders = ( + (eff_tp,) + if remote_tp_size == self.world_size + else get_moriio_read_holder_tp_ranks( + self.tp_rank, self.world_size, remote_tp_size + ) + ) + notify_ports = tuple( + str( + remote_notify_port + + get_port_offset(int(remote_dp_rank), tp_rank, remote_tp_size) + ) + for tp_rank in holders + ) if not any(local_block_ids): # A request aborted before decode allocation still carries the # producer's block table. There is no local destination or forward, # so release the producer blocks without building sessions or # posting RDMA reads. - self.moriio_wrapper.send_notify( - transfer_id, - remote_host, - str( - remote_notify_port - + get_port_offset( - int(remote_dp_rank), - eff_tp, - remote_tp_size, - ) - ), - message_type="release", - message_fields={"consumer_tp_size": self.world_size}, - ) + self._send_read_release(transfer_id, remote_host, notify_ports) return if self._has_mamba: @@ -3987,13 +4065,13 @@ def _read_blocks( group_idx = self.layer_to_group[layer_name] local_layer_blocks = local_block_ids[group_idx] remote_layer_blocks = remote_block_ids[group_idx] - offs = self._compute_block_transfer_offsets( + offs = self._compute_read_attention_offsets( layer_name, local_layer_blocks, remote_layer_blocks, remote_moriio_meta, - remote_tp_size=remote_tp_size, - remote_engine_id=remote_dp_engine_id, + remote_tp_size, + remote_dp_engine_id, ) statuses.append( self._post_read_with_backoff( @@ -4022,13 +4100,19 @@ def _read_blocks( self._reads_issued_this_step.extend(statuses) self._recving_transfers_callback_addr[request_id] = ( remote_host, - str( - remote_notify_port - + get_port_offset( - int(remote_dp_rank), - eff_tp, - remote_tp_size, - ) - ), + notify_ports, transfer_id, ) + + def _send_read_release( + self, transfer_id: str, host: str, ports: tuple[str, ...] + ) -> None: + """Release a READ's blocks on every prefill rank that holds them.""" + for port in ports: + self.moriio_wrapper.send_notify( + transfer_id, + host, + port, + message_type="release", + message_fields={"consumer_tp_size": self.world_size}, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 5e7200b430c7..f9d2bfe01134 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -418,6 +418,31 @@ def merge_contiguous_offsets( ) +def validate_transfer_block_ids( + local_block_ids: list[int], + remote_block_ids: list[int], + local_num_blocks: int, + remote_num_blocks: int, +) -> None: + """Check that a block transfer stays inside both sides' caches. + + A shorter or empty local list transfers only what decode allocated (a READ + after a partial or full prefix hit); a longer local list is a bug. + """ + if len(local_block_ids) > len(remote_block_ids): + raise ValueError( + "local_block_ids longer than remote_block_ids: " + f"{len(local_block_ids)} > {len(remote_block_ids)}" + ) + for side, block_ids, num_blocks in ( + ("local", local_block_ids, local_num_blocks), + ("remote", remote_block_ids, remote_num_blocks), + ): + invalid = next((b for b in block_ids if not 0 <= b < num_blocks), None) + if invalid is not None: + raise ValueError(f"{side} block id {invalid} is outside [0, {num_blocks})") + + def compute_block_transfer_offsets( layer_name: str, kv_cache: torch.Tensor, @@ -429,43 +454,15 @@ def compute_block_transfer_offsets( [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] ] = merge_contiguous_offsets, ) -> tuple[list[int], list[int], list[int]]: - # A shorter/empty local list transfers only what decode allocated (READ-mode - # partial or full-prefix-hit); a longer local list is a genuine bug. - if len(local_block_ids) > len(remote_block_ids): - raise ValueError( - "local_block_ids longer than remote_block_ids: " - f"{len(local_block_ids)} > {len(remote_block_ids)}" - ) geometry = cast( LayerTransferGeometry, get_layer_transfer_geometry( layer_name, kv_cache, layer_to_spec, remote_num_blocks ), ) - invalid_local = next( - ( - block_id - for block_id in local_block_ids - if not 0 <= block_id < geometry.num_blocks - ), - None, - ) - if invalid_local is not None: - raise ValueError( - f"local block id {invalid_local} is outside [0, {geometry.num_blocks})" - ) - invalid_remote = next( - ( - block_id - for block_id in remote_block_ids - if not 0 <= block_id < remote_num_blocks - ), - None, + validate_transfer_block_ids( + local_block_ids, remote_block_ids, geometry.num_blocks, remote_num_blocks ) - if invalid_remote is not None: - raise ValueError( - f"remote block id {invalid_remote} is outside [0, {remote_num_blocks})" - ) element_size = kv_cache.element_size() transfer_size_byte = geometry.block_len per_block = geometry.transfers_per_block From 139feb45348d3eb303107a896e21303db3a0344e Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:00:39 -0400 Subject: [PATCH 014/631] [ROCm][Perf][DSv4] Standalone mono decode layer for DeepSeek-V4-Pro on MI355X (#61482) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Codex Co-authored-by: vllm-agent --- .buildkite/test_areas/models_basic.yaml | 20 + tests/models/test_dsv4_mono_decode.py | 373 ++++ tests/models/test_dsv4_mono_layers.py | 355 ++++ tests/models/test_dsv4_mono_numerics.py | 776 ++++++++ vllm/config/kernel.py | 4 +- vllm/models/deepseek_v4/amd/model.py | 7 + vllm/models/deepseek_v4/amd/mono/README.md | 112 ++ vllm/models/deepseek_v4/amd/mono/__init__.py | 5 + .../amd/mono/attention/__init__.py | 4 + .../deepseek_v4/amd/mono/attention/back.py | 1021 ++++++++++ .../deepseek_v4/amd/mono/attention/device.py | 504 +++++ .../deepseek_v4/amd/mono/attention/front.py | 573 ++++++ .../deepseek_v4/amd/mono/attention/gemv.py | 261 +++ .../deepseek_v4/amd/mono/attention/plan.py | 102 + .../deepseek_v4/amd/mono/common/__init__.py | 9 + vllm/models/deepseek_v4/amd/mono/common/mx.py | 78 + .../models/deepseek_v4/amd/mono/common/ops.py | 301 +++ .../amd/mono/common/peer_memory.py | 104 + .../deepseek_v4/amd/mono/common/plan.py | 135 ++ .../deepseek_v4/amd/mono/common/sync.py | 177 ++ vllm/models/deepseek_v4/amd/mono/dims.py | 162 ++ vllm/models/deepseek_v4/amd/mono/layer.py | 927 +++++++++ vllm/models/deepseek_v4/amd/mono/runner.py | 463 +++++ vllm/models/deepseek_v4/amd/mono/sources.py | 8 + .../deepseek_v4/amd/mono/stages/__init__.py | 5 + .../deepseek_v4/amd/mono/stages/allreduce.py | 103 + .../deepseek_v4/amd/mono/stages/gemv.py | 294 +++ .../models/deepseek_v4/amd/mono/stages/moe.py | 1733 +++++++++++++++++ .../deepseek_v4/amd/mono/stages/moe_plan.py | 54 + .../deepseek_v4/amd/mono/stages/moe_shape.py | 262 +++ .../deepseek_v4/amd/mono/stages/router.py | 85 + .../deepseek_v4/amd/mono/stages/seam.py | 427 ++++ vllm/models/deepseek_v4/amd/mono_decode.py | 423 ++++ 33 files changed, 9865 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_dsv4_mono_decode.py create mode 100644 tests/models/test_dsv4_mono_layers.py create mode 100644 tests/models/test_dsv4_mono_numerics.py create mode 100644 vllm/models/deepseek_v4/amd/mono/README.md create mode 100644 vllm/models/deepseek_v4/amd/mono/__init__.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/__init__.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/back.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/device.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/front.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/gemv.py create mode 100644 vllm/models/deepseek_v4/amd/mono/attention/plan.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/__init__.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/mx.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/ops.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/peer_memory.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/plan.py create mode 100644 vllm/models/deepseek_v4/amd/mono/common/sync.py create mode 100644 vllm/models/deepseek_v4/amd/mono/dims.py create mode 100644 vllm/models/deepseek_v4/amd/mono/layer.py create mode 100644 vllm/models/deepseek_v4/amd/mono/runner.py create mode 100644 vllm/models/deepseek_v4/amd/mono/sources.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/__init__.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/allreduce.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/gemv.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/moe.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/moe_plan.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/moe_shape.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/router.py create mode 100644 vllm/models/deepseek_v4/amd/mono/stages/seam.py create mode 100644 vllm/models/deepseek_v4/amd/mono_decode.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index d4880005f4be..59983e0e45db 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -105,6 +105,7 @@ steps: - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_mono_decode.py - tests/models/test_muse_spark_rocm_loading.py + - tests/models/test_dsv4_mono_decode.py - tests/models/test_deepseek_v4_vl_stream_weights.py - tests/models/hy_v4/test_triton_ihc.py - vllm/models/muse_spark/ @@ -114,6 +115,7 @@ steps: - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/hy_v4/test_triton_ihc.py + - pytest -v -s models/test_dsv4_mono_decode.py - label: ":amd: (MI355) MiniMax M3 MonoKernel" key: minimax-m3-monokernel-mi355 @@ -347,6 +349,24 @@ steps: - pytest -v -s tests/models/test_dsv41_mono_numerics.py - pytest -v -s tests/models/test_dsv41_mono_layers.py +- label: ':amd: (MI355) DeepSeek-V4-Pro Mono Decode Numerics' + key: dsv4-mono-decode-numerics-amd + dind: false + device: mi355_8 + num_devices: 8 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - docker/Dockerfile.rocm_base + - vllm/models/deepseek_v4/amd/ + - tests/models/test_dsv4_mono_numerics.py + - tests/models/test_dsv4_mono_layers.py + commands: + - pytest -v -s tests/models/test_dsv4_mono_numerics.py + - pytest -v -s tests/models/test_dsv4_mono_layers.py + - label: ':amd: (MI300) DeepSeek-V4.1 gfx942 Mono Decode Layers' key: dsv41-mono-decode-layers-gfx942-amd dind: false diff --git a/tests/models/test_dsv4_mono_decode.py b/tests/models/test_dsv4_mono_decode.py new file mode 100644 index 000000000000..ee17ac9852f6 --- /dev/null +++ b/tests/models/test_dsv4_mono_decode.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4-Pro mono decode layer runs only where its kernels are exact: +decode-only steps of at most 64 rows with causal SWA windows, eager or in a +FULL graph. Every backbone layer takes the whole layer: a CSA layer through +its top-k (its compressor and indexer run between the launches), an HCA layer +through every compressed row (vLLM's c128a table), a hash-routed layer with +its tokens' experts. Deployments the kernels cannot run are an error.""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.deepseek_v4.amd import mono_decode as md +from vllm.models.deepseek_v4.amd.mono.dims import HIDDEN, MAX_TOKENS, WINDOW + +M = 6 +HCA_WIDTH = 512 +DEFAULT = object() + + +class _Runner: + def __init__(self): + self.calls = [] + + def forward(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return ("x", "residual", "post_mix", "res_mix", "extra") + + +def _cache(block): + return torch.zeros(4, block * md.RECORD, dtype=torch.uint8) + + +def _layer(ratio=md.CSA, tid2eid=None): + attn = SimpleNamespace( + prefix="comp", + swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=_cache(64)), + kv_cache=_cache(64), + compress_ratio=ratio, + topk_indices_buffer=torch.zeros(64, 1024, dtype=torch.int32), + ) + ffn = SimpleNamespace(gate=SimpleNamespace(tid2eid=tid2eid)) + return SimpleNamespace(attn=attn, ffn=ffn) + + +def _swa(**overrides): + fields = dict( + num_prefills=0, + num_decodes=1, + num_decode_tokens=M, + decode_swa_width=WINDOW, + decode_swa_indices=torch.zeros(64, 1, WINDOW, dtype=torch.int32), + decode_swa_lens=torch.zeros(64, dtype=torch.int32), + token_to_req_indices=torch.zeros(64, dtype=torch.int32), + slot_mapping=torch.zeros(64, dtype=torch.int64), + block_size=64, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _c128a(rows=M, width=128): + """VLLM's c128a decode table: its step's rows of a 64-row buffer as wide as + the longest context, cut to the step's active width.""" + buffer = torch.arange(64 * HCA_WIDTH, dtype=torch.int32).view(64, HCA_WIDTH) + return buffer[:rows, :width].view(rows, 1, -1) + + +def _comp(rows=M, **overrides): + fields = dict( + block_size=256, + block_table=torch.zeros(1, 4, dtype=torch.int32), + c128a_global_decode_topk_indices=_c128a(rows), + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _inputs(rows=M): + return dict( + x=torch.zeros(rows, 4, HIDDEN, dtype=torch.bfloat16), + positions=torch.zeros(rows, dtype=torch.int64), + input_ids=torch.arange(rows), + residual=torch.zeros(rows, 4, HIDDEN, dtype=torch.bfloat16), + post_mix=torch.zeros(rows, 4, 1), + res_mix=torch.zeros(rows, 4, 4), + ) + + +@pytest.fixture +def run(monkeypatch): + """Call the mono path for one layer under a forward context built from + ``mode`` and ``metadata``; returns (result, runner calls). ``run.built`` + records the runner's build arguments.""" + built: list = [] + + def _run(mode=CUDAGraphMode.FULL, metadata=DEFAULT, layer=None, **inputs): + if metadata is DEFAULT: + metadata = {"swa": _swa(), "comp": _comp()} + fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata) + runner = _Runner() + monkeypatch.setattr(md, "is_forward_context_available", lambda: True) + monkeypatch.setattr(md, "get_forward_context", lambda: fc) + + def build(*args): + built.append(args) + return runner + + monkeypatch.setattr(md, "_mono_runner", build) + monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) + mono = md.MonoDecodeLayer(HCA_WIDTH) + return mono(layer or _layer(), **{**_inputs(), **inputs}), runner.calls + + _run.built = built # type: ignore[attr-defined] + return _run + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_csa_step_runs_mono_with_its_topk(run, mode): + layer = _layer() + out, calls = run(mode, layer=layer) + assert out == ("x", "residual", "post_mix", "res_mix") and len(calls) == 1 + kwargs = calls[0][1] + assert kwargs["topk_indices"] is layer.attn.topk_indices_buffer + # the CSA cache: a 256-token page holds 64 compressed records + assert kwargs["comp_cache"].shape == (4, 64, md.RECORD) + assert kwargs["fold"] and kwargs["route_ids"] is None + assert "global_keys" not in kwargs + + +@pytest.mark.parametrize( + "ratio, indexed, expected", + [ + (md.CSA, True, ["compressor", "indexer", "topk"]), + # a short context's top-k is every row, already filled + (md.CSA, False, ["compressor", "indexer"]), + (md.HCA, False, ["compressor"]), + ], + ids=["csa", "csa-short-context", "hca"], +) +def test_select_runs_the_compressor_then_the_indexer(run, ratio, indexed, expected): + """Between K1 and K2 the layer's compressor writes its row and a CSA + layer's indexer selects the top-k from K1's normed input and qr.""" + log = [] + normed = torch.empty(M, HIDDEN, device="meta") + qr = torch.empty(M, 1536, device="meta") + q = torch.empty(M, 4, device="meta") if indexed else None + + class Compressor: + fused_wkv_wgate = SimpleNamespace(weight=torch.empty(8, HIDDEN, device="meta")) + + def __call__(self, score, positions, rope): + log.append("compressor") + + class Indexer: + compressor = Compressor() + + def weights_proj(self, x): + return x[:, :4], None + + def __call__(self, x, q_rank, score, weights, positions, rope): + assert x is normed and q_rank is qr + log.append("indexer") + return q, None, weights + + def indexer_op(self, x, index_q, _, weights): + assert x is normed and index_q is q + log.append("topk") + + layer = _layer(ratio) + layer.attn.compressor = Compressor() + layer.attn.indexer = Indexer() if ratio == md.CSA else None + layer.attn.rotary_emb = layer.attn.indexer_rotary_emb = None + _, calls = run(layer=layer) + calls[0][1]["select"](normed, qr) + assert log == expected + + +@pytest.mark.parametrize("decode_tokens", [M, M - 2], ids=["exact", "padded"]) +def test_hca_step_reads_every_compressed_row(run, decode_tokens): + """An HCA layer's keys are its c128a table's rows: the runner is built as + wide as vLLM's table, and the kernels read the table's buffer at its row + stride. A FULL graph padded past the step's decode tokens reads its pad + rows from the same buffer (their slots are -1).""" + comp = _comp(rows=decode_tokens) + metadata = {"swa": _swa(num_decode_tokens=decode_tokens), "comp": comp} + out, calls = run(metadata=metadata, layer=_layer(ratio=md.HCA)) + assert out is not None and len(calls) == 1 + assert [args[1] for args in run.built] == [HCA_WIDTH] + kwargs = calls[0][1] + assert kwargs["global_keys"] and kwargs["topk"] == HCA_WIDTH + keys = kwargs["topk_indices"] + table = comp.c128a_global_decode_topk_indices + assert keys.data_ptr() == table.data_ptr() and keys.stride(0) == HCA_WIDTH + assert keys.untyped_storage().nbytes() >= M * HCA_WIDTH * keys.element_size() + assert kwargs["comp_cache"].shape == (4, 2, md.RECORD) + + +def test_hash_layer_routes_its_tokens(run): + tid2eid = torch.arange(100 * 6).view(100, 6) + ids = torch.tensor([3, 1, 4, 1, 5, 9]) + out, calls = run(layer=_layer(tid2eid=tid2eid), input_ids=ids) + assert out is not None + assert torch.equal(calls[0][1]["route_ids"], tid2eid[ids]) + + +def test_first_layer_runs_mono_from_the_embedding(run): + """The first layer's residual is the embedding broadcast to the four + streams, which K1 takes as it is (no owed post to fold).""" + out, calls = run(residual=None, post_mix=None, res_mix=None) + assert out is not None and len(calls) == 1 + args, kwargs = calls[0] + assert args[2] is args[1] and not kwargs["fold"] + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict(metadata={"swa": _swa(num_prefills=1), "comp": _comp()}), + dict(metadata=None), # a profile / dummy run + dict(metadata={"swa": _swa()}), + dict( + metadata={ + "swa": _swa(), + "comp": _comp(c128a_global_decode_topk_indices=None), + }, + layer=_layer(ratio=md.HCA), + ), + dict(**_inputs(MAX_TOKENS + 1)), + dict(residual=torch.zeros(M, HIDDEN, 4, dtype=torch.bfloat16).mT), + dict(layer=_layer(tid2eid=torch.zeros(100, 6)), input_ids=None), + ], + ids=[ + "piecewise", + "prefill", + "no-metadata", + "no-comp-metadata", + "no-c128a-metadata", + "rows", + "non-contiguous", + "hash-without-ids", + ], +) +def test_other_steps_take_the_original_path(run, case): + out, calls = run(**case) + assert out is None and not calls + + +def _decoder_layer(layer_id=5, **attn): + attn = ( + dict( + prefix=f"model.layers.{layer_id}.attn", + compress_ratio=md.CSA, + indexer=object(), + kv_cache_dtype="fp8_ds_mla", + ) + | attn + ) + quant_method = SimpleNamespace(mxfp4_backend=Mxfp4MoeBackend.AITER_MXFP4_BF16) + ffn = SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(quant_method=quant_method) + ), + shared_experts=object(), + use_mega_moe=False, + fuse_heterogeneous_shared_expert=False, + is_fused_shared_expert_enabled=False, + gate=SimpleNamespace(tid2eid=None, bias_vl=None), + n_routed_experts=384, + n_activated_experts=6, + routed_scaling_factor=2.5, + swiglu_limit=10.0, + scoring_func="sqrtsoftplus", + renormalize=True, + ) + return SimpleNamespace( + attn=SimpleNamespace(**attn), + ffn=ffn, + use_fused_mhc=True, + fuse_mhc_rmsnorm=True, + ) + + +def _config(monokernel=True, max_model_len=8192): + parallel = dict( + enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1 + ) + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(num_hidden_layers=61), + max_model_len=max_model_len, + ), + parallel_config=SimpleNamespace(**parallel), + kernel_config=SimpleNamespace(monokernel=monokernel), + ) + + +def _deployment(monkeypatch, tp=8, aiter_linear=True): + """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" + import vllm.platforms.rocm as rocm + from vllm._aiter_ops import rocm_aiter_ops + + monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) + monkeypatch.delenv("VLLM_ROCM_FP8_PADDING", raising=False) + monkeypatch.setattr(rocm, "get_cdna_version", lambda: 4) + monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: tp) + monkeypatch.setattr(md, "_compute_units", lambda: 256) + monkeypatch.setattr( + rocm_aiter_ops, "is_linear_fp8_enabled", classmethod(lambda cls: aiter_linear) + ) + name = "vllm.models.deepseek_v4.amd.mono.runner" + runner = ModuleType(name) + runner.BLOCKS = 256 # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, name, runner) + + +@pytest.mark.parametrize( + "layer, eligible", + [ + (_decoder_layer(), True), + (_decoder_layer(compress_ratio=md.HCA, indexer=None), True), + (_decoder_layer(layer_id=61), False), + ], + ids=["csa", "hca", "draft"], +) +def test_backbone_layers_take_the_whole_layer(monkeypatch, layer, eligible): + _deployment(monkeypatch) + mono = md.MonoDecodeLayer.create(layer, _config()) + assert (mono is not None) == eligible + + +def test_off_without_the_flag(monkeypatch): + """No --monokernel: no mono layer, whatever the deployment.""" + _deployment(monkeypatch, tp=4) + assert ( + md.MonoDecodeLayer.create(_decoder_layer(), _config(monokernel=False)) is None + ) + + +@pytest.mark.parametrize( + "case, match", + [ + (dict(tp=4), "tensor parallel size 8"), + (dict(aiter_linear=False), "VLLM_ROCM_USE_AITER_LINEAR"), + (dict(padding="0"), "VLLM_ROCM_FP8_PADDING"), + ], + ids=["tp", "aiter-linear", "fp8-padding"], +) +def test_opt_in_outside_the_kernels_raises(monkeypatch, case, match): + """The kernels read TP8 shards of vLLM's preshuffled, padded aiter FP8 + weights: any other deployment is an error, not garbage.""" + padding = case.get("padding") + _deployment(monkeypatch, **{k: v for k, v in case.items() if k != "padding"}) + if padding is not None: + monkeypatch.setenv("VLLM_ROCM_FP8_PADDING", padding) + with pytest.raises(ValueError, match=match): + md.MonoDecodeLayer.create(_decoder_layer(), _config()) + + +def test_hca_key_table_is_vllms_c128a_width(monkeypatch): + """The runner's HCA key table is as wide as the c128a table vLLM's metadata + builder allocates: every compressed row of the longest context, in 128-row + multiples (sparse_mla).""" + _deployment(monkeypatch) + layer = _decoder_layer(compress_ratio=md.HCA, indexer=None) + mono = md.MonoDecodeLayer.create(layer, _config(max_model_len=1_000_000)) + assert mono is not None and mono.hca_width == 7936 diff --git a/tests/models/test_dsv4_mono_layers.py b/tests/models/test_dsv4_mono_layers.py new file mode 100644 index 000000000000..b83109a9e045 --- /dev/null +++ b/tests/models/test_dsv4_mono_layers.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4-Pro mono decode path through the real model: a five-layer +DeepSeek-V4-Pro at TP8 on CDNA4, served with ``--monokernel`` on and off on +the same random weights: every decoder layer's outputs at every decode step +must match. Unlike ``test_dsv4_mono_numerics.py``, which checks the kernels +against vLLM's ops one stage at a time, this runs +``DeepseekV4DecoderLayer.forward`` in the engine -- the model's hook, the +per-step path choice, the weights taken off the layer, the metadata, the +compressor and the indexer's top-k -- so a change to the decoder layer that +the kernels do not follow fails here. Scope: eager steps of one token a +request; speculative decoding's multi-token rows and graph replays are not +covered. + +The checkpoint's own first five layers (``compress_ratios`` 128, 128, 4, 128, +4; ``num_hash_layers`` 3) hold every variant the kernels build, each taking the +whole layer (K1 + K2): +- 0: hash-routed HCA, the first layer (K1 takes the embedding, no fold) +- 1: hash-routed HCA +- 2: hash-routed CSA (vLLM's indexer between K1 and K2) +- 3: bias-routed HCA +- 4: bias-routed CSA +A 32K ``max_model_len`` builds the HCA kernels 256 keys wide; the prompts +(under 8.3K) keep vLLM's c128a decode table at its 128-wide floor, so the +kernels read a table narrower than they are built for. +""" + +import zlib + +import pytest +import torch + +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.platforms import current_platform + +MODEL = "deepseek-ai/DeepSeek-V4-Pro" +TP = 8 +LOAD_FORMAT = "dsv4_mono_random" +TRUNCATED = { + "num_hidden_layers": 5, + # one entry a layer, the MTP layer's included + "compress_ratios": [128, 128, 4, 128, 4, 0], +} +LAYERS = (0, 1, 2, 3, 4) +STEPS = 5 # decode steps after the prefill's token + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +@torch.no_grad() +def _fill(name: str, t: torch.Tensor) -> None: + """A random value for every weight, seeded by its name (alike on every + rank and in every run), in the ranges the checkpoint's formats hold.""" + if t.device.type == "meta": + return # materialized and filled by the dummy loader's own path + g = torch.Generator(device=t.device).manual_seed(zlib.crc32(name.encode())) + + def randn(scale=1.0, shift=0.0): + return torch.randn(t.shape, generator=g, device=t.device) * scale + shift + + if t.dtype in (torch.uint8, torch.float8_e8m0fnu) and "scale" in name: + # E8M0 2^-12 .. 2^-8: one a 128 x 128 FP8 block, one a 32 MXFP4 + # values. Larger random layers make the stack chaotic: atomics' order + # alone moves vLLM's own logits run to run. + u8 = t.view(torch.uint8) + u8.copy_(torch.randint(115, 120, u8.shape, generator=g, device=t.device)) + elif t.dtype in (torch.uint8, torch.float4_e2m1fn_x2): + u8 = t.view(torch.uint8) + u8.copy_(torch.randint(0, 256, u8.shape, generator=g, device=t.device)) + elif t.dtype == torch.float8_e4m3fn: + t.copy_(randn().clamp(-6, 6)) + elif name.endswith("tid2eid"): + # a hash layer's table: distinct experts (of 384) a token + ids = torch.rand(t.shape[0], 384, generator=g, device=t.device) + t.copy_(ids.argsort(1)[:, : t.shape[1]]) + elif not torch.is_floating_point(t): + return + elif "norm" in name: + t.copy_(randn(0.1, 1.0)) + elif name.endswith(("hc_attn_scale", "hc_ffn_scale")): + t.copy_(randn(0.5).abs() + 0.5) + elif name.endswith(("hc_attn_base", "hc_ffn_base", "attn_sink")): + t.copy_(randn(0.1)) + elif name.endswith("embed_tokens.weight"): + t.copy_(randn()) + else: + t.copy_(randn(0.02)) + + +@register_model_loader(LOAD_FORMAT) +class RandomWeightLoader(DummyModelLoader): + """Dummy loading with weights in the checkpoint's ranges, before post-load + processing, so every processed copy follows (plain dummy weights leave the + E8M0 scales zero, and both paths would agree on zeros).""" + + def load_weights(self, model, model_config) -> None: + from vllm.distributed import get_tensor_model_parallel_world_size + + super().load_weights(model, model_config) + for name, t in model.state_dict().items(): + _fill(name, t) + # the routed experts' intermediate past this rank's share, if padded, + # is the loader's zero padding: the kernels skip it + inter = model_config.hf_text_config.moe_intermediate_size + inter //= get_tensor_model_parallel_world_size() + for m in model.modules(): + w13, w2 = getattr(m, "w13_weight", None), getattr(m, "w2_weight", None) + if w13 is None or w2 is None or w13.device.type == "meta": + continue + padded = w13.shape[1] // 2 + if padded > inter: + w13, w2 = w13.data.view(torch.uint8), w2.data.view(torch.uint8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + + +class MonoProbe: + """Worker extension (``worker_extension_cls``): it also brings this module, + and so the loader, into the workers.""" + + def _decoder_layers(self) -> dict: + from vllm.model_executor.models.utils import extract_layer_index + from vllm.models.deepseek_v4.amd.model import DeepseekV4DecoderLayer + + model = self.model_runner.model # type: ignore[attr-defined] + return { + extract_layer_index(m.attn.prefix): m + for m in model.modules() + if isinstance(m, DeepseekV4DecoderLayer) + } + + def capture_steps(self, lens: list[int]) -> None: + """Keep each decoder layer's outputs (x, residual, post_mix, res_mix) + for every decode token taken in a decode-only step (the steps the + mono path takes), keyed by (layer, prompt length, position): a + request's step and row can differ between engines, its prompt length + and position cannot. + ``lens``: the prompt lengths, more than ``STEPS`` apart; a row's + request is the longest prompt not past its position.""" + from vllm.forward_context import get_forward_context + + self._captured: dict = {} + self._gaps: dict = {} + lens = sorted(lens) + + def decode_only(layer) -> bool: + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return False + swa = md.get(layer.attn.swa_cache_layer.prefix) + return swa is not None and swa.num_prefills == 0 + + def keys(layer_id, pos): + for row, p in enumerate(pos.tolist()): + yield row, (layer_id, max(n for n in lens if n <= p), p) + + def hook(layer_id): + def keep(module, args, kwargs, out): + if not decode_only(module): + return + pos = kwargs.get("positions", args[1] if len(args) > 1 else None) + for row, k in keys(layer_id, pos): + kept = [t[row].float().cpu() for t in out[:4]] + self._captured.setdefault(k, kept) + + return keep + + def gap_hook(layer_id, layer): + def keep(ffn, args, kwargs): + # vLLM's path of a bias-routed layer only (the mono kernels + # route inside): each row's top-6 margin, sqrt(softplus(logits)) + # + bias as vLLM routes + if ffn.gate.tid2eid is not None or not decode_only(layer): + return + h = args[0] if args else kwargs["hidden_states"] + logits = h.float() @ ffn.gate.weight.float().T + score = torch.nn.functional.softplus(logits).sqrt() + score = score + ffn.gate.e_score_correction_bias.float() + top = score.topk(7, dim=-1).values + gap = (top[:, 5] - top[:, 6]).tolist() + for row, k in keys(layer_id, self._positions): + self._gaps.setdefault(k, gap[row]) + + return keep + + def pos_hook(module, args, kwargs): + self._positions = kwargs.get( + "positions", args[1] if len(args) > 1 else None + ) + + for i, layer in self._decoder_layers().items(): + layer.register_forward_hook(hook(i), with_kwargs=True) + layer.register_forward_pre_hook(pos_hook, with_kwargs=True) + gap = gap_hook(i, layer) + layer.ffn.register_forward_pre_hook(gap, with_kwargs=True) + + def save_capture(self, path: str) -> None: + rank = self.rank # type: ignore[attr-defined] + torch.save((self._captured, self._gaps), f"{path}.rank{rank}") + + def mono_layers(self) -> dict: + from vllm.models.deepseek_v4.amd import mono_decode + + layers = { + i: None if m.mono is None else m.mono._weights is not None + for i, m in self._decoder_layers().items() + } + runner = mono_decode._runner + return { + "layers": layers, + "epoch": -1 if runner is None else int(runner.epoch[0]), + } + + +def _decode(vllm_runner, monkeypatch, mono: bool, prompts, path: str): + """Greedy decode of ``prompts`` with the mono path on or off: (each + prompt's output, the workers' mono layers); each decoder layer's outputs + for the decode tokens of decode-only steps, and vLLM's routing margins, + saved under ``path``.""" + from vllm import SamplingParams + from vllm.inputs import TokensPrompt + + # the deployment the mono path takes (``MonoDecodeLayer.create``) + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0") + monkeypatch.setenv("VLLM_ROCM_FP8_PADDING", "1") + with vllm_runner( + MODEL, + tensor_parallel_size=TP, + load_format=LOAD_FORMAT, + hf_overrides=TRUNCATED, + worker_extension_cls=f"{__name__}.MonoProbe", + moe_backend="aiter", + monokernel=mono, + kv_cache_dtype="fp8", + # the serving page (vllm_runner's default 16 holds no C128 row) + block_size=256, + tokenizer_mode="deepseek_v4", + # eager: the hooks read every step (the kernels take eager steps as + # they take FULL graph replays) + enforce_eager=True, + max_model_len=32768, + max_num_seqs=len(prompts), + # the steps mixing prefill and decode take vLLM's path and are not + # captured (the runner turns chunked prefill off: a step's tokens + # must cover the model length) + max_num_batched_tokens=32768, + enable_prefix_caching=False, + gpu_memory_utilization=0.5, + ) as runner: + runner.llm.collective_rpc("capture_steps", args=([len(p) for p in prompts],)) + params = SamplingParams(temperature=0.0, max_tokens=STEPS + 1) + outs = runner.llm.generate( + [TokensPrompt(prompt_token_ids=p) for p in prompts], params + ) + runner.llm.collective_rpc("save_capture", args=(path,)) + probe = runner.llm.collective_rpc("mono_layers") + return [o.outputs[0] for o in outs], probe + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + assert a.shape == b.shape, (a.shape, b.shape) + return ((a - b).norm() / b.norm()).item() + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path): + if torch.accelerator.device_count() < TP: + pytest.skip(f"needs {TP} GPUs") + pytest.importorskip("flydsl") + g = torch.Generator().manual_seed(0) + # contexts past the CSA layers' 1024 top-k (4 tokens a compressed entry), + # so the indexer selects; 512 apart, so a row's prompt is unambiguous + jitter = torch.randint(256, (8,), generator=g).tolist() + lens = [4400 + 512 * i + r for i, r in enumerate(jitter)] + prompts = [torch.randint(1000, 60000, (n,), generator=g).tolist() for n in lens] + + on, probe_on = _decode(vllm_runner, monkeypatch, True, prompts, f"{tmp_path}/on") + off, probe_off = _decode( + vllm_runner, monkeypatch, False, prompts, f"{tmp_path}/off" + ) + + # every layer took the mono path on every rank, at every decode step + for rank in probe_on: + assert rank["layers"] == dict.fromkeys(LAYERS, True), rank + assert rank["epoch"] >= len(LAYERS) * STEPS, rank + assert rank["epoch"] % len(LAYERS) == 0, rank + assert all(set(r["layers"].values()) == {None} for r in probe_off) + # each decode token's input matches between the runs while the greedy + # tokens before it do: (prompt length, position) pairs to compare + same: set[tuple[int, int]] = set() + for prompt, a, b in zip(prompts, on, off): + for step in range(STEPS): + if a.token_ids[: step + 1] != b.token_ids[: step + 1]: + break + same.add((len(prompt), len(prompt) + step)) + + # each layer's outputs for those tokens, mono against vLLM's path, token by + # token. Layer 0 sees the same inputs on both paths, but its FFN input and + # output differ by up to test_dsv4_mono_numerics' 2% / 3% (MXFP8 + # activations against vLLM's FP8 a 128), which random weights amplify + # through the MoE: 5%. Each later layer takes the previous one's slightly + # different output: 10%. The hash layers' (0-2) experts are fixed by the + # token; a bias-routed layer's token (3, 4: past the hash layers, its input + # carries their differences) whose top-6 routing is a near-tie on vLLM's + # path may pick another expert. The margin: the FFN input's own MXFP8 + # difference (2e-2), plus twice the token's residual difference (a + # logit moves by its std, 1.7, times the input's relative difference, a + # top score by a quarter of that, both scores either way). Once it has, + # its residual stream differs and the later layers are exempt for it. + # The residual stream, set before the FFN: 4% (2.2% the most measured, + # layer 4). Its mixes for the next seam (post_mix in (0, 2), res_mix in + # [0, 1]; the last layer's feed the model's hc_post) are sigmoids of it, + # whose relative difference blows up where they saturate: 0.1 absolute + # (0.055 measured). With random weights about a third of the tokens + # reach a near-tie by layer 4 (every one within its margin); at most half + # may. A layer the kernels no longer follow is off by far more, on every + # token. + for rank in range(TP): + got, _ = torch.load(f"{tmp_path}/on.rank{rank}") + want, gaps = torch.load(f"{tmp_path}/off.rank{rank}") + keys: list[tuple[int, int]] = sorted( + same & {k[1:] for k in got} & {k[1:] for k in want} + ) + assert len(keys) >= len(prompts), keys + flipped: set = set() + for i in LAYERS: + for k in keys: + if k in flipped: + continue + (x, *stream), (x_ref, *stream_ref) = got[(i, *k)], want[(i, *k)] + res_err = _rel(stream[0], stream_ref[0]) + assert res_err < 0.04, (rank, i, k, "residual", res_err) + for name, a, b in zip( + ("post_mix", "res_mix"), stream[1:], stream_ref[1:] + ): + err = (a.float() - b.float()).abs().max().item() + assert err < 0.1, (rank, i, k, name, err) + tie = 2e-2 + 2 * res_err + err = _rel(x, x_ref) + if err >= (0.05 if i == 0 else 0.10): + gap = gaps.get((i, *k), float("inf")) # none: hash-routed + assert gap < tie, (rank, i, k, "x", err, gap, tie) + flipped.add(k) + assert len(flipped) <= len(keys) // 2, (rank, sorted(flipped)) diff --git a/tests/models/test_dsv4_mono_numerics.py b/tests/models/test_dsv4_mono_numerics.py new file mode 100644 index 000000000000..07607198092c --- /dev/null +++ b/tests/models/test_dsv4_mono_numerics.py @@ -0,0 +1,776 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Numerics of the DeepSeek-V4-Pro mono decode kernels at TP8 on CDNA4 against +the ops vLLM runs for the same decode step: random weights in the checkpoint's +formats (FP8 128 x 128 blocks, MXFP4 experts), processed as vLLM processes +them (aiter's preshuffled shared gate_up, its routed experts' layout). + +The whole layer (K1 + K2) against aiter's fused mHC seams, the ROCm attention +op chain (block-scaled FP8 GEMMs, the fused q norm / RoPE / KV insert, the +sparse MLA, the inverse-RoPE MXFP8 wo_a BMM), the TP all-reduce and +``DeepseekV4MoE``: on CSA layers (their top-k) and HCA ones (every compressed +row), bias-routed and hash-routed, at one- and two-hop widths. + +Each stage is compared on its own input, so a failure names the stage. The +kernels quantize activations as MXFP8 (E8M0 a 32), where vLLM's FP8 GEMMs take +fp32 scales a 128: where the two differ by more than rounding, each is +measured against the exact fp32 result and the kernels must be as close to it +as vLLM is.""" + +import pytest +import ray +import torch + +from vllm.platforms import current_platform + +from ..utils import ( + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_process_parallel, +) + +MODEL = "deepseek-ai/DeepSeek-V4-Pro" +TP = 8 +HIDDEN, HC = 7168, 4 +Q_RANK, HEAD_DIM, NOPE, ROPE = 1536, 512, 448, 64 +N_HEADS, O_GROUPS, O_RANK = 128, 16, 1024 +WINDOW, TOPK, SWA_BLOCK, CACHE_BLOCK = 128, 1024, 64, 256 +RECORD, DATA, ALIGN = 584, 576, 576 +REPLAYS = 200 + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.float().reshape(-1), b.float().reshape(-1) + return torch.nn.functional.cosine_similarity(a, b, dim=0).item() + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + """Relative L2 error of ``a`` against ``b``: unlike the cosine, it sees a + wrong scale.""" + a, b = a.float().reshape(-1), b.float().reshape(-1) + return ((a - b).norm() / b.norm()).item() + + +def _random_moe(moe, gd, gr, vocab: int) -> None: + """Random parameters in the loaded layout, before post-load processing: this + rank's shards (``gd``, a device generator) of the MXFP4 routed experts + (random nibbles, E8M0 scales) and the FP8 shared expert (one E8M0 a + 128 x 128 block); the replicated router from ``gr``, seeded alike on every + rank (a hash layer's table: distinct experts a token).""" + u8 = torch.uint8 + for lin in (moe.shared_experts.gate_up_proj, moe.shared_experts.down_proj): + w = torch.randn(lin.weight.shape, generator=gd, device=gd.device) + lin.weight.data.copy_(w.clamp(-6, 6).to(lin.weight.dtype)) + sc = lin.weight_scale_inv.data.view(u8) + sc.copy_( + torch.randint(116, 121, sc.shape, generator=gd, device=gd.device, dtype=u8) + ) + e = moe.experts.routed_experts + for name, lo, hi in (("weight", 0, 256), ("weight_scale", 118, 122)): + for w in ("w13", "w2"): + t = getattr(e, f"{w}_{name}").data.view(u8) + t.copy_( + torch.randint(lo, hi, t.shape, generator=gd, device=gd.device, dtype=u8) + ) + moe.gate.weight.data.copy_( + torch.randn(moe.gate.weight.shape, generator=gr, device=gr.device) * 0.02 + ) + if moe.gate.e_score_correction_bias is not None: + b = moe.gate.e_score_correction_bias + b.data.copy_(torch.randn(b.shape, generator=gr, device=gr.device) * 0.02) + if moe.gate.tid2eid is not None: + experts, topk = moe.n_routed_experts, moe.n_activated_experts + ids = torch.rand(vocab, experts, generator=gr, device=gr.device) + moe.gate.tid2eid.data.copy_(ids.argsort(1)[:, :topk].to(torch.int32)) + + +# ---------------------------------------------------------------- caches + + +def _paged_cache(num_blocks: int, block: int, device) -> torch.Tensor: + """[blocks, block, 584] u8 over pages of round_up(block * 584, 576) bytes, + the way vLLM binds a packed V4 cache.""" + page = -(-block * RECORD // ALIGN) * ALIGN + raw = torch.zeros(num_blocks * page, dtype=torch.uint8, device=device) + return torch.as_strided(raw, (num_blocks, block, RECORD), (page, RECORD, 1)) + + +def _clone_cache(cache: torch.Tensor) -> torch.Tensor: + raw = torch.as_strided(cache, (cache.shape[0] * cache.stride(0),), (1,)).clone() + return torch.as_strided(raw, cache.shape, cache.stride()) + + +def _rope(kv: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor): + """GPT-J RoPE of the last 64 dims (fp32 math) of rows [n, 512] or heads + [n, h, 512].""" + cs = cos_sin[positions.long()] + if kv.dim() == 3: + cs = cs[:, None] + cos, sin = cs[..., :32], cs[..., 32:] + x = kv.float() + e, o = x[..., NOPE::2].clone(), x[..., NOPE + 1 :: 2].clone() + x[..., NOPE::2] = e * cos - o * sin + x[..., NOPE + 1 :: 2] = e * sin + o * cos + return x + + +def _write_records(cache, block: int, slots, kv): + """fp8_ds_mla records of roped bf16 kv rows at ``slots``: the 448 nope dims + in e4m3 with a UE8M0 per 64, the rope dims in bf16, as vLLM's insert.""" + nope = kv.float()[:, :NOPE].view(-1, 7, 64) + exp = torch.ceil(torch.log2(nope.abs().amax(-1).clamp_min(1e-4) / 448.0)) + q = (nope * torch.exp2(-exp)[..., None]).clamp(-448, 448).reshape(-1, NOPE) + data = torch.cat( + [ + q.to(torch.float8_e4m3fn).view(torch.uint8), + kv[:, NOPE:].to(torch.bfloat16).contiguous().view(torch.uint8), + ], + 1, + ) + scales = torch.zeros(kv.shape[0], 8, dtype=torch.uint8, device=kv.device) + scales[:, :7] = (exp + 127).clamp(0, 255).to(torch.uint8) + flat, d0, s0 = _record_offsets(cache, block, slots) + ar = torch.arange(DATA, device=kv.device) + flat[(d0[:, None] + ar).reshape(-1)] = data.reshape(-1) + ar = torch.arange(8, device=kv.device) + flat[(s0[:, None] + ar).reshape(-1)] = scales.reshape(-1) + + +def _record_offsets(cache, block: int, slots): + n = cache.untyped_storage().nbytes() - cache.storage_offset() + flat = torch.as_strided(cache, (n,), (1,)) + slots = slots.to(cache.device).long() + b, o = slots // block, slots % block + return ( + flat, + b * cache.stride(0) + o * DATA, + b * cache.stride(0) + block * DATA + o * 8, + ) + + +def _read_records(cache, block: int, slots) -> torch.Tensor: + """The fp32 key rows of the records at ``slots``.""" + flat, d0, s0 = _record_offsets(cache, block, slots) + data = flat[(d0[:, None] + torch.arange(DATA, device=cache.device)).reshape(-1)] + data = data.view(-1, DATA) + sc = flat[(s0[:, None] + torch.arange(7, device=cache.device)).reshape(-1)] + nope = data[:, :NOPE].contiguous().view(torch.float8_e4m3fn).float().view(-1, 7, 64) + nope = nope * torch.exp2(sc.view(-1, 7).float() - 127)[..., None] + rope = data[:, NOPE:].contiguous().view(torch.bfloat16).float() + return torch.cat([nope.reshape(-1, NOPE), rope], 1) + + +def _normed_kv(n: int, kv_norm, g, device): + x = torch.randn(n, HEAD_DIM, generator=g).to(device) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)) * kv_norm.float() + return x.to(torch.bfloat16) + + +def _decode_step(ratio: int, R: int, T: int, kv_norm, cos_sin, seed: int): + """R requests x T tokens past random contexts: the step's metadata as + vLLM's builders lay it out, its window rows and the compressed cache + filled; a CSA layer's random top-k a token (request-local indices), an HCA + layer's every compressed row (slots, through the block table: vLLM's + c128a table).""" + g = torch.Generator().manual_seed(seed) + dev = kv_norm.device + ctx = [int(c) for c in torch.randint(3000, 40000, (R,), generator=g)] + M = R * T + spans = [ + (max(0, c - WINDOW - T) // SWA_BLOCK, (c + T - 1) // SWA_BLOCK) for c in ctx + ] + n_swa = 1 + sum(hi - lo + 1 for lo, hi in spans) + perm = (torch.randperm(n_swa - 1, generator=g) + 1).tolist() + bt = torch.zeros(R, max(hi for _, hi in spans) + 1, dtype=torch.int64) + for r, (lo, hi) in enumerate(spans): + for lb in range(lo, hi + 1): + bt[r, lb] = perm.pop() + swa = _paged_cache(n_swa, SWA_BLOCK, dev) + + def slot(r, p): + return int(bt[r, p // SWA_BLOCK]) * SWA_BLOCK + p % SWA_BLOCK + + pos = [ctx[r] + t for r in range(R) for t in range(T)] + t2r = [r for r in range(R) for _ in range(T)] + swa_idx = torch.full((M, WINDOW), -1, dtype=torch.int32) + swa_len = torch.zeros(M, dtype=torch.int32) + for i in range(M): + ps = range(max(pos[i] - WINDOW + 1, 0), pos[i] + 1) + swa_idx[i, : len(ps)] = torch.tensor([slot(t2r[i], p) for p in ps]) + swa_len[i] = len(ps) + hist = [ + (r, p) for r in range(R) for p in range(max(0, ctx[r] - WINDOW - T), ctx[r]) + ] + hp = torch.tensor([p for _, p in hist], device=dev) + kv = _rope(_normed_kv(len(hist), kv_norm, g, dev), hp, cos_sin) + _write_records(swa, SWA_BLOCK, torch.tensor([slot(r, p) for r, p in hist]), kv) + + entries = CACHE_BLOCK // ratio + lens = [(c + T) // ratio for c in ctx] + nblk = [-(-n // entries) for n in lens] + perm = (torch.randperm(sum(nblk), generator=g) + 1).tolist() + cbt = torch.zeros(R, max(nblk), dtype=torch.int32) + for r in range(R): + for b in range(nblk[r]): + cbt[r, b] = perm.pop() + comp = _paged_cache(1 + sum(nblk), entries, dev) + for r in range(R): + idx = torch.arange(lens[r]) + slots = cbt[r, idx // entries].long() * entries + idx % entries + kv = _rope(_normed_kv(lens[r], kv_norm, g, dev), (idx * ratio).to(dev), cos_sin) + _write_records(comp, entries, slots, kv) + if ratio == 4: + topk = torch.full((M, TOPK), -1, dtype=torch.int32) + for i in range(M): + cand = (pos[i] + 1) // ratio + n = min(cand, TOPK) + topk[i, :n] = torch.randperm(cand, generator=g)[:n].to(torch.int32) + else: + n_max = max((p + 1) // ratio for p in pos) + width = max(128, 1 << (n_max - 1).bit_length()) + topk = torch.full((M, width), -1, dtype=torch.int32) + for i in range(M): + j = torch.arange((pos[i] + 1) // ratio) + topk[i, : len(j)] = cbt[t2r[i], j // entries] * entries + j % entries + return dict( + M=M, + positions=torch.tensor(pos, dtype=torch.int64, device=dev), + slot_mapping=torch.tensor( + [slot(t2r[i], pos[i]) for i in range(M)], dtype=torch.int64, device=dev + ), + token_to_req=torch.tensor(t2r, dtype=torch.int32, device=dev), + swa_indices=swa_idx.to(dev), + swa_lens=swa_len.to(dev), + swa_cache=swa, + topk=topk.to(dev), + comp_cache=comp, + comp_block_table=cbt.to(dev), + comp_entries=entries, + ) + + +def _keys(s: dict, i: int, ratio: int) -> tuple[torch.Tensor, torch.Tensor]: + """Token i's (compressed slots, window slots) of the step.""" + top = s["topk"][i] + top = top[top >= 0].long() + if ratio == 4: + e, r = s["comp_entries"], s["token_to_req"][i] + top = s["comp_block_table"][r, top // e].long() * e + top % e + return top, s["swa_indices"][i, : s["swa_lens"][i]].long() + + +def _fp8_block(n: int, k: int, g, device): + """A random FP8 weight [n, k] e4m3 and its 128 x 128 block scales as vLLM + holds them after loading: the rows padded as on ROCm, fp32 powers of two.""" + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _maybe_pad_fp8_weight, + ) + + w = torch.randn(n, k, generator=g).clamp(-6, 6).to(torch.float8_e4m3fn) + s = torch.randint(116, 121, (n // 128, k // 128), generator=g).float() + return _maybe_pad_fp8_weight(w.to(device)), torch.exp2(s - 127).to(device) + + +def _deq(w: torch.Tensor, s: torch.Tensor) -> torch.Tensor: + return w.float() * s.repeat_interleave(128, 0).repeat_interleave(128, 1) + + +class _ExactMoE: + """A MoE layer's output on this rank in fp32 from its weights as loaded + (before post-load processing reshuffles them): MXFP4 / FP8 dequantized, + activations unquantized, vLLM's routing.""" + + FP4 = (0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6) + + def __init__(self, moe): + e = moe.experts.routed_experts + u8 = torch.uint8 + self.moe = moe + self.w13 = ( + e.w13_weight.data.view(u8).clone(), + e.w13_weight_scale.data.view(u8).clone(), + ) + self.w2 = ( + e.w2_weight.data.view(u8).clone(), + e.w2_weight_scale.data.view(u8).clone(), + ) + self.lut = torch.tensor(self.FP4, device=e.w13_weight.device) + + def shared(self, sh) -> None: + """The shared expert's weights, after their scales' upcast.""" + + def deq(lin): + w, sc = lin.weight.float(), lin.weight_scale_inv.float() + return w * sc.repeat_interleave(128, 0).repeat_interleave(128, 1) + + self.sgu, self.sd = deq(sh.gate_up_proj), deq(sh.down_proj) + + def _experts(self, raw, ids): + w, sc = raw[0][ids], raw[1][ids] + lo, hi = self.lut[(w & 15).long()], self.lut[(w >> 4).long()] + v = torch.stack([lo, hi], -1).flatten(-2) + return v * torch.exp2(sc.float() - 127).repeat_interleave(32, -1) + + @staticmethod + def _swiglu(gu): + n = gu.shape[-1] // 2 + g, u = gu[..., :n].clamp(max=10), gu[..., n:].clamp(-10, 10) + return torch.nn.functional.silu(g) * u + + def __call__(self, x: torch.Tensor, route: torch.Tensor | None) -> torch.Tensor: + moe, x = self.moe, x.float() + scores = torch.nn.functional.softplus(x @ moe.gate.weight.float().T).sqrt() + if route is None: + bias = moe.gate.e_score_correction_bias + ids = (scores + bias).topk(moe.n_activated_experts, dim=1).indices + else: + ids = route.long() + w = scores.gather(1, ids) + w = w / w.sum(1, keepdim=True) * moe.routed_scaling_factor + out = self._swiglu(x @ self.sgu.T) @ self.sd.T + for t in range(x.shape[0]): + w13, w2 = self._experts(self.w13, ids[t]), self._experts(self.w2, ids[t]) + act = self._swiglu(torch.einsum("knh,h->kn", w13, x[t])) + out[t] += (w[t, :, None] * torch.einsum("khn,kn->kh", w2, act)).sum(0) + return out + + +class _ExactAttention: + """A layer's attention on this rank in fp32 from its normed input: the + weights dequantized, the keys as a cache holds them (the window's rows of + the step written by the path under test), wo_b's partial out.""" + + def __init__(self, w: dict, eps: float, sink: torch.Tensor): + self.wqkv = _deq(w["wqkv"], w["wqkv_s"]) + self.wq_b = _deq(w["wq_b"], w["wq_b_s"]) + self.wo_a = _deq(w["wo_a"], w["wo_a_s"]) + self.wo_b = _deq(w["wo_b"], w["wo_b_s"]) + self.q_norm, self.kv_norm = w["q_norm"].float(), w["kv_norm"].float() + self.eps, self.sink = eps, sink.float() + + def _rms(self, x, w): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * w + + def kv(self, x, positions, cos_sin) -> torch.Tensor: + """The KV rows of a step's tokens: kv_norm, then RoPE.""" + kv = (x.float() @ self.wqkv.T)[:, Q_RANK:] + return _rope(self._rms(kv, self.kv_norm), positions, cos_sin) + + def __call__(self, x, s: dict, ratio: int, swa_cache, cos_sin) -> torch.Tensor: + M, pos = x.shape[0], s["positions"] + qr_kv = x.float() @ self.wqkv.T + q = (self._rms(qr_kv[:, :Q_RANK], self.q_norm) @ self.wq_b.T).view(M, -1, 512) + q = _rope(self._rms(q, 1.0), pos, cos_sin) + heads = q.shape[1] + o = torch.empty(M, heads, HEAD_DIM, device=x.device) + for i in range(M): + top, win = _keys(s, i, ratio) + k = torch.cat( + [ + _read_records(s["comp_cache"], s["comp_entries"], top), + _read_records(swa_cache, SWA_BLOCK, win), + ] + ) + sc = (q[i] @ k.T) * HEAD_DIM**-0.5 + m = torch.maximum(sc.amax(1), self.sink[:heads]) + p = torch.exp(sc - m[:, None]) + den = p.sum(1) + torch.exp(self.sink[:heads] - m) + o[i] = (p @ k) / den[:, None] + # inverse RoPE of each head's last 64 dims + cs = cos_sin[pos.long()][:, None] + cos, sin = cs[..., :32], cs[..., 32:] + e, od = o[..., NOPE::2].clone(), o[..., NOPE + 1 :: 2].clone() + o[..., NOPE::2] = e * cos + od * sin + o[..., NOPE + 1 :: 2] = od * cos - e * sin + groups = self.wo_a.shape[0] // O_RANK + og = o.view(M, groups, -1) + z = torch.einsum("mgk,gnk->mgn", og, self.wo_a.view(groups, O_RANK, -1)) + return z.reshape(M, -1) @ self.wo_b.T + + +def _attn_output(runner, epoch: int, M: int, rank: int) -> torch.Tensor: + """The wo_b partials every rank pushed to this rank's ATTN region in the + launch pair of ``epoch``, summed: [M, hidden] fp32. On two hops a rank + holds its own columns' partials only: the ranks' columns gathered.""" + from vllm.distributed import tensor_model_parallel_all_reduce + from vllm.models.deepseek_v4.amd.mono.layer import peer_half_bytes + + half = peer_half_bytes(runner.tp) + n = runner.tp * M * HIDDEN // 2 + base = (epoch & 1) * half + pairs = runner.peer.bytes[base : base + n * 8].view(torch.int32).view(n, 2) + parts = pairs[:, 0].contiguous().view(torch.bfloat16).view(runner.tp, M, HIDDEN) + out = parts.float().sum(0) + if runner.two_hop(M): + cols = slice(rank * HIDDEN // runner.tp, (rank + 1) * HIDDEN // runner.tp) + own = torch.zeros_like(out) + own[:, cols] = out[:, cols] + out = tensor_model_parallel_all_reduce(own) + return out + + +@ray.remote(num_gpus=1, max_calls=1) +def _mono_layer_numerics(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.worker.workspace import init_workspace_manager + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + m.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0") + rocm_aiter_ops.refresh_env_variables() + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + init_workspace_manager(device) + _check_layer(rank, tp_size, device) + + +def _check_layer(rank: int, tp: int, device) -> None: + from aiter.ops.batched_gemm_op_a8w8 import batched_gemm_a8w8_mxscale + from aiter.ops.inverse_rope_group_quant import inverse_rope_group_quant + + from vllm._aiter_ops import rocm_aiter_ops as ops + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce + from vllm.engine.arg_utils import EngineArgs + from vllm.forward_context import set_forward_context + from vllm.model_executor.kernels.mhc.torch import mhc_pre_torch + from vllm.model_executor.layers.mhc import MHCFusedPostPreOp, MHCPreOp + from vllm.model_executor.model_loader.utils import process_weights_after_loading + from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE + from vllm.models.deepseek_v4.amd.mono.dims import INDEX_TOPK + from vllm.models.deepseek_v4.amd.mono.layer import scratch_layout + from vllm.models.deepseek_v4.amd.mono.runner import ( + AttnWeights, + MonoLayerWeights, + MonoRunner, + ) + from vllm.models.deepseek_v4.amd.mono_decode import _hca_width, _mx_scales + from vllm.models.deepseek_v4.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) + from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope + from vllm.utils.torch_utils import set_default_torch_dtype + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_ragged_indices_from_dense, + ) + + vc = EngineArgs( + model=MODEL, + tensor_parallel_size=tp, + moe_backend="aiter", + load_format="dummy", + max_model_len=65536, + trust_remote_code=True, + ).create_engine_config() + cfg = vc.model_config.hf_config + eps = cfg.rms_norm_eps + bf = torch.bfloat16 + g = torch.Generator().manual_seed(1000) + gs = torch.Generator().manual_seed(1001 + rank) + H, G = N_HEADS // tp, O_GROUPS // tp + # the HCA key tables: the deployment's (64K context: 512, within the top-k) + # and a long context's (wider than the top-k: ``bounded``) + hca_dep, hca_long = _hca_width(vc), 2048 + assert hca_dep <= INDEX_TOPK < hca_long + runner = MonoRunner(tp, rank, get_tp_group().cpu_group, device, max_topk=hca_long) + with set_current_vllm_config(vc): + fused = MHCFusedPostPreOp() + pre = MHCPreOp() + + def random_moe(layer): + """A layer's MoE (layer 0's routes by hash) and its exact reference.""" + with set_current_vllm_config(vc), set_default_torch_dtype(bf), device: + moe = DeepseekV4MoE(vc, prefix=f"model.layers.{layer}.ffn") + _random_moe( + moe, + torch.Generator(device=device).manual_seed(2001 + rank + 10 * layer), + torch.Generator(device=device).manual_seed(2000 + layer), + cfg.vocab_size, + ) + exact = _ExactMoE(moe) # the weights as loaded, before processing + process_weights_after_loading(moe, vc.model_config, device) + exact.shared(moe.shared_experts) + moe.shared_experts.prepare_gateup_preshuffle() + assert moe.shared_experts._gateup_scale is not None + return moe, exact + + moes = {layer: random_moe(layer) for layer in (3, 0)} + + def f32(*shape, s=1.0, gen=g): + return (torch.randn(*shape, generator=gen) * s).to(device) + + mix = (2 + HC) * HC + seams = { + sub: dict( + fn=f32(mix, HC * HIDDEN, s=0.01), + scale=f32(3, s=0.5).abs() + 0.5, + base=f32(mix, s=0.1), + norm=(1 + f32(HIDDEN, s=0.1)).to(bf), + ) + for sub in ("attn", "ffn") + } + + def seam(sub, x, res, post, comb): + sm = seams[sub] + with set_current_vllm_config(vc): + return fused( + x, res, post, comb, sm["fn"], sm["scale"], sm["base"], eps, + cfg.hc_eps, cfg.hc_eps, 2.0, cfg.hc_sinkhorn_iters, + norm_weight=sm["norm"], norm_eps=eps, + ) # fmt: skip + + def mix_ref(sub, res): + sm = seams[sub] + post, comb, _ = mhc_pre_torch( + res, sm["fn"], sm["scale"], sm["base"], eps, cfg.hc_eps, cfg.hc_eps, + 2.0, cfg.hc_sinkhorn_iters, + ) # fmt: skip + return post, comb + + def moe_out(moe, x, ids): + with set_forward_context(None, vc): + return moe(x.clone(), ids) + + # the attention's weights, as vLLM holds them after loading and its + # preshuffle (wqkv and wo_b shuffled in place, scales fp32) + w = {} + w["wqkv"], w["wqkv_s"] = _fp8_block(Q_RANK + HEAD_DIM, HIDDEN, g, device) + w["wq_b"], w["wq_b_s"] = _fp8_block(H * HEAD_DIM, Q_RANK, gs, device) + w["wo_a"], w["wo_a_s"] = _fp8_block(G * O_RANK, H // G * HEAD_DIM, gs, device) + w["wo_b"], w["wo_b_s"] = _fp8_block(HIDDEN, G * O_RANK, gs, device) + w["q_norm"] = (1 + f32(Q_RANK, s=0.1)).to(bf) + w["kv_norm"] = (1 + f32(HEAD_DIM, s=0.1)).to(bf) + sink = f32(H, gen=gs) + wqkv_sh = ops.shuffle_weight(w["wqkv"], layout=(16, 16)) + wo_b_sh = ops.shuffle_weight(w["wo_b"], layout=(16, 16)) + wo_a_e8m0 = _mx_scales(w["wo_a_s"])[::4, ::4].contiguous() + exact_attn = _ExactAttention(w, eps, sink) + + steps = [] + # (compress ratio, requests, tokens a request, key table width, MoE layer) + for ratio, R, T, width, layer in ( + (4, 1, 7, 0, 3), + (128, 8, 7, hca_long, 0), + (128, 2, 7, hca_dep, 3), + (4, 4, 7, 0, 0), + ): + with set_current_vllm_config(VllmConfig()): + rope = build_deepseek_v4_rope( + cfg, head_dim=HEAD_DIM, rope_head_dim=ROPE, + max_position_embeddings=cfg.max_position_embeddings, + compress_ratio=ratio, + ) # fmt: skip + cos_sin = rope.cos_sin_cache.to(device=device, dtype=torch.float32) + cos_bf, sin_bf = (c.contiguous() for c in cos_sin.to(bf).chunk(2, dim=-1)) + s = _decode_step(ratio, R, T, w["kv_norm"], cos_sin, seed=7 + ratio + R) + M = s["M"] + tag = f"ratio {ratio}, {R} x {T}, width {width}, MoE layer {layer}" + hca = ratio == 128 + moe, exact_moe = moes[layer] + sh, e = moe.shared_experts, moe.experts.routed_experts + mw = MonoLayerWeights( + attn=AttnWeights( + layer_id=3, wqkv=wqkv_sh, wqkv_scale=_mx_scales(w["wqkv_s"]), + q_norm=w["q_norm"], kv_norm=w["kv_norm"], wq_b=w["wq_b"], + wq_b_scale=_mx_scales(w["wq_b_s"]), wo_a=w["wo_a"], + wo_a_scale=_mx_scales(w["wo_a_s"]), wo_b=wo_b_sh, + wo_b_scale=_mx_scales(w["wo_b_s"]), attn_sink=sink, + cos_sin=cos_sin, ratio=ratio, + ), + hc_attn_fn=seams["attn"]["fn"], hc_attn_scale=seams["attn"]["scale"], + hc_attn_base=seams["attn"]["base"], attn_norm=seams["attn"]["norm"], + hc_ffn_fn=seams["ffn"]["fn"], hc_ffn_scale=seams["ffn"]["scale"], + hc_ffn_base=seams["ffn"]["base"], ffn_norm=seams["ffn"]["norm"], + gate_w=moe.gate.weight, bias=moe.gate.e_score_correction_bias, + w13=e.w13_weight, w13_s=e.w13_weight_scale, w2=e.w2_weight, + w2_s=e.w2_weight_scale, sgu=sh.gate_up_proj.weight, + sgu_s=_mx_scales(sh._gateup_scale), sw2=sh.down_proj.weight, + sw2_s=_mx_scales(sh.down_proj.weight_scale_inv), + ) # fmt: skip + mw.attn.check(runner.d) + mw.check(runner.d) + + # the layer's inputs at the attention seam (the same on every rank) + gi = torch.Generator().manual_seed(11 + ratio) + res0 = torch.randn(M, HC, HIDDEN, generator=gi) + res0 = (res0 * torch.rsqrt(res0.pow(2).mean(-1, keepdim=True))).to(bf) + res0 = res0.to(device) + x0 = (torch.randn(M, HIDDEN, generator=gi) * 0.5).to(bf).to(device) + comb0 = torch.rand(M, HC, HC, generator=gi) + 0.1 + for _ in range(20): + comb0 = comb0 / comb0.sum(-1, keepdim=True) + comb0 = comb0 / comb0.sum(-2, keepdim=True) + comb0 = comb0.to(device) + post0 = (torch.sigmoid(torch.randn(M, HC, 1, generator=gi)) * 2).to(device) + ids = torch.randint(0, cfg.vocab_size, (M,), generator=gi).to(device) + hashed = moe.gate.tid2eid is not None + route = moe.gate.tid2eid[ids].contiguous() if hashed else None + swa_v, swa_m = s["swa_cache"], _clone_cache(s["swa_cache"]) + + # vLLM's decode chain + res, post, comb, xa = seam("attn", x0, res0, post0, comb0) + x8, xs = ops.group_fp8_quant(xa, transpose_scale=True) + qr_kv = ops.gemm_a8w8_blockscale_bpreshuffle(x8, wqkv_sh, xs, w["wqkv_s"], bf) + qr8, qrs, kvn = ops.fused_qk_rmsnorm_group_quant( + qr_kv[:, :Q_RANK], w["q_norm"], eps, qr_kv[:, Q_RANK:], w["kv_norm"], eps, + group_size=128, transpose_scale=False, + ) # fmt: skip + q = ops.gemm_a8w8_blockscale(qr8, w["wq_b"], qrs, w["wq_b_s"], [128, 128], bf) + swa2d = swa_v.view(swa_v.shape[0], -1) + q_roped = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q.view(M, H, HEAD_DIM), kvn, swa2d, s["slot_mapping"], s["positions"], + cos_sin, H, eps, SWA_BLOCK, True, False, True, False, + ) # fmt: skip + swa_ind, swa_ptr = build_ragged_indices_from_dense( + s["swa_indices"].reshape(M, -1), s["swa_lens"] + ) + if hca: + lens = (s["topk"] >= 0).sum(1).to(torch.int32) + top_ind, top_ptr = build_ragged_indices_from_dense(s["topk"], lens) + else: + valid = torch.ones(M, dtype=torch.bool, device=device) + top_ind, top_ptr, _ = compute_global_topk_ragged_indices_and_indptr( + s["topk"], s["token_to_req"], s["comp_block_table"], + s["comp_entries"], valid, + ) # fmt: skip + attn = torch.empty(M, H, HEAD_DIM, dtype=bf, device=device) + ops.triton_sparse_mla_fwd( + q_roped, swa_v, attn, HEAD_DIM**-0.5, swa_ptr, swa_ind, + attn_sink=sink, extra_kv_buffer=s["comp_cache"], + extra_kv_indptr=top_ptr, extra_kv_indices=top_ind, + ) # fmt: skip + o8, os_ = inverse_rope_group_quant( + attn, s["positions"], cos_bf, sin_bf, num_groups=G, quant_group_size=128 + ) + z = batched_gemm_a8w8_mxscale( + o8, w["wo_a"].view(G, O_RANK, -1), os_, + wo_a_e8m0.view(G, O_RANK // 128, -1), dtype=bf, + ).flatten(1) # fmt: skip + z8, zs = ops.group_fp8_quant(z.contiguous(), transpose_scale=True) + part = ops.gemm_a8w8_blockscale_bpreshuffle(z8, wo_b_sh, zs, w["wo_b_s"], bf) + a = tensor_model_parallel_all_reduce(part.clone()) + + # ---- the whole layer: K1 + K2 + epoch = int(runner.epoch[0]) + o = runner.forward( + mw, x0, res0, post0, comb0, s["positions"], s["slot_mapping"], + swa_m, s["swa_indices"], s["swa_lens"], s["token_to_req"], + topk_indices=s["topk"], comp_cache=s["comp_cache"], + select=lambda normed, qr: None, + comp_block_table=s["comp_block_table"], route_ids=route, + topk=width, global_keys=hca, + ) # fmt: skip + torch.accelerator.synchronize() + # the attention seam: the residual vLLM's, the mixes the fp32 reference's + assert _cos(runner.res_mid[:M], res) > 0.99999, tag + pm, cm = mix_ref("attn", runner.res_mid[:M]) + torch.testing.assert_close( + runner.post_a[:M].view_as(pm), pm, rtol=1e-4, atol=1e-5 + ) + torch.testing.assert_close( + runner.comb_a[:M].view_as(cm), cm, rtol=1e-4, atol=1e-5 + ) + normed = runner.normed[:M] + assert _cos(normed, xa) > 0.999 and _rel(normed, xa) < 0.02, tag + # the step's KV records, as each path wrote them: each against the + # exact KV rows of its own input + kv_m = _read_records(swa_m, SWA_BLOCK, s["slot_mapping"]) + kv_v = _read_records(swa_v, SWA_BLOCK, s["slot_mapping"]) + err_v = _rel(kv_v, exact_attn.kv(xa, s["positions"], cos_sin)) + err_m = _rel(kv_m, exact_attn.kv(normed, s["positions"], cos_sin)) + assert err_m < 1.2 * err_v + 5e-3, (tag, "kv", err_m, err_v) + assert _cos(kv_m, kv_v) > 0.998, (tag, _cos(kv_m, kv_v)) + # the attention output (wo_b's partials, reduced): both against the + # exact attention of their own input over their own cache + a_m = _attn_output(runner, epoch, M, rank) + ref_v = tensor_model_parallel_all_reduce( + exact_attn(xa, s, ratio, swa_v, cos_sin) + ) + ref_m = tensor_model_parallel_all_reduce( + exact_attn(normed, s, ratio, swa_m, cos_sin) + ) + err_v, err_m = _rel(a, ref_v), _rel(a_m, ref_m) + assert err_m < 1.2 * err_v + 5e-3, (tag, err_m, err_v) + assert _cos(a_m, a) > 0.995, (tag, _cos(a_m, a)) + # the FFN half on the kernels' own attention output + res2_m, _, _, _ = seam("ffn", a_m.to(bf), runner.res_mid[:M], + runner.post_a[:M].view(M, HC, 1), + runner.comb_a[:M]) # fmt: skip + torch.testing.assert_close(o[1], res2_m, rtol=2**-7, atol=1e-4) + pm, cm = mix_ref("ffn", o[1]) + torch.testing.assert_close(o[2].view_as(pm), pm, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(o[3].view_as(cm), cm, rtol=1e-4, atol=1e-5) + sm = seams["ffn"] + with set_current_vllm_config(vc): + _, _, xf_m = pre( + o[1], sm["fn"], sm["scale"], sm["base"], eps, + cfg.hc_eps, cfg.hc_eps, 2.0, cfg.hc_sinkhorn_iters, + norm_weight=sm["norm"], norm_eps=eps, + ) # fmt: skip + # the MoE on K2's own FFN input (where routing cannot differ) + off = scratch_layout(M, tp, width)["normed"][0] + xk = runner.scratch(M, width)[off : off + M * HIDDEN * 2] + xk = xk.view(bf).view(M, HIDDEN) + cos, rel = _cos(xk, xf_m), _rel(xk, xf_m) + assert cos > 0.999 and rel < 0.02, (tag, "ffn input", rank, cos, rel) + mo = moe_out(moe, xk, ids) + ref = tensor_model_parallel_all_reduce(exact_moe(xk, route)) + err_v, err_m = _rel(mo, ref), _rel(o[0], ref) + assert err_m < 1.1 * err_v + 2e-3, (tag, err_m, err_v) + assert _cos(o[0], mo) > 0.999, (tag, _cos(o[0], mo)) + + outs = runner._outs(M, res0) + meta = dict( + topk_indices=s["topk"], comp_cache=s["comp_cache"], + select=lambda n, q: None, comp_block_table=s["comp_block_table"], + route_ids=route, topk=width, global_keys=hca, + ) # fmt: skip + args = (s["positions"], s["slot_mapping"], swa_m, s["swa_indices"], + s["swa_lens"], s["token_to_req"]) # fmt: skip + + def step(mw=mw, x0=x0, res0=res0, post0=post0, comb0=comb0, + args=args, meta=meta, outs=outs): # fmt: skip + runner.forward(mw, x0, res0, post0, comb0, *args, **meta, outs=outs) + return list(outs) + + steps.append(step) + + # the hand-off protocol under back-to-back launches: the layers' widths and + # key tables alternating in one graph, every replay bit for bit the first's + outs = [t for step in steps for t in step()] + torch.accelerator.synchronize() + ref_outs = [t.clone() for t in outs] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for step in steps: + step() + for i in range(REPLAYS): + graph.replay() + if i % 50 == 49 or i == REPLAYS - 1: + torch.accelerator.synchronize() + bad = [ + j + for j, (t, r) in enumerate(zip(outs, ref_outs)) + if not torch.equal(t, r) + ] + assert not bad, f"replay {i}: outputs {bad} differ from the first run" + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch): + if torch.accelerator.device_count() < TP: + pytest.skip(f"needs {TP} GPUs") + pytest.importorskip("flydsl") + multi_process_parallel(monkeypatch, TP, 1, _mono_layer_numerics) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 8c4432d6942a..b50bc32e9a31 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -299,8 +299,8 @@ class KernelConfig: monokernel: bool = False """If True, run the eligible layers of decode steps on the model's MonoKernel: a persistent kernel that runs most or all of one decoder layer - in a single launch (DeepSeek-V4.1 on MI355X). A deployment that cannot run - it is refused at startup.""" + in a single launch (DeepSeek-V4.1 on MI355X, DeepSeek-V4-Pro at TP8 on + MI355X). A deployment that cannot run it is refused at startup.""" moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 47dab93ba2fb..367c8f2cec2c 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -80,6 +80,7 @@ DeepseekV4AiterMegaMoEExperts, finalize_mega_moe_weights, ) +from vllm.models.deepseek_v4.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950, on_gfx1250 @@ -946,6 +947,8 @@ def __init__( self.fuse_mhc_rmsnorm = self.hidden_size in _AITER_MHC_FUSED_RMSNORM_SIZES else: self.fuse_mhc_rmsnorm = HAS_TILELANG_MHC and self.use_fused_mhc + # Decode steps of an eligible layer may run as the mono decode layer. + self.mono = MonoDecodeLayer.create(self, vllm_config) def hc_pre( self, @@ -1092,6 +1095,10 @@ def forward( ) -> tuple[ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None ]: + if self.mono is not None: + out = self.mono(self, x, positions, input_ids, post_mix, res_mix, residual) + if out is not None: + return out if not self.use_fused_mhc: return self._forward_unfused_post_pre( x, positions, input_ids, post_mix, res_mix, residual diff --git a/vllm/models/deepseek_v4/amd/mono/README.md b/vllm/models/deepseek_v4/amd/mono/README.md new file mode 100644 index 000000000000..9cec8a459f50 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/README.md @@ -0,0 +1,112 @@ +# DeepSeek-V4-Pro mono decode layer (gfx950, FlyDSL) + +One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, +the attention, the attention all-reduce, the FFN seam, the MoE and the MoE +all-reduce, in two persistent launches a layer. This implementation builds on +the DeepSeek-V4.1-Flash mono decode layer (`models/deepseek_v41/amd/mono`). +[ROCm/ATOM's V4.1 mono decode](https://github.com/ROCm/ATOM/tree/97359d6df46c59736468168b7be8bc50b482a1a9/atom/models/deepseek_v41/mono) +is the reference design: its persistent execution framework, synchronization, +peer memory, mHC seam and MoE stages underpin this implementation. The adapted +files name their ATOM origins and retain attribution to Advanced Micro +Devices, Inc., under the MIT License. Here they are specialized to Pro's +widths and numerics; the attention stages (`attention/`) read vLLM's KV records +and weights as loaded. + +## Scope + +- Layers: the whole layer on all 61 backbone layers (CSA, compress ratio 4; + HCA, ratio 128; layers 0-2 hash-routed). K1 takes the residual settled, + without folding the previous output in, on layer 0 (the embedding broadcast + to the four streams). Between K1 and K2 every layer hands vLLM its normed + input and bf16 qr (`select`): vLLM's compressor runs there, and on a CSA + layer its indexer; K2 then builds the key table, the CSA top-k through the + block table or the HCA layer's every compressed row from vLLM's c128a table + (`global_keys`), as wide as the longest context (8192 at 1M): its rounds, + split chunks and combine cover a token's own keys only (`bounded`). The MTP / + DSpark draft layers keep vLLM's path; the paths interleave freely (same + tensors at the layer boundary). +- Steps: decode only (FULL graph or eager), M <= 64 rows (8 requests x 7 + DSpark tokens and spare), TP8, `fp8_ds_mla` KV records (ROCm). +- Contract: a drop-in for `DeepseekV4DecoderLayer.forward`: + in `x` (the previous FFN output, reduced, bf16 [M, 7168]), `residual` + (bf16 [M, 4, 7168]), `post_mix` [M, 4, 1], `res_mix` [M, 4, 4] (f32) + out the MoE output (reduced), the residual after the FFN seam, and that + seam's post / comb / pre mixes. + +## Launches + + K1 seam (stages/seam.py): slice (224 x 32 columns): fold x into the + residual (post x + sum comb r), partial mixes (fn . R') and sums of + squares -> gate (S): 24 mixes -> pre / post / comb (Sinkhorn 20) -> + norm (S): the collapse (pre . R'), attn_norm -> vLLM MXFP8 -> X8 + attention front (attention/front.py): wqkv, q / kv norms + KV insert, + wq_b + RoPE -> Q; the normed input and qr out for vLLM's compressor + and indexer + K2 key table (attention/front.py, stage_kt), from the top-k the indexer + selected after K1 + attention back (attention/back.py): split attention (each head's q + RMSNorm scaling its scores), combine + inverse RoPE, wo_a, wo_b; the + bf16 wo_b partial pushed to every rank, or to its columns' owner + FFN seam: reduce (the ranks' partials summed in rank order, fp32 -> + bf16: the attention output), slice, gate, norm -> the MoE input + MoE (stages/moe.py): xq, router (bf16 GEMV, f32 logits), route + (sqrtsoftplus + bias, top-6, renormalized x 2.5; a hash layer's experts + given), shared expert (MXFP8 GEMVs, clamped SiLU), ug (MXFP4 x MXFP8 per + expert of the step's union), down (+ routing weight, top-k order; + + shared) -> push, rank-order sum -> out + +## Execution model and hand-offs + +- Grid = 256 CTAs (one a CU, all resident), 512 threads a CTA, LDS a per-stage + union. A stage is a task table: CTA `b` runs tasks `b, b + 256, ...`, stages + in order; a consumer only waits on earlier stages, so no deadlock while the + grid is resident. A GEMV task issues its weight loads before it polls its + activations. +- In a launch: tagged pairs (value + tag in one 8 B store) or plain data + + flags, at device scope. The tag is the launch pair's epoch (K1 `2e + 1`, + K2 `2e + 2`), a device word K2's CTA 0 moves on once every CTA has marked + the launch done: no clearing between steps, no host step hook, every + argument a fixed device pointer (graph-safe). A step width and key-table + width have a scratch of their own: under another layout a pair's tag word + could be plain data. +- Across ranks: one symmetric uncached buffer a rank (aiter `UncachedIpcHeap`), + every rank holding every peer's address; pushes are system-scope stores of + (value, tag) pairs. Every rank runs the same launches, so the epochs agree; + peer regions alternate by epoch parity, and a rank cannot run two steps ahead + of a peer (each step's reduce waits on every rank's push). + +## Numerics + +- Seams: the gate is aiter's `mhc_fused_post_pre_delayed_rmsnorm` reduce (post + fold order, fn as bf16 hi + lo, products summed in order, Triton sigmoid / + exp / rcp), the collapse aiter's `mhc_pre` (the four streams in fp32 -> + bf16, rsq(sum / H + eps), (x w) rstd). The gate matches the fp32 reference; + aiter's fused seam is ~2e-3 off it. +- Attention: activations as MXFP8 (one E8M0 per 32 values, `code = + clamp(ceil(log2(amax / 448)) + 127, 0, 254)`), where vLLM's path quantizes + per 128 with fp32 scales: the kernels match the exact result no worse. + Weights as vLLM loads them: wqkv, wo_b and the shared gate_up in aiter's + (16, 16) shuffle, the others e4m3 [N, K] row-major with rows padded as + `VLLM_ROCM_FP8_PADDING` pads them; the checkpoint's 128 x 128 E8M0 block + scales repeated as [N/32, K/32] (exact). KV records as vLLM's insert writes + them (dims 0..447 e4m3 with a UE8M0 per 64, RoPE dims bf16, 584 B); softmax + in base 2 with the attention sink folded in. +- MoE: vLLM's router (bf16 gate GEMV, f32 logits, `topk_hash_softplus_sqrt`) + and AITER a8w4 `fused_moe` (interleaved gate / up, swiglu limit 10, weights + in aiter's (16, 16) shuffle and `shuffle_scale` order); the shared expert as + MXFP8 linears with `silu_and_mul_with_clamp`. +- All-reduces (`stages/allreduce.py`): exact, fp32 in rank order -> bf16. One + hop (every rank pushes its partial to every rank) or, from 42 rows, two (to + the column's owner, which sums and pushes the bf16 sum to the others: 2 / tp + of the bytes); bit-identical results. + +## Validation + +- Per layer against vLLM's own modules in a TP8 process group + (`tests/models/test_dsv4_mono_numerics.py`), the protocol under graph + replays; through the real model on random weights + (`tests/models/test_dsv4_mono_layers.py`); the docking's path choice + (`tests/models/test_dsv4_mono_decode.py`). +- Bench: layers of one type back to back in one HIP graph, cold weights, us + per layer; then e2e (GSM8K, GPQA-Diamond, AIME25, RULER NIAH, AgentX as + InferenceX runs it). diff --git a/vllm/models/deepseek_v4/amd/mono/__init__.py b/vllm/models/deepseek_v4/amd/mono/__init__.py new file mode 100644 index 000000000000..863cdad925de --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4-Pro mono decode layer: one decoder layer of a decode step in two +persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The host +side is ``runner``; vLLM docks it in ``..mono_decode``.""" diff --git a/vllm/models/deepseek_v4/amd/mono/attention/__init__.py b/vllm/models/deepseek_v4/amd/mono/attention/__init__.py new file mode 100644 index 000000000000..da1bf8c269a0 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's attention stages: ``front`` (K1: wqkv, q / kv norms and the +KV insert, wq_b) and ``back`` (K2: split attention, combine, wo_a, wo_b).""" diff --git a/vllm/models/deepseek_v4/amd/mono/attention/back.py b/vllm/models/deepseek_v4/amd/mono/attention/back.py new file mode 100644 index 000000000000..64a3c890c30c --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/back.py @@ -0,0 +1,1021 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2 ``back``: one attention layer from the rotated query to the rank's wo_b +partial. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S rows: + + split (S x chunks of CK keys): flash-decoding over one chunk of a token's + keys (K1's key table: top-k compressed rows, then window rows) for + every head of the rank. The token's q in LDS, each head's RMSNorm + factor (``q_head_rstd``) from it; each wave one + (head tile, 16-key tile): its fp8_ds_mla rows loaded once, + dequantized into its QK operands and kept raw in LDS; QK (bf16 + MFMA), the chunk's max / sum per head, bf16 p in LDS; PV over the + chunk's raw rows (wave w: dims 64 w ..) -> the chunk's partial + (bf16) and its (max, sum) per head, published + A wide key table (``bounded``) skips the chunks past a token's + keys, known from its position. + combine (S x H): a (token, head)'s chunks merged with the sink folded in + (``bounded``: its live chunks only, ``merge_live``), + bf16, inverse GPT-J RoPE of dims 448.., vLLM's MXFP8 of each 32 + -> XO (published) + wo_a (G x 1024 / 32, two 16-row tiles): the group's XO rows times wo_a + -> bf16 -> MXFP8 of the 32 rows -> X8B (published) + wo_b (448, 16 rows): X8B times wo_b -> bf16 -> ``out`` + +Hand-offs are plain device-scope data behind tagged flags (half the bytes of a +tagged pair per value). Adapted from ATOM's V4.1 ``attn_post`` (K2a), which +reads a bf16 pool and ends in a peer all-reduce. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T, as_ir_value + +from ..dims import ( + HALF, + HEAD_DIM, + HEAD_TILE, + NOPE, + NORM_EPS, + ROPE, + ROWS, + WINDOW, + WOA_ROWS, + WOB_TASKS, + Dims, + ld, + tile_rows, + token_tiles, +) +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + WAVES, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + clamp_fp8, + first_task, + fp8_pack4, + fp8x8_bf16, + fresh, + gload, + gstore, + hw_exp2, + hw_rsq, + lane_gather, + ld_f32, + ld_i32, + mfma_bf16, + mx_code, + mx_mul, + pow2, + rsrc, + spin_until, + stamp, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import DATA, back_scratch, cdiv, front_end + +_STREAM = fx.Stream(None) +NEG = -3.4028234663852886e38 +LOG2E = 1.4426950408889634 +REC = DATA + 8 # a raw record row in LDS: 576 data bytes + 8 scale bytes +Q_ROW = ( + HEAD_DIM // 2 + 4 +) # a q row in LDS (i32 words): 512 bf16 + 16 B, no bank aliasing +P_PAD = 2 # i32 words past a P row: 16 head rows read together hit distinct banks +# wo_a's K order: waves 0..3 / 4..7 take the two 16-row tiles, a contiguous K +# range a wave, summed in wave order +WOA_SPLIT_K = WAVES * ROWS // WOA_ROWS +BACK_POINTS = 16 # start, stage ends, split internals, GEMV internals + + +@dataclass(frozen=True) +class BackBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio: 4 (CSA) or 128 (HCA) + timeline: bool = False + topk: int = 0 # a token's compressed keys at most (``Dims.topk``) + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +def head_groups(d: Dims, s: int) -> int: + """Wave groups a split unit's head tiles are spread over: two (a tile each, + 64-key chunks, twice the units) while the units fit the grid, else one (a + wave's keys loaded and dequantized once for every head tile).""" + return 2 if d.head_tiles == 2 and s * cdiv(d.keys, 64) <= BLOCKS else 1 + + +def chunk_keys(d: Dims, s: int) -> int: + """A split unit's keys: a 16-key tile a wave of a head group.""" + return 16 * (WAVES // head_groups(d, s)) + + +def p_row(d: Dims, s: int) -> int: + """A P row (one head's chunk of bf16 p) in LDS, i32 words.""" + return chunk_keys(d, s) // 2 + P_PAD + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- keys + + +def key_rec(c, t, k): + """(data address, scale address, valid) of token t's key k from K1's key + table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for + none (it reads a zero record: its scores are masked, its values zeros).""" + a, ratio, keys = c["args"], c["ratio"], c["d"].keys + # a table this launch wrote (``krdy``) is read past the CU's cache + cm = CM_DEV if const_expr("krdy" in c) else 0 + val = ld_i32(a["kt"], t * keys + fx.min(fx.max(k, 0), keys - 1), cm) + ok = val != fx.Int32(-1) + slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) + if const_expr(ratio > 0): + is_c = ok & (val < 0) + blk = is_c.select(a["comp_block"], a["swa_block"]) + stride = is_c.select(a["comp_stride"], a["swa_stride"]) + base = is_c.select(a["comp"], a["swa"]) + else: + blk, stride, base = a["swa_block"], a["swa_stride"], a["swa"] + b = slot // blk + off = slot % blk + page = base + fx.Int64(b) * stride + z = a["zrec"] + data = ok.select(page + fx.Int64(off * DATA), z) + sc = ok.select(page + fx.Int64(blk * DATA + off * 8), z + fx.Int64(DATA)) + return data, sc, ok + + +def _scale_byte(scw, blk64): + return (scw[blk64 // 4] >> (blk64 % 4 * 8)) & 0xFF + + +def _fp8_pair_f32(word, scale, hi): + v = rocdl.cvt_scalef32_pk_f32_fp8( + fx.Vector.make_type(2, fx.Float32), + as_ir_value(fx.Int32(word)), + as_ir_value(fx.Float32(scale)), + hi, + ) + v = fx.Vector(v) + return [v[0], v[1]] + + +# ---------------------------------------------------------------- split + + +@traced +def stage_split(c, unit): + """Flash-decoding over chunk ``unit % nchunk`` of token ``unit // nchunk`` + (``split_chunk``). A build whose key table is wider than a token's top-k + (``bounded``: every compressed row of a long context) skips a chunk past + the token's keys at most, known from its position alone: no load of its q + or keys.""" + nchunk = c["nchunk"] + t = unit // nchunk + ch = unit % nchunk + if const_expr(c.get("bounded", False)): + d = c["d"] + pos = ld_i32(c["args"]["pos"], 2 * t) + cap = fx.min((pos + 1) // c["ratio"], fx.Int32(d.topk_width)) + WINDOW + if ch * chunk_keys(d, c["S"]) < cap: + split_chunk(c, t, ch) + else: + split_chunk(c, t, ch) + + +@traced +def q_head_rstd(c): + """Each head's weightless RMSNorm of the token's q in LDS: + rsq(mean(q^2) + eps) -> ``qn`` [head]. vLLM norms q before its RoPE, which + keeps every pair's norm: the scores scale by it after the dot product.""" + tid, d, sl = c["tid"], c["d"], c["sl"] + tph = THREADS // d.heads # threads a head + words = HEAD_DIM // 2 // tph + h = tid // tph + i = tid % tph + acc = fx.Float32(0.0) + for w in range_constexpr(words): + v = fx.Int32(fx.ptr_load(sl["q"] + (h * Q_ROW + i * words + w))) + lo, hi = bf_lo(v), bf_hi(v) + acc = acc + lo * lo + hi * hi + offs = tuple(o for o in (32, 16, 8, 4, 2, 1) if o < tph) + acc = butterfly(acc, offs) + if i == 0: + mean = acc * fx.Float32(1.0 / HEAD_DIM) + fx.ptr_store(hw_rsq(mean + fx.Float32(NORM_EPS)), sl["qn"] + h) + gpu.barrier() + + +@traced +def split_chunk(c, t, ch): + """Flash-decoding over chunk ``ch`` of token ``t``, every head of the rank + -> the chunk's partial and (max, sum) per head. Wave (head group hg, key + tile kt): keys 16 kt .. of the chunk (loaded and dequantized once) against + the group's head tiles; then dims 64 w .. of the PV over every key and + head.""" + tid, lane, wave, d, s = c["tid"], c["lane"], c["wave"], c["d"], c["S"] + a, sl = c["args"], c["sl"] + nchunk = c["nchunk"] + ck = chunk_keys(d, s) + ngrp = head_groups(d, s) + nkt = WAVES // ngrp # 16-key tiles a chunk + nht = d.head_tiles // ngrp # head tiles a wave + hgi = wave // nkt + kt = wave % nkt + g = lane // 16 + r16 = lane % 16 + k0 = ch * ck + if const_expr("krdy" in c): + await_flags(c, c["krdy"], 1, lambda u: t) + # the key count, this lane's key and q: independent loads, out together + # (the key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t, CM_DEV if const_expr("krdy" in c) else 0) + key = k0 + kt * 16 + r16 + data, sc, okv = key_rec(c, t, key) + qrow = d.heads * HEAD_DIM // 8 + qv = [] + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + qv.append( + fx.Vector( + bo.buffer_load( + rsrc(a["q"]), + (t * d.q_rows) // 2 + fx.min(u, qrow - 1) * 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + if k0 < kv_len: + # a key's 448 e4m3 as 7 x 16 B (lane g: dims 64 c + 16 g ..), its 64 + # bf16 rope dims as 2 x 16 B (32 c + 8 g ..), its 7 scale bytes + scw = gload(sc, words=2) + nope = [ + gload(data + fx.Int64(64 * c4 + 16 * g), words=4) + for c4 in range(NOPE // 64) + ] + rope = [ + gload(data + fx.Int64(NOPE + (cc * 32 + 8 * g) * 2), words=4) + for cc in range(ROPE // 32) + ] + # q of every head into LDS (bf16 rows of 512 + pad), 16 B a thread + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + if u < qrow: + fx.ptr_store( + qv[i], + sl["q"] + + (u // (HEAD_DIM // 8)) * Q_ROW + + (u % (HEAD_DIM // 8)) * 4, + ) + # the raw rows to LDS for the PV (head group 0's waves) + if hgi == 0: + rawb = sl["raw"] + (kt * 16 + r16) * (REC // 4) + for c4 in range_constexpr(NOPE // 64): + fx.ptr_store(nope[c4], rawb + (64 * c4 + 16 * g) // 4) + for cc in range_constexpr(ROPE // 32): + fx.ptr_store(rope[cc], rawb + (NOPE + (cc * 32 + 8 * g) * 2) // 4) + if g == 0: + fx.ptr_store(scw, rawb + DATA // 4) + gpu.barrier() + q_head_rstd(c) + stamp(c, 5) + # QK: A = the keys (dequantized once), B = q of a head tile (LDS); the + # contraction's 32-dim chunk 2 c + h of the nope dims is {64 c + 16 g + + # 8 h ..} (lane g's 16 B load), q read in the same order + scale = fx.Int32(a["qk_scale"]).bitcast(fx.Float32) + st = sl["st"] + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(nht)] + for c4 in range_constexpr(NOPE // 64): + ksc = pow2(_scale_byte(scw, c4)) + for h2 in range_constexpr(2): + kw = fp8x8_bf16(nope[c4][2 * h2], nope[c4][2 * h2 + 1], ksc).bitcast( + fx.BFloat16 + ) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + 32 * c4 + 8 * g + 4 * h2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + for cc in range_constexpr(ROPE // 32): + kw = rope[cc].bitcast(fx.BFloat16) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + (NOPE + cc * 32 + 8 * g) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + # scores of keys 4 g + i (C rows), head r16 of a tile (C column) + ok = (key < kv_len) & okv + okw = ok.select(fx.Int32(1), fx.Int32(0)) + valid = [lane_gather(okw, 4 * g + i) != 0 for i in range(4)] + svs = [] + for ht in range_constexpr(nht): + head = (hgi * nht + ht) * HEAD_TILE + r16 + hscale = scale * fx.ptr_load(sl["qn"] + head) + sv = [ + fx.Float32(valid[i].select(accs[ht][i] * hscale, fx.Float32(NEG))) + for i in range(4) + ] + mw = butterfly( + fx.max(fx.max(sv[0], sv[1]), fx.max(sv[2], sv[3])), (32, 16), fx.max + ) + if g == 0: + fx.ptr_store(mw, st + ((hgi * nht + ht) * nkt + kt) * HEAD_TILE + r16) + svs.append(sv) + gpu.barrier() + lst = st + d.heads * nkt + for ht in range_constexpr(nht): + hti = hgi * nht + ht + m = fx.ptr_load(st + (hti * nkt) * HEAD_TILE + r16) + for j in range_constexpr(1, nkt): + m = fx.max(m, fx.ptr_load(st + (hti * nkt + j) * HEAD_TILE + r16)) + sv = svs[ht] + p = [ + fx.Float32(valid[i].select(hw_exp2(sv[i] - m), fx.Float32(0.0))) + for i in range(4) + ] + lw = (p[0] + p[2]) + (p[1] + p[3]) + lw = lw + xshfl(lw, 32) + lw = lw + xshfl(lw, 16) + pw = ( + fx.Vector.from_elements(p, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) + ) + pbase = ( + sl["p"] + (hti * HEAD_TILE + r16) * p_row(d, s) + (kt * 16 + 4 * g) // 2 + ) + fx.ptr_store(fx.Vector.from_elements([pw[0], pw[1]], fx.Int32), pbase) + if g == 0: + fx.ptr_store(lw, lst + (hti * nkt + kt) * HEAD_TILE + r16) + gpu.barrier() + # chunk stats: (max, sum) per head + if tid < d.heads: + hts = tid // HEAD_TILE + hr = tid % HEAD_TILE + mm = fx.ptr_load(st + (hts * nkt) * HEAD_TILE + hr) + ll = fx.ptr_load(lst + (hts * nkt) * HEAD_TILE + hr) + for j in range_constexpr(1, nkt): + mm = fx.max(mm, fx.ptr_load(st + (hts * nkt + j) * HEAD_TILE + hr)) + ll = ll + fx.ptr_load(lst + (hts * nkt + j) * HEAD_TILE + hr) + bo.buffer_store( + fx.Vector.from_elements([mm, ll], fx.Float32), + rsrc(c["pstat"]), + ((t * nchunk + ch) * d.heads + tid) * 2, + cache_modifier=CM_DEV, + ) + # PV: wave w dims 64 w + 4 r16 + dt (dt = 0..3: one LDS word holds the + # four of a key); the rope wave (dims 448..) reads bf16 pairs. Lane + # (r16, g) takes keys 4 g + j of a 16-key tile: A[dim][key] + is_rope = wave == WAVES - 1 + dbase = 64 * wave + sh = (wave % 4) * 8 + accs = [ + [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4)] + for _ in range(d.head_tiles) + ] + for kt2 in range_constexpr(nkt): + rb0 = sl["raw"] + (kt2 * 16 + 4 * g) * (REC // 4) + vv = fx.Vector.filled(16, 0.0, fx.Float32) # [dt][j] + if is_rope: + rv = [] + for j in range_constexpr(4): + rw = fx.Vector( + fx.ptr_load( + rb0 + j * (REC // 4) + (NOPE + 8 * r16) // 4, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ) + rv.append([bf_lo(rw[0]), bf_hi(rw[0]), bf_lo(rw[1]), bf_hi(rw[1])]) + vv = fx.Vector.from_elements( + [rv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + else: + nv = [] + for j in range_constexpr(4): + rb = rb0 + j * (REC // 4) + w = fx.Int32(fx.ptr_load(rb + (dbase + 4 * r16) // 4)) + scb = fx.Int32(fx.ptr_load(rb + DATA // 4 + wave // 4)) + scl = pow2((scb >> sh) & 0xFF) + nv.append( + _fp8_pair_f32(w, scl, False) + _fp8_pair_f32(w, scl, True) + ) + vv = fx.Vector.from_elements( + [nv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + for ht2 in range_constexpr(d.head_tiles): + pb = fx.Vector( + fx.ptr_load( + sl["p"] + + (ht2 * HEAD_TILE + r16) * p_row(d, s) + + (kt2 * 16 + 4 * g) // 2, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ).bitcast(fx.Int16) + for dt in range_constexpr(4): + av = ( + fx.Vector.from_elements( + [vv[dt * 4 + j] for j in range(4)], fx.Float32 + ) + .to(fx.BFloat16) + .bitcast(fx.Int16) + ) + accs[ht2][dt] = fx.Vector( + rocdl.mfma_f32_16x16x16bf16_1k( + T.vec(4, T.f32), [av, pb, accs[ht2][dt]] + ) + ) + stamp(c, 6) + # partial[t][ch][head][dim] bf16: lane (r16, g) holds C rows 4 g + i -> + # dims 64 w + 16 g + 4 i + dt of head ht2 * 16 + r16: 32 contiguous + # bytes, two 16 B stores (4 lanes fill a 128 B line) + for ht2 in range_constexpr(d.head_tiles): + hh = ht2 * HEAD_TILE + r16 + words = [] + for i in range_constexpr(4): + v4 = fx.Vector.from_elements( + [accs[ht2][dt][i] for dt in range(4)], fx.Float32 + ) + w2 = v4.to(fx.BFloat16).bitcast(fx.Int32) + words += [w2[0], w2[1]] + off = (((t * nchunk + ch) * d.heads + hh) * HEAD_DIM + dbase + 16 * g) // 2 + for hlf in range_constexpr(2): + bo.buffer_store( + fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32), + rsrc(c["part"]), + off + 4 * hlf, + cache_modifier=CM_DEV, + ) + publish(c, c["srdy"], t * nchunk + ch) + stamp(c, 7) + + +# ---------------------------------------------------------------- combine + + +def _chunk_loads(c, t, h, d0, cc): + """Chunk cc's (max, sum) of head h and its partial's 8 dims from d0.""" + d, nchunk = c["d"], c["nchunk"] + stv = fx.Vector( + bo.buffer_load( + rsrc(c["pstat"]), + ((t * nchunk + cc) * d.heads + h) * 2, + vec_width=2, + dtype=T.f32, + ) + ) + pv = fx.Vector( + bo.buffer_load( + rsrc(c["part"]), + (((t * nchunk + cc) * d.heads + h) * HEAD_DIM + d0) // 2, + vec_width=4, + dtype=T.i32, + ) + ) + vals = [] + for e in range_constexpr(4): + vals += [bf_lo(pv[e]), bf_hi(pv[e])] + return stv, vals + + +def merge_all(c, t, h, d0, nlive, sink): + """The chunks' (sum, 8 dims) of (t, h) weighted by exp2(max - the max), + the sink's term first, in chunk order: every chunk loaded at once (a + past-``nlive`` chunk weighs nothing).""" + m = fx.Float32(NEG) + ms, ls, vs = [], [], [] + for cc in range_constexpr(c["nchunk"]): + live = cc < nlive + stv, vals = _chunk_loads(c, t, h, d0, cc) + mc = fx.Float32(live.select(stv[0], fx.Float32(NEG))) + ms.append(mc) + ls.append(fx.Float32(live.select(stv[1], fx.Float32(0.0)))) + vs.append([fx.Float32(live.select(v, fx.Float32(0.0))) for v in vals]) + m = fx.max(m, mc) + m_final = fx.max(m, sink) + lsum = hw_exp2(sink - m_final) + tot = [fx.Float32(0.0)] * 8 + for cc in range_constexpr(c["nchunk"]): + wgt = hw_exp2(ms[cc] - m_final) + lsum = lsum + ls[cc] * wgt + tot = [tot[e] + vs[cc][e] * wgt for e in range(8)] + return lsum, tot + + +MERGE_BATCH = 4 # a wide table's chunks loaded together (``merge_live``) + + +@traced +def merge_live(c, t, h, d0, nlive, sink): + """``merge_all`` over the live chunks only, ``MERGE_BATCH`` at a time (a + wide key table: most of its chunks past a token's keys).""" + m = fx.Float32(NEG) + for c0 in range(0, nlive, MERGE_BATCH): + for u in range_constexpr(MERGE_BATCH): + live = c0 + u < nlive + stv, _ = _chunk_loads(c, t, h, d0, fx.min(c0 + u, nlive - 1)) + m = fx.max(m, fx.Float32(live.select(stv[0], fx.Float32(NEG)))) + m_final = fx.max(m, sink) + lsum = hw_exp2(sink - m_final) + t0, t1, t2, t3, t4, t5, t6, t7 = [fx.Float32(0.0)] * 8 + for c0 in range(0, nlive, MERGE_BATCH): + got = [ + _chunk_loads(c, t, h, d0, fx.min(c0 + u, nlive - 1)) + for u in range(MERGE_BATCH) + ] + for u in range_constexpr(MERGE_BATCH): + live = c0 + u < nlive + stv, vals = got[u] + wgt = hw_exp2(fx.Float32(live.select(stv[0], fx.Float32(NEG))) - m_final) + lsum = lsum + fx.Float32(live.select(stv[1], fx.Float32(0.0))) * wgt + v = [fx.Float32(live.select(x, fx.Float32(0.0))) for x in vals] + t0, t1, t2, t3 = ( + t0 + v[0] * wgt, + t1 + v[1] * wgt, + t2 + v[2] * wgt, + t3 + v[3] * wgt, + ) + t4, t5, t6, t7 = ( + t4 + v[4] * wgt, + t5 + v[5] * wgt, + t6 + v[6] * wgt, + t7 + v[7] * wgt, + ) + return lsum, [t0, t1, t2, t3, t4, t5, t6, t7] + + +COMB_HEADS = WAVES # heads a combine unit: a wave each + + +@traced +def stage_combine(c, unit): + """(token, 8 heads) ``unit``, a wave a head, 8 dims a lane: the chunks + merged with the sink folded in, bf16, inverse RoPE, MXFP8 of each 32 -> XO, + published (a flag a (token, head group)).""" + lane, wave, d, s = c["lane"], c["wave"], c["d"], c["S"] + a = c["args"] + hg_n = d.heads // COMB_HEADS + t = unit // hg_n + hg = unit % hg_n + h = hg * COMB_HEADS + wave + nchunk = c["nchunk"] + ck = chunk_keys(d, s) + if const_expr("krdy" in c): + # a key table this launch builds: its count once published, past the + # CU's cache (a CTA without split work gets here first) + await_flags(c, c["krdy"], 1, lambda u: t) + kv_len = ld_i32(a["klen"], t, CM_DEV if const_expr("krdy" in c) else 0) + nlive = (kv_len + ck - 1) // ck + pos = ld_i32(a["pos"], 2 * t) + sink = ld_f32(a["sink"], h) * fx.Float32(LOG2E) + d0 = lane * 8 # this lane's 8 dims + rope_lane = d0 >= NOPE + jj = fx.max((d0 - NOPE) // 2, 0) # 4 pairs from here + cs = [ + rope_lane.select(ld_f32(a["cos_sin"], pos * ROPE + jj + i), fx.Float32(1.0)) + for i in range(4) + ] + sn = [ + rope_lane.select( + ld_f32(a["cos_sin"], pos * ROPE + HALF + jj + i), fx.Float32(0.0) + ) + for i in range(4) + ] + await_flags(c, c["srdy"], nlive, lambda u: t * nchunk + u) + if const_expr(c.get("bounded", False)): + lsum, tot = merge_live(c, t, h, d0, nlive, sink) + else: + lsum, tot = merge_all(c, t, h, d0, nlive, sink) + inv = fx.Float32(1.0) / lsum + o = [bf16_round(tot[e] * inv) for e in range(8)] + # inverse GPT-J RoPE of each pair: e' = e cos + o sin, o' = o cos - e sin + y = [] + for i in range_constexpr(4): + e, od = o[2 * i], o[2 * i + 1] + y += [e * cs[i] + od * sn[i], od * cs[i] - e * sn[i]] + amax = abs(y[0]) + for v in y[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) # 32 dims = 4 lanes + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(y[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(y[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["xo"]), + t * (d.q_rows // 4) + (h * HEAD_DIM + d0) // 4, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["xos"]), + t * (d.q_rows // 32) + (h * HEAD_DIM + d0) // 32, + cache_modifier=CM_DEV, + ) + publish(c, c["crdy"], t * hg_n + hg) + + +# ---------------------------------------------------------------- wo_a / wo_b + + +@traced +def stage_woa(c, task, ops, t0, n): + """Rows 32 task .. of wo_a -> bf16 -> vLLM's MXFP8 of the 32 -> X8B.""" + tid, red, d, lane = c["tid"], c["red"], c["d"], c["lane"] + split = WOA_SPLIT_K + gemv_mfmas(c, d.group_k, c["xl"], c["xsl"], red, ops, split, tiled=True, rows=n) + gpu.barrier() + mine = tid < WOA_ROWS * n + tl = fx.min(tid // WOA_ROWS, n - 1) + t = plus(t0, tl) + row = tid % WOA_ROWS + tile = row // ROWS + y = bf16_round( + row_sum(red, row % ROWS, tl, [tile * split + w for w in range(split)]) + ) + amax = butterfly(abs(y), (1, 2, 4, 8, 16), fx.max) + code = mx_code(amax) + q = clamp_fp8(y * mx_mul(code)) + w = fp8_pack4( + q, + *[ + lane_gather(q.bitcast(fx.Int32), fx.min(lane + k, 63)).bitcast(fx.Float32) + for k in (1, 2, 3) + ], + ) + if mine & (row % 4 == 0): + bo.buffer_store( + w, + rsrc(c["x8b"]), + t * (d.o_rows // 4) + (task * WOA_ROWS + row) // 4, + cache_modifier=CM_DEV, + ) + if mine & (row == 0): + bo.buffer_store( + code, rsrc(c["x8bs"]), t * (d.o_rows // 32) + task, cache_modifier=CM_DEV + ) + gpu.barrier() + + +def woa_loads(c, task): + a = c["args"] + return gemv_loads( + c, + a["woa"], + a["woa_s"], + c["d"].group_k, + task * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + ld=ld(c["d"].group_k), + ) + + +def wob_loads(c, task, pred=None): + a, d = c["args"], c["d"] + return gemv_loads( + c, + a["wob"], + a["wob_s"], + d.o_rows, + task * ROWS, + pred=pred, + shuffled=True, + ) + + +def _wob_store(a, t, col, v0, v1): + word = ( + fx.Vector.from_elements([v0, v1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["out"]), (t * a["out_stride"] + col) // 2) + + +@traced +def stage_wob(c, task, ops, task2, ops2, has2, t0, n, mark=False): + """Rows 16 task .. (and 16 task2 .., ``has2``) of wo_b -> bf16 -> out, or + (``c["wob_out"]``, the mono layer) the caller's hook with each bf16 pair.""" + tid, red = c["tid"], c["red"] + o_rows = c["d"].o_rows + a = c["args"] + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red, ops, rows=n) + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red + RED, ops2, rows=n) + gpu.barrier() + if const_expr(mark): + stamp(c, 15) + per = ROWS * n // 2 # bf16 pairs a task + k = tid // per + emit = c.get("wob_out") or (lambda t, col, v0, v1: _wob_store(a, t, col, v0, v1)) + if (tid < 2 * per) & ((k == 0) | has2): + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + rk = red + k * RED + col = (k == 0).select(task, task2) * ROWS + 2 * rp + v0 = row_sum(rk, 2 * rp, tl) + v1 = row_sum(rk, 2 * rp + 1, tl) + emit(t, col, v0, v1) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def back_plan(key): + """(nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob) of a build.""" + s, d = key.tokens, Dims(key.tp, key.topk) + ck = chunk_keys(d, s) + nchunk = cdiv(d.keys, ck) + n_split = s * nchunk + n_comb = s * (d.heads // COMB_HEADS) + COMB0 = n_split % BLOCKS + assert d.woa_tasks <= BLOCKS + assert BLOCKS < WOB_TASKS <= 2 * BLOCKS + # the GEMVs past the attention units, so CTAs without attention work have + # their weights in flight from the start + WOA0 = (n_split + n_comb) % BLOCKS + WOB0 = (WOA0 + d.woa_tasks) % BLOCKS + # wo_b's weights before the wo_a stage only while the wo_a CTAs have no + # attention work (loads complete in issue order: else they would hold up + # wo_a's activation rows) + early_wob = n_split + n_comb + d.woa_tasks <= BLOCKS + return nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob + + +def back_lds_members(s: int, d: Dims) -> dict: + """K2's attention LDS structs: the split stage's and the GEMVs' (a union's + members: one at a time).""" + rows = tile_rows(s) + ck = chunk_keys(d, s) + + @fx.struct + class SplitLds: + q: fx.Array[fx.Int32, d.heads * Q_ROW, 16] + raw: fx.Array[fx.Int32, ck * REC // 4, 16] + p: fx.Array[fx.Int32, d.heads * p_row(d, s), 16] + st: fx.Array[fx.Float32, 2 * d.heads * (WAVES // head_groups(d, s)), 16] + # each head's q norm (``q_head_rstd``) + qn: fx.Array[fx.Float32, d.heads, 16] + + @fx.struct + class GemvLds: + red: fx.Array[fx.Float32, 2 * RED, 16] + xl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 4), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 32), 16] + + return {"split": SplitLds, "gemv": GemvLds} + + +def back_lds(s: int, d: Dims): + """K2's attention LDS: the union of ``back_lds_members``.""" + m = back_lds_members(s, d) + SplitLds, GemvLds = m["split"], m["gemv"] + + @fx.union + class BackLds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + + return BackLds + + +def back_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The attention back's context: LDS views (``lds``: an allocated + ``back_lds``, or a dict of the members' views), mailbox (tag ``tag``), + scratch regions, arguments by name.""" + s, d = key.tokens, Dims(key.tp, key.topk) + tid = fx.thread_idx.x + if isinstance(lds, dict): + spl, gml = lds["split"], lds["gemv"] + else: + spl, gml = lds.split.peek(), lds.gemv.peek() + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": gml.red.ptr, + "xl": gml.xl.ptr, + "xsl": gml.xsl.ptr, + "d": d, + "ratio": key.ratio, + "mb": Mailbox(tag), + "nchunk": back_plan(key)[0], + "sl": { + "q": spl.q.ptr, + "raw": spl.raw.ptr, + "p": spl.p.ptr, + "st": spl.st.ptr, + "qn": spl.qn.ptr, + }, + "args": args, + } + for region, (off, _) in ( + layout or back_scratch(s, d, start=front_end(s, d)) + ).items(): + c[region] = scratch + fx.Int64(off) + return c + + +EPOCH_MARKS = 4 # word offset of the per-CTA marks in the epoch buffer + + +@traced +def epoch_begin(c, epoch, ep): + """This CTA has read the launch pair's epoch: its own mark word (no + contended atomic) := ep + 1, which ``epoch_end`` checks.""" + if c["tid"] == 0: + bo.buffer_store( + ep + 1, rsrc(epoch), EPOCH_MARKS + c["bid"], cache_modifier=CM_DEV + ) + + +@traced +def epoch_end(c, epoch, ep, reset=None): + """CTA 0, once every CTA marked this epoch (thread u checks CTA u's mark): + ``reset()`` (state later pairs must find cleared), then the epoch moved on. + The next pair's hand-offs carry a tag no earlier pair wrote; later launches' + kernels start after this one ends.""" + if c["bid"] == 0: + tid = c["tid"] + if tid < BLOCKS: + r = rsrc(epoch) + + def load_mark(): + return fx.Int32( + bo.buffer_load( + r, + EPOCH_MARKS + tid, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + + _ = spin_until(load_mark, lambda v: v != ep + 1) + gpu.barrier() + if tid == 0: + if const_expr(reset is not None): + reset() + gstore(epoch, ep + 1, words=1) + + +@traced +def run_back(c, key, bid): + """The attention back's stages on this CTA: split, combine, wo_a, wo_b.""" + s = key.tokens + dd = Dims(key.tp, key.topk) + # plain ints: a Python object read inside traced control flow is carried + woa_tasks, groups, heads, group_k = dd.woa_tasks, dd.groups, dd.heads, dd.group_k + q_rows, o_rows, wob_tasks = dd.q_rows, dd.o_rows, WOB_TASKS + a = c["args"] + tiles = token_tiles(s) + _, n_split, n_comb, COMB0, WOA0, WOB0, early_wob = back_plan(key) + stamp(c, 0) + for unit in range(first_task(bid, 0), n_split, BLOCKS): + stage_split(c, unit) + stamp(c, 1) + for unit in range(first_task(bid, COMB0), n_comb, BLOCKS): + stage_combine(c, unit) + gpu.barrier() + stamp(c, 2) + # the GEMV weights, every CTA right after its attention work: its wo_a + # task (if any), then its one or two wo_b tasks + woa_task = first_task(bid, WOA0) + has_woa = woa_task < woa_tasks + woa_t = fx.min(woa_task, woa_tasks - 1) + ops = gemv_loads( + c, + a["woa"], + a["woa_s"], + group_k, + woa_t * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + pred=has_woa, + ld=ld(group_k), + ) + wob_task = first_task(bid, WOB0) + wob_second = wob_task + BLOCKS + has_second = wob_second < wob_tasks + + def wob_weights(): + return ( + wob_loads(c, wob_task), + wob_loads(c, fx.min(wob_second, wob_tasks - 1), pred=has_second), + ) + + if const_expr(early_wob): + ops1, ops2 = wob_weights() + if const_expr(c.get("tl") is not None): + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + stamp(c, 8) + # wo_a: every token tile of its group's XO rows, the next tile's in flight + if has_woa: + grp = woa_t // (woa_tasks // groups) + await_flags(c, c["crdy"], s, lambda u: u * (heads // COMB_HEADS) + grp) + stamp(c, 9) + xo = ( + c["xo"], + c["xos"], + q_rows // 4, + grp * group_k // 4, + group_k // 4, + q_rows // 32, + grp * group_k // 32, + group_k // 32, + ) + regs = rows_loads(c, *xo, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, group_k // 4, group_k // 32, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *xo, *tiles[ti + 1]) + stage_woa(c, woa_t, ops, t0, n) + stamp(c, 10) + publish(c, c["ardy"], woa_t) + stamp(c, 3) + # wo_b + if const_expr(not early_wob): + ops1, ops2 = wob_weights() + stamp(c, 11) + await_flags(c, c["ardy"], woa_tasks, lambda u: u) + stamp(c, 12) + x8b = ( + c["x8b"], + c["x8bs"], + o_rows // 4, + 0, + o_rows // 4, + o_rows // 32, + 0, + o_rows // 32, + ) + regs = rows_loads(c, *x8b, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, o_rows // 4, o_rows // 32, c["xl"], c["xsl"], n) + if const_expr(ti == 0): + stamp(c, 14) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *x8b, *tiles[ti + 1]) + stage_wob( + c, + wob_task, + ops1, + fx.min(wob_second, wob_tasks - 1), + ops2, + has_second, + t0, + n, + mark=ti == 0, + ) + if const_expr(ti == 0): + stamp(c, 13) + stamp(c, 4) diff --git a/vllm/models/deepseek_v4/amd/mono/attention/device.py b/vllm/models/deepseek_v4/amd/mono/attention/device.py new file mode 100644 index 000000000000..f16bb4feaff9 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/device.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device primitives of the mono layer's attention stages. + +Wave reductions and shuffles, hardware math, bf16 / fp8 / MX packing, the +scaled MFMA, buffer and 64-bit global accesses, and the tagged-pair mailbox. +Adapted from ATOM's mono decode framework (``atom/mono/device``), itself ported +from the FlyDSL GLM-5 layer kernel. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +# The execution model: one CTA a CU of an MI355X, all resident. +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16, NT = 2. SC1 alone is +# device scope (past the per-XCD L2s); NT streams data read once. +CM_DEV = 16 +CM_NT = 2 + +FP8_MAX = 448.0 +FP8, FP4 = 0, 4 # f8f6f4 MFMA operand formats +UNIT_SCALE = 127 # E8M0 code of 1.0 +FLT_MIN = 1.1754943508222875e-38 + +POLL_MAX = 12 # mailbox pairs polled a batch + + +def traced(fn): + """FlyDSL's AST rewrite for a helper, so its ``if`` / ``while`` / ``for`` on + traced values lower to scf ops (kernels rewrite only their own body).""" + return ASTRewriter.transform(fn) + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``aiter::_``.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN5aiter{len(name)}{name}E" + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base``, then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ----------------------------------------------------------------------------- memory + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def ld_i32(ptr, i, cm=0): + return fx.Int32( + bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32, cache_modifier=cm) + ) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def ld_i64(ptr, i): + """Element i of an i64 buffer (two dwords).""" + w = fx.Vector(bo.buffer_load(rsrc(ptr), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(w[1]) << 32) | fx.Int64(fx.Uint32(w[0])) + + +def _gptr(addr): + return llvm.IntToPtrOp( + ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr)) + ).result + + +def gload(addr, words=4, nt=False): + """``words`` i32 (1, 2 or 4) at a 64-bit global address: a cache past 4 GiB is + reachable only this way (a buffer offset is 32 bits).""" + i32 = ir.IntegerType.get_signless(32) + ty = i32 if words == 1 else ir.VectorType.get([words], i32) + v = llvm.LoadOp(ty, _gptr(addr), alignment=4 * words, nontemporal=nt or None).result + return fx.Int32(v) if words == 1 else fx.Vector(v) + + +def atomic_add_dev(addr, value): + """Device-scope atomic add on the global i32 at ``addr``; the old value. + Relaxed: an acquire / release here would write back and invalidate L2.""" + return fx.Int32( + fx.atomic_add( + _gptr(addr), + fx.Int32(value), + syncscope="agent", + ordering=fx.AtomicOrdering.Monotonic, + ) + ) + + +def gstore(addr, value, words=4): + """Store ``words`` i32 (an Int32 or a Vector) at a 64-bit global address.""" + v = as_ir_value(fx.Int32(value) if words == 1 else value) + llvm.StoreOp(v, _gptr(addr), alignment=4 * words) + + +# ----------------------------------------------------------------------------- math + + +def fresh(v): + """``v`` through an opaque move, so the lane math built on it is computed at + its use instead of being hoisted out of a task loop and held in registers.""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def tile_map(tid, width): + """(tid // width, tid % width) for 0 <= tid < THREADS without a division: + shifts for a power of two, else a multiply-shift exact on that range.""" + if width & (width - 1) == 0: + sh = width.bit_length() - 1 + return tid >> sh, tid & (width - 1) + m = -(-(1 << 20) // width) + assert all((t * m) >> 20 == t // width for t in range(THREADS)), width + row = (tid * m) >> 20 + return row, tid - row * width + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def fma(a, b, c): + return fx.Float32(fmath.fma(fx.Float32(a), fx.Float32(b), fx.Float32(c))) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest (IEEE), given ``r = RN(1 / d)``.""" + q = x * r + return fma(fma(-q, d, x), r, q) + + +def permlane_swap(off, x, y): + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def lane_f32(v, src_lane): + return lane_gather(v.bitcast(fx.Int32), src_lane).bitcast(fx.Float32) + + +def readlane(v, src_lane): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def xshfl(v, off): + """Value of lane ``lane ^ off`` (DPP in a row, permlane swaps across).""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the xor ``offsets``, in that order.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +def row_shr(v, lane, k): + """Lane - k of the same 16-lane row, else 0 (DPP row_shr).""" + moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) + return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) + + +# ----------------------------------------------------------------------------- packing + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def fp8_pack4(a, b, c, d): + """Four f32 (in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +def fp8x8_bf16(w0, w1, scale): + """Eight E4M3 of two dwords times an f32 ``scale`` -> eight bf16 as a + 4-dword Vector (exact for a power-of-two scale in range).""" + out = [] + for w in (w0, w1): + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(scale)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32) + + +def clamp_fp8(v): + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def pow2(code): + """The fp32 power of two a biased exponent (1..254) stands for.""" + return (code << 23).bitcast(fx.Float32) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact for code in [1, 253].""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def ceil_exp(v): + """Biased exponent of the smallest power of two >= v (v > 0 finite, normal).""" + bits = v.bitcast(fx.Int32) + return ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + + +def mx_code(amax): + """VLLM's ROCm MXFP8 scale code of a 32-group: clamp(ceil(log2(amax / 448)) + + 127, 0, 254), amax floored at FLT_MIN.""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + return fx.min(fx.max(ceil_exp(q), fx.Int32(0)), fx.Int32(254)) + + +def mx_mul(code): + """2^(127 - code) as fp32, for code in [0, 254] (the quant multiplier).""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +# ----------------------------------------------------------------------------- MFMA + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8): + """16x16x128 scaled MFMA (gfx950 f8f6f4): lane l (row l % 16) holds K bytes + 16 (l // 16) .. +16 and 64 + 16 (l // 16) .. +16 of the 128-step, and the + E8M0 of K block l // 16.""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, 0, sa, 0, sb] + ) + ) + + +def mfma_bf16(a, b, c): + """16x16x32 bf16 MFMA: lane l (row / col l % 16) holds K 8 (l // 16) .. +8.""" + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +# ----------------------------------------------------------------------------- mailbox + + +@traced +def _spin(load_all, pending): + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +def spin_until(load, pending): + """``load()`` again (s_sleep between) while ``pending(value)``; the value.""" + return _spin(load, pending) + + +class Mailbox: + """Tagged-pair hand-off between the CTAs of one launch (device scope). + + Every 32-bit value is stored next to the launch's tag in one 8-byte store; a + consumer polls until every tag matches, so a hand-off is one round trip with + no separate flag. The tag is ``epoch * 64 + layer + 1`` (``epoch`` a device + counter bumped once a step): a pair from an earlier layer or step never + carries a live tag, and nothing is cleared between launches.""" + + def __init__(self, tag): + self.tag = tag + + def put(self, base, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def put_bf(self, base, i, vs): + """Elements i .. i + len(vs) (2 or 4, i aligned) as bf16 pairs, one store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(base), + i, + cache_modifier=CM_DEV, + ) + + def put_words(self, base, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. (fp8 words must come this way).""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of ``specs`` = [(base, pair index, npairs in {1, 2})]: all + re-loaded while any tag is stale. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + tag = self.tag + + def load_all(): + words = [] + for b, i, n in specs: + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +LOAD_BATCH = 4 + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over v = tid + THREADS i < n: ``batch`` rounds' loads + issued before their uses; a round past n loads element n - 1, uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def barrier(): + gpu.barrier() + + +# ----------------------------------------------------------------------------- timeline + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock.""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +@traced +def stamp(c, point): + """A timeline build's stamp: thread 0 of the CTA stores the clock at + ``point`` (``c["tl"]`` [BLOCKS, points] i64; None: not a timeline build).""" + if const_expr(c.get("tl") is None): + return + if c["tid"] == 0: + t = memrealtime() + w = fx.Vector.from_elements([fx.Int32(t), fx.Int32(t >> 32)], fx.Int32) + bo.buffer_store(w, rsrc(c["tl"]), (c["bid"] * c["tl_points"] + point) * 2) diff --git a/vllm/models/deepseek_v4/amd/mono/attention/front.py b/vllm/models/deepseek_v4/amd/mono/attention/front.py new file mode 100644 index 000000000000..23376ca4ff93 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/front.py @@ -0,0 +1,573 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K1 ``front``: one attention layer (the counts below are TP8's) from its +normed input row to the rotated query and the step's KV rows in the +sliding-window cache. + +One launch a layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, S rows. +Every CTA issues its weight loads (its wqkv part and its wq_b tasks) at the +kernel start, before it waits on anything; the stages: + + wqkv_a (128 row tiles x 2 K halves, 16 rows; a CTA each): the input's MXFP8 + half (X8, from the seam's norm) a token tile at a time (the next + tile's rows in flight) times wqkv's -> fp32 partials -> QKVP + qkv (S): the halves summed -> bf16; q_norm -> bf16 -> MXFP8 -> QX8 + (published, a flag a token); kv_norm -> bf16 -> GPT-J RoPE -> bf16 + -> the fp8_ds_mla record at the token's slot + wq_b (H x 512 / 16, 16 rows): QX8 times wq_b -> bf16 -> GPT-J RoPE of + each head's dims 448.. -> q (global, bf16) + kt (S, run by K2 after vLLM's indexer): the token's attention keys -- + top-k compressed slots through the kv source's block table (or, + ``global_keys``, every compressed slot of vLLM's c128a table), then + window slots -> KT + +Each elementwise step follows vLLM's ROCm path (``mxfp8_quantize``, +``FusedQKVRMSNorm``, ``fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert``); +the GEMV reductions differ from the Triton GEMMs' only in their order. +Adapted from ATOM's V4.1 ``attn_pre`` (K1), less its mHC seam. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from ..dims import ( + HALF, + HEAD_DIM, + HIDDEN, + KV_DIM, + NOPE, + NORM_EPS, + Q_RANK, + QKV_ROWS, + ROPE, + ROWS, + WINDOW, + WQKV_KPARTS, + WQKV_TASKS, + X_GROUPS, + X_WORDS, + Dims, + ld, + tile_rows, + token_tiles, +) +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + Mailbox, + bf16_round, + butterfly, + ceil_exp, + clamp_fp8, + div_rn, + first_task, + fp8_pack4, + fresh, + gstore, + hw_rsq, + ld_f32, + ld_i32, + mx_code, + mx_mul, + rsrc, + stamp, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import DATA, front_scratch + +_STREAM = fx.Stream(None) +QX_WORDS = Q_RANK // 4 +KSTEPS = HIDDEN // 128 // WQKV_KPARTS # a wqkv part's 128-steps +KWORDS = KSTEPS * 32 +KGROUPS = KSTEPS * 4 +XH = HIDDEN // WQKV_KPARTS # a K half's values +XH_COLS = XH // 8 // 2 # 16 B items (8 values) a half of a K half +XH_RPR = THREADS // XH_COLS # half-rows a round +XQ_ITEMS = XH // 8 # an xq unit's 16 B items (one a thread) +FRONT_POINTS = 6 # start, kt done, wqkv done, qkv done, wq_b done, (spare) + + +@dataclass(frozen=True) +class FrontBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (0: no top-k keys) + timeline: bool = False + topk: int = 0 # a token's compressed keys at most (``Dims.topk``) + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- kt + + +@traced +def _kt_round(c, t, k, ntopk, kv_len, width, req, cm): + """``stage_kt``'s key k of token t.""" + a, ratio, d = c["args"], c["ratio"], c["d"] + sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1) + s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx) + val = (k < kv_len).select(s_slot, fx.Int32(-1)) + if const_expr(ratio > 0): + local = ld_i32(a["topk"], t * width + fx.min(k, width - 1)) + if const_expr(c.get("global_keys", False)): + c_slot = local | fx.Int32(-2147483648) + else: + lc = fx.max(local, 0) + e = a["comp_block"] + blk = ld_i32(a["comp_bt"], req * a["bt_stride"] + lc // e) + c_slot = (blk * e + lc % e) | fx.Int32(-2147483648) + c_val = (local >= 0).select(c_slot, fx.Int32(-1)) + val = (k < ntopk).select(c_val, val) + if k < d.keys: + bo.buffer_store(val, rsrc(a["kt"]), t * d.keys + k, cache_modifier=cm) + + +@traced +def stage_kt(c, t): + """Token t's keys for K2: ``KT[t][k] = slot | 1 << 31`` for its top-k compressed + rows (k < ntopk), the window slot for k in [ntopk, kv_len), -1 for an + invalid or absent key; KLEN[t] = kv_len (0 for a pad row), KLEN[S + t] = + ntopk. The top-k table holds request-local compressed indices (through the + block table) or, ``global_keys``, the slots themselves (a row of the + runtime width ``bt_stride``: vLLM's every compressed key of a C128 layer).""" + tid = c["tid"] + a, ratio = c["args"], c["ratio"] + d, glob = c["d"], c.get("global_keys", False) + width = a["bt_stride"] if glob else fx.Int32(d.topk_width) + # device scope where this launch reads the table (K2's) + cm = CM_DEV if const_expr(c.get("dev", False)) else 0 + slot = ld_i32(a["slot"], 2 * t) + live = slot >= 0 + nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) + if const_expr(ratio > 0): + pos = ld_i32(a["pos"], 2 * t) + ntopk = live.select(fx.min((pos + 1) // ratio, width), fx.Int32(0)) + req = ld_i32(a["t2r"], t) + else: + ntopk = fx.Int32(0) + req = 0 + kv_len = ntopk + nswa + if const_expr(glob): + # a wide table: the rounds over the token's keys only (the split's + # chunks past them skipped, ``attention.back.stage_split``) + for i in range(0, (kv_len + THREADS - 1) // THREADS): + _kt_round(c, t, fresh(tid) + THREADS * i, ntopk, kv_len, width, req, cm) + else: + for i in range_constexpr((d.keys + THREADS - 1) // THREADS): + _kt_round(c, t, fresh(tid) + THREADS * i, ntopk, kv_len, width, req, cm) + if tid == 0: + bo.buffer_store(kv_len, rsrc(a["klen"]), t, cache_modifier=cm) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t, cache_modifier=cm) + + +# ---------------------------------------------------------------- wqkv_a + + +@traced +def stage_wqkv(c, task, part, ops, t0, n, live): + """Rows 16 task .. of wqkv over K half ``part`` against tokens t0 .. t0 + n + -> fp32 partials QKVP[part] (``live``: a CTA with a wqkv part).""" + tid, red = c["tid"], c["red"] + gemv_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, ops, rows=n, nsteps=KSTEPS) + gpu.barrier() + if (tid < ROWS * n) & live: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + v = row_sum(red, r, tl) + c["mb"].put(c["qkvp"], (part * c["S"] + t) * QKV_ROWS + task * ROWS + r, v) + gpu.barrier() + + +# ---------------------------------------------------------------- q / kv norms + + +@traced +def stage_qkv(c, t): + """Token t: the K halves summed -> bf16; q_norm -> bf16 -> MXFP8 (QX8); + kv_norm -> bf16 -> RoPE -> bf16 -> the fp8_ds_mla record at + slot_mapping[t]. 256 threads (8 q columns and 2 kv columns each); threads + 256.. mirror thread 255 and store nothing.""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, s = c["args"], c["S"] + tt = fx.min(tid, 255) + mine = tid < 256 + qc = [fx.min(tt * 8 + j, Q_RANK - 1) for j in range(8)] + # independent loads first: they overlap the partials' wait + qw = [c["ldbf"](a["qn_w"], qc[j]) for j in range(8)] + kw = [c["ldbf"](a["kvn_w"], tt * 2 + j) for j in range(2)] + pos = ld_i32(a["pos"], 2 * t) + slot = ld_i32(a["slot"], 2 * t) + jr = tt - NOPE // 2 + is_rope = jr >= 0 + jj = fx.max(jr, 0) + cs = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + jj), fx.Float32(1.0)) + sn = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + HALF + jj), fx.Float32(0.0)) + specs = [] + for p in range_constexpr(2): + specs += [(c["qkvp"], (p * s + t) * QKV_ROWS + qc[j], 1) for j in range(8)] + specs += [ + (c["qkvp"], (p * s + t) * QKV_ROWS + Q_RANK + tt * 2 + j, 1) + for j in range(2) + ] + got = c["mb"].poll(specs) + vals = [ + bf16_round(got[j][0].bitcast(fx.Float32) + got[10 + j][0].bitcast(fx.Float32)) + for j in range(10) + ] + xq = [((tt * 8 + j) < Q_RANK).select(vals[j], fx.Float32(0.0)) for j in range(8)] + xk = vals[8:10] + ssq = xq[0] * xq[0] + for j in range_constexpr(1, 8): + ssq = ssq + xq[j] * xq[j] + ssq = butterfly(ssq, (32, 16, 8, 4, 2, 1)) + ssk = butterfly(xk[0] * xk[0] + xk[1] * xk[1], (32, 16, 8, 4, 2, 1)) + if (lane == 0) & mine: + fx.ptr_store(ssq, red + wave) + fx.ptr_store(ssk, red + 4 + wave) + gpu.barrier() + tq = (fx.ptr_load(red + 0) + fx.ptr_load(red + 1)) + ( + fx.ptr_load(red + 2) + fx.ptr_load(red + 3) + ) + tk = (fx.ptr_load(red + 4) + fx.ptr_load(red + 5)) + ( + fx.ptr_load(red + 6) + fx.ptr_load(red + 7) + ) + rq = hw_rsq( + div_rn(tq, fx.Float32(float(Q_RANK)), fx.Float32(1.0 / Q_RANK)) + + fx.Float32(NORM_EPS) + ) + rk = hw_rsq( + div_rn(tk, fx.Float32(float(KV_DIM)), fx.Float32(1.0 / KV_DIM)) + + fx.Float32(NORM_EPS) + ) + # q: bf16 normed row, then vLLM's MXFP8 (a 32-group = 4 threads) + ys = [bf16_round((xq[j] * rq) * qw[j]) for j in range(8)] + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = mx_code(butterfly(amax, (1, 2), fx.max)) + mul = mx_mul(code) + # the bf16 qr, vLLM's indexer's input + if mine & (tt * 8 < Q_RANK): + words = [ + fx.Vector.from_elements([ys[2 * i], ys[2 * i + 1]], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + for i in range(4) + ] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a["qr_out"]), + t * (Q_RANK // 2) + 4 * tt, + ) + if mine & (tt * 8 < Q_RANK): + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["qx8"]), + t * QX_WORDS + 2 * tt, + cache_modifier=CM_DEV, + ) + if tt % 4 == 0: + bo.buffer_store( + code, + rsrc(c["qx8s"]), + t * (Q_RANK // 32) + tt // 4, + cache_modifier=CM_DEV, + ) + # kv: this thread's pair (2 tt, 2 tt + 1) + kn = [bf16_round((xk[j] * rk) * kw[j]) for j in range(2)] + e0 = bf16_round(kn[0] * cs - kn[1] * sn) + e1 = bf16_round(kn[0] * sn + kn[1] * cs) + # record quant of the nope dims: a 64-block = 32 threads (half a wave) + amax = butterfly(fx.max(abs(e0), abs(e1)), (1, 2, 4, 8, 16), fx.max) + rcode = ceil_exp( + div_rn( + fx.max(amax, fx.Float32(1e-4)), fx.Float32(448.0), fx.Float32(1.0 / 448.0) + ) + ) + inv = mx_mul(rcode) + q0 = clamp_fp8(e0 * inv) + q1 = clamp_fp8(e1 * inv) + n0 = xshfl(q0, 1) + n1 = xshfl(q1, 1) + word = fp8_pack4(q0, q1, n0, n1) + rope_word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + blk = a["swa_block"] + sblock = slot // blk + soff = slot % blk + base = a["swa"] + fx.Int64(sblock) * fx.Int64(a["swa_stride"]) + rec = base + fx.Int64(soff * DATA) + if mine & (tt % 32 == 0) & (tt < NOPE // 2): + fx.ptr_store(rcode.bitcast(fx.Float32), red + 8 + tt // 32) + gpu.barrier() + if mine & (slot >= 0): + if tt < NOPE // 2: + if tt % 2 == 0: + gstore(rec + fx.Int64(2 * tt), word, words=1) + else: + gstore(rec + fx.Int64(NOPE + 4 * (tt - NOPE // 2)), rope_word, words=1) + if tt == 0: + sb = [fx.ptr_load(red + 8 + b).bitcast(fx.Int32) for b in range(7)] + w0 = sb[0] | (sb[1] << 8) | (sb[2] << 16) | (sb[3] << 24) + w1 = sb[4] | (sb[5] << 8) | (sb[6] << 16) + sc = base + fx.Int64(blk * DATA + soff * 8) + gstore(sc, fx.Vector.from_elements([w0, w1], fx.Int32), words=2) + publish(c, c["qrdy"], t) + + +# ---------------------------------------------------------------- wq_b + + +def wqb_rope(c, task0, pos_t, n): + """(cos, sin) of this thread's pair in ``stage_wqb``'s epilogue over a tile + of n rows (``pos_t``: its token's position; every task of the CTA shares + the pair's dims: tasks BLOCKS apart are whole heads apart).""" + a = c["args"] + rp = c["tid"] % (ROWS // 2) + dim = (task0 * ROWS) % HEAD_DIM + 2 * rp + jj = fx.max((dim - NOPE) >> 1, 0) + return ld_f32(a["cos_sin"], pos_t * ROPE + jj), ld_f32( + a["cos_sin"], pos_t * ROPE + HALF + jj + ) + + +def wqb_pos(c, t0, n): + """This thread's token position for ``stage_wqb``'s epilogue of tile + t0 .. t0 + n.""" + per = ROWS * n // 2 + tl = fx.min((c["tid"] % per) // (ROWS // 2), n - 1) + return ld_i32(c["args"]["pos"], 2 * (t0 + tl)) + + +@traced +def stage_wqb(c, task0, bops, t0, n, rope): + """Rows 16 (task0 + BLOCKS k) .. of the rank's H x 512 query rows, every + task k of the CTA: GEMVs -> bf16 -> GPT-J RoPE of a head's dims 448.. + (fp32) -> bf16 -> q.""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + nt = len(bops) + for k in range_constexpr(nt): + gemv_mfmas(c, Q_RANK, c["qxl"], c["qxsl"], red + k * RED, bops[k], rows=n) + gpu.barrier() + per = ROWS * n // 2 # bf16 pairs a task + if tid < nt * per: + k = tid // per + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + r0 = (task0 + BLOCKS * k) * ROWS + 2 * rp + rk = red + k * RED + v0 = bf16_round(row_sum(rk, 2 * rp, tl)) + v1 = bf16_round(row_sum(rk, 2 * rp + 1, tl)) + is_rope = (task0 * ROWS) % HEAD_DIM + 2 * rp >= NOPE + cs, sn = rope + e0 = is_rope.select(v0 * cs - v1 * sn, v0) + e1 = is_rope.select(v0 * sn + v1 * cs, v1) + word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["q"]), (t * d.q_rows + r0) // 2) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def front_lds(s: int, d: Dims): + """K1's LDS (the attention front's stages).""" + rows = tile_rows(s) + wqb_per_cta = d.wqb_tasks // BLOCKS + + @fx.struct + class FrontLds: + red: fx.Array[fx.Float32, wqb_per_cta * RED + 16, 16] + qxl: fx.Array[fx.Int32, rows * lds_row(QX_WORDS), 16] + qxsl: fx.Array[fx.Int32, rows * lds_row(Q_RANK // 32), 16] + xl: fx.Array[fx.Int32, rows * lds_row(KWORDS), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(KGROUPS), 16] + + return FrontLds + + +def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The front stages' context: LDS views, mailbox (tag ``tag``), scratch + regions, arguments (``args``: the kernel's, by name).""" + s, d = key.tokens, Dims(key.tp, key.topk) + tid = fx.thread_idx.x + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "xl": lds.xl.ptr, + "xsl": lds.xsl.ptr, + "qxl": lds.qxl.ptr, + "qxsl": lds.qxsl.ptr, + "mb": Mailbox(tag), + "d": d, + "ratio": key.ratio, + "ldbf": lambda ptr, i: fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ), + "args": args, + } + for region, (off, _) in (layout or front_scratch(s, d)).items(): + c[region] = scratch + fx.Int64(off) + return c + + +@traced +def run_front(c, key, bid, before_wqkv=None, after_wqkv=None): + """The front's stages on this CTA, from the input's MXFP8 X8 / XRDY that a + stage before them (the mono layer's seam norm) publishes for every token. + ``before_wqkv()`` runs ahead of the wqkv weights' loads (a latency-critical + stage of the caller's), ``after_wqkv()`` right after. The step's key table + is K2's (``stage_kt``): an index source's top-k comes after this launch.""" + s = key.tokens + d = c["d"] + wqb_tasks = d.wqb_tasks + a = c["args"] + tiles = token_tiles(s) + wqb_per_cta = wqb_tasks // BLOCKS + assert wqb_tasks % BLOCKS == 0 + n_wqkv = WQKV_TASKS * WQKV_KPARTS # 256: every CTA's + assert n_wqkv == BLOCKS + QKV0 = 0 # on wqkv CTAs + WQB0 = 0 + stamp(c, 0) + + # ---- Loads complete in issue order, so every CTA issues its latency- + # critical loads before its later stages' weights: its wqkv part, its + # first input tile, the next tiles while it multiplies, then its wq_b + # tasks. The q / kv norms run on wqkv CTAs, their partials' wait behind + # the wq_b weights that are due by then anyway. + wqkv_unit = first_task(bid, 0) + has_wqkv = wqkv_unit < n_wqkv + task = fx.min(wqkv_unit, n_wqkv - 1) % WQKV_TASKS + part = fx.min(wqkv_unit, n_wqkv - 1) // WQKV_TASKS + wqb_task = first_task(bid, WQB0) + if const_expr(before_wqkv is not None): + before_wqkv() + kops = gemv_loads( + c, + a["wqkv"], + a["wqkv_s"], + HIDDEN, + task * ROWS, + kbeg=part * KSTEPS, + nsteps=KSTEPS, + pred=has_wqkv, + shuffled=True, + ) + if const_expr(after_wqkv is not None): + after_wqkv() + stamp(c, 1) + poss = [wqb_pos(c, t0, n) for t0, n in tiles] + + # ---- wqkv_a: one (row tile, K half) a CTA, every token tile; the wq_b + # weights issued behind the input rows' loads + def wqb_weights(): + return [ + gemv_loads( + c, + a["wqb"], + a["wqb_s"], + Q_RANK, + (wqb_task + BLOCKS * i) * ROWS, + ld=ld(Q_RANK), + ) + for i in range(wqb_per_cta) + ] + + await_flags(c, c["xrdy"], s, lambda u: WQKV_KPARTS * u + part) + xr = ( + c["x8"], + c["x8s"], + X_WORDS, + part * KWORDS, + KWORDS, + X_GROUPS, + part * KGROUPS, + KGROUPS, + ) + regs = rows_loads(c, *xr, tiles[0][0], tiles[0][1], pred=has_wqkv) + if const_expr(len(tiles) == 1): + bops = wqb_weights() + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, KWORDS, KGROUPS, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + t1, n1 = tiles[ti + 1] + regs = rows_loads(c, *xr, t1, n1, pred=has_wqkv) + if const_expr(ti + 2 == len(tiles)): + bops = wqb_weights() + stage_wqkv(c, task, part, kops, t0, n, has_wqkv) + stamp(c, 2) + + # ---- q / kv norms, KV insert + for t in range(first_task(bid, QKV0), s, BLOCKS): + stage_qkv(c, t) + stamp(c, 3) + + # ---- wq_b: every task's rows of a token tile at a time, the next + # tile's rows in flight + ropes = [wqb_rope(c, wqb_task, poss[ti], tiles[ti][1]) for ti in range(len(tiles))] + await_flags(c, c["qrdy"], s, lambda u: u) + qx = ( + c["qx8"], + c["qx8s"], + QX_WORDS, + 0, + QX_WORDS, + Q_RANK // 32, + 0, + Q_RANK // 32, + ) + regs = rows_loads(c, *qx, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, QX_WORDS, Q_RANK // 32, c["qxl"], c["qxsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *qx, *tiles[ti + 1]) + stage_wqb(c, wqb_task, bops, t0, n, ropes[ti]) + stamp(c, 4) diff --git a/vllm/models/deepseek_v4/amd/mono/attention/gemv.py b/vllm/models/deepseek_v4/amd/mono/attention/gemv.py new file mode 100644 index 000000000000..5cc9fd25a450 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/gemv.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 GEMV of the mega kernel: 16 rows of a row-major e4m3 weight [N, K] +(vLLM's loaded layout, E8M0 scales in 32x32 blocks [N / 32, K / 32]) against +every token of a tile, its MXFP8 rows in LDS, on the gfx950 scaled MFMA. + +A task may cover a window of K (``kbeg`` .. ``kbeg + nsteps`` 128-steps): a +split-K part, its activations in LDS holding just that window. The weight +operands are loaded first (all in flight, before the task waits on its +activations); wave w takes the window's steps w, w + 8, ... (or, ``split``, a +contiguous range of waves 0 .. split - 1), its 16 x 16 partial in ``red``, +summed in wave order by ``row_sum``. Adapted from ATOM's +``attn_pre.gemv_fp8_*``, which reads aiter's (16, 16)-preshuffled layout. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import T + +from .device import THREADS, WAVES, mfma_scaled, rsrc, tile_map, traced + +# a token's LDS row of MXFP8 words / E8M0 codes, padded by 16 B so 16 rows read +# by one MFMA operand do not share a bank group +LDS_PAD = 4 +RED = WAVES * 64 * 4 # one task's accumulators in ``red`` (f32) + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def _steps(c, nsteps, split, tiled): + """This wave's steps of the window: [(local step, live)].""" + wave = c["wave"] + if const_expr(split is None): + per = (nsteps + WAVES - 1) // WAVES + ragged = nsteps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < nsteps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and nsteps % split == 0 + per = nsteps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def _step_load(c, r_w, r_ws, k, row0, st, pred=None, shuffled=False, ld=None): + """Global step ``st``'s weight operand of rows row0 .. +16: two 16 B chunks a + lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16) + and the E8M0 byte of its K block. ``pred`` false: masked (out of range: no + memory traffic, zeros). ``shuffled``: the weight in aiter's (16, 16) + shuffle (a 16-row group's 32-column blocks of 16 x 32 bytes, each a 16-column + half after the other), where a lane's chunk is 16 contiguous bytes too. + ``ld``: a row-major weight's row stride (bytes), past ``k``.""" + lane = c["lane"] + j = lane // 16 + if const_expr(shuffled): + assert k % 64 == 0 + blocks = (row0 // 16) * (k // 64) + 2 * st + lo = fx.Vector( + bo.buffer_load(r_w, blocks * 256 + lane * 4, vec_width=4, dtype=T.i32) + ) + hi = fx.Vector( + bo.buffer_load(r_w, (blocks + 1) * 256 + lane * 4, vec_width=4, dtype=T.i32) + ) + else: + row = row0 + lane % 16 + base = row * (k if ld is None else ld) + 128 * st + 16 * j + lo = fx.Vector(bo.buffer_load(r_w, base // 4, vec_width=4, dtype=T.i32)) + hi = fx.Vector(bo.buffer_load(r_w, (base + 64) // 4, vec_width=4, dtype=T.i32)) + kg = k // 32 + wsi = (row0 // 32) * kg + st * 4 + j + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def gemv_loads( + c, + w, + ws, + k, + row0, + split=None, + tiled=False, + kbeg=0, + nsteps=None, + pred=None, + shuffled=False, + ld=None, +): + """The weight operands of rows row0 .. +16 (``tiled``: row0 + 16 (wave // + split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight + (``pred`` false: masked off, no memory traffic; ``shuffled``: the weight in + aiter's (16, 16) shuffle; ``ld``: a row-major weight's row stride).""" + nsteps = k // 128 if nsteps is None else nsteps + if const_expr(pred is None): + r_w, r_ws = rsrc(w), rsrc(ws) + else: + # a masked-off CTA reads through a zero-size buffer: zeros, no traffic + nb = pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + r_w, r_ws = rsrc(w, nb), rsrc(ws, nb) + if const_expr(tiled): + row0 = row0 + 16 * (c["wave"] // split) + return [ + _step_load( + c, + r_w, + r_ws, + k, + row0, + kbeg + (st if live is None else fx.min(st, nsteps - 1)), + pred, + shuffled, + ld, + ) + for st, live in _steps(c, nsteps, split, tiled) + ] + + +def _step_mfma(c, nsteps, xl, xsl, st, ops, acc, live, rows): + lane = c["lane"] + lo, hi, sa = ops + x_row, s_row = lds_row(nsteps * 32), lds_row(nsteps * 4) + j = lane // 16 + col = fx.min(lane % 16, rows - 1) + if const_expr(live is not None): + st = fx.min(st, nsteps - 1) + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [hi[d] for d in range(4)], fx.Int32 + ) + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + [xhi[d] for d in range(4)], fx.Int32 + ) + sb = fx.ptr_load(xsl + (col * s_row + st * 4 + j)) + if const_expr(live is not None): + # a step past the end adds zero at unit scale + av = fx.Vector.from_elements( + [live.select(av[d], fx.Int32(0)) for d in range(8)], fx.Int32 + ) + sa = live.select(sa, fx.Int32(127)) + sb = live.select(sb, fx.Int32(127)) + return mfma_scaled(av, xv, acc, sa, sb) + + +def gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None): + """``gemv_loads``' operands against ``rows`` tokens' MXFP8 rows of the + window in LDS (``xl`` words, ``xsl`` codes, rows ``lds_row`` of the window) + -> ``red[wave][lane]`` (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16).""" + lane, wave = c["lane"], c["wave"] + nsteps = k // 128 if nsteps is None else nsteps + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_steps(c, nsteps, split, tiled)): + acc = _step_mfma(c, nsteps, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' accumulators, summed in wave order.""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def span_items(width, n): + """A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.""" + wq = width // 4 + rpr = THREADS // wq + assert rpr >= 1, width + return wq, rpr, (n + rpr - 1) // rpr + + +def rows_loads( + c, + words_base, + codes_base, + stride_w, + off_w, + words, + stride_g, + off_g, + groups, + t0, + n, + pred=None, +): + """Tokens t0 .. t0 + n's published MXFP8 row slices -- ``words`` i32 from + ``off_w`` and ``groups`` codes (an i32 each) from ``off_g``, rows + ``stride_w`` / ``stride_g`` i32 apart -- loaded (issued, not waited on), + 16 B a lane-item: a thread one column, rows a round apart. Plain loads: + their flags are seen, and nothing in the launch read these lines before. + ``pred`` false: masked off (zeros, no memory traffic).""" + tid = c["tid"] + nb = None if pred is None else pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + out = [] + for base, stride, off, width in ( + (words_base, stride_w, off_w, words), + (codes_base, stride_g, off_g, groups), + ): + wq, rpr, rounds = span_items(width, n) + row0, col = tile_map(tid, wq) + a0 = col * 4 + off + t0 * stride + r = rsrc(base, nb) + regs = [] + for i in range(rounds): + row = row0 + rpr * i + if (i + 1) * rpr > n or rpr * wq < THREADS: + row = fx.min(row, n - 1) + regs.append( + fx.Vector( + bo.buffer_load(r, row * stride + a0, vec_width=4, dtype=T.i32) + ) + ) + out.append(regs) + return out + + +@traced +def _store_span(c, regs, width, dst, n): + tid = c["tid"] + wq, rpr, _ = span_items(width, n) + row0, col = tile_map(tid, wq) + lrow = lds_row(width) + base = dst + (row0 * lrow + col * 4) + for i in range_constexpr(len(regs)): + if const_expr((i + 1) * rpr <= n and rpr * wq == THREADS): + fx.ptr_store(regs[i], base + i * rpr * lrow) + else: + if (row0 + rpr * i < n) & (tid < rpr * wq): + fx.ptr_store(regs[i], base + i * rpr * lrow) + + +def rows_store(c, regs, words, groups, xl, xsl, n): + """``rows_loads``' registers into LDS rows 0 .. n (``lds_row`` strides).""" + _store_span(c, regs[0], words, xl, n) + _store_span(c, regs[1], groups, xsl, n) + gpu.barrier() diff --git a/vllm/models/deepseek_v4/amd/mono/attention/plan.py b/vllm/models/deepseek_v4/amd/mono/attention/plan.py new file mode 100644 index 000000000000..15f6b3583293 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/attention/plan.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Geometry of the attention stages (``front``, ``back``): the KV record, the +split attention's chunks and the scratch (mailbox) layout; the model's widths +and a TP rank's share of them are in ``..dims``.""" + +from vllm.models.deepseek_v4.amd.mono.dims import ( + HEAD_DIM, + HIDDEN, + Q_RANK, + QKV_ROWS, + WQKV_KPARTS, + X_GROUPS, + Dims, +) + +# KV record (V4 fp8_ds_mla): 576 B data rows (448 e4m3 | 64 bf16), then 8 B of +# scales a token (7 UE8M0 + pad) after the block's rows +DATA = 576 +SCALE_BYTES = 8 + +# attention: a split is 16 keys +BK = 16 +COLS = 32 # PV columns a group (one MX group of the output) +GROUPS = HEAD_DIM // COLS # 16 + +PAIR = 8 # a mailbox (value, tag) pair + + +def cdiv(a, b): + return -(-a // b) + + +def round_up(a, b): + return cdiv(a, b) * b + + +def pair_layout(items, start=0, align=256): + out, off = {}, start + for name, pairs in items: + off = round_up(off, align) + out[name] = (off, pairs * PAIR) + off += pairs * PAIR + return out + + +def byte_layout(items, start=0, align=256): + out, off = {}, start + for name, nbytes in items: + off = round_up(off, align) + out[name] = (off, nbytes) + off += nbytes + return out + + +def front_scratch(s: int, d: Dims) -> dict: + """K1's scratch regions -> (byte offset, bytes): published data (plain + words behind a flag) and mailbox pairs.""" + return byte_layout( + [ + ("x8", s * HIDDEN), # e4m3 words (S > 16: quantized once, shared) + ("x8s", s * X_GROUPS * 4), # an i32 code a group + ("xrdy", WQKV_KPARTS * s * PAIR), # a flag a (token, K half) + ("qkvp", WQKV_KPARTS * s * QKV_ROWS * PAIR), + ("qx8", s * Q_RANK), # e4m3 words + ("qx8s", s * Q_RANK // 32 * 4), # an i32 code a group + ("qrdy", s * PAIR), + ] + ) + + +def back_scratch(s: int, d: Dims, start: int = 0) -> dict: + """K2's scratch regions, laid out past ``start``: published data (plain) + and flag pairs.""" + ck = ( + 64 if d.head_tiles == 2 and s * cdiv(d.keys, 64) <= 256 else 128 + ) # back.chunk_keys + nchunk = cdiv(d.keys, ck) + return byte_layout( + [ + ("part", s * nchunk * d.heads * HEAD_DIM * 2), # bf16 partials + ("pstat", s * nchunk * d.heads * 8), # (max, sum) f32 + ("srdy", s * nchunk * PAIR), + ("xo", s * d.q_rows), # e4m3 words + ("xos", s * d.q_rows // 32 * 4), # an i32 code a group + ("crdy", s * d.heads // 8 * PAIR), + ("x8b", s * d.o_rows), + ("x8bs", s * d.o_rows // 32 * 4), + ("ardy", d.woa_tasks * PAIR), + ], + start=start, + ) + + +def front_end(s: int, d: Dims) -> int: + """Where K2's regions start: past K1's, which share the launch tag.""" + return round_up(max(o + n for o, n in front_scratch(s, d).values()), 256) + + +def scratch_bytes(s: int, d: Dims) -> int: + b = back_scratch(s, d, start=front_end(s, d)) + return round_up(max(o + n for o, n in b.values()), 4096) diff --git a/vllm/models/deepseek_v4/amd/mono/common/__init__.py b/vllm/models/deepseek_v4/amd/mono/common/__init__.py new file mode 100644 index 000000000000..397642c6b441 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP +group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX +arithmetic), ``plan`` (execution model, layouts, build keys) and +``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v4/amd/mono/common/mx.py b/vllm/models/deepseek_v4/amd/mono/common/mx.py new file mode 100644 index 000000000000..8f29d81ac0d2 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/mx.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/mx.py +"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale +codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight +tile's operand loads and the scaled MFMA that takes them. + +An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.ops import CM_NT + +FP8_MAX = 448.0 +# the f8f6f4 MFMA's operand formats (cbsz / blgp) +FP8, FP4 = 0, 4 +UNIT_SCALE = 127 # the E8M0 code of 1.0 + + +def clamp_fp8(v): + """``v`` clamped to the E4M3 range.""" + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def code_ceil(v): + """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), + clamped to [1, 254]: the codes a scale may take.""" + bits = v.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale + bit for bit, a multiply where the division is a dozen instructions.""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def fp4_tile_load(r_w, base, rg, k, st, at): + """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` + (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a + step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). + Nontemporal: a weight a step reads once.""" + return fx.Vector( + bo.buffer_load( + r_w, + base + rg * (2 * k) + st * 256 + at, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): + """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 + ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its + byte ``sb_byte``).""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] + ) + ) diff --git a/vllm/models/deepseek_v4/amd/mono/common/ops.py b/vllm/models/deepseek_v4/amd/mono/common/ops.py new file mode 100644 index 000000000000..6369b5856bca --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/ops.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ops.py +# atom/mono/device/ranks.py +# atom/mono/device/stamps.py +"""Device-side primitives of the mono kernels: wave reductions and shuffles, the +approximate hardware math, bf16 / fp8 packing, buffer resources, the +execution-model constants and the kernel symbol naming. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs +and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``. + +Also the TP group seen from inside a kernel (every rank's peer buffer, an +all-reduce's rank-ordered sum) and a timeline build's phase stamps.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +from vllm.models.deepseek_v4.amd.mono.common.plan import THREADS, WAVES + +# ---------------------------------------------------------------- primitives +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 + + +CM_SYS = 17 + + +# NT = 2: nontemporal. A weight a layer reads once streams past the caches it +# would only evict the step's reused data from (the MoE at 48 rows: 143 -> 133 us) +CM_NT = 2 + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def fresh(v): + """``v`` (a per-lane Int32) through an opaque move, here: the lane math built + on it is computed at its use, where LLVM would otherwise hoist it out of a task + loop or CSE it across stages into registers held for the whole kernel (spills, + and a kernel with scratch is refused: ``runtime.widths``).""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS). Powers of two only: any other + offset would take the lane ^ 1 path.""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python + int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. + A round whose ``use`` stores to global memory keeps the next round's loads + behind it (they may alias), so the rounds' loads would go out in series. A + round past n loads element n - 1 (in range) and uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` + + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in + order; all of them by default).""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def ballot(pred): + """The wave's lanes where ``pred`` holds, a bit a lane.""" + return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) + + +def popcount(x): + return fx.Int32(fx.ctpop(x)) + + +def lanes_below(bal, lane): + """How many of the lanes in ballot ``bal`` are below ``lane``.""" + return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + + +def load_ptr64(table, i): + """Entry ``i`` of an i64 table in memory, wave-uniform (each half through + ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" + pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + """The f32 of a packed bf16 pair word's low half.""" + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + """The f32 of a packed bf16 pair word's high half.""" + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return bf_lo(w), bf_hi(w) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +# ---------------------------------------------------------------- TP ranks +def peer_bases(peers, tp): + """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), + loaded once: a put's address must not wait on a load of its own.""" + return [load_ptr64(peers, p) for p in range(tp)] + + +def sum_partials(poll, own, pair_of, tp): + """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox + region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> + (lo sum, hi sum). Every rank sums in the same order, so they agree.""" + ws = poll([(own, pair_of(src), 1) for src in range(tp)]) + acc0, acc1 = bf2_f32(ws[0][0]) + for src in range_constexpr(1, tp): + lo, hi = bf2_f32(ws[src][0]) + acc0 = acc0 + lo + acc1 = acc1 + hi + return acc0, acc1 + + +# ---------------------------------------------------------------- timeline stamps +@traced +def stamp(on, tls, tid, k): + # compile-time gate outside, traced condition inside: they cannot be one `and` + if const_expr(on): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) diff --git a/vllm/models/deepseek_v4/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v4/amd/mono/common/peer_memory.py new file mode 100644 index 000000000000..58c95a567e07 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/peer_memory.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/consensus.py +# atom/mono/runtime/peer_memory.py +"""Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every +rank must reach before anything collective (a rank that turned the mono layer +off alone would leave the others spinning on the GPU).""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +# ---------------------------------------------------------------- TP agreement +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" + + +def tp_agree(ok: bool, group) -> bool: + """True when every rank of ``group`` passed ``ok`` true. Every rank must call + it at the same point; ``group`` is a CPU (gloo) process group, or None for a + single rank.""" + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +# ---------------------------------------------------------------- peer buffers +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a + non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no rank + overwrites a peer's unread data, is the model's layout's to say. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so a + runner reaches it only after ``tp_agree``. A failed mapping on any rank raises + ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + # a view, not an owner: close() frees the memory under it + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: # every rank must still reach the agreement + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent. Only once no + kernel of any rank can still touch it (the runner is being dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + self.bytes = None + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 diff --git a/vllm/models/deepseek_v4/amd/mono/common/plan.py b/vllm/models/deepseek_v4/amd/mono/common/plan.py new file mode 100644 index 000000000000..10f688ac3b6a --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/plan.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/plan/execution.py +# atom/mono/plan/layout.py +# atom/mono/plan/arith.py +# atom/mono/plan/shard.py +# atom/mono/plan/build_key.py +"""Plain-Python planning of the mono kernels: the execution model (one CTA a CU, +all resident), mailbox layouts, arithmetic that runs on ints and traced values +alike, TP shards, and build keys (every build parameter, and this package's +sources, in FlyDSL's JIT cache key).""" + +import enum +import functools +import hashlib +from dataclasses import dataclass, fields +from pathlib import Path + +# ---------------------------------------------------------------- execution model +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of + CTAs past 0: placements that add task counts run past the grid), then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ---------------------------------------------------------------- mailbox layouts +PAIR_BYTES = 8 # a (value, tag) pair + + +def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: + """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), + each region from an ``align``-byte boundary at or past ``start``.""" + out, off = {}, start + for name, pairs in items: + off = (off + align - 1) // align * align + out[name] = (off, pairs * PAIR_BYTES) + off += pairs * PAIR_BYTES + return out + + +# ---------------------------------------------------------------- arithmetic +def sel(pred, a, b): + """A if pred else b: a conditional on Python values, ``select`` on traced + ones (rewrapped in the traced operand's type).""" + if isinstance(pred, bool): + return a if pred else b + out = pred.select(a, b) + traced = [x for x in (a, b) if not isinstance(x, int)] + return type(traced[0])(out) if traced else out + + +def imin(a, b): + return sel(a < b, a, b) + + +def i32(x): + """``x``, asserted inside the device's Int32 range when it is a Python int.""" + if isinstance(x, int): + assert -(2**31) <= x < 2**31, x + return x + + +def cdiv(a, b): + """ceil(a / b) for a >= 0, b > 0.""" + return i32((a + b - 1) // b) + + +# ---------------------------------------------------------------- TP shards +class ShardError(ValueError): + """A width the TP size does not divide, or a TP size the model refuses.""" + + +@dataclass(frozen=True) +class Shard: + tp: int + + def __post_init__(self): + if self.tp < 1: + raise ShardError(f"TP {self.tp}") + + def split(self, full: int, what: str) -> int: + """``full`` / tp, refusing a remainder: a rank's share of ``what``.""" + if full % self.tp: + raise ShardError(f"{what} {full} is not divisible by TP {self.tp}") + return full // self.tp + + +def padded(n: int, align: int) -> int: + return cdiv(n, align) * align + + +# ---------------------------------------------------------------- build keys +_SCALARS = (int, float, bool, str, enum.Enum) +_ROOT = Path(__file__).resolve().parents[1] # the mono package + + +@functools.cache +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the mono package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these.""" + h = hashlib.sha256() + for rel in paths: + root = _ROOT / rel + files = [root] if root.is_file() else sorted(root.rglob("*.py")) + if not files: + raise FileNotFoundError(f"no kernel sources at {root}") + for f in files: + h.update(str(f.relative_to(_ROOT)).encode()) + h.update(f.read_bytes()) + return h.hexdigest()[:16] + + +def key_tuple(key, sources: str) -> tuple: + """``key``'s (field name, value) pairs, every value a scalar, then + ("sources", ``sources``): the package's ``source_digest``.""" + out = [] + for f in fields(key): + value = getattr(key, f.name) + if not isinstance(value, _SCALARS): + raise TypeError( + f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" + " a scalar to reach the JIT cache key" + ) + out.append((f.name, value)) + out.append(("sources", sources)) + return tuple(out) diff --git a/vllm/models/deepseek_v4/amd/mono/common/sync.py b/vllm/models/deepseek_v4/amd/mono/common/sync.py new file mode 100644 index 000000000000..2a167dd91c1d --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/common/sync.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/sync.py +"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and +between GPUs (system scope), every address a named region (``Addr``).""" + +from dataclasses import dataclass +from enum import Enum + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.ops import ( + CM_DEV, + CM_SYS, + bf16_pair, + rsrc, + traced, +) + +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +class Space(Enum): + SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs + PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs + + +@dataclass(frozen=True) +class Addr: + """A mailbox region's address: the region's name and space, and the traced + address. The space sets the hand-off's scope: device for SCRATCH, system for + PEER.""" + + region: str + space: Space + value: object + + +def sreg(base, offset, region): + """Scratch region ``region`` at byte ``offset`` of ``base``.""" + return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) + + +def preg(base, offset, region): + """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer + ``base``.""" + return Addr(region, Space.PEER, base + fx.Int64(offset)) + + +# the store / load cache policy of a hand-off in each space +_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The mono layer's launches tag their + hand-offs with their epoch (``layer.py``), so a pair from an earlier step + never carries a live tag. Every address is an ``Addr``: its space picks the + scope. + """ + + def __init__(self, tag): + self.tag = tag + + def put(self, a, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_bf(self, a, i, vs): + """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one + store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a.value), + i, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_words(self, a, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], + all of one space. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + spaces = {a.space for a, _, _ in specs} + assert len(spaces) == 1, f"one poll batch spans {spaces}" + cm = _SCOPE_CM[spaces.pop()] + bases = [a.value for a, _, _ in specs] + tag = self.tag + + def load_all(): + words = [] + for b, (_, i, n) in zip(bases, specs): + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +@traced +def publish(put, a, i, v, who): + """This CTA's plain stores made readable by every CTA, then (thread ``who``) + the mailbox pair (``a``, i) := v. The stores are drained first: a flag may + not overtake the data it announces, and a workgroup-scope release fence does + not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit + mode: the workgroup's waves share one CU, the readers here do not).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if who: + put(a, i, v) diff --git a/vllm/models/deepseek_v4/amd/mono/dims.py b/vllm/models/deepseek_v4/amd/mono/dims.py new file mode 100644 index 000000000000..92842056fac3 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/dims.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4-Pro's widths as the mono kernels read them, and ``Dims``, a TP +rank's share of them with every count the stages derive from it.""" + +from dataclasses import dataclass + +from vllm.models.deepseek_v4.amd.mono.common.plan import Shard, cdiv, padded + +HIDDEN = 7168 +Q_RANK = 1536 +N_HEADS = 128 +O_GROUPS = 16 +INDEX_TOPK = 1024 # a CSA token's top-k compressed keys +EXPERTS = 384 +TOPK = 6 +MOE_INTER = 3072 # a routed expert's intermediate width +SHARED_INTER = 3072 +ROUTED_SCALE = 2.5 +SWIGLU_LIMIT = 10.0 +NORM_EPS = 1e-6 # every RMSNorm's (and the mHC seams') +HC_SINKHORN = 20 +HC_EPS = 1e-6 +HC_POST_MULT = 2.0 +# a step's rows at most: 8 requests x (1 + 6 DSpark drafts) and spare +MAX_TOKENS = 64 +# the step width from which the all-reduces take two hops (``stages.allreduce``) +TWO_HOP_ROWS = 42 + +HC = 4 # the mHC seams' streams +HEAD_DIM = 512 +ROPE = 64 +NOPE = HEAD_DIM - ROPE +HALF = ROPE // 2 +KV_DIM = HEAD_DIM +O_RANK = 1024 +WINDOW = 128 +EXPERT_PAD = 128 # the loader pads a routed expert's rank share to this +ROWS = 16 # GEMV rows a task +TILE = 16 # tokens a GEMV pass (an MFMA's N) +HEAD_TILE = 16 # heads an attention MFMA takes +WOA_ROWS = 32 # a wo_a task's rows: one MX group of its output +UG_PART = 32 # intermediate columns an ug task part: one FP4 group +WQKV_KPARTS = 2 # wqkv's K split +QKV_ROWS = Q_RANK + KV_DIM +X_WORDS = HIDDEN // 4 # an MXFP8 hidden row's i32 words +X_GROUPS = HIDDEN // 32 # an MXFP8 hidden row's E8M0 groups +WQKV_TASKS = QKV_ROWS // ROWS +WOB_TASKS = HIDDEN // ROWS +DOWN_TASKS = HIDDEN // ROWS + + +def ld(k: int) -> int: + """The row stride of a row-major FP8 weight of K = ``k`` (bytes): vLLM pads + the rows of a multiple of 512 bytes by 256 (``VLLM_ROCM_FP8_PADDING``).""" + return k + 256 if k % 512 == 0 else k + + +@dataclass(frozen=True) +class Dims: + """A TP rank's share of the model's widths.""" + + tp: int + # a token's compressed keys at most (0: the CSA top-k): an HCA layer's + # attention over every compressed row of a long context reads more + topk: int = 0 + + def __post_init__(self): + shard = Shard(self.tp) + shard.split(N_HEADS, "attention heads") + shard.split(O_GROUPS, "output groups") + shard.split(MOE_INTER, "routed intermediate") + shard.split(SHARED_INTER, "shared intermediate") + + # ------------------------------------------------------------ attention + @property + def heads(self) -> int: + return N_HEADS // self.tp + + @property + def head_tiles(self) -> int: + return self.heads // HEAD_TILE + + @property + def groups(self) -> int: + return O_GROUPS // self.tp + + @property + def group_k(self) -> int: + """wo_a's K a group: its heads' outputs.""" + return self.heads // self.groups * HEAD_DIM + + @property + def o_rows(self) -> int: + return self.groups * O_RANK + + @property + def q_rows(self) -> int: + return self.heads * HEAD_DIM + + @property + def wqb_tasks(self) -> int: + return self.q_rows // ROWS + + @property + def woa_tasks(self) -> int: + return self.o_rows // WOA_ROWS + + @property + def topk_width(self) -> int: + """A token's compressed keys at most.""" + return self.topk or INDEX_TOPK + + @property + def keys(self) -> int: + """A token's attention keys at most: its top-k, then its window.""" + return self.topk_width + WINDOW + + # ------------------------------------------------------------ MoE + @property + def inter_real(self) -> int: + """A routed expert's real intermediate width on this rank.""" + return MOE_INTER // self.tp + + @property + def inter(self) -> int: + """... padded as the loader pads it (the rest zero weights).""" + return padded(self.inter_real, EXPERT_PAD) + + @property + def sh_inter(self) -> int: + """The shared expert's intermediate width on this rank.""" + return SHARED_INTER // self.tp + + @property + def shared_tasks(self) -> int: + return self.sh_inter // ROWS + + @property + def mid_words(self) -> int: + """MXFP8 words of a pick's routed intermediate.""" + return self.inter // 4 + + @property + def down_scale_cols(self) -> int: + """w2_s's e8m0 columns: inter / 32, padded to a multiple of 8 + (aiter ``shuffle_scale``).""" + return padded(self.inter // 32, 8) + + +def token_tiles(s): + """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its + weights once and runs every pass on them, an x tile in LDS at a time.""" + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + """The LDS rows an x tile holds.""" + return min(s, TILE) + + +__all__ = ["Dims", "cdiv", "ld", "tile_rows", "token_tiles"] diff --git a/vllm/models/deepseek_v4/amd/mono/layer.py b/vllm/models/deepseek_v4/amd/mono/layer.py new file mode 100644 index 000000000000..ed90fc4536e8 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/layer.py @@ -0,0 +1,927 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono decode layer's two launches (README): K1 = the attention seam +(``stages.seam``: slice, gate, the norm into vLLM's MXFP8) + the attention front +(``attention.front``); K2 = the attention back (``attention.back``, wo_b pushed +to every rank) + the FFN seam (the TP sum folded in) + the MoE (``stages.moe``, +its all-reduce in-kernel too). + +Tags: K1's hand-offs carry 2 e + 1, K2's 2 e + 2 (e: the launch pair's epoch, +moved on by K2's CTA 0 at its end), so K2's seam may reuse K1's seam regions. +Peer regions alternate by e's parity; every rank runs the same launches, so +the epochs agree across ranks. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Int32, Int64, T + +from .attention import back as mb_back +from .attention import front as mb_front +from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc +from .attention.plan import back_scratch, front_scratch +from .common.mx import clamp_fp8 +from .common.ops import butterfly, fp8_pack4, memrealtime, peer_bases, traced +from .common.plan import WAVES, first_task, key_tuple +from .common.sync import Mailbox, publish, sreg +from .dims import HIDDEN, INDEX_TOPK, MAX_TOKENS, X_GROUPS, X_WORDS, Dims +from .sources import SOURCES +from .stages import allreduce, moe, seam +from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape + +_STREAM = fx.Stream(None) +# K1's MXFP8 x and the MoE's: plain words in both, the same size, so the two +# launches share the MoE's (placed later); any other name is one region, so a +# word is never a tag to one launch and data to another +SHARED_REGIONS = ("x8", "x8s") +COUNTER_WORDS = 2 * 256 # the MoE's ug queue and down counts, a word a slot +COUNTERS = ("ugq", "dq") +# the epoch buffer, in words: the epoch, a mark a CTA (back.EPOCH_MARKS), then on +# lines of their own the MoE's counters +EPOCH_COUNTERS = -(-(mb_back.EPOCH_MARKS + BLOCKS) // 64) * 64 +EPOCH_WORDS = EPOCH_COUNTERS + COUNTER_WORDS +# a timeline build's K2 stamps a CTA: the attention back's, then the MoE's +K2_POINTS = mb_back.BACK_POINTS + moe.TL_POINTS + + +@dataclass(frozen=True) +class MonoBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (its attention's key sources) + timeline: bool = False + # K1's seam folds the previous layer's output into the residual; off for a + # residual that arrives settled (the first layer, after an Engram) + fold: bool = True + # a hash-routed MoE (``moe_shape.RouteShape.hash_route``) + hash_route: bool = False + # a token's compressed keys at most (0: the model's top-k; ``Dims.topk``) + topk: int = 0 + # the all-reduces in two hops (``stages.allreduce``) + two_hop: bool = False + # the key table's compressed keys are slots, not indices through the block + # table (``front.stage_kt``) + global_keys: bool = False + + +def _variant(key: MonoBuild) -> dict: + """A build's variant flags past the original's, for its kernel symbol.""" + out = {"h": 1} if key.hash_route else {} + if key.topk: + out["k"] = key.topk + if key.global_keys: + out["g"] = 1 + if key.two_hop: + out["a"] = 2 + return out + + +def _dims(key: MonoBuild) -> Dims: + return Dims(key.tp, key.topk) + + +def _moe_key(key: MonoBuild) -> MoeBuild: + return MoeBuild( + key.tokens, + key.tp, + timeline=key.timeline, + hash_route=key.hash_route, + two_hop=key.two_hop, + ) + + +def _front_key(key: MonoBuild): + return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline, key.topk) + + +def _back_key(key: MonoBuild): + return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline, key.topk) + + +def scratch_layout(s: int, tp: int, topk: int = 0) -> dict: + """Every region of step width s's launches -> (byte offset, bytes), disjoint: + K1's front and seam, K2's back, the MoE's regions, the normed rows and their + flags. Each width has a scratch of its own (``MonoRunner.scratch``): + under another width's layout a mailbox pair's tag word could be plain data, + which a poll may take for a current tag. The MoE's counters, whose slots + outlive a step width, are in the epoch buffer.""" + d = Dims(tp, topk) + out: dict = {} + off = 0 + + def place(regions): + nonlocal off + base = off + for name, (o, n) in regions.items(): + assert name not in out or name in SHARED_REGIONS, name + out[name] = (base + o, n) + off = max(off, base + max(o + n for o, n in regions.values())) + off = -(-off // 256) * 256 + + place(front_scratch(s, d)) + place(seam.scratch_layout(s)) + place(back_scratch(s, d, start=0)) + regions = moe.scratch_layout(MoeBuild(s, tp)).items() + mr = {n: v for n, v in regions if n not in COUNTERS} + lo = min(o for o, _ in mr.values()) + place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) + row = s * HIDDEN * 2 + place({"normed": (0, row), "xrdy_moe": (row + 256, s * 8)}) + place({"krdy": (0, s * 8)}) + return out + + +def scratch_bytes(s: int, tp: int, topk: int = 0) -> int: + return max(o + n for o, n in scratch_layout(s, tp, topk).values()) + + +def peer_half_bytes(tp: int) -> int: + """One parity's peer regions: the attention's partials, then the MoE's.""" + d, s = Dims(tp), MAX_TOKENS + return seam.attn_peer_bytes(s, d) + moe.peer_bytes(s, d) + + +def _epoch(epoch): + return fx.Int32(bo.buffer_load(rsrc(epoch), 0, vec_width=1, dtype=T.i32)) + + +# ---------------------------------------------------------------- K1 + + +@traced +def store_x8(c, t, col, ys, live): + """The attention seam's norm output (``seam.stage_norm``'s 8 columns at + ``col``) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the + front's X8 / X8S, plain at device scope.""" + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = moe.vllm_mx_code(butterfly(amax, (1, 2), fx.max)) + mul = ((254 - code) << 23).bitcast(fx.Float32) + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + if live: + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + col // 4, + cache_modifier=CM_DEV, + ) + if c["tid"] % 4 == 0: + bo.buffer_store( + code, + rsrc(c["x8s"]), + t * X_GROUPS + col // 32, + cache_modifier=CM_DEV, + ) + + +@traced +def publish_x8(c, t): + """Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] < 2: + c["mb"].put(c["xrdy"], 2 * t + c["tid"], fx.Int32(1)) + gpu.barrier() + + +def build_mono_k1(key: MonoBuild): + s, tp, d = key.tokens, key.tp, _dims(key) + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp, key.topk) + fkey = _front_key(key) + FrontLds = mb_front.front_lds(s, d) + n_slices = seam.SLICES + # a norm collapses with its token's gate: the gates first + GATE0 = n_slices + NORM0 = n_slices + s + # past the grid a CTA runs its stages in order: a gate never waits on a norm + assert n_slices <= BLOCKS + + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + + @fx.union + class K1Lds: + seam: SeamLds + front: FrontLds # type: ignore[valid-type] + + name = kernel_symbol( + "dsv4_mono_k1", + s=s, + tp=tp, + r=key.ratio, + tl=key.timeline, + f=key.fold, + **_variant(key), + ) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k1( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + # unread: the seam collapses with its own pre-mix; dropping it moves + # the kernels' register allocation + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + normed_out: Int64, + qr_out: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K1Lds) + sl = lds.seam.peek() + ep = _epoch(epoch) + tag = (ep << 1) + 1 + fargs = { + "wqkv": wqkv, + "wqkv_s": wqkv_s, + "qn_w": qn_w, + "kvn_w": kvn_w, + "wqb": wqb, + "wqb_s": wqb_s, + "cos_sin": cos_sin, + "pos": pos, + "slot": slot, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "q": q, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "topk": topk, + "t2r": t2r, + "kt": kt, + "klen": klen, + "qr_out": qr_out, + } + flayout = {n: layout[n] for n in front_scratch(s, d)} + c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + amb = Mailbox(tag) + tid = fx.thread_idx.x + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "args": { + "res_in": res_in, + "pend": pend, + "post_in": post_in, + "comb_in": comb_in, + "pre_in": pre_in, + "hc_fn": hc_fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "norm_w": attn_w, + "res_out": res_out, + "post_out": post_out, + "comb_out": comb_out, + "pre_out": pre_out, + }, + "d": d, + "fold": key.fold, + } + for region in seam.REGIONS: + ca[region] = sreg(scratch, layout[region][0], region) + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_front.FRONT_POINTS + + def seam_slice(): + for task in range(first_task(bid, 0), n_slices, BLOCKS): + seam.stage_slice(ca, task) + gpu.barrier() + + def store(*v): + store_x8(c, *v) + store_normed(normed_out, HIDDEN, *v) + + def norms(): + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, store) + publish_x8(c, t) + + def gates(): + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + + def seam_rest(): + gates() + norms() + gpu.barrier() + + mb_front.run_front(c, fkey, bid, before_wqkv=seam_slice, after_wqkv=seam_rest) + + @flyc.jit + def launch( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + normed_out: Int64, + qr_out: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k1( + res_in, + pend, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + attn_w, + res_out, + post_out, + comb_out, + pre_out, + wqkv, + wqkv_s, + qn_w, + kvn_w, + wqb, + wqb_s, + cos_sin, + pos, + slot, + swa, + swa_stride, + swa_block, + q, + swa_idx, + swa_lens, + comp_block, + comp_bt, + bt_stride, + topk, + t2r, + kt, + klen, + normed_out, + qr_out, + scratch, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- K2 + + +@traced +def store_normed(normed, hidden, t, col, ys, live): + """The FFN seam's norm output -> the MoE input row (bf16, device scope: + this launch reads it).""" + if live: + for j in range_constexpr(8): + bo.buffer_store( + ys[j].to(fx.BFloat16), + rsrc(normed), + t * hidden + col + j, + cache_modifier=CM_DEV, + ) + + +def _seam_lds(s): + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + pl: fx.Array[fx.Float32, s * seam.COLS, 16] + + return SeamLds + + +@traced +def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep): + """K2's FFN half on this CTA: the FFN seam (the attention's TP partials + summed and folded into the residual), the MoE and its all-reduce into + ``a["out"]``, then the launch pair's epoch moved on.""" + s, tp, d = key.tokens, key.tp, _dims(key) + layout = scratch_layout(s, tp, key.topk) + mkey = _moe_key(key) + rs = route_shape(mkey) + n_slices = seam.SLICES + GATE0 = n_slices + NORM0 = n_slices + s + bid = fx.block_idx.x + tid = fx.thread_idx.x + sl = lds.seam.peek() + scratch = a["scratch"] + seam_args = ( + "res_in", + "post_in", + "comb_in", + "pre_in", + "hc_fn", + "hc_scale", + "hc_base", + "res_out", + "post_out", + "comb_out", + "pre_out", + ) + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "pend_lds": sl.pl.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": own, + "args": {n: a[n] for n in seam_args} | {"norm_w": a["ffn_w"]}, + "d": d, + "two_hop": key.two_hop, + } + normed = scratch + fx.Int64(layout["normed"][0]) + for region in seam.REGIONS: + ca[region] = sreg(scratch, layout[region][0], region) + mlayout = {n: layout[n] for n in moe.scratch_layout(mkey) if n not in COUNTERS} + mlayout["xrdy"] = layout["xrdy_moe"] + margs = moe.moe_args( + normed, + a["gate_w"], + a["bias"], + a["w13"], + a["w13_s"], + a["w2"], + a["w2_s"], + a["sgu"], + a["sgu_s"], + a["sw2"], + a["sw2_s"], + a["out"], + ) + cb = moe.moe_context( + s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own + ) + cb["x_cm"], cb["x_ready"] = CM_DEV, True + cb["two_hop"] = key.two_hop + cb["tag"] = ep & 255 + for i, region in enumerate(COUNTERS): + cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region) + for task in range(first_task(bid, 0), n_slices, BLOCKS): + seam.stage_reduce(ca, task) + seam.stage_slice(ca, task) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_normed(normed, HIDDEN, *v)) + publish(cb["put"], cb["xrdy"], t, 1, tid == 0) + gpu.barrier() + + # ---- the MoE, its all-reduce into ``out`` + if const_expr(key.timeline): + # the seam's end, then the MoE's stages' (stamped in LDS), after the + # attention back's points + tl = a["tl"] + tls = lds.moe.peek().tls.ptr + if tid == 0: + fx.ptr_store(memrealtime(), tls + 0) + moe.run_moe(cb, mkey, bid, tls, 0) + if tid < moe.TL_POINTS: + t64 = fx.Int64(fx.ptr_load(tls + tid)) + w = fx.Vector.from_elements([fx.Int32(t64), fx.Int32(t64 >> 32)], fx.Int32) + bo.buffer_store( + w, rsrc(tl), (bid * K2_POINTS + mb_back.BACK_POINTS + tid) * 2 + ) + else: + moe.run_moe(cb, mkey, bid, 0, 0) + + def reset(): + # the MoE counter slot 128 launch pairs ahead: its last use long + # done, its next use far off + slot = (ep + 128) & 255 + for region in COUNTERS: + gstore(cb[region].value + fx.Int64(slot * 4), fx.Int32(0), words=1) + + mb_back.epoch_end({"bid": bid, "tid": tid}, a["epoch"], ep, reset) + + +def build_mono_k2(key: MonoBuild): + s, tp, d = key.tokens, key.tp, _dims(key) + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp, key.topk) + bkey = _back_key(key) + mkey = _moe_key(key) + rs = route_shape(mkey) + assert rs.experts // 64 in SORT_NETS + BackM = mb_back.back_lds_members(s, d) + SplitLds, GemvLds = BackM["split"], BackM["gemv"] + MoeLds = moe.moe_smem(s, rs, key.timeline) + half = peer_half_bytes(tp) + assert seam.SLICES <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class K2Lds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol( + "dsv4_mono_k2", + s=s, + tp=tp, + r=key.ratio, + tl=key.timeline, + **_variant(key), + ) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k2( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, # unread, as K1's + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K2Lds) + ep = _epoch(epoch) + tag = (ep << 1) + 2 + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + own = sym + par + bargs = { + "q": q, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "comp": comp, + "comp_stride": comp_stride, + "comp_block": comp_block, + "kt": kt, + "klen": klen, + "pos": pos, + "sink": sink, + "qk_scale": qk_scale, + "cos_sin": cos_sin, + "woa": woa, + "woa_s": woa_s, + "wob": wob, + "wob_s": wob_s, + "out": 0, + "out_stride": 0, + "zrec": zrec, + } + blayout = {n: layout[n] for n in back_scratch(s, d, start=0)} + views = {"split": lds.split.peek(), "gemv": lds.gemv.peek()} + c = mb_back.back_context(bkey, views, bargs, scratch, tag, blayout) + # a key table wider than a token's top-k: chunks past a token's keys + # skipped from its position + c["bounded"] = d.topk_width > INDEX_TOPK + amb = Mailbox(tag) + + def push(t, col, v0, v1): + # this rank's bf16 partial pair to the ATTN region of every rank, + # or of the column's owner (``allreduce``) + allreduce.push( + lambda dst, pair, vals: amb.put_bf(dst, 2 * pair, vals), + lambda p: bases[p], + 0, + "attn", + rank, + t, + s, + col, + [v0, v1], + d, + key.two_hop, + ) + + c["wob_out"] = push + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, K2_POINTS + mb_back.epoch_begin(c, epoch, ep) + # the key table, from the top-k vLLM's indexer selected after K1, + # on the last CTAs: each split task waits on its token's + ckt = { + "S": s, + "d": d, + "tid": c["tid"], + "dev": True, + "ratio": key.ratio, + "global_keys": key.global_keys, + "args": { + "slot": slot, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "pos": pos, + "t2r": t2r, + "topk": topk, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "kt": kt, + "klen": klen, + }, + } + c["krdy"] = sreg(scratch, layout["krdy"][0], "krdy") + for t in range(first_task(bid, BLOCKS - s), s, BLOCKS): + mb_front.stage_kt(ckt, t) + mb_back.publish(c, c["krdy"], t) + mb_back.run_back(c, bkey, bid) + gpu.barrier() + + ffn = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + tl=tl, + ) + run_ffn(key, lds, ffn, amb, bases, own, rank, ep) + + @flyc.jit + def launch( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k2( + q, + swa, + swa_stride, + swa_block, + comp, + comp_stride, + comp_block, + kt, + klen, + slot, + swa_idx, + swa_lens, + t2r, + topk, + comp_bt, + bt_stride, + pos, + sink, + qk_scale, + cos_sin, + woa, + woa_s, + wob, + wob_s, + zrec, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch diff --git a/vllm/models/deepseek_v4/amd/mono/runner.py b/vllm/models/deepseek_v4/amd/mono/runner.py new file mode 100644 index 000000000000..47bea49be826 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/runner.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4-Pro mono decode layer: host side. + +One ``MonoRunner`` a TP rank runs every mono layer of a decode step (at most +``MAX_TOKENS`` rows): two persistent launches a layer (``layer``) from the +layer's inputs at the attention seam to its outputs at the next one -- the +MoE's reduced output, the residual after the FFN seam and that seam's mixes -- +with both TP all-reduces inside the kernels (symmetric peer memory). Every +argument is a device pointer: the launches can be captured in a HIP graph. The +kernels move their mailbox epoch on themselves; every rank runs the same +launches, so the ranks' epochs agree. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +import torch + +from .common.plan import BLOCKS +from .dims import ( + EXPERTS, + HC, + HEAD_DIM, + HIDDEN, + INDEX_TOPK, + MAX_TOKENS, + Q_RANK, + QKV_ROWS, + TWO_HOP_ROWS, + Dims, + ld, +) +from .layer import ( + EPOCH_WORDS, + MonoBuild, + build_mono_k1, + build_mono_k2, + peer_half_bytes, + scratch_bytes, +) + +__all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "MonoLayerWeights", "MonoRunner"] + +_LOG2E = 1.4426950408889634 + + +def _ensure_writable_flydsl_cache() -> None: + """Aiter points FlyDSL at its bundled, read-only cache; new kernels need a + writable one.""" + cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR") + if cur and os.access(cur, os.W_OK): + return + path = Path.home() / ".flydsl" / "cache" + path.mkdir(parents=True, exist_ok=True) + os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path) + + +def _check_tensors(owner, tag: str, want: dict, lds: dict | None = None) -> None: + """Each named tensor of ``owner`` has the (shape, dtype) the kernels read, + contiguous -- or a 2-D one's rows ``lds[name]`` apart: their buffer loads + are unbounded, so a mismatch would read garbage rather than fault.""" + for name, (shape, dtype) in want.items(): + t = getattr(owner, name) + ld = (lds or {}).get(name) + laid = t.is_contiguous() if ld is None else t.stride() == (ld, 1) + assert tuple(t.shape) == shape and t.dtype == dtype and laid, ( + f"{tag} {name}: {tuple(t.shape)} {t.dtype} {t.stride()}, want {shape} " + f"{dtype}" + ("" if ld is None else f" rows {ld} apart") + ) + + +@dataclass +class AttnWeights: + """One layer's attention tensors on this rank, in vLLM's loaded layout.""" + + layer_id: int + # [2048, 7168] e4m3: [wq_a; wkv], aiter (16, 16) shuffle + wqkv: torch.Tensor + wqkv_scale: torch.Tensor # [64, 224] uint8 E8M0 (32 x 32 blocks) + q_norm: torch.Tensor # [1536] bf16 + kv_norm: torch.Tensor # [512] bf16 + wq_b: torch.Tensor # [H * 512, 1536] e4m3, padded rows (``dims.ld``) + wq_b_scale: torch.Tensor # [H * 16, 48] uint8 + wo_a: torch.Tensor # [G * 1024, group_k] e4m3, padded rows + wo_a_scale: torch.Tensor # [G * 32, group_k / 32] uint8 + wo_b: torch.Tensor # [7168, G * 1024] e4m3, aiter (16, 16) shuffle + wo_b_scale: torch.Tensor # [hidden / 32, G * 32] uint8 + attn_sink: torch.Tensor # [>= H] f32 + cos_sin: torch.Tensor # [positions, 64] f32 (the layer's rope table) + ratio: int # the layer's compress ratio + + def check(self, d: Dims) -> None: + h, g, k, q = d.heads, d.groups, HIDDEN, Q_RANK + want = { + "wqkv": ((QKV_ROWS, k), torch.float8_e4m3fn), + "wqkv_scale": ((QKV_ROWS // 32, k // 32), torch.uint8), + "wq_b": ((h * HEAD_DIM, q), torch.float8_e4m3fn), + "wq_b_scale": ((h * HEAD_DIM // 32, q // 32), torch.uint8), + "wo_a": ((g * 1024, d.group_k), torch.float8_e4m3fn), + "wo_a_scale": ((g * 32, d.group_k // 32), torch.uint8), + "wo_b": ((k, g * 1024), torch.float8_e4m3fn), + "wo_b_scale": ((k // 32, g * 32), torch.uint8), + } + lds = {"wq_b": ld(q), "wo_a": ld(d.group_k)} + _check_tensors(self, f"layer {self.layer_id}", want, lds) + assert self.cos_sin.dtype == torch.float32 and self.cos_sin.shape[-1] == 64 + + +@dataclass +class MonoLayerWeights: + """One layer's tensors on this rank, in vLLM's loaded layout.""" + + attn: AttnWeights + hc_attn_fn: torch.Tensor # [24, 4 * hidden] f32 + hc_attn_scale: torch.Tensor # [3] f32 + hc_attn_base: torch.Tensor # [24] f32 + attn_norm: torch.Tensor # [hidden] bf16 + hc_ffn_fn: torch.Tensor + hc_ffn_scale: torch.Tensor + hc_ffn_base: torch.Tensor + ffn_norm: torch.Tensor + gate_w: torch.Tensor # [experts, hidden] bf16 + # [experts] f32 (e_score_correction_bias); None for a hash layer, whose + # tokens' experts each step gives (``route_ids``) + bias: torch.Tensor | None + w13: torch.Tensor # [experts, 2 inter, hidden / 2] fp4x2, aiter (16, 16) shuffle + w13_s: torch.Tensor # its scales, aiter shuffle_scale order + w2: torch.Tensor # [experts, hidden, inter / 2] fp4x2 + w2_s: torch.Tensor + # shared gate_up [2 inter, hidden] e4m3, aiter (16, 16) shuffle + sgu: torch.Tensor + sgu_s: torch.Tensor # [2 inter / 32, hidden / 32] E8M0 + sw2: torch.Tensor # shared down [hidden, inter] e4m3 + sw2_s: torch.Tensor + + def check(self, d: Dims) -> None: + """The MoE's tensors: every routed expert, TP-sharded intermediates (the + routed one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 + E8M0 blocks for the shared expert.""" + fp4, e4m3, u8 = torch.float4_e2m1fn_x2, torch.float8_e4m3fn, torch.uint8 + e, k = EXPERTS, HIDDEN + want = { + "gate_w": ((e, k), torch.bfloat16), + "w13": ((e, 2 * d.inter, k // 2), fp4), + "w13_s": ((e * 2 * d.inter, k // 32), u8), + "w2": ((e, k, d.inter // 2), fp4), + "w2_s": ((e * k, d.down_scale_cols), u8), + "sgu": ((2 * d.sh_inter, k), e4m3), + "sgu_s": ((2 * d.sh_inter // 32, k // 32), u8), + "sw2": ((k, d.sh_inter), e4m3), + "sw2_s": ((k // 32, d.sh_inter // 32), u8), + } + if self.bias is not None: + want["bias"] = ((e,), torch.float32) + _check_tensors(self, "MoE", want) + + +class MonoRunner: + """The mono decode runner of one TP rank (``group``: its TP group, for the + peer-memory handle exchange; a gloo / CPU group).""" + + def __init__( + self, + tp: int, + rank: int, + group, + device: torch.device | str = "cuda", + max_topk: int = 0, + timeline: bool = False, + ): + """``max_topk``: the widest key table a layer's build takes (``Dims.topk``), + past the CSA top-k. ``timeline``: the kernels' stage stamps into + ``tl`` (K1's, K2's: [BLOCKS, points] of the device clock), a profile.""" + _ensure_writable_flydsl_cache() + from .common.peer_memory import PeerBuffer + + self.tp, self.rank = tp, rank + self.d = d = Dims(tp) + keys = Dims(tp, max(max_topk, INDEX_TOPK)).keys + rows = MAX_TOKENS + dev = self.device = torch.device(device) + self._scratch: dict[tuple[int, int], torch.Tensor] = {} + # [epoch, -, -, -, a mark per CTA, the MoE's counters] + self.epoch = torch.zeros(EPOCH_WORDS, dtype=torch.int32, device=dev) + self.q = torch.zeros(rows, d.heads, HEAD_DIM, dtype=torch.bfloat16, device=dev) + self.kt = torch.zeros(rows * keys, dtype=torch.int32, device=dev) + self.klen = torch.zeros(2 * rows, dtype=torch.int32, device=dev) + self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) + self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # an index source's K1 outputs, vLLM's indexer's inputs + self.normed = torch.zeros(rows, HIDDEN, dtype=torch.bfloat16, device=dev) + self.qr = torch.zeros(rows, Q_RANK, dtype=torch.bfloat16, device=dev) + # the attention seam's outputs, K1 -> K2 (its pre-mix unread: the FFN + # seam collapses with its own) + self.res_mid = torch.zeros(rows, HC, HIDDEN, dtype=torch.bfloat16, device=dev) + self.post_a = torch.zeros(rows, HC, dtype=torch.float32, device=dev) + self.comb_a = torch.zeros(rows, HC, HC, dtype=torch.float32, device=dev) + self.pre_a = torch.zeros(rows, HC, dtype=torch.float32, device=dev) + self.peer = PeerBuffer(2 * peer_half_bytes(tp), group, rank, tp, dev) + self.peer.bytes.zero_() + self._qk_scale = ( + torch.tensor([HEAD_DIM**-0.5 * _LOG2E], dtype=torch.float32) + .view(torch.int32) + .item() + ) + self._kernels: dict = {} + self.timeline = timeline + self.tl = ( + torch.zeros(2, BLOCKS, 64, dtype=torch.int64, device=dev) + if timeline + else None + ) + + def scratch(self, tokens: int, topk: int = 0) -> torch.Tensor: + """Step width ``tokens``'s scratch (of the builds of key table width + ``topk``): a layout has memory of its own (``layer.scratch_layout``). + Allocated at the width's first step, which is eager: vLLM runs every + graph's batch eagerly before capturing it.""" + buf = self._scratch.get((tokens, topk)) + if buf is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"Mono decode: step width {tokens} first reached inside " + "a CUDA graph capture" + ) + buf = torch.zeros( + scratch_bytes(tokens, self.tp, topk), + dtype=torch.uint8, + device=self.device, + ) + self._scratch[(tokens, topk)] = buf + return buf + + def supports(self, tokens: int) -> bool: + return 1 <= tokens <= MAX_TOKENS + + @staticmethod + def two_hop(tokens: int) -> bool: + """A step width's all-reduces in two hops (``stages.allreduce``).""" + return tokens >= TWO_HOP_ROWS + + def kernels( + self, + tokens: int, + ratio: int, + fold: bool = True, + hash_route: bool = False, + topk: int = 0, + global_keys: bool = False, + ): + key = (tokens, ratio, fold, hash_route, topk, global_keys) + if key not in self._kernels: + b = MonoBuild( + tokens, + self.tp, + ratio, + fold=fold, + hash_route=hash_route, + topk=topk, + global_keys=global_keys, + timeline=self.timeline, + two_hop=self.two_hop(tokens), + ) + self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) + return self._kernels[key] + + @staticmethod + def _route(w: MonoLayerWeights, route_ids: torch.Tensor | None) -> torch.Tensor: + """The MoE's bias argument: the layer's bias, or a hash layer's step + experts ([M, topk] int32).""" + if route_ids is None: + assert w.bias is not None + return w.bias + assert route_ids.dtype == torch.int32 and route_ids.is_contiguous() + return route_ids + + def _outs(self, M: int, residual: torch.Tensor) -> tuple: + dev = residual.device + return ( + torch.empty(M, HIDDEN, dtype=torch.bfloat16, device=dev), + torch.empty_like(residual), + torch.empty(M, HC, 1, dtype=torch.float32, device=dev), + torch.empty(M, HC, HC, dtype=torch.float32, device=dev), + torch.empty(M, HC, dtype=torch.float32, device=dev), + ) + + def forward( + self, + w: MonoLayerWeights, + x: torch.Tensor, # [M, hidden] bf16: the previous FFN output (reduced) + residual: torch.Tensor, # [M, 4, hidden] bf16 + post_mix: torch.Tensor, # [M, 4, 1] f32 + res_mix: torch.Tensor, # [M, 4, 4] f32 + positions: torch.Tensor, # [M] int64 + slot_mapping: torch.Tensor, # [M] int64 + swa_cache: torch.Tensor, # [blocks, block, 584] uint8 + swa_indices: torch.Tensor, # [>= M, 1, 128] int32 + swa_lens: torch.Tensor, # [>= M] int32 + token_to_req: torch.Tensor, # [>= M] int32 + topk_indices: torch.Tensor, + comp_cache: torch.Tensor, # [blocks, block / ratio, 584] uint8 + select: Callable[[torch.Tensor, torch.Tensor], None], + comp_block_table: torch.Tensor | None = None, + outs: tuple | None = None, + fold: bool = True, + route_ids: torch.Tensor | None = None, + topk: int = 0, + global_keys: bool = False, + ): + """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs + (``outs``: buffers to write them into). ``select(normed, qr)``: vLLM's + compressor and, on a CSA layer, its top-k selection into + ``topk_indices``, run between the kernels on K1's normed input and qr. + ``route_ids``: a hash layer's experts of the step's tokens. ``topk``: + the build's key table width (``Dims.topk``); ``global_keys``: + ``topk_indices`` holds compressed slots (a row of its own width), not + indices through ``comp_block_table``.""" + M = x.shape[0] + assert self.supports(M), M + a = w.attn + ratio = a.ratio + assert ratio > 0 + k1, k2 = self.kernels(M, ratio, fold, route_ids is not None, topk, global_keys) + scratch = self.scratch(M, topk) + st = torch.cuda.current_stream() + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + for t_ in (x, residual, post_mix, res_mix, positions, slot_mapping): + assert t_.is_contiguous() + assert positions.dtype == slot_mapping.dtype == torch.int64 + assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32 + assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8 + assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584 + swa_block = swa_cache.shape[1] + assert topk_indices.dtype == torch.int32 + assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584 + comp, comp_stride, comp_block = ( + comp_cache, + comp_cache.stride(0), + comp_cache.shape[1], + ) + if global_keys: + # the table's own width where the block table's stride was + width = topk_indices.stride(0) + assert width <= Dims(self.tp, topk).topk_width + bt, bt_stride, topk_t = self._dummy, width, topk_indices + else: + assert comp_block_table is not None + assert topk_indices.stride(0) == self.d.topk_width + assert comp_block_table.dtype == torch.int32 + bt, bt_stride, topk_t = ( + comp_block_table, + comp_block_table.stride(0), + topk_indices, + ) + k1( + residual.data_ptr(), + x.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + self._dummy.data_ptr(), # pre_in, unread + w.hc_attn_fn.data_ptr(), + w.hc_attn_scale.data_ptr(), + w.hc_attn_base.data_ptr(), + w.attn_norm.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + a.wqkv.data_ptr(), + a.wqkv_scale.data_ptr(), + a.q_norm.data_ptr(), + a.kv_norm.data_ptr(), + a.wq_b.data_ptr(), + a.wq_b_scale.data_ptr(), + a.cos_sin.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + self.q.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + comp_block, + bt.data_ptr(), + bt_stride, + topk_t.data_ptr(), + token_to_req.data_ptr(), + self.kt.data_ptr(), + self.klen.data_ptr(), + self.normed.data_ptr(), + self.qr.data_ptr(), + scratch.data_ptr(), + self.epoch.data_ptr(), + self.tl[0].data_ptr() if self.tl is not None else 0, + stream=st, + ) + select(self.normed[:M], self.qr[:M]) + k2( + self.q.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + comp.data_ptr(), + comp_stride, + comp_block, + self.kt.data_ptr(), + self.klen.data_ptr(), + slot_mapping.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + token_to_req.data_ptr(), + topk_t.data_ptr(), + bt.data_ptr(), + bt_stride, + positions.data_ptr(), + a.attn_sink.data_ptr(), + self._qk_scale, + a.cos_sin.data_ptr(), + a.wo_a.data_ptr(), + a.wo_a_scale.data_ptr(), + a.wo_b.data_ptr(), + a.wo_b_scale.data_ptr(), + self._zero_rec.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self._dummy.data_ptr(), # pre_in, unread + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + self._route(w, route_ids).data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + scratch.data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + self.tl[1].data_ptr() if self.tl is not None else 0, + stream=st, + ) + return outs diff --git a/vllm/models/deepseek_v4/amd/mono/sources.py b/vllm/models/deepseek_v4/amd/mono/sources.py new file mode 100644 index 000000000000..f2a5b6940851 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/sources.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's kernel sources, in every build's JIT cache key +(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" + +from vllm.models.deepseek_v4.amd.mono.common.plan import source_digest + +SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v4/amd/mono/stages/__init__.py b/vllm/models/deepseek_v4/amd/mono/stages/__init__.py new file mode 100644 index 000000000000..c9926b5adf4a --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The stages the mono layer adapts from ROCm/ATOM: the mHC seams (``seam``), +the MoE (``moe``, ``router``) and the GEMV / LDS helpers they share (``gemv``); +the in-kernel all-reduce of both (``allreduce``).""" diff --git a/vllm/models/deepseek_v4/amd/mono/stages/allreduce.py b/vllm/models/deepseek_v4/amd/mono/stages/allreduce.py new file mode 100644 index 000000000000..030e35b07897 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/allreduce.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's in-kernel all-reduce of a bf16 [S, hidden] partial over the +TP ranks (the attention output, the MoE output), into a peer region of each +rank's symmetric buffer, one of two ways: + + one hop every rank pushes its partial to every rank (PART: [source rank] + [token][hidden / 2] bf16 pairs); a reader sums the ranks' pairs + in rank order + two hops a rank pushes each column to its owner (``owner``: hidden / tp + columns a rank); the owner sums them in rank order and pushes + the sums to every other rank (SUM: [token][hidden / 2] pairs), + where the readers take them + +Two hops carry 2 / tp of one hop's bytes a rank (a quarter at TP8) for one more +crossing: the wide steps of a wide TP group take them (``MonoBuild.two_hop``). +Both sum in rank order in fp32, so they agree bit for bit. Every value travels +as a tagged pair (``common.sync``).""" + +from flydsl.expr import range_constexpr + +from vllm.models.deepseek_v4.amd.mono.common.ops import ( + bf16_round, + bf_hi, + bf_lo, + sum_partials, + traced, +) +from vllm.models.deepseek_v4.amd.mono.common.sync import preg +from vllm.models.deepseek_v4.amd.mono.dims import HIDDEN, Dims + +PAIR_BYTES = 8 + + +def region_bytes(tokens: int, d: Dims) -> int: + """A reduction's peer region: PART, then SUM (two hops).""" + return (d.tp + 1) * tokens * HIDDEN // 2 * PAIR_BYTES + + +def part_pair(d: Dims, src, t, s, col): + """PART's pair of rank ``src``, token t, columns col, col + 1.""" + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +def sum_pair(d: Dims, t, s, col): + """SUM's pair of token t, columns col, col + 1 (past PART).""" + return (d.tp * s + t) * (HIDDEN // 2) + col // 2 + + +def owner(d: Dims, col): + """The rank that sums column ``col`` on two hops.""" + return col // (HIDDEN // d.tp) + + +def peer_at(bases, p): + """Peer ``p``'s base (``p`` traced): the select of the peers' bases.""" + base = bases[0] + for q in range_constexpr(1, len(bases)): + base = (p == q).select(bases[q], base) + return base + + +def push(put, addr, off, name, rank, t, s, col, vals, d: Dims, two_hop: bool): + """This rank's partial pair at (t, col) -- ``vals``, ``put(region, pair, + vals)`` -- to every rank's PART (one hop) or to the column's owner's (two + hops); ``addr(p)``: peer p's base, ``off``: the region's offset in it.""" + if two_hop: + bases = [addr(p) for p in range(d.tp)] + dst = preg(peer_at(bases, owner(d, col)), off, name) + put(dst, part_pair(d, rank, t, s, col), vals) + else: + for p in range_constexpr(d.tp): + dst = preg(addr(p), off, name) + put(dst, part_pair(d, rank, t, s, col), vals) + + +def reduce_pair(c, own, t, s, col, d: Dims): + """One hop: the ranks' PART pairs at (t, col) summed in rank order, fp32 + (lo, hi). Two hops, the owner's sum (``reduce_owned``).""" + return sum_partials(c["poll"], own, lambda src: part_pair(d, src, t, s, col), d.tp) + + +def reduce_owned(c, own, off, name, t, s, col, d: Dims): + """Two hops, on the column's owner: the ranks' sum (``reduce_pair``), its + bf16 pushed to every other rank's SUM -> (lo, hi).""" + acc0, acc1 = reduce_pair(c, own, t, s, col, d) + lo, hi = bf16_round(acc0), bf16_round(acc1) + at = 2 * sum_pair(d, t, s, col) + for p in range_constexpr(d.tp): + _put_other(c, p, preg(c["peer_addr"](p), off, name), at, [lo, hi]) + return acc0, acc1 + + +@traced +def _put_other(c, p, region, at, vals): + if p != c["rank"]: + c["put_bf"](region, at, vals) + + +def reduce_gathered(c, own, t, s, col, d: Dims): + """Two hops, off the column's owner: its sum from SUM -> (lo, hi).""" + w = c["poll"]([(own, sum_pair(d, t, s, col), 1)])[0][0] + return bf_lo(w), bf_hi(w) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/gemv.py b/vllm/models/deepseek_v4/amd/mono/stages/gemv.py new file mode 100644 index 000000000000..8913cf55d1cf --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/gemv.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (its GEMV and LDS helpers). +"""The FP8 GEMV tiles of the seam and MoE stages and the helpers they share: +token tiles, LDS rows of MXFP8 words, scalar loads and the mailbox-to-LDS +copy.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.mx import ( + UNIT_SCALE, + mfma_scaled, +) +from vllm.models.deepseek_v4.amd.mono.common.ops import ( + CM_NT, + fresh, + rsrc, + traced, +) +from vllm.models.deepseek_v4.amd.mono.common.plan import WAVES +from vllm.models.deepseek_v4.amd.mono.common.sync import POLL_MAX + +ROWS = 16 # GEMV rows a task +TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def token_tiles(s): + """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its + weights once and runs every pass on them, an x tile in LDS at a time.""" + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + """The LDS rows an x tile holds.""" + return min(s, TILE) + + +# a token's LDS row of MXFP8 words / E8M0 codes, 16 of which a MFMA B operand +# gathers, is padded by 16 B: at a multiple of 64 words they share a bank group +LDS_PAD = 4 + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def lds_tail(k): + """Words a K = ``k`` row's last K step reads past its ``lds_row`` (selected + out): the last row's buffer must hold them.""" + return max(0, 32 * -(-k // 128) - lds_row(k // 4)) + + +def plus(t0, v): + """Token ``v`` of the tile from ``t0`` (a first tile adds nothing; ``t0`` + traced: a tile a CTA picks at run time).""" + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +# ---------------------------------------------------------------- GEMV +def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False): + """K step st's global operands of rows 16 rg ..: the two 64-column weight + chunks (the upper clamped into the row) and the raw scale word. The (16, 16) + shuffle lays a row group out in 32-column blocks of 16 rows, a chunk two of + them (lanes 0-31, 32-63); ``k`` % 64 = 32 (a last quarter step): the lanes + of a block past the row read its last block. ``row_major`` (vLLM's dense + layout, k % 64 == 0): the same registers from a row-major [N, K] weight -- + lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64.""" + lane = c["lane"] + j = lane // 16 + kg = k // 32 + if const_expr(row_major): + assert k % 64 == 0 + row = rg * 16 + lane % 16 + lo_at = row * (k // 4) + st * 32 + j * 4 + hi_at = row * (k // 4) + fx.min(2 * st + 1, k // 64 - 1) * 16 + j * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + if const_expr(k % 64 == 0): + lo_at = (rg * (k // 64) + 2 * st) * 256 + lane * 4 + else: + lo_at = rg * (k * 4) + fx.min(4 * st + lane // 32, kg - 1) * 128 + lane % 32 * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + if const_expr(k % 64 == 0): + hi_at = (rg * (k // 64) + fx.min(2 * st + 1, k // 64 - 1)) * 256 + lane * 4 + else: + hi_at = rg * (k * 4) + fx.min(2 * st + 1, k // 64 - 1) * 256 + lane * 4 + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def _fp8_step_mfma(c, k, xl, xsl, st, ops, acc, live=None, rows=None): + """``acc`` through K step st's scaled MFMA: ``_fp8_step_load``'s operands + against every token's MXFP8 row in LDS. ``live`` false: a step past the end + (its operands those of the last step), zero at unit scale. ``k`` % 64 = 32: + a lane's lower block past the row is zero on both sides.""" + s, lane = c["S"] if rows is None else rows, c["lane"] + lo, hi, sa = ops + kg = k // 32 + x_row, s_row = lds_row(k // 4), lds_row(kg) + j = lane // 16 + col = fx.min(lane % 16, s - 1) + if const_expr(live is not None): + st = fx.min(st, (k + 127) // 128 - 1) + lo = [live.select(lo[d], fx.Int32(0)) for d in range(4)] + past = (2 * st + 1) * 64 >= k + if const_expr(live is not None): + past = past | ~live + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [past.select(fx.Int32(0), hi[d]) for d in range(4)], + fx.Int32, + ) + grp = st * 4 + j + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + if const_expr(k % 64 != 0): + over = 4 * st + j // 2 >= kg + av = fx.Vector.from_elements( + [over.select(fx.Int32(0), av[d]) for d in range(4)] + + [av[4 + d] for d in range(4)], + fx.Int32, + ) + xlo = [over.select(fx.Int32(0), xlo[d]) for d in range(4)] + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + + [past.select(fx.Int32(0), xhi[d]) for d in range(4)], + fx.Int32, + ) + sb = fx.ptr_load(xsl + (col * s_row + fx.min(grp, kg - 1))) + unit = grp >= kg + if const_expr(live is not None): + unit = unit | ~live + sa = unit.select(fx.Int32(UNIT_SCALE), sa) + sb = unit.select(fx.Int32(UNIT_SCALE), sb) + return mfma_scaled(av, xv, acc, sa, sb) + + +def _gemv_steps(c, k, split, tiled=False): + """This wave's K steps (a step past the end marked, ``split`` waves past + split repeat the last one's; ``tiled``: wave w takes part w % split of tile + w / split): [(step, live)].""" + wave = c["wave"] + steps = (k + 127) // 128 + if const_expr(split is None): + per = (steps + WAVES - 1) // WAVES + ragged = steps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < steps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and steps % split == 0 + per = steps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def gemv_fp8_loads(c, w, ws, k, rg, split=None, tiled=False, row_major=False): + """The weight operands of ``gemv_fp8``'s rows for this wave, every one in + flight at once (a step's latency each would put them in series): loads + only. A step past the end loads the last step. ``tiled`` (split x tiles = + WAVES): row groups rg .., one a group of ``split`` waves.""" + steps = (k + 127) // 128 + r_w, r_ws = rsrc(w), rsrc(ws) + if const_expr(tiled): + assert WAVES % split == 0 + rg = rg + c["wave"] // split + return [ + _fp8_step_load( + c, + r_w, + r_ws, + k, + rg, + st if live is None else fx.min(st, steps - 1), + row_major, + ) + for st, live in _gemv_steps(c, k, split, tiled) + ] + + +def gemv_fp8_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=None): + """``gemv_fp8_loads``' operands against the rows in LDS (``rows``: a token + tile's, every token's by default) -> + ``red[wave][lane]``; the scaled 16x16x128 MFMA, chunk kc in dwords 0-3 and kc + + 1 in 4-7. ``k`` a multiple of 32: a last half step has its upper chunk zeroed + on both sides at unit scale (a NaN byte past the row must not reach the + MFMA), its loads clamped into the row: the last row group's would read past + the tensor; a last quarter step its lower chunk's upper block too. A step + past the end adds zero at unit scale. + + The K order is the original kernel's, which ``row_sum`` completes: + wave w takes steps w, w + 8, ... (``split`` None: aiter's preshuffled group32 + GEMM, P0.5 microbenchmark 1), or waves 0 .. split - 1 one contiguous K range + each, the other waves zero (aiter's split-K bmm, whose last split sums the + partials in split order).""" + lane, wave = c["lane"], c["wave"] + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_gemv_steps(c, k, split, tiled)): + acc = _fp8_step_mfma(c, k, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +# ---------------------------------------------------------------- mailbox -> LDS +@traced +def poll_copy(c, idx, width, spans, live=None): + """Mailbox spans [(region, pair_of, n, dst, row)] into LDS: element u (pair + ``pair_of(u)``) to dst + u, or ``lds_at(u, row)`` for a span of rows + ``row`` words wide. Thread ``idx`` of ``width`` takes u = idx, idx + + width, ..., POLL_MAX a round trip, each batch stored before the next is + polled (holding them all spills past S = 16); a thread past a span's end + re-reads its last element and stores nothing. + + ``live``: a span's elements the reader needs (traced, <= n), a span each; + a batch past them is not polled and an element past them neither polled + nor written (one past a step's rows is never written).""" + idx = fresh(idx) # its elements' indices not hoisted out of a task loop + if const_expr(live is None): + _poll_store_batches( + c, + [ + (region, pair_of, n, (dst, row), idx + width * i) + for region, pair_of, n, dst, row in spans + for i in range((n + width - 1) // width) + ], + ) + return + for (region, pair_of, n, dst, row), need in zip(spans, live): + chunks = [ + (region, pair_of, need, (dst, row), idx + width * i) + for i in range((n + width - 1) // width) + ] + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + if chunks[b0][4] < need: + _poll_store_batches(c, chunks[b0 : b0 + POLL_MAX]) + + +@traced +def _poll_store_batches(c, chunks): + """``poll_copy``'s chunks [(region, pair_of, n, (dst, row), u)], POLL_MAX a + round trip, each batch stored before the next is polled.""" + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + batch = chunks[b0 : b0 + POLL_MAX] + got = c["poll"]([(r, p(fx.min(u, n - 1)), 1) for r, p, n, _, u in batch]) + for (_, _, n, (dst, row), u), v in zip(batch, got): + at = u if row is None else lds_at(u, row) + if u < n: + fx.ptr_store(v[0], dst + at) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/moe.py b/vllm/models/deepseek_v4/amd/mono/stages/moe.py new file mode 100644 index 000000000000..f694b6e77bed --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/moe.py @@ -0,0 +1,1733 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe.py +"""The mono layer's MoE stages: one layer's FFN, from the FFN seam's normed rows +to the reduced output (its TP all-reduce in-kernel, ``allreduce``). + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S = ``tokens`` rows, +E = 384 routed experts, top-6. The rank holds every routed expert's 1/TP of the +intermediate width (TP-sharded MoE) and the shared expert's 1/TP (``Dims``). +Stages, with TP8's task counts: + + xq (S): the row as MXFP8 for the routed experts (aiter + ``fused_mx_quant_moe_sort``: amax seeded 1e-10) and for the shared + one (ATOM ``quantize_fp8``: amax floored 1e-4) -> X8R, X8 (plain + words, then the token's flag X8RDY) + router (E / 16 x 4, 16 rows x a quarter of K): bf16 GEMV -> the part's + f32 sum -> LOGIT + route (S): the parts summed in order -> bf16 logits; aiter + ``topk_gating_kernel_opt`` (sqrtsoftplus + bias, top-k, renormalized + x ``ROUTED_SCALE``; a hash layer's experts given), one wave a + token -> ROUTE + shared (24, 16 gate + 16 up rows): FP8 GEMV -> bf16 -> aiter silu_and_mul + (gate <= 10, up in [-10, 10], ocml exp) -> bf16 -> SMID + ug (|U| x 6, 2 x (32 gate + 32 up) rows of one expert of the union U of + every token's top-k; w13 gate / up interleaved a 16-row tile): the + MXFP4 x MXFP8 GEMV, a wave the gate and up 16-row tiles of one + 32-row scale block over a half of K (each B operand from LDS read + once for both), the slices summed pairwise; then its epilogue: + clamp, silu x up, bf16, + and the stage-2 input quant to MXFP8 (amax seeded 1e-10, the ceil + code) -> MID (plain words at device scope, then the task's flag, UGF) + down (448, 16 rows; a CTA's two run together): per expert of U the MXFP4 + x MXFP8 GEMV of its intermediate (x its routing weight -> bf16), + summed in each token's top-k order, rounded to bf16 at every add + (aiter's stage 2 adds bf16 atomically, in no fixed order); the + shared expert's FP8 down projection (quantize_fp8 of SMID) -> bf16; + routed + shared -> bf16 -> pushed to every rank's FFN region, or + to its columns' owner (``allreduce``) + reduce (224, 32 columns): the TP partials in rank order 0 .. TP - 1 + (``sum_partials``), fp32 -> bf16 -> ``out`` +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.mx import ( + FP4, + FP8, + FP8_MAX, + UNIT_SCALE, + clamp_fp8, + code_ceil, + fp4_tile_load, + mfma_scaled, + rcp_pow2, + scale_index, +) +from vllm.models.deepseek_v4.amd.mono.common.ops import ( + CM_DEV, + CM_NT, + _hw_f32, + ballot, + bf16_round, + bf_hi, + bf_lo, + butterfly, + div_rn, + fp8_pack4, + fresh, + hw_exp2, + hw_rcp, + lane_gather, + lanes_below, + popcount, + row_sum, + rsrc, + stamp, + traced, +) +from vllm.models.deepseek_v4.amd.mono.common.plan import ( + BLOCKS, + THREADS, + WAVES, + first_task, + pair_layout, +) +from vllm.models.deepseek_v4.amd.mono.common.sync import preg, publish, sreg +from vllm.models.deepseek_v4.amd.mono.dims import ( + DOWN_TASKS, + HIDDEN, + ROUTED_SCALE, + SWIGLU_LIMIT, + UG_PART, + X_GROUPS, + X_WORDS, + Dims, +) +from vllm.models.deepseek_v4.amd.mono.stages import allreduce +from vllm.models.deepseek_v4.amd.mono.stages import moe_plan as mplan + +# helpers the kernel calls are imported by name: FlyDSL keys a build by the +# source of the same-directory functions it names, never through a module +from vllm.models.deepseek_v4.amd.mono.stages.gemv import ( + ROWS, + TILE, + gemv_fp8_loads, + gemv_fp8_mfmas, + ld_f32, + lds_row, + lds_tail, + plus, + poll_copy, + tile_rows, + token_tiles, +) +from vllm.models.deepseek_v4.amd.mono.stages.moe_shape import ( + MASK_WORDS, + MoeBuild, + route_shape, +) +from vllm.models.deepseek_v4.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.deepseek_v4.amd.mono.stages.router import ( + router_tile_mfma, + router_weight_loads, + router_x_loads, +) +from vllm.models.deepseek_v4.amd.mono.stages.seam import attn_peer_bytes + +# a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other +assert mplan.TILE == TILE +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +# (task, slot) scale blocks a down wave holds a round (``stage_down``): 3 took +# no less time, 4 spills +DOWN_R_MAX = 2 +# a down slot's K steps on this many independent MFMA accumulators +DOWN_CHAINS = 2 +# a reduce task's columns: at 128 a thread polled 6 pairs at 48 rows, each +# pass a round trip in series (32: 1-2 passes, S=48 -1 us) +RED_COLS = 32 +# CTA b polls route copy b % ROUTE_COPIES: every CTA on one small region +# serializes on its lines (S=24: -0.5 us) +ROUTE_COPIES = 8 +TAGS = 256 # the ug queue's / down count's slots, a word each: ``c["tag"]`` picks one + +XQ0 = 0 +DOWN0 = 0 +RED0 = 0 + + +# ``timeline`` stamps: kernel start, then the end of each stage's task loop +TL_STAGES = ("xq", "router", "route", "shared", "ug", "down", "reduce") +TL_POINTS = 1 + len(TL_STAGES) + + +# the MoE stages' kernel arguments, in launch order +ARGS = ( + "x", + "gate_w", + "bias", + "w13", + "w13_s", + "w2", + "w2_s", + "sgu", + "sgu_s", + "sw2", + "sw2_s", + "out", +) + + +def scratch_layout(key: MoeBuild) -> dict[str, tuple[int, int]]: + rs = route_shape(key) + tokens, experts, topk = key.tokens, rs.experts, rs.topk + d = rs.dims + pairs = { + # the routed experts' MXFP8 x (aiter's fused_mx_quant_moe_sort) and the + # shared expert's (ATOM quantize_fp8): their amax floors differ + "x8r": tokens * X_WORDS // 2, + "x8rs": tokens * X_GROUPS // 2, + "x8": tokens * X_WORDS // 2, + "x8s": tokens * X_GROUPS // 2, + # a token's four MXFP8 regions written (plain words): ``stage_xq`` + "x8rdy": tokens, + "logit": tokens * experts * ROUTER_PARTS, + "route": ROUTE_COPIES * tokens * 2 * topk, + "smid": tokens * d.sh_inter, + # SMID quantize_fp8'd (``smid_quant``): the down stage's MXFP8 rows + "smq": tokens * d.sh_inter // 4, + "smc": tokens * d.sh_inter // 32, + # a pick a row, in the slot order (``RouteShape.order0``) + "mid": rs.picks * d.mid_words, + "mids": rs.picks * d.inter // 32, + # an ug task's done flag (``stage_ug``) + "ugf": (rs.max_tiles if rs.multi else rs.max_u) * rs.ug_parts, + # the ug task queue: a word a mailbox tag (a layer), counting up from the + # step's clear (``take_ug_task``); plain words, 2 a pair + "ugq": TAGS // 2, + # the down stage's arrival count, a word a tag (``take_down_rank``) + "dq": TAGS // 2, + } + return pair_layout(pairs.items()) + + +def peer_bytes(tokens: int, d: Dims) -> int: + """FFN: the MoE output's all-reduce (``allreduce``).""" + return allreduce.region_bytes(tokens, d) + + +def ffn_offset(tokens: int, d: Dims) -> int: + """FFN's byte offset in the peer buffer: past the ATTN region.""" + return attn_peer_bytes(tokens, d) + + +def red_tasks(d: Dims) -> int: + return HIDDEN // RED_COLS + + +# ---------------------------------------------------------------- xq +@traced +def stage_xq(c, t): + """Token t, a group of 32 a thread.""" + tid = c["tid"] + a, hidden = c["args"], HIDDEN + # thread g < H / 32 owns the 32 values of group g; the rest mirror the last + g = fx.min(tid, hidden // 32 - 1) + mine = tid < hidden // 32 + # 32 bf16 as four 16-byte loads (a load an element made the stage issue bound) + words = [ + fx.Vector( + bo.buffer_load( + rsrc(a["x"]), + (t * hidden + g * 32) // 2 + 4 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=c["x_cm"], + ) + ) + for q in range(4) + ] + v = [(bf_lo if i % 2 == 0 else bf_hi)(words[i // 8][i % 8 // 2]) for i in range(32)] + amax0 = abs(v[0]) + for y in v[1:]: + amax0 = fx.max(amax0, abs(y)) + # the routed experts' (aiter fused_mx_quant_moe_sort): amax floored 1e-10; + # the shared expert's: vLLM's mxfp8_e4m3_quantize (its MXFP8 linear) + if mine: + put_mxfp8_group(c, "x8r", "x8rs", t, g, v, fx.max(amax0, fx.Float32(1e-10))) + put_mxfp8_group_vllm(c, "x8", "x8s", t, g, v, amax0) + + +def put_mxfp8_group(c, words_mb, codes_mb, t, g, v, amax): + """Token t's group g of 32 values ``v`` as MXFP8 (the ceil code of + ``amax`` / 448, x / scale clamped, RNE) -> scratch ``words_mb`` / + ``codes_mb`` (plain words: the token's X8RDY announces them).""" + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + inv = rcp_pow2(code) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * X_GROUPS + g) + + +FLT_MIN = 1.1754943508222875e-38 + + +def vllm_mx_code(amax): + """VLLM's ``mxfp8_e4m3_quantize`` scale code of a 32-group: the smallest + power of two >= RN(max(amax, FLT_MIN) / 448), clamped to [0, 254].""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + bits = q.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.min(fx.max(e, fx.Int32(0)), fx.Int32(254)) + + +def put_mxfp8_group_vllm(c, words_mb, codes_mb, t, g, v, amax): + """``put_mxfp8_group`` with vLLM's rule (``vllm_mx_code``; x * 2^(127 - code), + which code 0 keeps finite).""" + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * X_GROUPS + g) + + +# ---------------------------------------------------------------- router +@traced +def stage_router(c, task): + """Router task (token group, row tile, K part): K part ``part`` of rows + 16 ``tile`` .. of the bf16 gate weight [experts, hidden] times the group's + tokens (``RouteShape.router_groups``), a token tile at a time -- the + decode router's tile (``router``), the + same tiles in the same order -> the part's f32 sum (``route`` sums the + parts in order, rounds to bf16). Every tile's x loads issue first: a task + over every token's tiles in series read their x serially (48 rows: a router + CTA 9 us, 8 of them x).""" + s, tid, red, lane, wave = c["S"], c["tid"], c["red"], c["lane"], c["wave"] + a = c["args"] + experts = c["rs"].experts + groups, per = c["rs"].router_groups + per_group = experts // ROWS * ROUTER_PARTS + tile, part = task % per_group // ROUTER_PARTS, task % ROUTER_PARTS + g0 = 0 if groups == 1 else task // per_group * per + # the group's token tiles: (first token, tokens), tokens traced past one group + tiles = [ + ( + plus(g0, j0), + min(TILE, s - j0) + if groups == 1 + else fx.min(fx.Int32(min(TILE, per - j0)), s - plus(g0, j0)), + ) + for j0 in range(0, per, TILE) + ] + k = HIDDEN + wvs = router_weight_loads(a["gate_w"], k, tile, part, lane, wave, CM_NT) + xs = [ + router_x_loads(a["x"], k, part, n, lane, wave, c["x_cm"], t0) for t0, n in tiles + ] + for (t0, n), xvs in zip(tiles, xs): + router_tile_mfma(wvs, xvs, lane, wave, red) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + v = row_sum(red, r, tl) + c["put"]( + c["logit"], + (plus(t0, tl) * experts + tile * ROWS + r) * ROUTER_PARTS + part, + v, + ) + gpu.barrier() + + +# ---------------------------------------------------------------- route +def _score(x): + """sqrt(log2(1 + exp2(x log2e)) ln2), x itself past 20 (topk_gating).""" + sp = fx.Float32( + (x > fx.Float32(20.0)).select( + x, + _hw_f32( + "llvm.amdgcn.log.f32", fx.Float32(1.0) + hw_exp2(x * fx.Float32(LOG2E)) + ) + * fx.Float32(LN2), + ) + ) + return fx.Float32(fmath.sqrt(sp)) + + +@traced +def stage_route(c, t): + """VLLM's topk_hash_softplus_sqrt for token t on wave 0 (f32 logits).""" + lane, wave = c["lane"], c["wave"] + experts, topk = c["rs"].experts, c["rs"].topk + if wave == 0: + ept = experts // 64 + # each logit's router part sums, added in part order (vLLM's gate: f32 out) + assert ROUTER_PARTS == 4 + got = c["poll"]( + [ + (c["logit"], (t * experts + lane + 64 * i) * ROUTER_PARTS + h, 2) + for i in range(ept) + for h in (0, 2) + ] + ) + logits = [] + for i in range_constexpr(ept): + parts = got[2 * i] + got[2 * i + 1] + tot = fx.Float32(0.0) + for p in range_constexpr(ROUTER_PARTS): + tot = tot + parts[p].bitcast(fx.Float32) + logits.append(tot) + if const_expr(c["rs"].hash_route): + my_id, my_w, total = _hash_picks(c, t, logits) + else: + my_id, my_w, total = _topk_picks(c, logits) + scale = fx.Float32(ROUTED_SCALE) / fx.max(total, fx.Float32(1e-20)) + if lane < topk: + for r in range_constexpr(ROUTE_COPIES): + at = (r * c["S"] + t) * 2 * topk + lane + c["put"](c["route"], at, my_id) + c["put"](c["route"], at + topk, my_w * scale) + gpu.barrier() + + +def _hash_picks(c, t, logits): + """A hash layer's picks of token t: lane k < topk's expert, given (the + bias argument holds the step's [S, topk] ids), and its score from the + logits (expert e's on lane e % 64, ``logits[e // 64]``); the scores' sum + in top-k order.""" + lane, topk = c["lane"], c["rs"].topk + eid = fx.Int32( + bo.buffer_load( + rsrc(c["args"]["bias"]), + t * topk + fx.min(lane, topk - 1), + vec_width=1, + dtype=T.i32, + ) + ) + logit = fx.Float32(0.0) + for i in range(len(logits)): + v = lane_gather(logits[i].bitcast(fx.Int32), eid & 63).bitcast(fx.Float32) + logit = (eid >> 6 == i).select(v, logit) + sc = _score(logit) + w = fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc)) + total = fx.Float32(0.0) + bits = w.bitcast(fx.Int32).ir_value() + for k in range(topk): + wk = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(k).ir_value())) + total = total + wk.bitcast(fx.Float32) + return eid, w, total + + +def _topk_picks(c, logits): + """``stage_route``'s top-k of a token from its f32 logits (lane l holds + experts l + 64 i): lane k < topk's expert, its score, the scores' sum.""" + lane, a = c["lane"], c["args"] + experts, topk = c["rs"].experts, c["rs"].topk + ept = experts // 64 + vals, orig, idxs = [], [], [] + for i in range_constexpr(ept): + e = lane + 64 * i + sc = _score(logits[i]) + # aiter's NaN guards: a NaN score weighs 0 and never wins + orig.append(fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc))) + vals.append(fx.max(sc + ld_f32(a["bias"], e), fx.Float32(float("-inf")))) + idxs.append(fx.Int32(e)) + for p, q in c["rs"].sort_net: + sw = vals[p] < vals[q] + vals[p], vals[q] = sw.select(vals[q], vals[p]), sw.select(vals[p], vals[q]) + orig[p], orig[q] = sw.select(orig[q], orig[p]), sw.select(orig[p], orig[q]) + idxs[p], idxs[q] = sw.select(idxs[q], idxs[p]), sw.select(idxs[p], idxs[q]) + cursor = fx.Int32(0) + total = fx.Float32(0.0) + my_w, my_id = fx.Float32(0.0), fx.Int32(0) + for k in range_constexpr(topk): + cv, ci, co = fx.Float32(float("-inf")), fx.Int32(0), fx.Float32(0.0) + for i in range_constexpr(ept): + at = cursor == i + cv, ci, co = ( + at.select(vals[i], cv), + at.select(idxs[i], ci), + at.select(orig[i], co), + ) + best = butterfly(cv, (32, 16, 8, 4, 2, 1), fx.max) + ball = ballot(cv == best) + win = fx.Int32( + (ball == fx.Int64(0)).select(fx.Int32(0), fx.Int32(fmath.cttz(ball))) + ) + win_id = fx.Int32(rocdl.readlane(T.i32, ci.ir_value(), win.ir_value())) + won = (cursor < ept) & (ci == win_id) + mine_o = fx.Float32(won.select(co, fx.Float32(0.0))) + w = fx.Int32( + rocdl.readlane( + T.i32, mine_o.bitcast(fx.Int32).ir_value(), (win_id & 63).ir_value() + ) + ).bitcast(fx.Float32) + cursor = cursor + won.select(1, 0) + total = total + w + my_w = (lane == k).select(w, my_w) + my_id = (lane == k).select(win_id, my_id) + return my_id, my_w, total + + +@traced +def load_route(c, tab): + """Every token's (ids, weights) into LDS ``tab``, then the union U of the ids + (ascending): tab[2 S topk ..] = |U|, the experts of U, each (t, k)'s slot in + U, each expert's slot, the picks in slot order (a slot's ascending) and each + slot's first position (``RouteShape``). + + A pick a thread, no scan in series: each pick ORs its token's bit into its + expert's token bitmap, so an expert's pick count is its bitmap's popcount, + a slot's first position the counts of the experts below it (a wave's prefix + sum; slots follow the experts ascending), and a pick's rank in its slot the + bitmap's tokens below its own. The bitmaps are the same whatever order the + ORs land in, so every CTA derives the same pick order.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rs = c["rs"] + topk, max_u, experts = rs.topk, rs.max_u, rs.experts + n = s * topk + assert s <= 32 * MASK_WORDS + # a thread's route words in one poll batch: a pass each was a round trip + # each in series past 256 picks (48 rows: 5 us from the last route to the + # table); a word past 2n re-reads the last and is not stored + idxs = [plus(pas, tid) for pas in range(0, 2 * n, THREADS)] + copy_at = c["bid"] % ROUTE_COPIES * (2 * n) + got = c["poll"]([(c["route"], copy_at + fx.min(idx, 2 * n - 1), 1) for idx in idxs]) + for i in range_constexpr(len(idxs)): + if idxs[i] < 2 * n: + fx.ptr_store(got[i][0], tab + idxs[i]) + for pas in range_constexpr(0, MASK_WORDS * experts, THREADS): + idx = plus(pas, tid) + if idx < MASK_WORDS * experts: + fx.ptr_store(fx.Int32(0), tab + (rs.mask0 + idx)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + fx.llvm.atomic_or( + tab + (rs.mask0 + eid * MASK_WORDS + t // 32), + fx.Int32(1) << (t % 32), + syncscope="workgroup", + ) + gpu.barrier() + base = 2 * n + # a chunk of 64 experts a wave (wave 0 alone through 6 chunks in series + # was 3.3 us of the table's 5 at 48 rows): each chunk's own counts and + # prefixes, its totals through LDS ``red`` (free here), then each chunk + # past the totals of the chunks below it -- the serial walk's values + chunks = experts // 64 + assert chunks <= WAVES + sums = c["red"] + e = lane + 64 * fx.min(wave, chunks - 1) + count = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + e * MASK_WORDS + w)) + count = count + popcount(word) + ball = ballot(count != 0) + # this chunk's inclusive prefixes, lanes ascending: picks, and pick tiles + incl = lane_prefix(count, lane) + tiles = mplan.pick_tiles(count, rs.x_tile) + tincl = lane_prefix(tiles, lane) if rs.multi else tiles + totals = [popcount(ball), lane_gather(incl, 63), lane_gather(tincl, 63)] + if (wave < chunks) & (lane == 0): + for k in range_constexpr(3): + fx.ptr_store(totals[k].bitcast(fx.Float32), sums + (wave * 3 + k)) + gpu.barrier() + below_c = [fx.Int32(0), fx.Int32(0), fx.Int32(0)] # n_u, first, tile0 + for j in range_constexpr(chunks - 1): + for k in range_constexpr(3): + v = fx.ptr_load(sums + (j * 3 + k)).bitcast(fx.Int32) + below_c[k] = below_c[k] + (j < wave).select(v, fx.Int32(0)) + n_u, first, tile0 = below_c + slot = n_u + lanes_below(ball, lane) + # e afresh: hoisted, its address held a VGPR across the prefix sums (spill) + e_at = fresh(e) + if wave < chunks: + fx.ptr_store(slot, tab + (rs.expert_slot0 + e_at)) + if count != 0: + fx.ptr_store(fx.Int32(e), tab + (base + 1 + slot)) + fx.ptr_store(first + incl - count, tab + (rs.first0 + slot)) + if const_expr(rs.multi): + # the slot's pick tiles: their first index, and each one's slot + tf = tile0 + tincl - tiles + if count != 0: + fx.ptr_store(tf, tab + (rs.tile_first0 + slot)) + for j in range_constexpr(rs.slot_tiles): + if j < tiles: + fx.ptr_store(slot, tab + (rs.tile_slot0 + tf + j)) + if (wave == chunks - 1) & (lane == 0): + n_all = n_u + totals[0] + fx.ptr_store(n_all, tab + base) + fx.ptr_store(fx.Int32(n), tab + (rs.first0 + n_all)) + if const_expr(rs.multi): + fx.ptr_store(tile0 + totals[2], tab + (rs.tile_first0 + n_all)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + slot = fx.ptr_load(tab + (rs.expert_slot0 + eid)) + fx.ptr_store(slot, tab + (base + 1 + max_u + tid)) + # the tokens below t that picked this expert: the picks before tid in it + rank = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + eid * MASK_WORDS + w)) + lo = t - 32 * w + keep = (lo >= 32).select( + fx.Int32(-1), + (lo > 0).select((fx.Int32(1) << fx.max(lo, 0)) - 1, fx.Int32(0)), + ) + rank = rank + popcount(word & keep) + pos = fx.ptr_load(tab + (rs.first0 + slot)) + rank + fx.ptr_store(tid, tab + (rs.order0 + pos)) + gpu.barrier() + + +def lane_prefix(v, lane): + """The wave's inclusive prefix sum of v, lanes ascending.""" + incl = v + for d in (1, 2, 4, 8, 16, 32): + moved = lane_gather(incl, fx.max(lane - d, 0)) + incl = incl + (lane >= d).select(moved, fx.Int32(0)) + return incl + + +# ---------------------------------------------------------------- shared ug +def _silu_mul(g, u, limit): + """VLLM's silu_and_mul_with_clamp (the shared expert's): gate min limit, + up clamped; silu = g / (1 + expf(-g)) rounded to bf16, then times up.""" + g = bf16_round(fx.min(g, fx.Float32(limit))) + u = fx.max(fx.min(u, fx.Float32(limit)), fx.Float32(-limit)) + act = bf16_round(g / (fx.Float32(1.0) + fx.Float32(fmath.exp(-g)))) + return act * u + + +@traced +def stage_shared(c, task, gops, t0, n): + """Gate rows 16 task .. (weights ``gops``), up rows 576 + 16 task .. against + tokens t0 .. t0 + n (their x in LDS; ``n`` traced: a tile a CTA picks at run + time) -> bf16 -> silu x up -> SMID; an odd task then quantizes its group of + 32 (``smid_quant``).""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + rows = None if isinstance(n, int) else n + hidden = HIDDEN + gemv_fp8_mfmas(c, hidden, c["xl"], c["xsl"], red, gops, rows=rows) + gpu.barrier() + gate = fx.Float32(0.0) + if tid < ROWS * n: + gate = row_sum(red, tid % ROWS, tid // ROWS) + gpu.barrier() + uops = gemv_fp8_loads( + c, + a["sgu"], + a["sgu_s"], + hidden, + d.shared_tasks + task, + ) + gemv_fp8_mfmas(c, hidden, c["xl"], c["xsl"], red, uops, rows=rows) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + t = plus(t0, tid // ROWS) + up = row_sum(red, r, tid // ROWS) + v = fx.Float32( + _silu_mul(bf16_round(gate), bf16_round(up), SWIGLU_LIMIT).to(fx.BFloat16) + ) + c["put"](c["smid"], t * d.sh_inter + task * ROWS + r, v) + fx.ptr_store(v, c["pair"] + tid) + gpu.barrier() + if task % 2 == 1: + smid_quant(c, task, t0, n) + gpu.barrier() + + +@traced +def smid_quant(c, task, t0, n): + """Odd shared task ``task``: its 16 columns (LDS ``pair``) and the even + task's (SMID) are a quantize_fp8 group of 32 (amax floored 1e-4, the ceil + code) -> SMQ / SMC, a thread a token. Once here, not in each of the 256 CTAs + whose down stage reads them: that took a down CTA 14 us at 48 rows.""" + tid, sh = c["tid"], c["d"].sh_inter + assert c["d"].shared_tasks % 2 == 0 # tasks pair into groups + if tid < n: + t = plus(t0, tid) + got = c["poll"]( + [(c["smid"], t * sh + (task - 1) * ROWS + r, 1) for r in range(ROWS)] + ) + vals = [got[r][0].bitcast(fx.Float32) for r in range(ROWS)] + [ + fx.ptr_load(c["pair"] + (tid * ROWS + r)) for r in range(ROWS) + ] + amax = abs(vals[0]) + for y in vals[1:]: + amax = fx.max(amax, abs(y)) + # vLLM's down_proj input quant (mxfp8_e4m3_quantize) + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + g = task // 2 + for w in range_constexpr(8): + word = fp8_pack4(*[clamp_fp8(vals[4 * w + i] * inv) for i in range(4)]) + c["put_words"](c["smq"], t * (sh // 4) + g * 8 + w, [word]) + c["put_words"](c["smc"], t * (sh // 32) + g, [code]) + + +# ---------------------------------------------------------------- routed ug +@traced +def gemv_fp4_pair(c, w, ws, e, n_rows, k, rg0, ks, xl, xsl, scale_cols, col): + """Tiles rg0 and rg0 + 1 (16 rows each, one 32-row scale block) of expert + e's (16, 16)-preshuffled MXFP4 weight [n_rows, k] times 16 MXFP8 rows in + LDS, over this wave's K slice ``ks`` (``ug_rounds``) -> both tiles' partial + C, a Vector of 8 (rows 4 (lane / 16) .., column lane % 16, whose LDS row is + ``col``). Each B operand read from LDS feeds both tiles' MFMAs: the waves of + a task share x, so a wave a tile over all of K read it once a tile. + """ + lane, wave = c["lane"], c["wave"] + sl = c["wsl"] + wave * wsl_words(c["d"]) + slices = c["rs"].ug_k_slices + rounds = ug_rounds(k, ks, slices) + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + # rounds whose loads are in flight past the one computing: a quarter's + # 2-step rounds two + ahead = slices // 2 + ops = [ + fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[q]) + for q in range(min(ahead, len(rounds))) + ] + for r in range_constexpr(len(rounds)): + sc0, st0, n, sc_steps = rounds[r] + sv, avs = ops[r] + fx.ptr_store(sv, sl + lane * (sc_steps // 2)) + sas = [ + [ + lds_u8( + sl, + scale_index( + t * 16 + lane % 16, (st0 + i) * 4 + lane // 16, scale_cols + ) + - 128 * sc0, + ) + for i in range(n) + ] + for t in range(2) + ] + accs = fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col) + if const_expr(r + ahead < len(rounds)): + ops.append( + fp4_pair_loads( + c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[r + ahead] + ) + ) + return fx.Vector.from_elements( + [accs[t][q] for t in range(2) for q in range(4)], fx.Float32 + ) + + +def ug_rounds(k, ks, slices): + """K slice ``ks`` of ``slices``' rounds (first step of the scales loaded, + first step, steps, steps of scales loaded); ``ks`` traced: the steps are. + Halves: the 8-step scale blocks 2 r + ks, the odd block left split 4 + 4, + so both run k / 256 steps. Quarters: steps 8 i + 2 ks, + 1 of each block + i, a round loading only their scales (the shuffled layout keeps 2 steps' + together: a dword a lane, where the whole block was 4).""" + blocks = k // 128 // 8 + if slices == 4: + return [(fx.Int32(8 * i) + ks * 2,) * 2 + (2, 2) for i in range(blocks)] + assert slices == 2 + rounds = [ + (fx.Int32(16 * r) + ks * 8, fx.Int32(16 * r) + ks * 8, 8, 8) + for r in range(blocks // 2) + ] + if blocks % 2: + last = (blocks - 1) * 8 + rounds.append((fx.Int32(last), fx.Int32(last) + ks * 4, 4, 8)) + return rounds + + +def fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, sc0, st0, n, sc_steps): + """``gemv_fp4_pair``'s loads of a round: the scales of steps sc0 .. sc0 + + sc_steps (of the tiles' 32-row block: 128 bytes a step, sc_steps / 2 dwords + a lane) and both tiles' weight operands of steps st0 .. st0 + n. Loads + only.""" + lane = c["lane"] + block = (e * n_rows + rg0 * 16) // 32 * (scale_cols * 8) + words = sc_steps // 2 + sv = bo.buffer_load( + rsrc(ws), + block + 32 * sc0 + lane * words, + vec_width=words, + dtype=T.i32, + cache_modifier=CM_NT, + ) + sv = fx.Vector(sv) if words > 1 else fx.Int32(sv) + return sv, [fp4_tile_weights(c, w, e, n_rows, k, rg0 + t, st0, n) for t in range(2)] + + +def scale_block_loads(c, ws, first_row, scale_cols): + """The loads of the 32-row e8m0 block holding ``first_row``: in aiter's + shuffle_scale order a block is contiguous (both 16-row halves, all of K), so + a few 16-byte loads a lane fetch it, where a byte load a K step doubled the + stage's load instructions and bound it by their issue (2026-09-29, ug). A + lane past the block re-reads its last 16 bytes.""" + lane = c["lane"] + words = scale_cols * 8 + block = first_row // 32 * words + return [ + fx.Vector( + bo.buffer_load( + rsrc(ws), + block + fx.min(256 * i + lane * 4, words - 4), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for i in range((words + 255) // 256) + ] + + +def down_scale_words(d) -> int: + """A down tile's 32-row scale block in LDS: ``scale_block_store`` writes + whole 256-word rounds.""" + return down_lds_words(d.down_scale_cols * 8) + + +def wsl_words(d) -> int: + """A wave's LDS scale words: an ug batch's 256, or a down round's blocks -- + a (task, slot) each, ``DOWN_R_MAX`` at most (``stage_down``).""" + return max(256, DOWN_R_MAX * down_scale_words(d)) + + +def scale_block_store(c, svs, sl): + """``scale_block_loads``' words -> this wave's LDS block ``sl`` + (``down_scale_words``; a lane past the block writes past it, inside it).""" + lane = c["lane"] + for i in range_constexpr(len(svs)): + fx.ptr_store(svs[i], sl + (256 * i + lane * 4)) + + +def lds_scales(sl, local_row, lane, st0, blk, scale_cols): + """K steps st0 .. st0 + blk's scale bytes of this lane's row from a 32-row + scale block in LDS (``local_row`` its row in the block). An even step and + the next are bytes 2 apart of one word (``scale_index``'s d): a read a + pair, not a step (down: 3 reads a slot, not 5).""" + out, word = [], None + for i in range(blk): + idx = scale_index(local_row, (st0 + i) * 4 + lane // 16, scale_cols) + if (st0 + i) % 2 == 0 or word is None: + word = fx.ptr_load(sl + idx // 4) + out.append((word >> (idx % 4 * 8)) & 0xFF) + return out + + +def lds_u8(ptr, i): + return (fx.ptr_load(ptr + i // 4) >> (i % 4 * 8)) & 0xFF + + +def fp4_tile_weights(c, w, e, n_rows, k, rg, st0, blk, k_real=None): + """The weight operands of expert e's 16-row tile rg, K steps st0 .. st0 + + blk: loads only, nothing waits on them here. ``k_real``: a lane whose 32 + columns lie past it (zero weights) reads its row group's first lane's + address instead (no extra traffic) and holds zero.""" + lane = c["lane"] + base = e * n_rows * (k // 2) // 4 # dwords: expert e + out = [] + for i in range_constexpr(blk): + past = k_real is not None and ((st0 + i) * 128 + 96) >= k_real + if const_expr(past): + real = (st0 + i) * 128 + lane // 16 * 32 < k_real + at = real.select(lane * 4, lane % 16 * 4) + else: + at = lane * 4 + v = fp4_tile_load(rsrc(w), base, rg, k, st0 + i, at) + if const_expr(past): + v = fx.Vector.from_elements( + [real.select(v[q], fx.Int32(0)) for q in range(4)], fx.Int32 + ) + out.append(v) + return out + + +def fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col): + """``accs`` (two tiles') chained through ``n`` K steps from st0 against MXFP8 + rows in LDS (column lane % 16: row ``col``): a step's B operand loaded once + for both tiles' MFMAs.""" + lane = c["lane"] + x_row = lds_row(k // 4) + s_row = lds_row(k // 32) + j = lane // 16 + accs = list(accs) + for i in range_constexpr(n): + bv = fp8_operand( + [ + fx.ptr_load( + xl + (col * x_row + (st0 + i) * 32 + 16 * h + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + for h in range(2) + ] + ) + sb = fx.ptr_load(xsl + (col * s_row + (st0 + i) * 4 + j)) + for t in range_constexpr(2): + accs[t] = mfma_scaled(avs[t][i], bv, accs[t], sas[t][i], sb, FP4, FP8) + return accs + + +def fp8_operand(halves): + """A lane's MXFP8 MFMA operand of a K step: its two 16-byte halves (K + 16 j .. and 64 + 16 j .. of the step, j = lane / 16).""" + return fx.Vector.from_elements( + [fx.Vector(h).bitcast(fx.Int32)[d] for h in halves for d in range(4)], + fx.Int32, + ) + + +def down_lds_words(n): + """A down wave's LDS block of n words: ``scale_block_store`` writes whole + 256-word rounds.""" + return -(-n // 256) * 256 + + +def st_dev(v, region, i): + """Word i of scratch ``region`` (plain words, not tagged pairs) at device + scope: past the per-XCD L2, so another CTA reads it after the flag.""" + bo.buffer_store(v, rsrc(region.value), i, cache_modifier=CM_DEV) + + +@traced +def take_ug_task(c): + """This CTA's next ug task past the first round, from the launch's queue + (``ugq``, a word a tag, counting from 0): the CTAs that ran a shared, + router or route task come to it late and so take fewer, where a fixed + placement gave every CTA the same count and down waited on theirs. A + task's result is the same whichever CTA runs it.""" + if c["tid"] == 0: + at = c["ugq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + task = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(task), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +@traced +def stage_ug(c, task): + """One token tile: slot task // ug_parts of U (``RouteShape.ug_parts``), the + ``ug_groups`` FP4 groups of intermediate columns from 32 ug_groups (task % + ug_parts) (``ug_wave_pair``); then silu x up, rounded to bf16, and the + MXFP8 requant of each token's 32 values of a group -> MID (the stage-1 + epilogue and the stage-2 input quant). Past one: ``stage_ug_tile``.""" + s, tid, lane, tab = c["S"], c["tid"], c["lane"], c["tab"] + rs = c["rs"] + ug_parts = rs.ug_parts + u = task // ug_parts + if u < c["rs"].n_union(tab, s): + e = c["rs"].union_expert(tab, s, u) + # the slot's picks, a column each (a column past them repeats the last); + # every token's row is in LDS: a column reads its pick's token's + first = c["rs"].slot_first(tab, u) + count = c["rs"].slot_first(tab, u + 1) - first + pick = c["rs"].pick_at(tab, first + fx.min(lane % 16, count - 1)) + ug_pass(c, e, task % ug_parts * rs.ug_groups, first, count, pick // rs.topk) + publish(c["put"], c["ugf"], task, 1, tid == 0) + gpu.barrier() + + +def ug_wave_pair(c, part0): + """This wave's share of an ug task from group ``part0``: group p of its + ``ug_groups`` (wave / (2 slices)), its 16-column half bb = wave % 2 -- w13's + gate tile 2 b and up tile 2 b + 1 (aiter GUGU interleaves them a 16-row + tile; b = 2 (part0 + p) + bb), one 32-row scale block -- over K slice ks = + wave / 2 % slices (``gemv_fp4_pair``): (p, 2 b, ks, whether the group is in + the real width).""" + wave, slices = c["wave"], c["rs"].ug_k_slices + p = wave // (2 * slices) + b = (part0 + p) * 2 + wave % 2 + return p, 2 * b, wave // 2 % slices, part0 + p < c["d"].inter_real // UG_PART + + +def ug_red_slot(rs, p, bb, ks, t): + """RED's 64-lane slot of an ug task's partial C: group p, half bb, K slice + ks, gate (t 0) or up (t 1).""" + return ((p * 2 + bb) * rs.ug_k_slices + ks) * 2 + t + + +def pairwise_sum(vs): + """Vs summed as a balanced tree, neighbours first: (v0 + v1) + (v2 + v3).""" + while len(vs) > 1: + vs = [vs[i] + vs[i + 1] for i in range(0, len(vs), 2)] + return vs[0] + + +def ug_tile_task(c, task): + """Past one token tile, ug task ``task``: its slot's expert, its first + group, and its pick tile's first pick and the picks from it.""" + tab, rs, s = c["tab"], c["rs"], c["S"] + ug_parts = rs.ug_parts + tile = task // ug_parts + u = rs.tile_slot(tab, tile) + first = mplan.tile_pick0( + rs.slot_first(tab, u), tile, rs.tile_first(tab, u), rs.x_tile + ) + count = rs.slot_first(tab, u + 1) - first + return rs.union_expert(tab, s, u), task % ug_parts * rs.ug_groups, first, count + + +@traced +def stage_ug_tile(c, task, total): + """Past one token tile: ug task ``task`` of ``total`` (a pick tile of its + slot), then the next task from the queue (``take_ug_task``). No next task's + weights prefetched under this one: held across its epilogue and gather they + spilled the merged K2 (S = 24 / 48), and without them it ran faster.""" + tid, lane = c["tid"], c["lane"] + e, part0, first, count = ug_tile_task(c, task) + gather_x8r(c, first, count) + ug_pass(c, e, part0, first, count, fx.min(lane % 16, count - 1)) + nxt = (total > BLOCKS).select(take_ug_task(c) + BLOCKS, total) + # the epilogue's MID stores done before the flag + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid == 0: + c["put"](c["ugf"], task, 1) + return nxt + + +@traced +def ug_pass(c, e, part0, first, count, col): + """An ug task's GEMV and epilogue for the slot's picks first .. first + + min(count, 16) (column lane % 16's x in LDS row ``col``).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, d, rs = c["args"], c["d"], c["rs"] + p, rg0, ks, real = ug_wave_pair(c, part0) + # a group past the real width (the padding to inter) has zero weights: + # its GEMV is zero unread, and the epilogue still writes its zero MID + acc = fx.Vector.filled(8, 0.0, fx.Float32) + if real: + acc = gemv_fp4_pair( + c, + a["w13"], + a["w13_s"], + e, + 2 * d.inter, + HIDDEN, + rg0, + ks, + c["xl"], + c["xsl"], + X_GROUPS, + col, + ) + for t in range_constexpr(2): + slot = ug_red_slot(rs, p, wave % 2, ks, t) + fx.ptr_store( + fx.Vector.from_elements([acc[4 * t + q] for q in range(4)], fx.Float32), + red + (slot * 64 + lane) * 4, + ) + gpu.barrier() + # thread (p, j, i): group p, the slot's pick j, intermediate column + # 32 part + i, a pass of THREADS at a time; a group's 32 threads stay + # lane-aligned + cols = rs.tile_picks + for pas in range_constexpr(0, rs.ug_groups * cols * UG_PART, THREADS): + ug_epilogue(c, part0, pas + tid, first, count) + + +@traced +def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live): + """MXFP8 rows ``token_of(j)``, j < ``live`` (<= the x tile; traced or int), of + scratch ``words_mb`` / ``codes_mb`` into LDS rows j of ``xl`` / ``xsl``: + the first ``live`` lanes wait on their token's X8RDY, then plain 16-byte + loads (half the bytes of the tagged pairs a poll would read).""" + tid, rows = c["tid"], c["rs"].x_tile + if tid < rows: + c["poll"]([(c["x8rdy"], token_of(fx.min(tid, live - 1)), 1)]) + gpu.barrier() + for mb, width, dst in ((words_mb, X_WORDS, xl), (codes_mb, X_GROUPS, xsl)): + quads = width // 4 + for i in range_constexpr(-(-rows * quads // THREADS)): + q = tid + THREADS * i + row = q // quads + if (q < rows * quads) & (row < live): + v = fx.Vector( + bo.buffer_load( + rsrc(c[mb].value), + token_of(row) * width + q % quads * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + fx.ptr_store(v, dst + (row * lds_row(width) + q % quads * 4)) + gpu.barrier() + + +def gather_x8r(c, first, count): + """The X8R rows and codes of the slot's picks first .. first + min(count, 16) + into LDS rows 0 .. (only those rows: an ug pass reads a pick's row, the + columns past them repeat the last one's).""" + tab, topk = c["tab"], c["rs"].topk + + def token(j): + return c["rs"].pick_at(tab, first + fx.min(j, count - 1)) // topk + + live = fx.min(count, c["rs"].x_tile) + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], token, live) + + +@traced +def ug_epilogue(c, part0, idx, first, count): + """Thread ``idx``'s value of an ug task: silu x up of its (group, pick j, + column), the group's MXFP8 scale, its byte of MID (the pick's row, ``first`` + + j in the pick order).""" + lane, red, d, rs = c["lane"], c["red"], c["d"], c["rs"] + cols = rs.tile_picks + if idx < rs.ug_groups * cols * UG_PART: + p = idx // (cols * UG_PART) + part = part0 + p + t = idx % (cols * UG_PART) // UG_PART + i = idx % UG_PART + tile = i // 16 + r = i % 16 + at = (16 * (r // 4) + t) * 4 + r % 4 + g, up = [ + pairwise_sum( + [ + fx.ptr_load(red + ug_red_slot(rs, p, tile, ks, t) * 64 * 4 + at) + for ks in range(rs.ug_k_slices) + ] + ) + for t in range(2) + ] + limit = SWIGLU_LIMIT + gc = fx.min(g, fx.Float32(limit)) + uc = fx.max(fx.min(up, fx.Float32(limit)), fx.Float32(-limit)) + sig = hw_rcp(fx.Float32(1.0) + hw_exp2(gc * fx.Float32(-LOG2E))) + # the stage-1 kernel stores the intermediate as bf16; the stage-2 input + # quant (fused_mx_quant_moe_sort) makes it MXFP8: amax floored 1e-10, + # the ceil code, x / scale clamped, RNE + v = bf16_round(gc * sig * uc) + amax = fx.max(butterfly(abs(v), (1, 2, 4, 8, 16), fx.max), fx.Float32(1e-10)) + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + q8 = clamp_fp8(v * rcp_pow2(code)) + nb = [ + lane_gather(q8.bitcast(fx.Int32), fx.min(lane + q, 63)).bitcast(fx.Float32) + for q in range(1, 4) + ] + # 4 threads' bytes -> one word + word = fp8_pack4(q8, *nb) + row = first + t + if (i % 4 == 0) & (t < count): + st_dev(word, c["mid"], row * d.mid_words + part * 8 + i // 4) + if (i == 0) & (t < count): + st_dev(code, c["mids"], mid_code_index(d, row, part)) + + +# ---------------------------------------------------------------- down +@traced +def stage_down(c, tasks): + """Rows 16 task .. of the hidden output for each of ``tasks`` (a CTA's one or + two, run together: a second pass would double the CTA's time): routed + experts of U (a wave each, round robin), the shared expert, their sum -> + pushed to every rank.""" + s, wave, tab, d = c["S"], c["wave"], c["tab"], c["d"] + n = len(tasks) + nu = c["rs"].n_union(tab, s) + sl = c["wsl"] + wave * wsl_words(d) + sw = down_scale_words(d) + # slots of U a round: the loads of a round stay the same + slots = DOWN_R_MAX // n + per_round = WAVES * slots + # the shared expert's down projection first: its input is ready long + # before the routed results, which a CTA mostly waits on here; after the + # rounds it sat on the path past the last ug task + down_shared(c, tasks) + for rnd in range(0, (nu + per_round - 1) // per_round, 1): + # the round's slots of U (wave, wave + WAVES, ...), every load of them + # and of every task in flight before the first wait: the weights and + # their scale blocks (no ug result: before the round's ug flags), then + # the slots' picks' intermediate rows, each lane its pick's straight + # from MID (``mid_operands``). A slot past U loads the last slot again + # (L2 hits) and contributes nothing + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + ucs = [fx.min(u, nu - 1) for u in us] + firsts = [c["rs"].slot_first(tab, u) for u in ucs] + svs, avs = down_weights(c, tasks, rnd, slots, nu) + # this round's slots' ug tasks done, not every slot of the wave: the + # round starts while later slots' ug run + c["poll"]([(c["ugf"], f, 1) for f in down_ug_flags(c, ucs)]) + counts = [c["rs"].slot_first(tab, u + 1) - f for u, f in zip(ucs, firsts)] + mids = [mid_operands(c, f, k) for f, k in zip(firsts, counts)] + for q in range_constexpr(len(svs)): + scale_block_store(c, svs[q], sl + q * sw) + for i in range_constexpr(slots): + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + mids[i], + firsts[i], + counts[i], + avs[j * slots + i], + us[i] < nu, + ) + # a slot of more than 16 picks: its next pick tile's rows + if const_expr(s > TILE): + for j0 in range_constexpr(TILE, s, TILE): + if j0 < counts[i]: + more = mid_operands(c, firsts[i] + j0, counts[i] - j0) + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + more, + firsts[i] + j0, + counts[i] - j0, + avs[j * slots + i], + us[i] < nu, + ) + gpu.barrier() + for t0, nn in token_tiles(s): + for j in range_constexpr(n): + down_combine(c, tasks[j], j, t0, nn) + + +@traced +def take_down_rank(c): + """This CTA's arrival rank at the down stage (DQ, a word a tag): which + CTAs take the extra tiles. A tile's result is the same whichever CTA + computes it.""" + if c["tid"] == 0: + at = c["dq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + rank = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(rank), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +def down_weights(c, tasks, rnd, slots, nu): + """Round ``rnd``'s weights of ``stage_down``: each (task, slot)'s scale block + and weight operands, this wave's slots of U (a slot past U: the last's). + Loads only.""" + s, wave, tab, a, d = c["S"], c["wave"], c["tab"], c["args"], c["d"] + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + es = [c["rs"].union_expert(tab, s, fx.min(u, nu - 1)) for u in us] + svs = [ + scale_block_loads(c, a["w2_s"], e * HIDDEN + task * 16, d.down_scale_cols) + for task in tasks + for e in es + ] + avs = [ + fp4_tile_weights( + c, a["w2"], e, HIDDEN, d.inter, task, 0, d.inter // 128, d.inter_real + ) + for task in tasks + for e in es + ] + return svs, avs + + +def mid_operands(c, first, count): + """The MFMA B operands of picks first .. first + min(count, 16) (column + lane % 16 the pick, past them the last): each K step's two 16-byte halves + of the lane's MID row and its group's code, loaded at device scope straight into + registers -- no LDS block a slot, so a round holds more slots.""" + lane, d = c["lane"], c["d"] + j = lane // 16 + row = first + fx.min(lane % 16, count - 1) + steps = d.inter // 128 + # the step's constant after an opaque lane word folds into the immediate + # offset; summed with j * 4 it was a hoisted VGPR a step (spill) + lane_word = fresh(row * d.mid_words + j * 4) + bvs = [] + for st in range_constexpr(steps): + halves = [] + for h in range_constexpr(2): + v = fx.Vector( + bo.buffer_load( + rsrc(c["mid"].value), + lane_word + (st * 32 + 16 * h), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + if const_expr(st * 128 + 64 * h + 48 >= d.inter_real): + # columns past the rank's real intermediate (the loader pads + # it with zero weights) are never written: stale bytes, which + # a NaN pattern would carry through the zero weights + past = st * 128 + 64 * h + 16 * j >= d.inter_real + v = fx.Vector.from_elements( + [past.select(fx.Int32(0), v[q]) for q in range(4)], fx.Int32 + ) + halves.append(v) + bvs.append(fp8_operand(halves)) + # this lane's codes, one a K step: contiguous in MIDS (``mid_code_index``), + # 4 a load -- a load a step was 5 of a slot's 16 load instructions, and the + # stage is load-issue bound (probe ``harness/down_probe.py``: -8 %) + base = mid_code_index(d, row, j) + sbs = [] + for w0 in range_constexpr(0, steps, 4): + n = min(4, steps - w0) + v = bo.buffer_load( + rsrc(c["mids"].value), + base + w0, + vec_width=n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + if const_expr(n == 1): + sbs.append(fx.Int32(v)) + else: + v = fx.Vector(v) + sbs += [fx.Int32(v[k]) for k in range(n)] + for st in range_constexpr(steps): + if const_expr((st + 1) * 128 > d.inter_real): + # a group past the real intermediate has no code written either + past = (4 * st + j) * 32 >= d.inter_real + sbs[st] = past.select(fx.Int32(UNIT_SCALE), sbs[st]) + return bvs, sbs + + +def mid_code_index(d, row, group): + """MIDS word of pick row ``row``'s MXFP8 group ``group`` (32 columns): a K + step's 4 groups apart, so the steps of one group % 4 lie together + (``mid_operands``' lane j reads groups j, j + 4, ..).""" + steps = d.inter // 128 + return row * (d.inter // 32) + group % 4 * steps + group // 4 + + +def down_ug_flags(c, ucs): + """This lane's ug flags of a down round's slots ``ucs``: flag f a (slot, + pick tile past one token tile, part), a tile past the slot's repeating its + last; lane l takes flags l, l + 64, .. (TP2's 18 parts a tile, 48 rows: + 108 flags a round).""" + rs = c["rs"] + ug_parts = rs.ug_parts + per_slot = rs.slot_tiles * ug_parts if rs.multi else ug_parts + n = len(ucs) * per_slot + return [ + _down_ug_flag(c, ucs, fx.min(c["lane"] + 64 * j, n - 1), per_slot) + for j in range(-(-n // 64)) + ] + + +def _down_ug_flag(c, ucs, f, per_slot): + rs, tab = c["rs"], c["tab"] + ug_parts = rs.ug_parts + mine = ucs[0] + for i in range_constexpr(1, len(ucs)): + mine = (f // per_slot == i).select(ucs[i], mine) + if const_expr(not rs.multi): + return mplan.ug_flag(mine, f % ug_parts, ug_parts) + tile0 = rs.tile_first(tab, mine) + tiles = rs.tile_first(tab, mine + 1) - tile0 + return mplan.ug_flag( + mplan.polled_tile(tile0, tiles, f % per_slot // ug_parts), + f % ug_parts, + ug_parts, + ) + + +@traced +def down_shared(c, tasks): + """Each of ``tasks``' rows of the shared expert's down projection + (its MXFP8 intermediate, ``load_smid_tile``) for every + token, bf16 -> LDS ``sdl[task][token][16 rows]`` (``down_combine``): + every task's weights in flight first.""" + s, tid, red, d = c["S"], c["tid"], c["red"], c["d"] + a = c["args"] + shops = [ + gemv_fp8_loads(c, a["sw2"], a["sw2_s"], d.sh_inter, task, row_major=True) + for task in tasks + ] + for t0, nn in token_tiles(s): + load_smid_tile(c, t0, nn, d.sh_inter // 32) + for j in range_constexpr(len(tasks)): + gemv_fp8_mfmas( + c, d.sh_inter, c["shared_xl"], c["shared_xsl"], red, shops[j], rows=nn + ) + gpu.barrier() + if tid < ROWS * nn: + v = row_sum(red, tid % ROWS, tid // ROWS) + at = (j * s + plus(t0, tid // ROWS)) * ROWS + tid % ROWS + fx.ptr_store(bf16_round(v).bitcast(fx.Int32), c["sdl"] + at) + gpu.barrier() + + +def push_ffn(c, t, col, vals): + """A partial bf16 pair of the MoE output at (t, col) into the all-reduce + (``allreduce.push``).""" + d = c["d"] + allreduce.push( + lambda dst, pair, v: c["put_bf"](dst, 2 * pair, v), + c["peer_addr"], + c["ffn_off"], + "ffn", + c["rank"], + t, + c["S"], + col, + vals, + d, + c["two_hop"], + ) + + +@traced +def down_combine(c, task, j, t0, n): + """Down task ``task`` (the CTA's j-th) for tokens t0 .. t0 + n: the routed + contributions summed in bf16 in each token's top-k order, plus the shared expert's + (``down_shared``), the sum pushed to every rank.""" + s, tid, tab = c["S"], c["tid"], c["tab"] + topk = c["rs"].topk + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + shared = fx.ptr_load(c["sdl"] + ((j * s + t) * ROWS + r)).bitcast(fx.Float32) + # aiter's stage 2 adds each bf16(v * w) into a bf16 buffer: a rounding + # at every add (its order is the atomics'; here top-k order) + routed = fx.Float32(0.0) + for k in range_constexpr(topk): + w = c["rs"].route_w(tab, s, t, k) + v = contrib_load(c, j, (t * topk + k) * 16 + r) + routed = bf16_round(routed + bf16_round(v * w)) + fx.ptr_store(bf16_round(routed + shared), c["pair"] + tid) + gpu.barrier() + if tid < ROWS * n // 2: + rp = tid % (ROWS // 2) + tl = tid // (ROWS // 2) + t = plus(t0, tl) + v0 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp)) + v1 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp + 1)) + col = task * ROWS + 2 * rp + push_ffn(c, t, col, [v0, v1]) + gpu.barrier() + + +@traced +def down_mfma(c, j, task, sl, mid, first, count, av, live): + """Down task ``task`` (the CTA's j-th) on a slot of U: the routed MFMA of its + picks first .. first + min(count, 16) (their B operands ``mid``: + ``mid_operands``; the scale block in ``sl``), each pick's contribution.""" + lane, tab, d = c["lane"], c["tab"], c["d"] + sas = lds_scales( + sl, task % 2 * 16 + lane % 16, lane, 0, d.inter // 128, d.down_scale_cols + ) + jj = lane % 16 + col = fx.min(jj, count - 1) + bvs, sbs = mid + # the K steps on DOWN_CHAINS accumulators in turn, summed in order: one + # chain put every FP8-rate MFMA's latency in series + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(DOWN_CHAINS)] + for i in range_constexpr(len(av)): + h = i % DOWN_CHAINS + accs[h] = mfma_scaled(av[i], bvs[i], accs[h], sas[i], sbs[i], FP4, FP8) + acc = accs[0] + for h in range_constexpr(1, DOWN_CHAINS): + acc = acc + accs[h] + pick = c["rs"].pick_at(tab, first + col) + for q in range_constexpr(4): + r = 4 * (lane // 16) + q + if live & (jj < count): + contrib_store(c, j, pick * 16 + r, acc[q]) + + +def contrib_store(c, j, i, v): + """Task j's routed contribution word i (``[pick (t, k)][16 rows]`` f32 bits).""" + fx.ptr_store(v.bitcast(fx.Int32), c["contrib"] + (j * c["rs"].picks * 16 + i)) + + +def contrib_load(c, j, i): + return fx.ptr_load(c["contrib"] + (j * c["rs"].picks * 16 + i)).bitcast(fx.Float32) + + +# ---------------------------------------------------------------- loaders + + +@traced +def load_smid_tile(c, t0, n, groups): + """The shared expert's MXFP8 intermediate of tokens t0 .. t0 + n (SMQ / + SMC, ``smid_quant``) into LDS ``shared_xl`` / ``shared_xsl`` rows 0 .., by every + thread.""" + tid, words = c["tid"], c["d"].sh_inter // 4 + poll_copy( + c, + tid, + THREADS, + [ + (c["smq"], lambda u: plus(t0 * words, u), n * words, c["shared_xl"], words), + ( + c["smc"], + lambda u: plus(t0 * groups, u), + n * groups, + c["shared_xsl"], + groups, + ), + ], + ) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """Columns 128 task ..: the TP ranks' partials, this rank's first, then + rank + 1, + 2, .., in fp32 -> bf16 -> ``out``.""" + s, tid = c["S"], c["tid"] + if const_expr(c["two_hop"]): + # this task's columns' owner sums them for every rank (``allreduce``) + if allreduce.owner(c["d"], task * RED_COLS) == c["rank"]: + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid, "owned") + else: + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid, "gathered") + else: + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid) + + +@traced +def reduce_pair(c, task, idx, how="sum"): + """Thread ``idx``'s column pair of token idx / 64 (``how``: the + ``allreduce`` way).""" + s, d = c["S"], c["d"] + a = c["args"] + if idx < s * RED_COLS // 2: + t = idx // (RED_COLS // 2) + col = task * RED_COLS + 2 * (idx % (RED_COLS // 2)) + own = preg(c["sym"], c["ffn_off"], "ffn") + if const_expr(how == "owned"): + acc0, acc1 = allreduce.reduce_owned( + c, own, c["ffn_off"], "ffn", t, s, col, d + ) + elif const_expr(how == "gathered"): + acc0, acc1 = allreduce.reduce_gathered(c, own, t, s, col, d) + else: + acc0, acc1 = allreduce.reduce_pair(c, own, t, s, col, d) + bo.buffer_store(acc0.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col) + bo.buffer_store(acc1.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col + 1) + + +def moe_smem(s, rs, timeline): + """The MoE stages' LDS (``timeline``: with its own stamp record).""" + + @fx.struct + class Smem: + # an ug task's partial C: 2 tiles a wave (``ug_red_slot``) + red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16] + tab: fx.Array[fx.Int32, rs.tab_words, 16] + pair: fx.Array[fx.Float32, ROWS * s, 16] + # one stage's buffers at a time (``stage_lds``) + stage: fx.Array[fx.Int32, sum(stage_lds(s, rs)["words"].values()), 16] + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] + + return Smem + + +def stage_lds(s, rs): + """Word offsets into ``Smem.stage`` of the buffers only one stage reads: the + shared stage's MXFP8 x tile; the waves' scale blocks (``wsl``, the ug and + down stages'), then the ug stage's routed MXFP8 x tile, or the down stage's + token tile of the shared expert's MXFP8 intermediate (``load_smid_tile``), + each down task's routed contributions (``contrib_store``) and its shared + expert rows (``down_shared``). A barrier ends every stage.""" + d = rs.dims + # the x tile the shared / ug stages hold (every token's up to one tile) + rows = min(s, rs.x_tile) + x8 = rows * lds_row(X_WORDS) + groups = rows * lds_row(X_GROUPS) + wsl = WAVES * wsl_words(d) + # the down stage's token tile of the shared intermediate (``token_tiles``) + srows = tile_rows(s) + shared_xl = srows * lds_row(d.sh_inter // 4) + lds_tail(d.sh_inter) + shared_xsl = srows * lds_row(d.sh_inter // 32) + contrib = 2 * rs.picks * 16 # two down tasks' at most + offsets = { + "x8l": 0, + "x8sl": x8, + "wsl": 0, + "x8rl": wsl, + "x8rsl": wsl + x8, + "shared_xl": wsl, + "shared_xsl": wsl + shared_xl, + "contrib": wsl + shared_xl + shared_xsl, + "sdl": wsl + shared_xl + shared_xsl + contrib, + } + words = { + "shared": x8 + groups, + "ug": wsl + x8 + groups, + "down": wsl + shared_xl + shared_xsl + contrib + 2 * s * ROWS, + } + return {"offsets": offsets, "words": {"all": max(words.values())}} + + +def moe_context(s, rs, lds, mb, bases, layout, scratch, args, rank, sym): + """The MoE stages' context: LDS, mailbox, scratch regions, arguments. Its lane + math starts from a ``fresh`` thread id: in K2 it is not the attention and + seam stages', held from theirs on.""" + tid = fresh(fx.thread_idx.x) + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "tab": lds.tab.ptr, + "pair": lds.pair.ptr, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "tag": mb.tag, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": sym, + "ffn_off": ffn_offset(s, rs.dims), + "two_hop": rs.two_hop, + "rs": rs, + "d": rs.dims, + "args": args, + # x (``ffn_normed``) before this launch: plain loads, no readiness wait + "x_cm": 0, + "x_ready": False, + } + for name, off in stage_lds(s, rs)["offsets"].items(): + c[name] = lds.stage.ptr + off + for region, (off, _) in layout.items(): + c[region] = sreg(scratch, off, region) + return c + + +@traced +def run_moe(c, key, bid, tls, tl0): + """The MoE stages, in order; ``timeline`` stamps from point ``tl0`` + 1.""" + s, timeline, tid = key.tokens, key.timeline, c["tid"] + router_tasks = c["rs"].router_tasks + route0, shared0 = c["rs"].route0, c["rs"].shared0 + shared_tasks = c["d"].shared_tasks + ug_tasks = c["rs"].max_u * c["rs"].ug_parts + for t in range(first_task(bid, XQ0), s, BLOCKS): + await_x(c, [t]) + stage_xq(c, t) + publish(c["put"], c["x8rdy"], t, 1, tid == 0) + stamp(timeline, tls, tid, tl0 + 1) + multi = c["rs"].multi + for task in range(first_task(bid, c["rs"].router0), router_tasks, BLOCKS): + await_x(c, list(range(s))) + stage_router(c, task) + stamp(timeline, tls, tid, tl0 + 2) + for t in range(first_task(bid, route0), s, BLOCKS): + stage_route(c, t) + stamp(timeline, tls, tid, tl0 + 3) + c["xl"], c["xsl"] = c["x8l"], c["x8sl"] + # a CTA's one shared unit (task, token tile) at most: straight-line code, + # whose tile loads LLVM does not hoist into one live set as it would out of + # a task loop (spills). A task's tiles in series on one CTA put three tiles' + # latency on the path down waits on (the weights read again a tile: L2 hits) + assert c["rs"].shared_units <= BLOCKS + unit = first_task(bid, shared0) + if unit < c["rs"].shared_units: + # the gate weights in flight before x's wait + a = c["args"] + gops = gemv_fp8_loads( + c, + a["sgu"], + a["sgu_s"], + HIDDEN, + unit % shared_tasks, + ) + if const_expr(not multi): + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: j, s) + stage_shared(c, unit, gops, 0, s) + else: + x_tile = c["rs"].x_tile + t0 = unit // shared_tasks * x_tile + n = fx.min(fx.Int32(x_tile), s - t0) + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: t0 + j, n) + stage_shared(c, unit % shared_tasks, gops, t0, n) + stamp(timeline, tls, tid, tl0 + 4) + c["xl"], c["xsl"] = c["x8rl"], c["x8rsl"] + if const_expr(multi): + # past one token tile a task gathers its pick tile's rows itself. The + # first round placed (the shared units' CTAs take its last tasks); the + # rest, past a round, from the queue: a CTA a shared or router task + # delayed takes fewer of them + load_route(c, c["tab"]) + n_tiles = c["rs"].tile_first(c["tab"], c["rs"].n_union(c["tab"], s)) + ug_total = mplan.ug_tasks(n_tiles, c["rs"].ug_parts) + task = first_task(bid, c["rs"].ug0) + while task < ug_total: + task = stage_ug_tile(c, task, ug_total) + else: + # one token tile: every task placed (its x in LDS once a CTA) + loaded = fx.Int32(0) + for task in range(first_task(bid, c["rs"].ug0), ug_tasks, BLOCKS): + if loaded == 0: + # X8R (ready long before the routing) first: after the routing + # lands only its own poll remains + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], lambda j: j, s) + load_route(c, c["tab"]) + loaded = fx.Int32(1) + stage_ug(c, task) + stamp(timeline, tls, tid, tl0 + 5) + down_tasks = DOWN_TASKS + assert BLOCKS < down_tasks <= 2 * BLOCKS + # the route table is in LDS already past one token tile (every CTA loads it + # for the ug queue) or when every CTA ran an ug iteration (more ug task + # slots than CTAs), whose first loads it; else load it + if const_expr(not multi and ug_tasks <= BLOCKS): + load_route(c, c["tab"]) + # every CTA its tile; the first down_tasks - BLOCKS to get here an extra + # one too: the CTAs whose ug ended first (by bid, those ending last held + # two tiles and ended the stage) + down_task = first_task(bid, DOWN0) + rank = take_down_rank(c) + if rank < down_tasks - BLOCKS: + stage_down(c, [down_task, BLOCKS + rank]) + else: + stage_down(c, [down_task]) + stamp(timeline, tls, tid, tl0 + 6) + for task in range(first_task(bid, RED0), red_tasks(c["d"]), BLOCKS): + stage_reduce(c, task) + stamp(timeline, tls, tid, tl0 + 7) + + +def await_x(c, tokens): + """In a launch that also writes x (K2: the FFN seam's norm): tokens' rows + ready.""" + if const_expr(c["x_ready"]): + c["poll"]([(c["xrdy"], t, 1) for t in tokens]) + + +def moe_args(*values) -> dict: + """The MoE stages' arguments by name, in ``ARGS`` order.""" + assert len(values) == len(ARGS) + return dict(zip(ARGS, values)) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/moe_plan.py b/vllm/models/deepseek_v4/amd/mono/stages/moe_plan.py new file mode 100644 index 000000000000..e08c0c9f93b3 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/moe_plan.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/moe_plan.py +"""Every per-step decision of which MoE ug flags (UGF) and intermediate rows +(MID) the ug stage writes and the down stage reads, defined once. + +The route table (``stage_route``) gives each expert with picks a slot of U, its +picks a run from ``slot_first``, and, past one token tile, its picks tiles of +``TILE`` from ``tile_first`` (each tile's slot ``tile_slot``). An ug task covers +a unit -- a slot, or past one token tile a pick tile -- and one of ``ug_parts`` parts +of the intermediate; its flag is ``ug_flag``. A down round polls each of its +slots' flags (``polled_tile`` its tiles) before it reads their MID rows. Called +by the kernel on traced Int32 and by the CPU property test on Python ints. +""" + +from vllm.models.deepseek_v4.amd.mono.common.plan import cdiv, imin + +TILE = 16 # picks a pick tile: an MFMA's N + + +def pick_tiles(count, tile=TILE): + """A slot's pick tiles: its ``count`` picks, ``tile`` a tile.""" + return cdiv(count, tile) + + +def ug_parts(groups, ug_groups): + """A unit's ug parts: the expert's ``groups`` FP4 groups of its real width, + ``ug_groups`` a part. A part wholly in the padding is not run: MID's padded + columns stay the step's zeros, against zero weights in down.""" + return cdiv(groups, ug_groups) + + +def ug_flag(unit, part, ug_parts): + """The UGF pair of ug task (``unit``, ``part``), ``ug_parts`` parts a unit.""" + return unit * ug_parts + part + + +def ug_tasks(units, ug_parts): + """A step's ug tasks: every part of each of its ``units``.""" + return units * ug_parts + + +def tile_pick0(slot_first, tile, tile_first, size=TILE): + """Pick tile ``tile``'s first pick, its slot's picks from ``slot_first`` and + tiles (``size`` picks each) from ``tile_first``.""" + return slot_first + (tile - tile_first) * size + + +def polled_tile(tile_first, tiles, j): + """The pick tile a down round polls as a slot's ``j``-th, the slot's + ``tiles`` tiles from ``tile_first``: the last past them.""" + return tile_first + imin(j, tiles - 1) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/moe_shape.py b/vllm/models/deepseek_v4/amd/mono/stages/moe_shape.py new file mode 100644 index 000000000000..3d9d2d8d6147 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/moe_shape.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe_shape.py +"""The MoE stages' build key and route shape: what they are built for (``MoeBuild``), +and the task counts, placements and LDS route table that follow from its +routing (``RouteShape``), with the table's readers. Plain Python but for the +readers; the kernel's stages are in ``moe``.""" + +from dataclasses import dataclass, field + +import flydsl.expr as fx + +from vllm.models.deepseek_v4.amd.mono.common.plan import BLOCKS, WAVES +from vllm.models.deepseek_v4.amd.mono.dims import ( + EXPERTS, + HIDDEN, + TOPK, + UG_PART, + X_GROUPS, + Dims, +) +from vllm.models.deepseek_v4.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v4.amd.mono.stages.gemv import ROWS, TILE, token_tiles +from vllm.models.deepseek_v4.amd.mono.stages.router import PARTS as ROUTER_PARTS + +ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) +# the LDS an x tile's MXFP8 rows may take (``RouteShape.x_tile``): beside them +# the ug stage holds its scale blocks, the route table and the partials +X_TILE_BYTES = 96 * 1024 +MASK_WORDS = 2 # an expert's token bitmap in the route table: tokens <= 64 +# a route lane's experts (experts / 64) sorted descending, ties kept in index +# order: the network by their count +SORT_NETS = { + 6: [ + (0, 1), + (2, 3), + (4, 5), + (0, 2), + (1, 4), + (3, 5), + (0, 1), + (2, 3), + (4, 5), + (1, 2), + (3, 4), + (2, 3), + ], +} + + +@dataclass(frozen=True) +class MoeBuild: + tokens: int + # the TP size: every width a rank holds (``Dims``) + tp: int + timeline: bool = field(default=False, metadata={"sym": "tl"}) + # FP4 groups an ug task (``RouteShape.ug_groups``) + ug_groups: int = 2 + # a hash layer's routing: each token's experts given (``RouteShape``) + hash_route: bool = False + # the all-reduce in two hops (``stages.allreduce``) + two_hop: bool = False + + +@dataclass(frozen=True) +class RouteShape: + """A build's routing: ``experts`` routed experts, ``topk`` a token of + ``tokens``, the task counts and placements that follow, and the route table + ``load_route`` lays out in LDS: every token's (ids, weights), |U|, the + experts of U, each (t, k)'s slot in U, each expert's slot, the picks in slot + order and each slot's first position in it (past one token tile also each + slot's first pick tile and each pick tile's slot).""" + + tokens: int + tp: int + # FP4 groups (32 intermediate columns) an ug task, its waves splitting K in + # halves (2) or quarters (1: twice the tasks) + ug_groups: int = 2 + # a hash layer: each token's top-k experts are given ([S, topk] int32 in + # the bias argument: the layer's tid2eid rows of the step's tokens), the + # router's scores only weigh them + hash_route: bool = False + # the all-reduce in two hops (``stages.allreduce``) + two_hop: bool = False + + @property + def experts(self) -> int: + return EXPERTS + + @property + def topk(self) -> int: + return TOPK + + @property + def dims(self) -> Dims: + return Dims(self.tp) + + @property + def ug_parts(self) -> int: + """An expert's ug tasks (``moe_plan.ug_parts``).""" + return mplan.ug_parts(self.dims.inter_real // UG_PART, self.ug_groups) + + @property + def ug_k_slices(self) -> int: + """The K slices of an ug task's waves: a wave a (group, 16-column half, + slice).""" + return WAVES // (2 * self.ug_groups) + + @property + def picks(self) -> int: + return self.tokens * self.topk + + @property + def max_u(self) -> int: + return min(self.picks, self.experts) + + @property + def router_groups(self) -> tuple[int, int]: + """``stage_router``'s token groups and their tokens: a group a token + tile, as many as one round of CTAs holds (a CTA running two serialized + them: the route waited on the 32 that did at 48 rows).""" + per_group = self.experts // ROWS * ROUTER_PARTS + groups = max(1, min(len(token_tiles(self.tokens)), BLOCKS // per_group)) + tokens = -(-self.tokens // groups) + assert (groups - 1) * tokens < self.tokens + return groups, tokens + + @property + def router_tasks(self) -> int: + # (token group, 16-row tile, K part): ``stage_router`` + return self.router_groups[0] * self.experts // ROWS * ROUTER_PARTS + + @property + def router0(self) -> int: + """Past the xq tasks (a token each from CTA 0).""" + return max(ROUTER0, self.tokens) + + @property + def route0(self) -> int: + return self.router0 + self.router_tasks + + @property + def shared0(self) -> int: + return self.route0 + max(8, self.tokens) + + @property + def x_tile(self) -> int: + """The rows of an x tile in LDS -- the shared stage's token tile, the + ug stage's pick tile: an MFMA's N (16) while they fit + ``X_TILE_BYTES``, else 8 (a wide hidden: the MFMA's columns past + them repeat the last).""" + row = HIDDEN + X_GROUPS + return TILE if TILE * row <= X_TILE_BYTES else TILE // 2 + + @property + def x_tiles(self) -> int: + """A step's x tiles.""" + return -(-self.tokens // self.x_tile) + + @property + def shared_units(self) -> int: + """The shared expert's (task, token tile) units (``stage_shared``).""" + return self.dims.shared_tasks * self.x_tiles + + @property + def ug0(self) -> int: + """The first ug round's placement: the shared units' CTAs last.""" + return self.shared0 + self.shared_units + + @property + def sort_net(self): + return SORT_NETS[self.experts // 64] + + @property + def tab_words(self) -> int: + return self.tile_slot0 + self.max_tiles if self.multi else self.queue0 + 1 + + @property + def multi(self) -> bool: + """More than one x tile: an ug task is a (tile of a slot's picks, part) + (``stage_ug``).""" + return self.tokens > self.x_tile + + @property + def slot_tiles(self) -> int: + """A slot's pick tiles at most (a token picks an expert once).""" + return self.x_tiles + + @property + def max_tiles(self) -> int: + """The pick tiles of U at most: a slot's partial tile each, then the + whole tiles, and never more than a tile a pick.""" + return min(self.picks, self.max_u + self.picks // self.x_tile) + + @property + def tile_first0(self) -> int: + """Past one token tile: each slot's first pick tile, then their count.""" + return self.queue0 + 1 + + @property + def tile_slot0(self) -> int: + """Past one token tile: each pick tile's slot of U.""" + return self.tile_first0 + self.max_u + 1 + + @property + def queue0(self) -> int: + """The ug task this CTA took from the queue (``take_ug_task``).""" + return self.mask0 + MASK_WORDS * self.experts + + @property + def expert_slot0(self) -> int: + """The route table's per-expert slots (a slot where the expert is in U).""" + return 3 * self.picks + 1 + self.max_u + + @property + def order0(self) -> int: + """The picks (t topk + k) in slot order, a slot's ascending.""" + return self.expert_slot0 + self.experts + + @property + def first0(self) -> int: + """Each slot's first position in the pick order, then the pick count.""" + return self.order0 + self.picks + + @property + def mask0(self) -> int: + """Each expert's tokens that picked it, a bit a token (``MASK_WORDS``).""" + return self.first0 + self.max_u + 1 + + @property + def tile_picks(self) -> int: + """A pick tile's picks at most: a token picks an expert once.""" + return min(self.tokens, self.x_tile) + + def route_w(self, tab, s, t, k): + return fx.ptr_load(tab + (t * 2 * self.topk + self.topk + k)).bitcast( + fx.Float32 + ) + + def n_union(self, tab, s): + return fx.ptr_load(tab + 2 * s * self.topk) + + def union_expert(self, tab, s, u): + return fx.ptr_load(tab + (2 * s * self.topk + 1 + u)) + + def slot_first(self, tab, u): + return fx.ptr_load(tab + (self.first0 + u)) + + def pick_at(self, tab, i): + return fx.ptr_load(tab + (self.order0 + i)) + + def tile_first(self, tab, u): + return fx.ptr_load(tab + (self.tile_first0 + u)) + + def tile_slot(self, tab, i): + return fx.ptr_load(tab + (self.tile_slot0 + i)) + + +def route_shape(key: MoeBuild) -> RouteShape: + return RouteShape(key.tokens, key.tp, key.ug_groups, key.hash_route, key.two_hop) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/router.py b/vllm/models/deepseek_v4/amd/mono/stages/router.py new file mode 100644 index 000000000000..bce9d67e717e --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/router.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/model_ops/deepseek_v41/router.py +"""The router GEMV tile for decode: bf16 x [T, K] times the gate weight +[experts, K] -> f32 logits, in one fixed summation order (K the hidden width). + +Rows 16 a tile, K in ``PARTS`` contiguous parts, 16x16x32 bf16 MFMA: in part +p wave w takes the part's K steps w, w + 8, ... (32 columns each) in that +order, and the eight wave partials are summed in wave order +(``common.ops.row_sum``); the MoE's route stage sums the parts in part +order. The MoE stages spread a tile's parts over PARTS CTAs, and past 16 tokens +a token tile a CTA too, which the order allows: a CTA alone cannot read a +tile, or every token's x, fast enough. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.ops import mfma_bf16, rsrc +from vllm.models.deepseek_v4.amd.mono.common.plan import WAVES + +ROWS = 16 # a tile's rows: the MFMA's M +PARTS = 4 # contiguous K parts of a tile, summed in order + + +def _k0s(k, part, lane, wave): + """The K offsets of this lane's loads in part ``part`` (8 bf16 each).""" + per_wave = k // 32 // WAVES // PARTS + assert per_wave * 32 * WAVES * PARTS == k + return [ + (part * per_wave * WAVES + wave + WAVES * i) * 32 + 8 * (lane // 16) + for i in range(per_wave) + ] + + +def router_weight_loads(w, k, task, part, lane, wave, w_cm=0): + """``router_tile``'s weight operands (rows 16 task .., K part ``part``). + ``w_cm``: their load cache policy (nontemporal where nothing else of the + step reuses them).""" + row = task * ROWS + lane % 16 + return [ + fx.Vector( + bo.buffer_load( + rsrc(w), + (row * k + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=w_cm, + ) + ) + for k0 in _k0s(k, part, lane, wave) + ] + + +def router_x_loads(x, k, part, s, lane, wave, x_cm=0, t0=0): + """``router_tile``'s x operands: the s tokens from ``t0`` (column + lane % 16; ``t0`` traced or an int). ``x_cm``: their load cache policy + (device scope when another CTA of the same launch wrote x).""" + tok = fx.min(lane % 16, s - 1) + if const_expr(not isinstance(t0, int) or t0 != 0): + tok = tok + t0 + return [ + fx.Vector( + bo.buffer_load( + rsrc(x), + (tok * k + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=x_cm, + ) + ) + for k0 in _k0s(k, part, lane, wave) + ] + + +def router_tile_mfma(wvs, xvs, lane, wave, red): + """``router_tile``'s MFMAs, K steps in order -> ``red[wave][lane]``.""" + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(len(wvs)): + acc = mfma_bf16(wvs[i].bitcast(fx.BFloat16), xvs[i].bitcast(fx.BFloat16), acc) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) diff --git a/vllm/models/deepseek_v4/amd/mono/stages/seam.py b/vllm/models/deepseek_v4/amd/mono/stages/seam.py new file mode 100644 index 000000000000..6a9cda403b41 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono/stages/seam.py @@ -0,0 +1,427 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (slice, gate, norm) and +# atom/models/deepseek_v41/mono/kernels/attn_post.py (the attention's TP reduce). +"""The mono layer's mHC seam stages, each a task table of a persistent launch +handing off through the tagged mailbox: + + slice (224, 32 hidden columns x 4 streams each): the owed post folded into + the residual (R_new, bf16 out at device scope) and the partial mix + projection / sum of squares (mailbox PMIX) + gate (S): the 224 partials -> post / comb (Sinkhorn) and the pre-mix, + also to the mailbox PREM + norm (S): PREM, then R_new's four streams -> the collapse -> the RMSNorm + -> the bf16 row, handed to the caller's store + reduce (224): the FFN seam's attention output, the TP ranks' partials + (pushed to the ATTN region) summed in the all-reduce's order -> + ``pend_lds``, where its slice reads it + +The collapse takes the gate's own pre-mix (aiter's ``mhc_pre``: the four streams +by pre in fp32, then bf16). Every other elementwise step follows aiter's +``mhc_fused_post_pre_delayed_rmsnorm`` rounding; the mix projection and the sums +of squares, split-K reductions, differ only in their split and order. +""" + +import math + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v4.amd.mono.common.ops import ( + CM_DEV, + batched_rounds, + bf16_round, + bf_hi, + bf_lo, + butterfly, + div_rn, + hw_exp2, + hw_rcp, + hw_rsq, + lane_gather, + rsrc, + traced, +) +from vllm.models.deepseek_v4.amd.mono.common.plan import ( + THREADS, + WAVES, + cdiv, + pair_layout, +) +from vllm.models.deepseek_v4.amd.mono.common.sync import preg +from vllm.models.deepseek_v4.amd.mono.dims import ( + HC, + HC_EPS, + HC_POST_MULT, + HC_SINKHORN, + HIDDEN, + NORM_EPS, + Dims, +) +from vllm.models.deepseek_v4.amd.mono.stages import allreduce +from vllm.models.deepseek_v4.amd.mono.stages.gemv import ( + ld_bf, + ld_f32, + plus, + token_tiles, +) + +MIX = HC * (HC + 2) # 24 projections: pre 4 | post 4 | comb 16 +LOG2E = 1.0 / math.log(2.0) + +COLS = 32 # hidden columns a slice task +SLICES = HIDDEN // COLS # 224 +KT = HC * COLS # projection columns a slice task +PART = MIX + 1 # a slice's partials a token: the mixes, then the sum of squares +# slices a gate thread sums: its groups' partials take at most THREADS threads +GATE_BLOCK = 14 +NORM_CHUNK = 2048 # a norm pass's columns: 256 threads x 8 +assert SLICES % GATE_BLOCK == 0 and SLICES // GATE_BLOCK * PART <= THREADS +# the seam's mailbox regions +REGIONS = ("pmix", "prem") + + +def scratch_layout(tokens: int) -> dict[str, tuple[int, int]]: + """Mailbox region -> (byte offset, bytes) of a seam; 8 B a pair.""" + pairs = { + "pmix": SLICES * tokens * PART, + "prem": tokens * HC, # the gate's pre-mix, the norm's collapse + } + return pair_layout(pairs.items()) + + +def attn_peer_bytes(tokens: int, d: Dims) -> int: + """ATTN: the attention output's all-reduce (``allreduce``).""" + return allreduce.region_bytes(tokens, d) + + +def fma(a, b, c): + return fx.Float32(fmath.fma(a, b, c)) + + +def sigmoid(z): + """Triton's ``tl.sigmoid``: 1 / (1 + exp2(-z log2(e))).""" + return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) + + +def lane_f32(v, src): + return lane_gather(v.bitcast(fx.Int32), src).bitcast(fx.Float32) + + +def pend_value(c, t, col, cc): + """The owed post's sublayer output at (token t, column col): read from + ``pend``, or where a caller staged it (``pend_lds``, ``[token][32]`` of this + task's columns: a reduction it did first).""" + if const_expr(c.get("pend_lds") is None): + return ld_bf(c["args"]["pend"], t * HIDDEN + col) + return fx.ptr_load(c["pend_lds"] + (t * COLS + cc)) + + +# ---------------------------------------------------------------- slice +def slice_at(c0, e): + """Element e of a slice task at column c0: (token, stream, hidden column).""" + return e // KT, (e % KT) // COLS, c0 + e % COLS + + +def slice_loads(c, c0, e): + """Element e's loads of ``stage_slice`` (e clamped into the step's).""" + a, hidden = c["args"], HIDDEN + t, o, col = slice_at(c0, e) + if const_expr(not c.get("fold", True)): + return (ld_bf(a["res_in"], (t * HC + o) * hidden + col),) + r = [ld_bf(a["res_in"], (t * HC + h) * hidden + col) for h in range(HC)] + cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] + return r, cm, ld_f32(a["post_in"], t * HC + o) + + +def slice_row(c, c0, e, got): + """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" + a = c["args"] + t, o, col = slice_at(c0, e) + if const_expr(not c.get("fold", True)): + # a settled residual (no owed post): stream o as it arrives + v = got[0] + else: + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + # the norms' collapse reads R_new back from res_out in other CTAs: device + # scope + bo.buffer_store( + v.to(fx.BFloat16), + rsrc(a["res_out"]), + (t * HC + o) * HIDDEN + col, + cache_modifier=CM_DEV, + ) + fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) + + +@traced +def stage_slice(c, task): + """R_new and the partial projections of hidden columns 32 task .. +32, all + four streams, every token.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl = c["rl"], c["fl"] + a, hidden = c["args"], HIDDEN + c0 = task * COLS + batched_rounds( + tid, + s * KT, + lambda e: slice_loads(c, c0, e), + lambda e, got: slice_row(c, c0, e, got), + ) + for i in range_constexpr((MIX * KT + THREADS - 1) // THREADS): + e = tid + THREADS * i + if e < MIX * KT: + j = e // KT + o = (e % KT) // COLS + # the original's two bf16 MFMAs take fn as hi = bf16(fn) and lo = + # bf16(fn - hi): their sum, exact in fp32, is the weight it applies + f = ld_f32(a["hc_fn"], j * HC * hidden + o * hidden + c0 + e % COLS) + hi = bf16_round(f) + fx.ptr_store(hi + bf16_round(f - hi), fl + e) + # R_new landed device-wide before any partial announces it + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + # sum of squares: token t on wave t % WAVES, two columns a lane + for t in range_constexpr(s): + if wave == t % WAVES: + x0 = fx.ptr_load(rl + (t * KT + 2 * lane)) + x1 = fx.ptr_load(rl + (t * KT + 2 * lane + 1)) + sq = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) + if lane == 0: + c["put"](c["pmix"], (task * s + t) * PART + MIX, sq) + # mixes: fp32 MFMA 16x16x4, A = fn rows (two 16-row tiles), B = R_new^T, + # a token tile at a time + for t0, n in token_tiles(s): + slice_mix(c, task, t0, n) + + +@traced +def slice_mix(c, task, t0, n): + """The mix projection partials of tokens t0 .. t0 + n (``stage_slice``).""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl, red = c["rl"], c["fl"], c["red"] + tok = plus(t0, fx.min(lane % 16, n - 1)) + kw = KT // WAVES + cs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + for kk in range_constexpr(kw // 4): + colk = wave * kw + lane // 16 + 4 * kk + b = fx.ptr_load(rl + (tok * KT + colk)) + for h in range_constexpr(2): + row = h * 16 + lane % 16 + av = fx.ptr_load(fl + (fx.min(row, MIX - 1) * KT + colk)) + av = (row < MIX).select(av, fx.Float32(0.0)) + ops = rocdl._split_mfma_operands([av, b, cs[h], 0, 0, 0]) + cs[h] = fx.Vector(rocdl.mfma_f32_16x16x4f32(T.vec(4, T.f32), *ops).result) + for h in range_constexpr(2): + fx.ptr_store(cs[h], red + ((wave * 2 + h) * 64 + lane) * 4) + gpu.barrier() + if tid < MIX * n: + tl = tid // MIX + t = plus(t0, tl) + j = tid % MIX + h = j // 16 + ln = 16 * ((j % 16) // 4) + tl + tot = fx.Float32(0.0) + for w in range_constexpr(WAVES): + tot = tot + fx.ptr_load(red + (((w * 2 + h) * 64 + ln) * 4 + j % 4)) + c["put"](c["pmix"], (task * s + t) * PART + j, tot) + gpu.barrier() + + +# ---------------------------------------------------------------- gate +def _row(v, lane, q): + """Comb lanes 8..23 hold ``comb[r][c]`` at lane 8 + 4 r + c: element q of + the row.""" + return lane_f32(v, (lane - 8) // 4 * 4 + 8 + q) + + +def _col(v, lane, q): + return lane_f32(v, 8 + 4 * q + (lane - 8) % 4) + + +@traced +def stage_gate(c, t): + """Token t's 224 slices, in slice order, -> post / comb / pre.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + red = c["red"] + a = c["args"] + blk = GATE_BLOCK + groups = SLICES // blk # 16 + if tid < groups * PART: + g = tid // PART + j = tid % PART + vals = c["poll"]( + [(c["pmix"], ((g * blk + k) * s + t) * PART + j, 1) for k in range(blk)] + ) + acc = vals[0][0].bitcast(fx.Float32) + for k in range_constexpr(1, blk): + acc = acc + vals[k][0].bitcast(fx.Float32) + fx.ptr_store(acc, red + (g * PART + j)) + gpu.barrier() + if wave == 0: + j = fx.min(lane, MIX - 1) + mix = fx.ptr_load(red + j) + sq = fx.ptr_load(red + MIX) + for g in range_constexpr(1, groups): + mix = mix + fx.ptr_load(red + (g * PART + j)) + sq = sq + fx.ptr_load(red + (g * PART + MIX)) + sc = [ld_f32(a["hc_scale"], i) for i in range(3)] + base = ld_f32(a["hc_base"], j) + # aiter's mhc_fused_post_pre_delayed_rmsnorm reduce: one rstd = + # rsq(fma(sum, 1/K, eps)), then each gate fma(mix rstd, scale, base) + v = mix * hw_rsq(fma(sq, fx.Float32(1.0 / (HC * HIDDEN)), fx.Float32(NORM_EPS))) + if lane < HC: + pre = sigmoid(fma(v, sc[0], base)) + fx.Float32(HC_EPS) + bo.buffer_store(pre, rsrc(a["pre_out"]), t * HC + lane) + c["put"](c["prem"], t * HC + lane, pre) + if (lane >= HC) & (lane < 2 * HC): + post = sigmoid(fma(v, sc[1], base)) + bo.buffer_store( + post * fx.Float32(HC_POST_MULT), + rsrc(a["post_out"]), + t * HC + lane - HC, + ) + cv = fma(v, sc[2], base) + row = [_row(cv, lane, q) for q in range(4)] + # Triton's exp: exp2 of (x - max) log2(e) + m = fx.max(fx.max(row[0], row[1]), fx.max(row[2], row[3])) + cv = hw_exp2((cv - m) * fx.Float32(LOG2E)) + # fast_dividef is x rcp(y); its first "+ eps" contracts into the fma + row = [_row(cv, lane, q) for q in range(4)] + eps = fx.Float32(HC_EPS) + cv = fma(cv, hw_rcp((row[0] + row[1]) + (row[2] + row[3])), eps) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((col[0] + col[1]) + (col[2] + col[3])) + eps) + for _ in range_constexpr(HC_SINKHORN - 1): + row = [_row(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((row[0] + row[1]) + (row[2] + row[3])) + eps) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((col[0] + col[1]) + (col[2] + col[3])) + eps) + if (lane >= 2 * HC) & (lane < MIX): + bo.buffer_store(cv, rsrc(a["comb_out"]), t * HC * HC + lane - 2 * HC) + gpu.barrier() + + +# ---------------------------------------------------------------- norm +@traced +def stage_norm(c, t, store): + """Token t's input collapse (``collapse_chunk``) and its RMSNorm, aiter's + ``mhc_pre``: rstd = rsq(sum / H + eps) of the fp32 collapse, bf16((x w) + rstd); 256 threads, chunks of 8 at 8 t + 2048 c, each chunk's bf16 values + handed to ``store(t, col, ys, live)`` by every thread (``live``: a thread's + chunk in the row; threads 256.. mirror thread 255).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, hidden = c["args"], HIDDEN + nch = cdiv(hidden, NORM_CHUNK) + tt = fx.min(tid, 255) + mine = tid < 256 + cols = [tt * 8 + NORM_CHUNK * ch for ch in range(nch)] + xs = [] + acc = fx.Float32(0.0) + got = c["poll"]([(c["prem"], t * HC + h, 1) for h in range(HC)]) + pre = [got[h][0].bitcast(fx.Float32) for h in range(HC)] + for ch in range_constexpr(nch): + col = fx.min(cols[ch], hidden - 8) + xf = collapse_chunk(c, t, col, pre) + xf = [(cols[ch] < hidden).select(x, fx.Float32(0.0)) for x in xf] + for x in xf: + acc = acc + x * x + xs.append([bf16_round(x) for x in xf]) + acc = butterfly(acc, (1, 2, 4, 8, 16, 32)) + if (lane == 0) & mine: + fx.ptr_store(acc, red + wave) + gpu.barrier() + tot = butterfly(fx.ptr_load(red + lane % 4), (1, 2)) + mean = div_rn(tot, fx.Float32(float(hidden)), fx.Float32(1.0 / hidden)) + r = hw_rsq(mean + fx.Float32(NORM_EPS)) + for ch in range_constexpr(nch): + col = fx.min(cols[ch], hidden - 8) + ys = [ + bf16_round((xs[ch][j] * ld_bf(a["norm_w"], col + j)) * r) for j in range(8) + ] + store(t, col, ys, mine & (cols[ch] < hidden)) + gpu.barrier() + + +def collapse_chunk(c, t, col, pre): + """Token t's layer input at columns col .. col + 8, fp32: R_new's four + streams (res_out, written device-wide by this launch's slices, whose + partials the gate waited on) by the pre-mix, summed pairwise as aiter's + ``mhc_pre`` reduces across its stream lanes.""" + a, hidden = c["args"], HIDDEN + rs = [] + for h in range(HC): + w = fx.Vector( + bo.buffer_load( + rsrc(a["res_out"]), + ((t * HC + h) * hidden + col) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + v = [] + for k in range(4): + v += [bf_lo(w[k]), bf_hi(w[k])] + rs.append(v) + return [ + (pre[0] * rs[0][j] + pre[1] * rs[1][j]) + + (pre[2] * rs[2][j] + pre[3] * rs[3][j]) + for j in range(8) + ] + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """The attention output at this slice's 32 columns: the TP ranks' partials + in the all-reduce's order (``sum_partials``), in fp32 -> bf16 -> + ``pend_lds``.""" + s, tid = c["S"], c["tid"] + if const_expr(c.get("two_hop", False)): + # this slice's columns' owner sums them for every rank (``allreduce``) + if allreduce.owner(c["d"], task * COLS) == c["rank"]: + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid), "owned") + else: + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid), "gathered") + else: + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid)) + gpu.barrier() + + +@traced +def reduce_attn_pair(c, task, idx, how="sum"): + """Thread ``idx``'s column pair of token idx / 16 (``stage_reduce``; + ``how``: the ``allreduce`` way, ``reduce_pair`` / ``_owned`` / + ``_gathered``).""" + s, d = c["S"], c["d"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = c0 + 2 * (idx % (COLS // 2)) + own = preg(c["sym"], 0, "attn") + if const_expr(how == "owned"): + acc0, acc1 = allreduce.reduce_owned(c, own, 0, "attn", t, s, col, d) + elif const_expr(how == "gathered"): + acc0, acc1 = allreduce.reduce_gathered(c, own, t, s, col, d) + else: + acc0, acc1 = allreduce.reduce_pair(c, own, t, s, col, d) + lds = c["pend_lds"] + (t * COLS + col - c0) + fx.ptr_store(bf16_round(acc0), lds) + fx.ptr_store(bf16_round(acc1), lds + 1) diff --git a/vllm/models/deepseek_v4/amd/mono_decode.py b/vllm/models/deepseek_v4/amd/mono_decode.py new file mode 100644 index 000000000000..b9eb6d11fd5c --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mono_decode.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4-Pro mono decode layer on ROCm CDNA4 (``--monokernel``). + +A decode step's backbone layer -- the attention seam, the attention, its TP +all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent +FlyDSL launches (``mono``, this directory) on the layer's loaded weights, the +TP reductions inside the kernels. Every layer takes the whole layer: between +its two launches vLLM's compressor runs on K1's normed input (and on a CSA +layer its indexer, from K1's qr, selects the top-k), then K2 builds the key +table -- the CSA layer's top-k through the block table, the HCA layer's every +compressed row (its metadata's slots). The first three layers route by hash: +their tokens' experts are gathered here and handed to the MoE. + +A step takes the kernels when it is decode only, at most ``MAX_TOKENS`` rows +(eager or a FULL CUDA graph: the model is not torch.compiled), with causal SWA +windows. Anything else runs the layer as before; the paths share every tensor +at the layer boundary, so they interleave freely. Every TP rank must take the +same path at every layer (the kernels wait on each other's pushes): the +conditions read only metadata every rank holds alike. +""" + +from typing import TYPE_CHECKING, Any + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v4.amd.mono.dims import ( + EXPERTS, + MAX_TOKENS, + ROUTED_SCALE, + SWIGLU_LIMIT, + TOPK, + WINDOW, + Dims, +) +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v4.amd.model import DeepseekV4DecoderLayer + +logger = init_logger(__name__) + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +TP_SIZE = 8 +CSA, HCA = 4, 128 # the layers' compress ratios: top-k / every compressed row +BLOCK = 128 # the checkpoint's FP8 block +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSION = 4 + +_runner = None + + +def _mono_runner(device: torch.device, max_topk: int): + """The process's mono runner (one a TP rank, shared by every layer; + ``max_topk``: its widest key table). Its peer-memory handshake is collective + over the TP group: every rank builds it at the same layer of the same + step.""" + global _runner + if _runner is None: + if torch.cuda.is_current_stream_capturing(): + # its peer memory is allocated and exchanged eagerly; vLLM warms + # every graph up eagerly before capturing it + raise RuntimeError( + "--monokernel was first reached inside a CUDA graph " + "capture, before an eager decode step." + ) + from vllm.models.deepseek_v4.amd.mono.runner import MonoRunner + + _runner = MonoRunner( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + get_tp_group().cpu_group, + device, + max_topk=max_topk, + ) + return _runner + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + """E8M0 scales as their bytes: a float8_e8m0fnu / uint8 tensor's own, or the + exponents of a power-of-two fp32 one (vLLM's aiter linears upcast them).""" + if scale.dtype == torch.float32: + bits = scale.contiguous().view(torch.int32) + code = (bits >> 23).to(torch.uint8) + assert torch.equal(bits, code.to(torch.int32) << 23), "not powers of two" + return code + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _mx_scales(scale: torch.Tensor) -> torch.Tensor: + """The checkpoint's 128 x 128 E8M0 weight scales as the kernels' 32 x 32 + ones: each code repeated over its block (exact: E8M0 is a power of two).""" + r = BLOCK // 32 + codes = _scale_bytes(scale) + return codes.repeat_interleave(r, 0).repeat_interleave(r, 1).contiguous() + + +def _block_scale(lin) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_fp8_block_weight_scale, + ) + + scale = get_fp8_block_weight_scale(lin) + assert scale is not None, f"{type(lin).__name__} without block scales" + return _mx_scales(scale) + + +def _hca_width(vllm_config) -> int: + """The HCA layers' key table width: vLLM's c128a table's, every compressed + row of the longest context (``sparse_mla``).""" + from vllm.models.deepseek_v4.sparse_mla import _C128A_TOPK_ALIGNMENT + + n = -(-vllm_config.model_config.max_model_len // HCA) + return -(-n // _C128A_TOPK_ALIGNMENT) * _C128A_TOPK_ALIGNMENT + + +def _compute_units() -> int: + return current_platform.num_compute_units(torch.accelerator.current_device_index()) + + +def _deployment_error(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None: + """Why this deployment cannot run the kernels at all, or None (an explicit + opt-in that cannot run is an error). They read every routed expert on each + rank, TP-sharded, in AITER's A8W4 MXFP4 layout, and vLLM's preshuffled FP8 + linears with padded rows.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + from vllm.platforms.rocm import get_cdna_version + + if get_cdna_version() != CDNA_VERSION: + return f"needs CDNA{CDNA_VERSION}" + if get_tensor_model_parallel_world_size() != TP_SIZE: + return f"needs tensor parallel size {TP_SIZE}" + pc = vllm_config.parallel_config + if pc.enable_expert_parallel or pc.enable_eplb or pc.data_parallel_size > 1: + return "needs tensor parallelism only: no expert or data parallelism" + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + return "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + experts = getattr(getattr(layer.ffn, "experts", None), "routed_experts", None) + if experts is not None: # else not a routed MoE: the layer checks decline it + backend = getattr(experts.quant_method, "mxfp4_backend", None) + if backend != Mxfp4MoeBackend.AITER_MXFP4_BF16: + return f"needs the AITER_MXFP4_BF16 MoE backend, not {backend}" + if not rocm_aiter_ops.is_linear_fp8_enabled(): + return "needs AITER's FP8 linears (VLLM_ROCM_USE_AITER_LINEAR)" + if not envs.VLLM_ROCM_FP8_PADDING: + return "needs vLLM's padded FP8 weight rows (VLLM_ROCM_FP8_PADDING)" + try: + from vllm.models.deepseek_v4.amd.mono.runner import BLOCKS + except ImportError as err: + return f"needs FlyDSL and AITER's FlyDSL helpers ({err})" + # every CTA of a launch stays resident: a partitioned GPU deadlocks + cus = _compute_units() + if cus < BLOCKS: + return f"needs {BLOCKS} compute units, the GPU has {cus}" + return None + + +def _select(layer: "DeepseekV4DecoderLayer", positions, normed, qr) -> None: + """The layer's compressor -- and a CSA layer's indexer, its top-k into the + shared top-k buffer -- from K1's normed input and bf16 qr, serial.""" + attn = layer.attn + comp = attn.compressor + assert comp is not None + score = torch.mm(normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32) + comp(score, positions, attn.rotary_emb) + indexer = attn.indexer + if indexer is None: + return + score = torch.mm( + normed, indexer.compressor.fused_wkv_wgate.weight.T, out_dtype=torch.float32 + ) + weights, _ = indexer.weights_proj(normed) + q, q_scale, weights = indexer( + normed, qr, score, weights, positions, attn.indexer_rotary_emb + ) + if q is not None: # else the short context's top-k is every row, filled + assert weights is not None + indexer.indexer_op( + normed, (q, q_scale) if q_scale is not None else q, None, weights + ) + + +class MonoDecodeLayer: + """One DeepSeek-V4-Pro decoder layer's mono path (``create``: None when the + layer is not eligible). Its weights are taken from the layer at the first + call, after loading and vLLM's preshuffle.""" + + def __init__(self, hca_width: int) -> None: + self.hca_width = hca_width + self._weights: Any = None + + @classmethod + def create( + cls, layer: "DeepseekV4DecoderLayer", vllm_config + ) -> "MonoDecodeLayer | None": + if not vllm_config.kernel_config.monokernel: + return None + why = _deployment_error(layer, vllm_config) + if why is not None: + raise ValueError(f"--monokernel {why}.") + ffn, attn = layer.ffn, layer.attn + config = vllm_config.model_config.hf_config + if extract_layer_index(attn.prefix) >= config.num_hidden_layers: + why = "draft layer" + elif not (layer.use_fused_mhc and layer.fuse_mhc_rmsnorm): + why = "needs aiter's fused mHC seam with its norm" + elif ( + ffn.use_mega_moe + or ffn.shared_experts is None + or ffn.fuse_heterogeneous_shared_expert + or ffn.is_fused_shared_expert_enabled + or ffn.gate.bias_vl is not None + ): + why = "needs the separate FP8 shared expert and the AITER MXFP4 MoE" + elif ( + ffn.n_routed_experts != EXPERTS + or ffn.n_activated_experts != TOPK + or ffn.routed_scaling_factor != ROUTED_SCALE + or ffn.swiglu_limit != SWIGLU_LIMIT + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x2.5)" + elif attn.compress_ratio not in (CSA, HCA) or ( + (attn.compress_ratio == CSA) != (attn.indexer is not None) + ): + why = f"compress ratio {attn.compress_ratio}" + elif attn.kv_cache_dtype != "fp8_ds_mla": + why = f"a {attn.kv_cache_dtype} KV cache" + if why is not None: + logger.info_once("DSv4 mono decode off for some layers: %s", why) + return None + logger.info_once("DSv4 mono decode: the whole backbone layers") + return cls(_hca_width(vllm_config)) + + def weights(self, layer: "DeepseekV4DecoderLayer"): + if self._weights is not None: + return self._weights + from vllm.models.deepseek_v4.amd.mono.runner import ( + AttnWeights, + MonoLayerWeights, + ) + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + # the kernels read vLLM's preshuffled wqkv / wo_b / shared gate_up + assert a._wqa_wkv_scale is not None and a._wo_b_scale is not None + assert sh is not None and sh._gateup_scale is not None + d = Dims(get_tensor_model_parallel_world_size()) + attn = AttnWeights( + layer_id=extract_layer_index(a.prefix), + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_mx_scales(a._wqa_wkv_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_block_scale(a.wq_b), + wo_a=a.wo_a.weight, + wo_a_scale=_block_scale(a.wo_a), + wo_b=a.wo_b.weight, + wo_b_scale=_mx_scales(a._wo_b_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(d) + self._weights = MonoLayerWeights( + attn=attn, + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_mx_scales(sh._gateup_scale), + sw2=sh.down_proj.weight, + sw2_s=_block_scale(sh.down_proj), + ) + self._weights.check(d) + return self._weights + + @staticmethod + def _skip(why: str) -> None: + logger.debug_once("DSv4 mono decode skipped a step: %s", why) + return None + + @staticmethod + def _decode_step(layer: "DeepseekV4DecoderLayer", rows: int): + """(why this step takes the original path -- None for a decode-only step + the kernels take -- the layer's SWA metadata, its compressed cache's).""" + if not 1 <= rows <= MAX_TOKENS: + return f"{rows} rows", None, None + if not is_forward_context_available(): + return "no forward context", None, None + fc = get_forward_context() + md = fc.attn_metadata + attn = layer.attn + swa: Any = None + comp: Any = None + if isinstance(md, dict): + swa = md.get(attn.swa_cache_layer.prefix) + comp = md.get(attn.prefix) + why = None + if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + # a PIECEWISE (breakable) capture is replayed for mixed batches too + why = f"{fc.cudagraph_runtime_mode} step" + elif not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + why = "no attention metadata" + elif swa is None or comp is None: + why = "layer metadata missing" + elif swa.num_prefills != 0 or swa.num_decodes == 0: + why = "not a decode-only step" + elif swa.num_decode_tokens > rows: + why = "more decode tokens than rows" + elif swa.decode_swa_width != WINDOW: + why = f"SWA width {swa.decode_swa_width}" + elif ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + why = "SWA decode metadata missing" + elif ( + attn.compress_ratio == HCA + and getattr(comp, "c128a_global_decode_topk_indices", None) is None + ): + why = "C128A decode metadata missing" + return why, swa, comp + + def __call__( + self, + layer: "DeepseekV4DecoderLayer", + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs (x, residual, post_mix, res_mix) by the mono + kernels, or None when this step takes the original path.""" + attn, M = layer.attn, x.shape[0] + first = residual is None + why, swa, comp = self._decode_step(layer, M) + tensors = (x,) if first else (x, residual, post_mix, res_mix) + contiguous = all(t is not None and t.is_contiguous() for t in tensors) + if why is None and not contiguous: + why = "non-contiguous layer inputs" + hash_route = layer.ffn.gate.tid2eid is not None + if why is None and hash_route and input_ids is None: + why = "a hash layer without its input ids" + if why is not None: + return self._skip(why) + assert swa is not None and comp is not None + runner = _mono_runner(x.device, self.hca_width) + logger.info_once( + "DSv4 mono decode layer active (decode steps of <= %d rows)", MAX_TOKENS + ) + if first: + # the embedding broadcast to the four streams: settled, no post owed + # (the mixes unread without the fold) + residual, post_mix, res_mix = x, x, x + assert residual is not None and post_mix is not None and res_mix is not None + route = None + if hash_route: + assert input_ids is not None + route = layer.ffn.gate.tid2eid[input_ids[:M]].contiguous() + if attn.compress_ratio == HCA: + topk = comp.c128a_global_decode_topk_indices[:M, 0] + kw: dict[str, Any] = dict(topk=self.hca_width, global_keys=True) + else: + topk = attn.topk_indices_buffer + kw = dict(comp_block_table=comp.block_table) + out = runner.forward( + self.weights(layer), + x, + residual, + post_mix, + res_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + topk_indices=topk, + comp_cache=_record_view( + attn.kv_cache, comp.block_size // attn.compress_ratio + ), + select=lambda normed, qr: _select(layer, positions, normed, qr), + fold=not first, + route_ids=route, + **kw, + ) + return out[:4] From cba669d45d127f56ed3e634074b620d2d16d05cd Mon Sep 17 00:00:00 2001 From: ukannika Date: Sat, 24 Oct 2026 05:03:34 -0700 Subject: [PATCH 015/631] [ROCm][PERF] Enable MiniMax-M3 Indexer CP (#57909) Signed-off-by: ukannika Signed-off-by: ukannika_amdeng Co-authored-by: Cursor Co-authored-by: Micah Williamson --- tests/kernels/attention/test_indexer_msa.py | 424 ++++++++ tests/kernels/attention/test_minimax_m3.py | 62 +- .../attention/test_minimax_m3_indexer_cp.py | 18 +- vllm/config/attention.py | 23 +- vllm/model_executor/layers/linear.py | 30 +- vllm/models/minimax_m3/amd/indexer_aiter.py | 652 ------------ vllm/models/minimax_m3/amd/indexer_cp.py | 87 ++ vllm/models/minimax_m3/amd/indexer_msa.py | 931 ++++++++++++++++++ vllm/models/minimax_m3/amd/model.py | 168 ++-- vllm/models/minimax_m3/amd/mtp.py | 109 +- 10 files changed, 1719 insertions(+), 785 deletions(-) create mode 100644 tests/kernels/attention/test_indexer_msa.py delete mode 100644 vllm/models/minimax_m3/amd/indexer_aiter.py create mode 100644 vllm/models/minimax_m3/amd/indexer_cp.py create mode 100644 vllm/models/minimax_m3/amd/indexer_msa.py diff --git a/tests/kernels/attention/test_indexer_msa.py b/tests/kernels/attention/test_indexer_msa.py new file mode 100644 index 000000000000..915c6d9afc29 --- /dev/null +++ b/tests/kernels/attention/test_indexer_msa.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import time + +import multiprocess as mp +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("the MiniMax-M3 MSA indexer is AMD-only", allow_module_level=True) + +import vllm.platforms.rocm as rocm_platform # noqa: E402 +from tests.utils import ensure_current_vllm_config # noqa: E402 +from vllm.distributed.parallel_state import ( # noqa: E402 + ensure_model_parallel_initialized, + init_distributed_environment, +) +from vllm.models.minimax_m3.amd.indexer_msa import ( # noqa: E402 + MAX_SUPPORTED_BLOCKS, + MSA_INDEX_HEAD_DIM, + MSA_SCORE_TYPE, + MSA_SPARSE_BLOCK_SIZE, + MSA_TOPK_BLOCKS, + SLOTS_MAX, + WAVE_SIZE, +) +from vllm.models.minimax_m3.amd.ops.sparse_pa import ( # noqa: E402 + ASM_PAGE_SIZE, + PAGE16_SIDES_PER_BLOCK, +) +from vllm.utils.network_utils import get_open_port # noqa: E402 +from vllm.utils.system_utils import update_environment_variables # noqa: E402 + +mp.set_start_method("spawn", force=True) + +# No head count: the gate takes one only to describe a CP shape, and the +# default cp_world of 1 is the unsharded indexer, which does not turn on one. +_ADMITTED = dict( + topk_blocks=MSA_TOPK_BLOCKS, + sparse_block_size=MSA_SPARSE_BLOCK_SIZE, + index_head_dim=MSA_INDEX_HEAD_DIM, + indexer_kv_dtype="fp8", + max_model_len=8192, + score_type=MSA_SCORE_TYPE, +) + +# The same config sharded, at MiniMax-M3's own shape: 4 index heads over 4 +# ranks is one apiece, and 4 x top-16 fills the merge exactly. +_ADMITTED_CP = dict( + _ADMITTED, + cp_world=4, + num_index_heads=4, + num_kv_heads=1, + max_query_len=1, + dcp_size=1, +) + +# The longest context the top-k has register slots for. +_MAX_LEN = MAX_SUPPORTED_BLOCKS * MSA_SPARSE_BLOCK_SIZE + + +@pytest.fixture +def aiter_msa_indexer_gate(monkeypatch): + import vllm.models.minimax_m3.amd.indexer_msa as indexer_msa_mod + + def probe(rocm=True, gfx950=True, aiter_attend=True, base=_ADMITTED, **overrides): + monkeypatch.setattr(indexer_msa_mod.current_platform, "is_rocm", lambda: rocm) + monkeypatch.setattr( + indexer_msa_mod, + "_minimax_m3_aiter_sparse_pa_requested", + lambda: aiter_attend, + ) + monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: gfx950) + return indexer_msa_mod.msa_indexer_unsupported_reason(**{**base, **overrides}) + + return probe + + +@pytest.mark.parametrize( + ("override", "expected"), + [ + ({}, None), + # Two ranks owning two index heads each, matching the model's kv split. + ({"cp_world": 2, "num_kv_heads": 2}, None), + # Sharding does not excuse the unsharded contract: the same checks run + # first, so one gate cannot admit at one width what it refuses at the + # other. + ({"indexer_kv_dtype": "bf16"}, "fp8 e4m3 index cache"), + ({"gfx950": False}, "gfx950"), + ({"max_model_len": _MAX_LEN + 1}, "the top-k is compiled for"), + # And then the conditions only sharding runs into. + ({"dcp_size": 2}, "decode_context_parallel_size"), + ({"cp_world": 3}, "divisible"), + ({"cp_world": 8, "num_index_heads": 8}, "candidates the merge holds"), + ({"max_query_len": 5}, "MFMA columns"), + ({"num_kv_heads": 2}, "1:1"), + ], +) +def test_msa_indexer_gate_rejects_unshardable_configs( + aiter_msa_indexer_gate, override, expected +): + reason = aiter_msa_indexer_gate(base=_ADMITTED_CP, **override) + if expected is None: + assert reason is None, f"{override} should be admitted, got {reason!r}" + else: + assert reason is not None, f"{override} should be refused" + assert expected in reason, f"{override} refused for the wrong reason: {reason}" + + +def test_msa_slot_cap_matches_what_aiter_compiled(): + msa_block_select = pytest.importorskip("aiter.ops.msa_block_select") + + assert SLOTS_MAX == msa_block_select.TOPK_MAX_SLOTS + assert WAVE_SIZE == msa_block_select.WAVE_SIZE + + def aiter_accepts(blocks: int) -> bool: + strips = -(-blocks // WAVE_SIZE) + return 1 << (strips - 1).bit_length() <= SLOTS_MAX + + assert aiter_accepts(MAX_SUPPORTED_BLOCKS) + assert not aiter_accepts(MAX_SUPPORTED_BLOCKS + 1) + + +_INDEX_HEADS = 4 # MiniMax-M3's own count, which is also the CP ceiling +_QUERY_LEN = 1 +_INIT_BLOCKS = 1 +_LOCAL_BLOCKS = 2 +_PAGES_PER_BLOCK = MSA_SPARSE_BLOCK_SIZE // ASM_PAGE_SIZE +_BATCH_SIZES = (32, 64) +_CONTEXT_LENS = (16 * 1024, 64 * 1024, 128 * 1024) +_SHAPES = [(batch, ctx) for batch in _BATCH_SIZES for ctx in _CONTEXT_LENS] +_MAX_ROWS = max(_BATCH_SIZES) * _QUERY_LEN +_FP8 = torch.float8_e4m3fn +_UNWRITTEN = -7 + + +def _shape_id(shape) -> str: + batch, ctx = shape + return f"b{batch}-{ctx // 1024}k" + + +def _cdiv(a: int, b: int) -> int: + return -(-a // b) + + +def _score_cols(max_seq_len: int, world: int = 1) -> int: + """Block-axis width both kernels demand of a shard of ``world``. + + Their contract, not a copy of the indexer's helper: lanes read whole + wave-wide strips and hold a power-of-two count of them. + """ + blocks = _cdiv(_cdiv(max_seq_len, MSA_SPARSE_BLOCK_SIZE), world) + strips = _cdiv(blocks, WAVE_SIZE) + return (1 << (strips - 1).bit_length()) * WAVE_SIZE + + +def _require_msa_kernels(): + if not rocm_platform.on_gfx950(): + pytest.skip("the MSA indexer's kernels are gfx950-only") + pytest.importorskip("aiter.ops.msa_attention") + + +def _decode_inputs(batch: int, ctx: int, device: torch.device): + gen = torch.Generator().manual_seed(0x5EED) + blocks = _cdiv(ctx, MSA_SPARSE_BLOCK_SIZE) + rows = batch * _QUERY_LEN + + query = torch.randn(rows, _INDEX_HEADS, MSA_INDEX_HEAD_DIM, generator=gen) + cache = torch.randn( + blocks, MSA_SPARSE_BLOCK_SIZE, MSA_INDEX_HEAD_DIM, generator=gen + ) + block_table = torch.stack( + [torch.randperm(blocks, generator=gen) for _ in range(batch)] + ) + # Long enough that the forced blocks and the top-k never cover the whole + # row, or every selection would be the trivial one. + floor = MSA_SPARSE_BLOCK_SIZE * (MSA_TOPK_BLOCKS + _INIT_BLOCKS + _LOCAL_BLOCKS) + seq_lens = torch.randint(floor, ctx + 1, (batch,), generator=gen) + seq_lens[0] = ctx # one row pinned to the full context + + return ( + query.to(device).to(_FP8), + cache.to(device).to(_FP8), + block_table.to(device).to(torch.int32), + seq_lens.to(device).to(torch.int32), + ) + + +@pytest.mark.parametrize("cp_world", [2, 4]) +@pytest.mark.parametrize("shape", _SHAPES, ids=_shape_id) +def test_cp_score_shard_holds_the_unsharded_scores(cp_world, shape): + """Each shard's scores are the unsharded ones at its round-robin columns. + + One GPU and no exchange: this is the half of the chain that can be + checked without peers, and the multi-rank test below leans on it. + """ + _require_msa_kernels() + from aiter.ops.msa_attention import pa_sparse_block_score_decode + + batch, ctx = shape + device = torch.device("cuda") + query, cache, block_table, seq_lens = _decode_inputs(batch, ctx, device) + blocks = _cdiv(ctx, MSA_SPARSE_BLOCK_SIZE) + rows = batch * _QUERY_LEN + + def score(cols: int, **shard) -> torch.Tensor: + # Blocks past a row's causal reach are left untouched, so the fill is + # what makes the two sides comparable there as well. + out = torch.full( + (_INDEX_HEADS, rows, cols), + float("-inf"), + device=device, + dtype=torch.float32, + ) + pa_sparse_block_score_decode( + query, + cache, + out, + block_table, + seq_lens, + init_blocks=_INIT_BLOCKS, + local_blocks=_LOCAL_BLOCKS, + query_len=_QUERY_LEN, + max_seq_len=ctx, + **shard, + ) + return out + + unsharded = score(_score_cols(ctx)) + # Row 0 spans the whole context, so every block of it is in reach and the + # comparison below cannot be passing on a buffer of untouched -inf. + assert torch.isfinite(unsharded[:, 0, :blocks]).all() + + for rank in range(cp_world): + shard = score(_score_cols(ctx, cp_world), rank=rank, world=cp_world) + owned = _cdiv(blocks - rank, cp_world) + assert torch.equal( + shard[:, :, :owned], + unsharded[:, :, rank : rank + owned * cp_world : cp_world], + ), f"rank {rank}/{cp_world} does not hold its round-robin columns" + + +_JOIN_TIMEOUT_S = 300 + + +def _check_cp_selection(shape, rank: int, world: int, device, peers) -> None: + """One shape's CP decode against an unsharded decode of the same inputs.""" + from aiter.ops.msa_attention import ( + pa_sparse_block_score_decode, + pa_sparse_block_topk, + pa_sparse_block_topk_cp, + ) + + batch, ctx = shape + query, cache, block_table, seq_lens = _decode_inputs(batch, ctx, device) + page16 = block_table * PAGE16_SIDES_PER_BLOCK + rows = batch * _QUERY_LEN + owned = _INDEX_HEADS // world + lo = rank * owned + + def outputs(): + return ( + torch.full( + (owned, rows, MSA_TOPK_BLOCKS), + _UNWRITTEN, + dtype=torch.int32, + device=device, + ), + torch.full( + (rows * owned, MSA_TOPK_BLOCKS * _PAGES_PER_BLOCK), + _UNWRITTEN, + dtype=torch.int32, + device=device, + ), + torch.full((rows * owned,), _UNWRITTEN, dtype=torch.int32, device=device), + ) + + shard = torch.full( + (_INDEX_HEADS, rows, _score_cols(ctx, world)), + float("-inf"), + device=device, + dtype=torch.float32, + ) + pa_sparse_block_score_decode( + query, + cache, + shard, + block_table, + seq_lens, + init_blocks=_INIT_BLOCKS, + local_blocks=_LOCAL_BLOCKS, + query_len=_QUERY_LEN, + max_seq_len=ctx, + rank=rank, + world=world, + ) + cp_topk, cp_bt, cp_ctx = outputs() + pa_sparse_block_topk_cp( + shard, + peers.cand_ptrs, + peers.cp_gen, + cp_topk, + seq_lens, + page16, + cp_bt, + cp_ctx, + max_seq_len=ctx, + block_size=MSA_SPARSE_BLOCK_SIZE, + cand_rows=_MAX_ROWS, + query_len=_QUERY_LEN, + init_blocks=_INIT_BLOCKS, + local_blocks=_LOCAL_BLOCKS, + num_kv_heads=owned, + pages_per_block=_PAGES_PER_BLOCK, + rank=rank, + world=world, + ) + + unsharded = torch.full( + (_INDEX_HEADS, rows, _score_cols(ctx)), + float("-inf"), + device=device, + dtype=torch.float32, + ) + pa_sparse_block_score_decode( + query, + cache, + unsharded, + block_table, + seq_lens, + init_blocks=_INIT_BLOCKS, + local_blocks=_LOCAL_BLOCKS, + query_len=_QUERY_LEN, + max_seq_len=ctx, + ) + ref_topk, ref_bt, ref_ctx = outputs() + pa_sparse_block_topk( + unsharded[lo : lo + owned].contiguous(), + ref_topk, + page16, + seq_lens, + sparse_bt=ref_bt, + sparse_ctx=ref_ctx, + max_seq_len=ctx, + block_size=MSA_SPARSE_BLOCK_SIZE, + query_len=_QUERY_LEN, + num_kv_heads=owned, + pages_per_block=_PAGES_PER_BLOCK, + ) + + torch.accelerator.synchronize() + where = f"rank {rank}/{world} at {_shape_id(shape)}" + assert (cp_topk != _UNWRITTEN).any(), f"{where}: the merge wrote no selection" + assert torch.equal(cp_topk, ref_topk), f"{where}: selected different blocks" + assert torch.equal(cp_bt, ref_bt), f"{where}: emitted a different page table" + assert torch.equal(cp_ctx, ref_ctx), f"{where}: emitted different page counts" + + +def _cp_worker(env) -> None: + update_environment_variables(env) + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") + torch.accelerator.set_device_index(device) + + from vllm.models.minimax_m3.amd.indexer_cp import get_indexer_cp_peers + + with ensure_current_vllm_config(): + init_distributed_environment() + ensure_model_parallel_initialized(world, 1) + + peers = get_indexer_cp_peers(_INDEX_HEADS // world, _MAX_ROWS, MSA_TOPK_BLOCKS) + for shape in _SHAPES: + _check_cp_selection(shape, rank, world, device, peers) + + +def _distributed_run(fn, world_size: int) -> None: + port = str(get_open_port()) + procs = [] + for rank in range(world_size): + env = { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(world_size), + "LOCAL_WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": port, + } + proc = mp.Process(target=fn, args=(env,)) + proc.start() + procs.append(proc) + + deadline = time.monotonic() + _JOIN_TIMEOUT_S + for proc in procs: + proc.join(max(0.0, deadline - time.monotonic())) + hung = [i for i, proc in enumerate(procs) if proc.is_alive()] + for proc in procs: + if proc.is_alive(): + proc.kill() + proc.join() + + assert not hung, f"ranks {hung} did not finish within {_JOIN_TIMEOUT_S}s" + failed = {i: proc.exitcode for i, proc in enumerate(procs) if proc.exitcode != 0} + assert not failed, f"ranks exited non-zero: {failed}" + + +@pytest.mark.parametrize("cp_world", [2, 4]) +def test_indexer_cp_selects_what_an_unsharded_indexer_selects(cp_world): + """The real thing: ``cp_world`` processes, real IPC, real exchange. + + Each rank runs the production CP decode over its block shard and compares + the selection against an unsharded decode of the same inputs over the + heads it owns. The gating lives here rather than in the worker, where a + skip reads back to the parent as a failed exit. + """ + _require_msa_kernels() + if torch.accelerator.device_count() < cp_world: + pytest.skip(f"needs {cp_world} GPUs") + + _distributed_run(_cp_worker, cp_world) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 918e4562cea3..a821a2f27b1d 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -1796,72 +1796,40 @@ def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch): assert backend.get_builder_cls().__name__ == "MiniMaxM3SparseAiterPAMetadataBuilder" -def test_aiter_indexer_requires_the_aiter_attend(monkeypatch): +def test_msa_indexer_requires_the_aiter_attend(monkeypatch): """The top-k emits the attend's page table, so it needs that attend. Selected without it, the indexer would reach its top-k with no table to resolve through; falling back instead surfaces the fp8 index cache as the configuration error it is. """ - import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod + import vllm.models.minimax_m3.amd.indexer_msa as indexer_msa_mod - monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True) - - def unsupported_reason() -> str | None: - return indexer_aiter_mod.aiter_indexer_unsupported_reason( - topk_blocks=TOPK, - sparse_block_size=BLOCK_SIZE, - num_index_heads=1, - index_head_dim=HEAD_DIM, - # The implementation also accepts this legacy fp8 alias. - indexer_kv_dtype="fp8_e4m3", - max_model_len=8192, - ) + monkeypatch.setattr(indexer_msa_mod.current_platform, "is_rocm", lambda: True) + kwargs = dict( + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + num_index_heads=1, + index_head_dim=HEAD_DIM, + indexer_kv_dtype="fp8_e4m3", + max_model_len=8192, + ) monkeypatch.setattr( - indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False + indexer_msa_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False ) - reason = unsupported_reason() + reason = indexer_msa_mod.msa_indexer_unsupported_reason(**kwargs) assert reason is not None and "AITER sparse PA attend" in reason # With the attend asked for, the gate moves on to the kernel-contract # checks rather than stopping here. monkeypatch.setattr( - indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True + indexer_msa_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True ) - reason = unsupported_reason() + reason = indexer_msa_mod.msa_indexer_unsupported_reason(**kwargs) assert reason is None or "AITER sparse PA attend" not in reason -def test_aiter_indexer_rejects_unbuilt_index_heads(monkeypatch): - """AITER builds the MSA score kernels for 1 or 2 index heads only. - - MiniMax-M3 has 4, so at attention TP=1 the indexer has to fall back at load - time rather than fail at the first forward. - """ - import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod - - monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True) - monkeypatch.setattr( - indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True - ) - - def unsupported_reason(num_index_heads: int) -> str | None: - return indexer_aiter_mod.aiter_indexer_unsupported_reason( - topk_blocks=TOPK, - sparse_block_size=BLOCK_SIZE, - num_index_heads=num_index_heads, - index_head_dim=HEAD_DIM, - indexer_kv_dtype="fp8", - max_model_len=8192, - ) - - if reason := unsupported_reason(2): - pytest.skip(f"AITER MSA indexer unavailable: {reason}") - reason = unsupported_reason(4) - assert reason is not None and "num_index_heads" in reason - - def test_aiter_consolidated_qknorm_enabled(monkeypatch): from vllm._aiter_ops import rocm_aiter_ops diff --git a/tests/kernels/attention/test_minimax_m3_indexer_cp.py b/tests/kernels/attention/test_minimax_m3_indexer_cp.py index fe3b8dbe90a9..2bc8b68d813b 100644 --- a/tests/kernels/attention/test_minimax_m3_indexer_cp.py +++ b/tests/kernels/attention/test_minimax_m3_indexer_cp.py @@ -319,9 +319,21 @@ def test_local_topk_merge_matches_full_topk(): assert torch.equal(got, ref) -def test_indexer_cp_config_rejects_fp8(): - """CP scores the index cache with no scale; fp8 must not enable it.""" +@pytest.mark.parametrize("indexer_kv_dtype", ["auto", "bf16", "fp8", "fp8_e4m3"]) +def test_indexer_cp_config_accepts_bf16_and_fp8(indexer_kv_dtype): + """bf16 enables this CP indexer, fp8 the AITER MSA one; both are valid.""" + from vllm.config.attention import AttentionConfig + + config = AttentionConfig( + minimax_m3_indexer_cp=True, indexer_kv_dtype=indexer_kv_dtype + ) + assert config.minimax_m3_indexer_cp + + +@pytest.mark.parametrize("indexer_kv_dtype", ["mxfp4", "nvfp4"]) +def test_indexer_cp_config_rejects_dtypes_with_no_cp_indexer(indexer_kv_dtype): + """Neither CP score kernel reads a block-scaled cache.""" from vllm.config.attention import AttentionConfig with pytest.raises(ValueError, match="minimax_m3_indexer_cp requires a bf16"): - AttentionConfig(minimax_m3_indexer_cp=True, indexer_kv_dtype="fp8") + AttentionConfig(minimax_m3_indexer_cp=True, indexer_kv_dtype=indexer_kv_dtype) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 7eb817b14f26..a01ddf97ded8 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -117,7 +117,10 @@ class AttentionConfig: indexer context-parallel at decode. Each rank reads 1/TP of the index cache, scores every index head on it, and per-head top-k candidates are exchanged and merged, so the selected blocks match the default path. - Requires a bf16 index cache (`indexer_kv_dtype` "auto" or "bf16").""" + `indexer_kv_dtype` picks the kernels that do it: a bf16 index cache ("auto" + or "bf16") runs the Triton CP indexer, an fp8 e4m3 one the AITER MSA CP + indexer on gfx950. This shards indexer *work* and not the KV cache, so it + is not `decode_context_parallel_size` and requires that to be 1.""" hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental @@ -175,13 +178,19 @@ def __post_init__(self) -> None: # The alias selects only MiniMax's sparse decode kernel. Dense # layers still use the platform's normal automatic backend. self.backend = None - # Checked here rather than at indexer selection: an fp8 index cache is - # served by the AITER indexer, which never reaches it. This also - # excludes the fused decode path, which needs fp8. - if self.minimax_m3_indexer_cp and self.indexer_kv_dtype not in ("auto", "bf16"): + # Checked here rather than at indexer selection, which the two CP + # backends reach separately: bf16 goes to the Triton indexer and fp8 + # to the AITER MSA one, so neither ever sees the other's dtype and + # neither can speak for the dtypes that have no CP indexer at all. + if self.minimax_m3_indexer_cp and self.indexer_kv_dtype not in ( + "auto", + "bf16", + "fp8", + "fp8_e4m3", + ): raise ValueError( - "minimax_m3_indexer_cp requires a bf16 index cache " - "(indexer_kv_dtype 'auto' or 'bf16'), got " + "minimax_m3_indexer_cp requires a bf16 or fp8 e4m3 index cache " + "(indexer_kv_dtype 'auto', 'bf16', 'fp8' or 'fp8_e4m3'), got " f"{self.indexer_kv_dtype!r}" ) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 7a9732107482..d3195472d3cd 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1381,6 +1381,11 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): *replication* path when ``tp_size > total_num_kv_heads`` (this is what makes a TP size greater than the KV-head count work). ``index_k`` is a single shared head, replicated to every rank. + + ``replicate_index_q`` moves ``index_q`` onto that replicated path too, so + every rank projects all the index heads rather than its own. That is what + indexer context parallelism needs: it inverts the split, scoring every head + against a shard of the blocks instead of one head against all of them. """ def __init__( @@ -1394,6 +1399,7 @@ def __init__( bias: bool = False, quant_config: QuantizationConfig | None = None, prefix: str = "", + replicate_index_q: bool = False, ) -> None: # index_q rides the KV-head sharding/replication path, so its head count # must match the KV heads. @@ -1408,6 +1414,7 @@ def __init__( self.total_num_kv_heads = total_num_kv_heads self.total_num_index_heads = total_num_index_heads self.index_head_size = index_head_size + self.replicate_index_q = replicate_index_q tp_size = get_tensor_model_parallel_world_size() self.num_heads = divide(self.total_num_heads, tp_size) @@ -1417,8 +1424,11 @@ def __init__( else: self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) self.num_kv_head_replicas = 1 - # index_q shards identically to the KV heads. - self.num_index_heads = self.num_kv_heads + # index_q shards identically to the KV heads, unless CP has asked for + # every head on every rank. + self.num_index_heads = ( + self.total_num_index_heads if replicate_index_q else self.num_kv_heads + ) # Global per-group sizes (replicated groups counted x tp_size, matching # the QKVParallelLinear convention). index_k is a single replicated head. @@ -1495,12 +1505,14 @@ def weight_loader_v2( weight_block_size, shard_size, shard_offset ) - # index_k is fully replicated: num_heads == tp_size makes - # load_qkv_weight pick shard_id_int == 0 on every rank. q/k/v/index_q ride - # the KV-head replication factor. - num_heads = ( - self.tp_size if loaded_shard_id == "index_k" else self.num_kv_head_replicas + # A replicated group takes num_heads == tp_size, which makes + # load_qkv_weight pick shard_id_int == 0 on every rank and so copy the + # whole group everywhere. index_k always is; index_q is too under CP. + # The rest ride the KV-head replication factor. + replicated = loaded_shard_id == "index_k" or ( + loaded_shard_id == "index_q" and self.replicate_index_q ) + num_heads = self.tp_size if replicated else self.num_kv_head_replicas param.load_qkv_weight( loaded_weight=loaded_weight, num_heads=num_heads, @@ -1534,7 +1546,9 @@ def weight_loader( param_data = param.data.narrow(output_dim, shard_offset, shard_size) if loaded_shard_id == "q": shard_rank = self.tp_rank - elif loaded_shard_id == "index_k": + elif loaded_shard_id == "index_k" or ( + loaded_shard_id == "index_q" and self.replicate_index_q + ): shard_rank = 0 # replicated to every rank else: shard_rank = self.tp_rank // self.num_kv_head_replicas diff --git a/vllm/models/minimax_m3/amd/indexer_aiter.py b/vllm/models/minimax_m3/amd/indexer_aiter.py deleted file mode 100644 index 1efed4b06131..000000000000 --- a/vllm/models/minimax_m3/amd/indexer_aiter.py +++ /dev/null @@ -1,652 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""AITER (ROCm) indexer impl for MiniMax M3. - -Scores index blocks and selects the top-k with AITER's fp8 MFMA kernels on both -sides of the batch: ``pa_sparse_block_score_decode`` for the uniform-query-length -decode rows and ``pa_sparse_block_score_prefill`` for the ragged prefill rows, -then ``pa_sparse_block_topk`` for each. The two scoring passes share one tile -body in AITER, so they agree block for block, and both encode the forced -init/local blocks as the same sentinel scores the Triton indexer uses. - -The top-k also emits the attend's page table. The winners are already in its -workgroup's LDS, so resolving them through the block table there costs one wave -and replaces the separate Triton pass in ``ops.sparse_pa``; the rows it writes -are one per (token, kv head) with the page ids folded head-minor, which is the -layout ``pa_decode_gluon`` reads after it flattens the cache. - -The score kernels are built on ``v_mfma_f32_16x16x32_fp8_fp8``, so this impl -requires an fp8 (e4m3) index cache and an fp8 index query -- the fused -QK-norm/RoPE kernel emits both directly when the index cache is e4m3. See -``aiter_indexer_unsupported_reason`` for the full set of limits; -``select_aiter_indexer_impl_cls`` refuses to pick this impl unless they all -hold, and the model falls back to the platform-neutral ``MiniMaxM3Indexer``. -""" - -import math -from dataclasses import dataclass -from functools import cache -from typing import ClassVar - -import torch -from torch import nn - -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.config.attention import IndexerKVDType -from vllm.forward_context import get_forward_context -from vllm.logger import init_logger -from vllm.models.minimax_m3.amd.ops.sparse_pa import ASM_PAGE_SIZE -from vllm.models.minimax_m3.common.indexer import ( - MiniMaxM3IndexerBackend, - MiniMaxM3IndexerCache, - MiniMaxM3IndexerDecodeMetadata, - MiniMaxM3IndexerImpl, - MiniMaxM3IndexerMetadata, - MiniMaxM3IndexerMetadataBuilder, - MiniMaxM3IndexerPrefillMetadata, -) - -# The single source of truth for whether the AITER attend was asked for; this -# indexer is only usable alongside it. -from vllm.models.minimax_m3.common.sparse_attention import ( - _minimax_m3_aiter_sparse_pa_requested, -) -from vllm.platforms import current_platform -from vllm.v1.attention.backend import ( - AttentionBackend, - CommonAttentionMetadata, - max_decode_query_len, -) -from vllm.v1.attention.backends.utils import split_decodes_and_prefills -from vllm.v1.kv_cache_interface import AttentionSpec - -logger = init_logger(__name__) - -# Compiled AITER MiniMax-M3 score/top-k contract -MSA_TOPK_BLOCKS = 16 -MSA_SCORE_TYPE = "max" -MSA_SPARSE_BLOCK_SIZE = 128 -MSA_INDEX_HEAD_DIM = 128 -# Index head counts the score kernels are instantiated for. -MSA_NUM_INDEX_HEADS = (1, 2) -# Wave width the top-k is written against: it gives one lane per output slot and -# reads the score row in wave-wide strips. -WAVE_SIZE = 64 -# Per-lane register slots the top-k can hold, which is what caps the context: a -# row may span at most SLOTS_MAX * WAVE_SIZE blocks. -SLOTS_MAX = 128 -MAX_SUPPORTED_BLOCKS = SLOTS_MAX * WAVE_SIZE -# MFMA columns the score pass has, one per (query token, index head) pair. -MFMA_COLS = 16 -# The fp8 MFMA the score kernels are built on exists on these targets only. -SUPPORTED_ARCHS = ("gfx950",) - - -def _pow2_ceil(n: int) -> int: - return 1 << (n - 1).bit_length() if n >= 1 else 1 - - -def score_block_width(max_seq_len: int, block_size: int) -> int: - """Block-axis width the top-k requires of the score buffer. - - Lanes read whole wave-wide strips with no tail guard and each holds a - power-of-two count of them, so the axis is padded past the block count. - """ - max_blk = math.ceil(max(max_seq_len, 1) / block_size) - return _pow2_ceil(math.ceil(max_blk / WAVE_SIZE)) * WAVE_SIZE - - -def aiter_indexer_max_decode_query_len(vllm_config: VllmConfig) -> int: - """Longest query a decode row can carry, which spec decode is what sets. - - Mirrors ``_init_reorder_batch_threshold(1, supports_spec_as_decode=True)``, - since that is what the builder splits the batch on and therefore what the - decode kernel will actually be handed. - """ - return max_decode_query_len(vllm_config) - - -@cache -def aiter_msa_kernels_unavailable_reason() -> str | None: - """Return why the AITER MSA score/top-k ops cannot be imported, or None. - - They are a recent addition, so an AITER that predates them imports fine - while these three names do not exist, and the failure would otherwise - surface as an ImportError from the middle of a forward. ``compile_ops`` - binds them lazily, so this costs the module import only -- the kernel build - itself still happens on the first call. - """ - try: - from aiter.ops.msa_attention import ( # noqa: F401 - pa_sparse_block_score_decode, - pa_sparse_block_score_prefill, - pa_sparse_block_topk, - ) - except ImportError as exc: - return f"AITER cannot supply the MSA score/top-k kernels ({exc})" - return None - - -def aiter_indexer_unsupported_reason( - *, - topk_blocks: int, - sparse_block_size: int, - num_index_heads: int, - index_head_dim: int, - indexer_kv_dtype: IndexerKVDType | str, - max_model_len: int, - max_decode_query_len: int = 1, - score_type: str = "max", -) -> str | None: - """Return why this config cannot use the AITER indexer, or None if it can. - - Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache - dtype, the compiled score/top-k contract, MFMA column limits, max context - in blocks, and whether AITER exposes the MSA kernels. - ``select_aiter_indexer_impl_cls`` logs the string and falls back when it is - not None. - """ - if not current_platform.is_rocm(): - return ( - "needs ROCm for the AITER fp8 MFMA score/top-k, " - f"got platform={current_platform.device_type!r}" - ) - if not _minimax_m3_aiter_sparse_pa_requested(): - # The top-k emits the attend's page table in page-16 numbering, which - # only addresses the interleaved cache the AITER attend reads. Paired - # with any other attend there is nowhere valid to write it, so this - # indexer is not usable on its own. - return ( - "needs the AITER sparse PA attend, whose page table its top-k " - "emits (rocm_aiter_ops + shuffle KV cache layout)" - ) - if indexer_kv_dtype not in ("fp8", "fp8_e4m3"): - # The score kernels are fp8 MFMA; there is no bf16 instantiation. - return ( - f"needs an fp8 e4m3 index cache, got indexer_kv_dtype={indexer_kv_dtype!r}" - ) - from vllm.platforms.rocm import on_gfx950 - - if not on_gfx950(): - return f"needs {' or '.join(SUPPORTED_ARCHS)} for the fp8 MFMA" - if score_type != MSA_SCORE_TYPE: - return f"needs score_type={MSA_SCORE_TYPE!r}, got score_type={score_type!r}" - if topk_blocks != MSA_TOPK_BLOCKS: - return f"needs topk_blocks={MSA_TOPK_BLOCKS}, got topk_blocks={topk_blocks}" - if sparse_block_size != MSA_SPARSE_BLOCK_SIZE: - return ( - f"needs sparse_block_size={MSA_SPARSE_BLOCK_SIZE}, " - f"got sparse_block_size={sparse_block_size}" - ) - if index_head_dim != MSA_INDEX_HEAD_DIM: - return ( - f"needs index_head_dim={MSA_INDEX_HEAD_DIM}, " - f"got index_head_dim={index_head_dim}" - ) - if num_index_heads not in MSA_NUM_INDEX_HEADS: - return ( - f"needs num_index_heads in {MSA_NUM_INDEX_HEADS}, " - f"got num_index_heads={num_index_heads}" - ) - # A decode row's whole query shares one MFMA tile, one column per (token, - # head) pair, so spec decode trades columns against the head count. - if num_index_heads * max_decode_query_len > MFMA_COLS: - return ( - f"num_index_heads={num_index_heads} x max_decode_query_len=" - f"{max_decode_query_len} exceeds the {MFMA_COLS} MFMA columns" - ) - max_blocks = math.ceil(max_model_len / sparse_block_size) - if max_blocks > MAX_SUPPORTED_BLOCKS: - return ( - f"max_model_len={max_model_len} needs {max_blocks} blocks per row, " - f"more than the top-k's {MAX_SUPPORTED_BLOCKS}" - ) - return aiter_msa_kernels_unavailable_reason() - - -class MiniMaxM3IndexerAiterBackend(MiniMaxM3IndexerBackend): - """Indexer side-cache backend selecting the AITER builder.""" - - @staticmethod - def get_builder_cls() -> type["MiniMaxM3IndexerAiterMetadataBuilder"]: - return MiniMaxM3IndexerAiterMetadataBuilder - - -@dataclass -class MiniMaxM3IndexerAiterMetadata(MiniMaxM3IndexerMetadata): - """Adds the per-row shape the ragged top-k needs. - - The uniform decode rows are recovered inside the kernel from ``seq_lens`` and - the shared query length, which also clamps cudagraph padding rows to nothing. - Prefill rows have no such shape, so it is materialized here once per forward - and shared by every layer -- which is also where the emitted page table's - tail block comes from, since a block count alone does not say how many tokens - the last block holds. - """ - - # [num_prefill_tokens] int32, cdiv(position + 1, sparse_block_size) per prefill row. - prefill_num_valid_pages: torch.Tensor | None = None - # [num_prefill_tokens] int32, the request each prefill row belongs to. - prefill_row_req_id: torch.Tensor | None = None - # [num_prefill_tokens] int32, causal token count (position + 1) per row. - prefill_kv_lens: torch.Tensor | None = None - - -class MiniMaxM3IndexerAiterMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): - """The Triton indexer's metadata plus the prefill rows' causal shape.""" - - 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) - hf_config = vllm_config.model_config.hf_config - text_config = getattr(hf_config, "text_config", hf_config) - self.sparse_block_size = int( - text_config.sparse_attention_config["sparse_block_size"] - ) - max_tokens = vllm_config.scheduler_config.max_num_batched_tokens - # Companions to the base's num_valid_pages_buffer, for the two vectors - # only the emitted table needs. - self.row_req_id_buffer = torch.empty( - max_tokens, dtype=torch.int32, device=device - ) - self.kv_lens_buffer = torch.empty(max_tokens, dtype=torch.int32, device=device) - - def build( - self, - common_prefix_len: int, - common_attn_metadata: CommonAttentionMetadata, - fast_build: bool = False, - ) -> MiniMaxM3IndexerAiterMetadata: - num_reqs = common_attn_metadata.num_reqs - num_tokens = common_attn_metadata.num_actual_tokens - query_start_loc = common_attn_metadata.query_start_loc - seq_lens = common_attn_metadata.seq_lens - block_table = common_attn_metadata.block_table_tensor - - 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=True, - ) - ) - assert num_decodes + num_prefills == num_reqs - assert num_decode_tokens + num_prefill_tokens == num_tokens - - # Decode-first batch: context lengths into the stable cudagraph buffer. - context_lens = self.context_len_buffer[:num_reqs] - context_lens.copy_( - common_attn_metadata.compute_num_computed_tokens(), non_blocking=True - ) - - prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None - prefill_num_valid_pages: torch.Tensor | None = None - prefill_row_req_id: torch.Tensor | None = None - prefill_kv_lens: torch.Tensor | None = None - if num_prefills > 0: - cu_seqlens_q = (query_start_loc[num_decodes:] - num_decode_tokens).to( - torch.int32 - ) - prefill_metadata = MiniMaxM3IndexerPrefillMetadata( - cu_seqlens_q=cu_seqlens_q, - seq_lens=seq_lens[num_decodes:], - context_lens=context_lens[num_decodes:], - block_table=block_table[num_decodes:], - max_query_len=common_attn_metadata.max_query_len, - max_seq_len=common_attn_metadata.max_seq_len, - ) - # A prefill row sees its own position, so its causal length and block - # count both follow from that alone; the request it belongs to comes - # from the query offsets. Prefill batches are never captured, so the - # stable buffers are only being reused here, not required. - positions = common_attn_metadata.positions - assert positions is not None - row_positions = positions[num_decode_tokens:num_tokens] - prefill_num_valid_pages = self.num_valid_pages_buffer[ - num_decode_tokens:num_tokens - ] - prefill_num_valid_pages.copy_( - row_positions // self.sparse_block_size + 1, non_blocking=True - ) - prefill_kv_lens = self.kv_lens_buffer[num_decode_tokens:num_tokens] - prefill_kv_lens.copy_(row_positions + 1, non_blocking=True) - prefill_row_req_id = self.row_req_id_buffer[num_decode_tokens:num_tokens] - prefill_row_req_id.copy_( - torch.searchsorted( - cu_seqlens_q[1:].contiguous(), - torch.arange( - num_prefill_tokens, - dtype=torch.int32, - device=cu_seqlens_q.device, - ), - right=True, - ), - non_blocking=True, - ) - - decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None - if num_decodes > 0: - qsl_cpu = common_attn_metadata.query_start_loc_cpu - query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] - decode_query_len = int(query_lens_cpu[0].item()) - assert decode_query_len > 0 - assert torch.all( - (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) - ) - assert num_decode_tokens == num_decodes * decode_query_len - decode_metadata = MiniMaxM3IndexerDecodeMetadata( - seq_lens=seq_lens[:num_decodes], - block_table=block_table[:num_decodes], - max_seq_len=common_attn_metadata.max_seq_len, - decode_query_len=decode_query_len, - max_decode_query_len=self.max_decode_query_len, - ) - - return MiniMaxM3IndexerAiterMetadata( - seq_lens=seq_lens, - max_seq_len=common_attn_metadata.max_seq_len, - slot_mapping=common_attn_metadata.slot_mapping, - num_actual_tokens=num_tokens, - num_decodes=num_decodes, - num_decode_tokens=num_decode_tokens, - num_prefills=num_prefills, - num_prefill_tokens=num_prefill_tokens, - prefill=prefill_metadata, - decode=decode_metadata, - prefill_num_valid_pages=prefill_num_valid_pages, - prefill_row_req_id=prefill_row_req_id, - prefill_kv_lens=prefill_kv_lens, - ) - - -class MiniMaxM3IndexerAiterImpl(MiniMaxM3IndexerImpl): - """AITER fp8 score + top-k for both prefill and decode.""" - - indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerAiterBackend - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - # Both passes are v_mfma_f32_16x16x32_fp8_fp8 with no bf16 instantiation, - # so an index cache of any other dtype would be read as e4m3 bytes. - # The selector will not get here, but nothing else may either. - if self.indexer_kv_dtype not in ("fp8", "fp8_e4m3"): - raise ValueError( - "The AITER indexer requires an fp8 e4m3 index cache, got " - f"indexer_kv_dtype={self.indexer_kv_dtype!r}" - ) - # Shared, stable-address page table + per-row context bound the top-k - # emits for the attend. Owned by the model so one allocation serves - # every layer; left None when it reserved none, in which case the - # attend rebuilds the table itself. - self.sparse_bt_buffer: torch.Tensor | None = None - self.sparse_ctx_buffer: torch.Tensor | None = None - - @property - def pages_per_block(self) -> int: - """Physical pages one selected block expands into for the attend.""" - return self.block_size // ASM_PAGE_SIZE - - def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]: - """The page table and context rows covering score rows ``[lo, hi)``. - - One table row per (token, kv head), head minor, which is the order - ``pa_decode_gluon`` reads once it flattens the cache. Slices of the - shared buffers rather than copies, so the attend sees the writes; a - model that reserved no buffers gets throwaway ones, and the attend - rebuilds the table itself in that case. - """ - kv_heads = self.num_kv_heads - if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None: - rows = (hi - lo) * kv_heads - device = self.index_cache.kv_cache.device - return ( - torch.empty( - (rows, self.topk_blocks * self.pages_per_block), - dtype=torch.int32, - device=device, - ), - torch.empty(rows, dtype=torch.int32, device=device), - ) - return ( - self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads], - self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads], - ) - - def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor: - """Score buffer for ``rows`` query rows. - - Left uninitialized on purpose: the score pass writes every block up to - the longest row it covers, and the top-k reads only the blocks its own - row can see, so nothing downstream observes the padded tail. Filling it - would cost a write over the whole [H, rows, width] extent, which at long - context is the largest tensor in the indexer. - """ - return torch.empty( - ( - self.num_index_heads, - rows, - score_block_width(max_seq_len, self.block_size), - ), - dtype=torch.float32, - device=self.index_cache.kv_cache.device, - ) - - def forward( - self, - index_query: torch.Tensor, - *, - decode_page16_block_table: torch.Tensor | None = None, - prefill_page16_block_table: torch.Tensor | None = None, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - from aiter.ops.msa_attention import ( - pa_sparse_block_score_decode, - pa_sparse_block_score_prefill, - pa_sparse_block_topk, - ) - - attn_metadata = get_forward_context().attn_metadata - if not isinstance(attn_metadata, dict): - return None, None # profiling run; caches unbound - md = attn_metadata[self.index_cache.prefix] - assert isinstance(md, MiniMaxM3IndexerAiterMetadata) - # The emitted page table addresses the main cache, whose blocks are a - # different group's than the index cache's, so the top-k resolves the - # selection through the attend's block table -- only the score pass reads - # the index cache and takes the indexer's. It has to be the page-16 - # rebase of that table, which the attend's own metadata builder does - # once per step; the indexer's metadata cannot reach another group's - # blocks, so the layer reads it off the attend and hands it in. - num_tokens = md.num_actual_tokens - nd = md.num_decode_tokens - iq = index_query[:num_tokens].view( - -1, self.num_index_heads, self.index_head_dim - ) - kv = self.index_cache.kv_cache - - # Both sides write into the single shared persistent buffer (decode at - # [:, :nd], prefill at [:, nd:]) and return views into it. The top-k - # takes the head/row strides, so these slices need no copy back. - buf = self.topk_indices_buffer - if buf is None: - buf = torch.empty( - (self.num_index_heads, num_tokens, self.topk_blocks), - dtype=torch.int32, - device=iq.device, - ) - - decode_topk: torch.Tensor | None = None - prefill_topk: torch.Tensor | None = None - - if md.num_decodes > 0: - d = md.decode - assert d is not None - assert decode_page16_block_table is not None, ( - "the AITER indexer's top-k emits the attend's page table and " - "needs the page-16 rebase of the attend's decode block table" - ) - score = self._new_score(nd, d.max_seq_len) - pa_sparse_block_score_decode( - iq[:nd], - kv, - score, - d.block_table, - d.seq_lens, - init_blocks=self.init_blocks, - local_blocks=self.local_blocks, - query_len=d.decode_query_len, - max_seq_len=d.max_seq_len, - ) - decode_topk = buf[:, :nd, :] - sparse_bt, sparse_ctx = self._table_rows(0, nd) - # A decode row's causal length is seq_len - query_len + token + 1, - # which the kernel derives itself, so the emitted table covers - # speculative rows without any extra per-row shape. - pa_sparse_block_topk( - score, - decode_topk, - decode_page16_block_table, - d.seq_lens, - sparse_bt, - sparse_ctx, - max_seq_len=d.max_seq_len, - block_size=self.block_size, - query_len=d.decode_query_len, - num_kv_heads=self.num_kv_heads, - pages_per_block=self.pages_per_block, - ) - - if md.num_prefills > 0: - p = md.prefill - assert p is not None - assert prefill_page16_block_table is not None, ( - "the AITER indexer's top-k emits the attend's page table and " - "needs the page-16 rebase of the attend's prefill block table" - ) - assert md.prefill_num_valid_pages is not None - assert md.prefill_row_req_id is not None - assert md.prefill_kv_lens is not None - score = self._new_score(num_tokens - nd, p.max_seq_len) - pa_sparse_block_score_prefill( - iq[nd:], - kv, - score, - p.block_table, - p.cu_seqlens_q, - p.seq_lens, - init_blocks=self.init_blocks, - local_blocks=self.local_blocks, - max_query_len=p.max_query_len, - max_seq_len=p.max_seq_len, - ) - prefill_topk = buf[:, nd:num_tokens, :] - sparse_bt, sparse_ctx = self._table_rows(nd, num_tokens) - # Ragged rows carry their request and causal length explicitly: the - # block count alone cannot place the tail block the table ends on. - pa_sparse_block_topk( - score, - prefill_topk, - prefill_page16_block_table, - p.seq_lens, - sparse_bt, - sparse_ctx, - max_seq_len=p.max_seq_len, - block_size=self.block_size, - num_valid_pages=md.prefill_num_valid_pages, - row_req_id=md.prefill_row_req_id, - kv_lens=md.prefill_kv_lens, - num_kv_heads=self.num_kv_heads, - pages_per_block=self.pages_per_block, - ) - - return decode_topk, prefill_topk - - -def select_aiter_indexer_impl_cls( - *, - topk_blocks: int, - sparse_block_size: int, - num_index_heads: int, - index_head_dim: int, - indexer_kv_dtype: IndexerKVDType | str, - score_type: str = "max", -) -> type[MiniMaxM3IndexerAiterImpl] | None: - """The AITER indexer impl if this config can use it, else None. - - ``None`` sends the caller to the platform-neutral ``MiniMaxM3Indexer``, - which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8 - index cache that lands here has nowhere to go. - """ - reason = aiter_indexer_unsupported_reason( - topk_blocks=topk_blocks, - sparse_block_size=sparse_block_size, - num_index_heads=num_index_heads, - index_head_dim=index_head_dim, - indexer_kv_dtype=indexer_kv_dtype, - max_model_len=get_current_vllm_config().model_config.max_model_len, - max_decode_query_len=aiter_indexer_max_decode_query_len( - get_current_vllm_config() - ), - score_type=score_type, - ) - if reason is not None: - logger.info_once("MiniMax M3 indexer: AITER unavailable (%s)", reason) - return None - logger.info_once( - "MiniMax M3 indexer: selected AITER (fp8 MFMA score + top-k) " - "[topk_blocks=%d, indexer_kv_dtype=%s]", - topk_blocks, - indexer_kv_dtype, - ) - return MiniMaxM3IndexerAiterImpl - - -class MiniMaxM3AiterIndexer(nn.Module): - """``MiniMaxM3Indexer``'s surface over the AITER impl. - - The platform-neutral wrapper picks its impl through ``common``'s selector - and forwards a Triton-only set of fused-table arguments, neither of which - can reach this impl without editing ``common``. This holds the same three - members the attention layer uses -- ``impl``, ``index_cache``, - ``num_index_heads`` -- and forwards the one argument AITER needs instead. - """ - - def __init__( - self, - *, - impl_cls: type[MiniMaxM3IndexerAiterImpl], - sparse_bt_buffer: torch.Tensor | None = None, - sparse_ctx_buffer: torch.Tensor | None = None, - **impl_kwargs, - ) -> None: - super().__init__() - self.impl = impl_cls(**impl_kwargs) - # Assigned rather than passed: the impl's base ``__init__`` lives in - # common and takes no table buffers. - self.impl.sparse_bt_buffer = sparse_bt_buffer - self.impl.sparse_ctx_buffer = sparse_ctx_buffer - - @property - def index_cache(self) -> MiniMaxM3IndexerCache: - return self.impl.index_cache - - @property - def num_index_heads(self) -> int: - return self.impl.num_index_heads - - def forward( - self, - index_query: torch.Tensor, - *, - decode_page16_block_table: torch.Tensor | None = None, - prefill_page16_block_table: torch.Tensor | None = None, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - return self.impl( - index_query, - decode_page16_block_table=decode_page16_block_table, - prefill_page16_block_table=prefill_page16_block_table, - ) diff --git a/vllm/models/minimax_m3/amd/indexer_cp.py b/vllm/models/minimax_m3/amd/indexer_cp.py new file mode 100644 index 000000000000..a5799bae792c --- /dev/null +++ b/vllm/models/minimax_m3/amd/indexer_cp.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Indexer-only context parallelism for MiniMax M3 (AMD MSA indexer). + +Every rank projects all H index heads and scores 1/P of the context blocks, then +one exchange gives it each shard's candidates for the heads it owns, cutting the +replicated index-cache read to 1/P. Shards cut to the full k so selection stays +exact; candidates carry a generation tag, so all ranks must call this in lockstep. +""" + +from dataclasses import dataclass + +import torch + +from vllm.distributed import get_tp_group +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def get_indexer_cp_group(): + """The group the candidate exchange runs over.""" + return get_tp_group() + + +# Granularity an uncached allocation under 2MB has to be a whole number of +# for the IPC export to succeed. At or above 2MB it is backed coarse-grained +# and exports regardless, but rounding up is free. +_HIP_PAGE = 4096 + + +def _ipc_alloc(nbytes: int, group) -> torch.Tensor: + """Allocate a zeroed, IPC-shared ``nbytes`` on every rank. + + Returns a ``[world]`` int64 CPU tensor of their device addresses, never freed. + """ + from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + nbytes = (nbytes + _HIP_PAGE - 1) // _HIP_PAGE * _HIP_PAGE + peers = CustomAllreduce.create_shared_buffer(nbytes, group=group.cpu_group) + return torch.tensor(peers, dtype=torch.int64, device="cpu") + + +@dataclass(frozen=True) +class IndexerCpPeers: + """The buffers the CP selector exchanges through. + + ``cand_ptrs`` is every rank's candidate buffer; ``cp_gen`` is this rank's own. + """ + + cand_ptrs: torch.Tensor + cp_gen: torch.Tensor + + +_PEERS: dict[tuple[int, int, int, int], IndexerCpPeers] = {} + + +def get_indexer_cp_peers(owned_heads: int, max_rows: int, topk: int) -> IndexerCpPeers: + """The buffers the CP selector exchanges through, allocated once per key. + + Cached and shared by every layer, sized to ``max_rows`` so captured graphs + always address the same buffer. Every rank must reach this with the same key. + """ + from aiter.ops.msa_block_select import CP_MAX_BLOCKS, topk_cp_candidate_numel + + group = get_indexer_cp_group() + key = (group.world_size, owned_heads, max_rows, topk) + peers = _PEERS.get(key) + if peers is not None: + return peers + + numel = topk_cp_candidate_numel(group.world_size, owned_heads, max_rows, topk) + peers = IndexerCpPeers( + cand_ptrs=_ipc_alloc(numel * 8, group), + cp_gen=torch.zeros(CP_MAX_BLOCKS, dtype=torch.int32, device="cuda"), + ) + _PEERS[key] = peers + logger.info_once( + "MiniMax M3 indexer CP: %d KiB candidate buffer per rank mapped over " + "%d ranks [owned_heads=%d, rows=%d, topk=%d, generations=2]", + numel * 8 // 1024, + group.world_size, + owned_heads, + max_rows, + topk, + ) + return peers diff --git a/vllm/models/minimax_m3/amd/indexer_msa.py b/vllm/models/minimax_m3/amd/indexer_msa.py new file mode 100644 index 000000000000..8d9b9032d522 --- /dev/null +++ b/vllm/models/minimax_m3/amd/indexer_msa.py @@ -0,0 +1,931 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MSA (gfx950/CDNA4) indexer impl for MiniMax M3. + +Scores index blocks and selects the top-k with AITER's fp8 MFMA kernels. Prefill +and decode differ only in the scorer -- ragged rows against uniform ones -- and +share one ``pa_sparse_block_topk``, which also emits the attend's page table in +one numbering for both. Under indexer CP, decode instead scores this rank's +stride of the blocks, cuts it to the full top-k, exchanges the candidates and +merges them, over those same kernels at ``world`` > 1. +""" + +import math +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.config.attention import IndexerKVDType +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.models.minimax_m3.amd.indexer_cp import IndexerCpPeers, get_indexer_cp_peers +from vllm.models.minimax_m3.amd.ops.sparse_pa import ASM_PAGE_SIZE +from vllm.models.minimax_m3.common.indexer import ( + MiniMaxM3IndexerBackend, + MiniMaxM3IndexerCache, + MiniMaxM3IndexerDecodeMetadata, + MiniMaxM3IndexerImpl, + MiniMaxM3IndexerMetadata, + MiniMaxM3IndexerMetadataBuilder, + MiniMaxM3IndexerPrefillMetadata, +) + +# The single source of truth for whether the AITER attend was asked for; this +# indexer is only usable alongside it. +from vllm.models.minimax_m3.common.sparse_attention import ( + _minimax_m3_aiter_sparse_pa_requested, +) +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv, next_power_of_2 +from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata +from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + +# MiniMax-M3 score/top-k shape contract +MSA_TOPK_BLOCKS = 16 +MSA_SCORE_TYPE = "max" +MSA_SPARSE_BLOCK_SIZE = 128 +MSA_INDEX_HEAD_DIM = 128 +# The fp8 MFMA the score kernels are built on exists on these targets only. +SUPPORTED_ARCHS = ("gfx950",) + +# Wave width the top-k is written against: it gives one lane per output slot +# and reads the score row in wave-wide strips. +WAVE_SIZE = 64 +# Per-lane register slots the top-k can hold, which is what caps the context: +# a row may span at most SLOTS_MAX * WAVE_SIZE blocks. +SLOTS_MAX = 128 +MAX_SUPPORTED_BLOCKS = SLOTS_MAX * WAVE_SIZE + + +def _score_width(max_seq_len: int, block_size: int, world: int = 1) -> int: + """Blocks one score row spans, padded as the top-k requires. + + Lanes read whole wave-wide strips with no tail guard and each holds a + power-of-two count of them, so the block axis is padded past the block + count the context actually needs. + + ``world`` is the block shard the row holds: a context-parallel row covers + only this rank's ``1/world`` of the blocks, so the padding is reached from + that count and not from the context's. + """ + blocks = cdiv(cdiv(max(max_seq_len, 1), block_size), world) + return next_power_of_2(cdiv(blocks, WAVE_SIZE)) * WAVE_SIZE + + +def _score_buffer( + heads: int, + rows: int, + max_seq_len: int, + block_size: int, + device: torch.device, + world: int = 1, +) -> torch.Tensor: + """Score buffer for ``rows`` query rows, padded as the top-k requires. + + Left uninitialized on purpose: the score pass writes every block up to the + longest row it covers and the top-k reads only the blocks its own row can + see, so nothing downstream observes the padded tail. Filling it would cost + a write over what is, at long context, the largest tensor in the indexer. + """ + width = _score_width(max_seq_len, block_size, world) + return torch.empty((heads, rows, width), dtype=torch.float32, device=device) + + +_CP_SCORE: dict[tuple[int, int, int, int, int, str, int], torch.Tensor] = {} + + +def _cp_score_buffer( + heads: int, + rows: int, + max_seq_len: int, + block_size: int, + device: torch.device, + world: int, +) -> torch.Tensor: + """The CP score buffer, allocated once and shared by every layer.""" + key = (heads, rows, max_seq_len, block_size, world, device.type, device.index or 0) + if key not in _CP_SCORE: + _CP_SCORE[key] = _score_buffer( + heads, rows, max_seq_len, block_size, device, world=world + ) + return _CP_SCORE[key] + + +class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend): + """Indexer side-cache backend selecting the MSA builder.""" + + @staticmethod + def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]: + return MiniMaxM3IndexerMSAMetadataBuilder + + +@dataclass +class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): + """Adds the per-row shape the ragged top-k needs. + + The uniform decode rows are recovered inside the kernel from ``seq_lens`` and + the shared query length, which also clamps cudagraph padding rows to nothing. + Prefill rows have no such shape, so it is materialized here once per forward + and shared by every layer -- which is also where the emitted page table's + tail block comes from, since a block count alone does not say how many tokens + the last block holds. + """ + + # [num_prefill_tokens] int32, causal block count per row. + prefill_num_valid_pages: torch.Tensor | None = None + # [num_prefill_tokens] int32, the request each prefill row belongs to. + prefill_row_req_id: torch.Tensor | None = None + # [num_prefill_tokens] int32, causal token count (position + 1) per row. + prefill_kv_lens: torch.Tensor | None = None + + +class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): + """The Triton indexer's metadata plus the prefill rows' causal shape.""" + + 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) + hf_config = vllm_config.model_config.hf_config + text_config = getattr(hf_config, "text_config", hf_config) + self.sparse_block_size = int( + text_config.sparse_attention_config["sparse_block_size"] + ) + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + # Companions to the base's num_valid_pages_buffer, for the two vectors + # only the emitted table needs. + self.row_req_id_buffer = torch.empty( + max_tokens, dtype=torch.int32, device=device + ) + self.kv_lens_buffer = torch.empty(max_tokens, dtype=torch.int32, device=device) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> MiniMaxM3IndexerMSAMetadata: + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + query_start_loc = common_attn_metadata.query_start_loc + seq_lens = common_attn_metadata.seq_lens + block_table = common_attn_metadata.block_table_tensor + + 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=True, + ) + ) + assert num_decodes + num_prefills == num_reqs + assert num_decode_tokens + num_prefill_tokens == num_tokens + + # Decode-first batch: context lengths into the stable cudagraph buffer. + context_lens = self.context_len_buffer[:num_reqs] + context_lens.copy_( + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True + ) + + prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None + prefill_num_valid_pages: torch.Tensor | None = None + prefill_row_req_id: torch.Tensor | None = None + prefill_kv_lens: torch.Tensor | None = None + if num_prefills > 0: + cu_seqlens_q = (query_start_loc[num_decodes:] - num_decode_tokens).to( + torch.int32 + ) + prefill_metadata = MiniMaxM3IndexerPrefillMetadata( + cu_seqlens_q=cu_seqlens_q, + seq_lens=seq_lens[num_decodes:], + context_lens=context_lens[num_decodes:], + block_table=block_table[num_decodes:], + max_query_len=common_attn_metadata.max_query_len, + max_seq_len=common_attn_metadata.max_seq_len, + ) + # A prefill row sees its own position, so its causal length and block + # count both follow from that alone; the request it belongs to comes + # from the query offsets. Prefill batches are never captured, so the + # stable buffers are only being reused here, not required. + positions = common_attn_metadata.positions + assert positions is not None + row_positions = positions[num_decode_tokens:num_tokens] + prefill_num_valid_pages = self.num_valid_pages_buffer[ + num_decode_tokens:num_tokens + ] + prefill_num_valid_pages.copy_( + row_positions // self.sparse_block_size + 1, non_blocking=True + ) + prefill_kv_lens = self.kv_lens_buffer[num_decode_tokens:num_tokens] + prefill_kv_lens.copy_(row_positions + 1, non_blocking=True) + prefill_row_req_id = self.row_req_id_buffer[num_decode_tokens:num_tokens] + prefill_row_req_id.copy_( + torch.searchsorted( + cu_seqlens_q[1:].contiguous(), + torch.arange( + num_prefill_tokens, + dtype=torch.int32, + device=cu_seqlens_q.device, + ), + right=True, + ), + non_blocking=True, + ) + + decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None + if num_decodes > 0: + qsl_cpu = common_attn_metadata.query_start_loc_cpu + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] + decode_query_len = int(query_lens_cpu[0].item()) + assert decode_query_len > 0 + assert torch.all( + (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) + ) + assert num_decode_tokens == num_decodes * decode_query_len + decode_metadata = MiniMaxM3IndexerDecodeMetadata( + seq_lens=seq_lens[:num_decodes], + block_table=block_table[:num_decodes], + max_seq_len=common_attn_metadata.max_seq_len, + decode_query_len=decode_query_len, + max_decode_query_len=self.max_decode_query_len, + ) + + return MiniMaxM3IndexerMSAMetadata( + seq_lens=seq_lens, + max_seq_len=common_attn_metadata.max_seq_len, + slot_mapping=common_attn_metadata.slot_mapping, + num_actual_tokens=num_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + prefill=prefill_metadata, + decode=decode_metadata, + prefill_num_valid_pages=prefill_num_valid_pages, + prefill_row_req_id=prefill_row_req_id, + prefill_kv_lens=prefill_kv_lens, + ) + + +class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): + """Fp8 score + top-k for both prefill and decode.""" + + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + + if self.indexer_kv_dtype not in ("fp8", "fp8_e4m3"): + raise ValueError( + "The MSA indexer requires an fp8 e4m3 index cache, got " + f"indexer_kv_dtype={self.indexer_kv_dtype!r}" + ) + + self.sparse_bt_buffer: torch.Tensor | None = None + self.sparse_ctx_buffer: torch.Tensor | None = None + + self.indexer_cp = False + self.cp_world = 1 + self.cp_rank = 0 + + self.cp_max_seq_len = 0 + + self.cp_cand_rows = 0 + self._cp_score: torch.Tensor | None = None + self._cp_peers: IndexerCpPeers | None = None + + def init_cp(self) -> None: + """Resolve the CP decode pass and allocate its buffers. + + Separate from ``__init__`` because the topology is assigned after it, + and persistent because the decode path is captured: a replay reuses the + addresses it recorded, so the chain has to hand back the same ones. + + Every shape is taken from ``max_model_len`` rather than the batch, for + the same reason. Capture goes through a dummy batch one token long, so a + width derived from that batch's own ``max_seq_len`` would bake a + one-block shard into the graph and then replay it at full context. The + per-request causal bound clips the real work at runtime instead. + """ + config = get_current_vllm_config() + self.cp_max_seq_len = config.model_config.max_model_len + spec = getattr(config, "speculative_config", None) + n_spec = int(getattr(spec, "num_speculative_tokens", 0) or 0) if spec else 0 + # A verification step's draft tokens are extra rows of one decode pass: + # they score against the same blocks, so the indexer sees them together. + max_query_len = n_spec + 1 + + comp = config.compilation_config + max_capture = getattr(comp, "max_cudagraph_capture_size", 0) or 0 + # Decode rows the chain may be handed, which a captured batch pads up to. + max_rows = ( + max(config.scheduler_config.max_num_seqs, max_capture) * max_query_len + ) + device = self.index_cache.kv_cache.device + # Every rank scores every index head over its own shard, so the score + # buffer carries the full head count; only what survives the merge is + # this rank's own run. + self._cp_score = _cp_score_buffer( + self.num_index_heads, + max_rows, + self.cp_max_seq_len, + self.block_size, + device, + self.cp_world, + ) + # The candidates never surface as a tensor here: the selector writes + # them into a peer's buffer and reads its own back within one launch, + # so what this layer needs is the mapping and not the storage. Shared + # with every other layer, which the generation each candidate carries + # is what makes safe. + self.cp_cand_rows = max_rows + self._cp_peers = get_indexer_cp_peers( + self.num_owned_index_heads, max_rows, self.topk_blocks + ) + logger.info_once( + "MiniMax M3 indexer CP: fused decode selector over %d ranks " + "[index_heads=%d, shard=%d blocks of %d]", + self.cp_world, + self.num_index_heads, + cdiv(cdiv(self.cp_max_seq_len, self.block_size), self.cp_world), + cdiv(self.cp_max_seq_len, self.block_size), + ) + + @property + def num_owned_index_heads(self) -> int: + """Index heads whose selection this rank ends up holding. + + Under CP the scoring pass runs every head on every rank, but the merge + leaves each rank only the run it owns -- so this, not + ``num_index_heads``, is the width of the top-k the attend reads and of + everything downstream of the exchange. It equals this rank's KV head + count, which is what the emitted page table is laid out against. + """ + return self.num_index_heads // self.cp_world + + @property + def pages_per_block(self) -> int: + """Physical pages one selected block expands into for the attend.""" + return self.block_size // ASM_PAGE_SIZE + + def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]: + """The page table and context rows covering score rows ``[lo, hi)``. + + One table row per (token, kv head), head minor, which is the order + ``pa_decode_gluon`` reads once it flattens the cache. Slices of the + shared buffers rather than copies, so the attend sees the writes; a + model that reserved no buffers gets throwaway ones, and the attend + rebuilds the table itself in that case. + """ + kv_heads = self.num_kv_heads + if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None: + rows = (hi - lo) * kv_heads + device = self.index_cache.kv_cache.device + return ( + torch.empty( + (rows, self.topk_blocks * self.pages_per_block), + dtype=torch.int32, + device=device, + ), + torch.empty(rows, dtype=torch.int32, device=device), + ) + return ( + self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads], + self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads], + ) + + def _decode_cp_aiter( + self, + index_query: torch.Tensor, # [num_decode_tokens, num_index_heads, D] + kv: torch.Tensor, + d: MiniMaxM3IndexerDecodeMetadata, + decode_topk: torch.Tensor, + decode_page16_block_table: torch.Tensor, + sparse_bt: torch.Tensor, + sparse_ctx: torch.Tensor, + ) -> None: + """Score this rank's block shard, then select across the shards. + + Two kernels: the selector nominates, exchanges over the IPC mapping and + merges in one launch, with no collective between the halves. Both bound + the context by ``cp_max_seq_len`` so capture cannot specialize a shard + width. Every rank must reach this with the same row count or it hangs. + """ + from aiter.ops.msa_attention import ( + pa_sparse_block_score_decode, + pa_sparse_block_topk_cp, + ) + + assert self._cp_score is not None + assert self._cp_peers is not None + rows = index_query.shape[0] + score = self._cp_score[:, :rows] + + pa_sparse_block_score_decode( + index_query, + kv, + score, + d.block_table, + d.seq_lens, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + query_len=d.decode_query_len, + max_seq_len=self.cp_max_seq_len, + rank=self.cp_rank, + world=self.cp_world, + ) + # Each rank nominates the full top-k of its shard and never + # topk/world, since the global winners may lie entirely inside one + # shard. Forced blocks need no pinning on the way in -- the round-robin + # puts each of them in exactly one rank's shard, and the score pass + # already wrote its sentinel there -- but the merge re-pins them, so + # the pin survives the round trip through a score. + pa_sparse_block_topk_cp( + score, + self._cp_peers.cand_ptrs, + self._cp_peers.cp_gen, + decode_topk, + d.seq_lens, + decode_page16_block_table, + sparse_bt, + sparse_ctx, + max_seq_len=self.cp_max_seq_len, + block_size=self.block_size, + cand_rows=self.cp_cand_rows, + query_len=d.decode_query_len, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + num_kv_heads=self.num_kv_heads, + pages_per_block=self.pages_per_block, + rank=self.cp_rank, + world=self.cp_world, + ) + + def _decode( + self, + index_query: torch.Tensor, # [num_decode_tokens, num_index_heads, D] + kv: torch.Tensor, + d: MiniMaxM3IndexerDecodeMetadata, + decode_topk: torch.Tensor, + decode_page16_block_table: torch.Tensor, + sparse_bt: torch.Tensor, + sparse_ctx: torch.Tensor, + ) -> None: + """Score the blocks, take the top-k, emit the attend's page table. + + The unsharded pair does both in one rank's worth of work, which is the + thing CP exists not to do, so the sharded case takes the fused + cross-rank selector instead. + """ + if self.indexer_cp: + self._decode_cp_aiter( + index_query, + kv, + d, + decode_topk, + decode_page16_block_table, + sparse_bt, + sparse_ctx, + ) + return + + from aiter.ops.msa_attention import ( + pa_sparse_block_score_decode, + pa_sparse_block_topk, + ) + + heads, tokens = decode_topk.shape[:2] + score = _score_buffer( + heads, tokens, d.max_seq_len, self.block_size, index_query.device + ) + pa_sparse_block_score_decode( + index_query, + kv, + score, + d.block_table, + d.seq_lens, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + query_len=d.decode_query_len, + max_seq_len=d.max_seq_len, + ) + pa_sparse_block_topk( + score, + decode_topk, + decode_page16_block_table, + d.seq_lens, + sparse_bt=sparse_bt, + sparse_ctx=sparse_ctx, + max_seq_len=d.max_seq_len, + block_size=self.block_size, + query_len=d.decode_query_len, + num_kv_heads=self.num_kv_heads, + pages_per_block=self.pages_per_block, + ) + + def _prefill( + self, + index_query: torch.Tensor, # [num_prefill_tokens, num_index_heads, D] + kv: torch.Tensor, + p: MiniMaxM3IndexerPrefillMetadata, + prefill_topk: torch.Tensor, + prefill_page16_block_table: torch.Tensor, + sparse_bt: torch.Tensor, + sparse_ctx: torch.Tensor, + num_valid_pages: torch.Tensor, + row_req_id: torch.Tensor, + kv_lens: torch.Tensor, + ) -> None: + """Score the blocks, take the top-k, emit the attend's page table. + + AITER both sides, CP or not: prefill has a whole chunk of query rows to + spread over the machine and nothing to gain from sharding the block + axis on top of that, so the sharded chain the decode path takes buys it + only an exchange it would otherwise not pay for. + """ + from aiter.ops.msa_attention import ( + pa_sparse_block_score_prefill, + pa_sparse_block_topk, + ) + + if self.indexer_cp: + # The scorer indexes the head axis directly rather than taking its + # stride, so this rank's run of the replicated projection has to be + # compacted. Tensor-parallel rows arrive contiguous already. + owned = self.num_owned_index_heads + lo = self.cp_rank * owned + index_query = index_query[:, lo : lo + owned, :].contiguous() + + heads, total_q = prefill_topk.shape[:2] + score = _score_buffer( + heads, total_q, p.max_seq_len, self.block_size, index_query.device + ) + pa_sparse_block_score_prefill( + index_query, + kv, + score, + p.block_table, + p.cu_seqlens_q, + p.seq_lens, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + max_query_len=p.max_query_len, + max_seq_len=p.max_seq_len, + ) + pa_sparse_block_topk( + score, + prefill_topk, + prefill_page16_block_table, + p.seq_lens, + sparse_bt=sparse_bt, + sparse_ctx=sparse_ctx, + max_seq_len=p.max_seq_len, + block_size=self.block_size, + num_valid_pages=num_valid_pages, + row_req_id=row_req_id, + kv_lens=kv_lens, + num_kv_heads=self.num_kv_heads, + pages_per_block=self.pages_per_block, + ) + + def forward( + self, + index_query: torch.Tensor, + *, + decode_page16_block_table: torch.Tensor | None = None, + prefill_page16_block_table: torch.Tensor | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return None, None # profiling run; caches unbound + md = attn_metadata[self.index_cache.prefix] + assert isinstance(md, MiniMaxM3IndexerMSAMetadata) + # The emitted page table addresses the main cache, whose blocks are a + # different group's than the index cache's, so the top-k resolves the + # selection through the attend's block table -- only the score pass reads + # the index cache and takes the indexer's. It has to be the page-16 + # rebase of that table, which the attend's own metadata builder does + # once per step; the indexer's metadata cannot reach another group's + # blocks, so the layer reads it off the attend and hands it in. + num_tokens = md.num_actual_tokens + nd = md.num_decode_tokens + iq = index_query[:num_tokens].view( + -1, self.num_index_heads, self.index_head_dim + ) + kv = self.index_cache.kv_cache + + # Both sides write into the single shared persistent buffer (decode at + # [:, :nd], prefill at [:, nd:]) and return views into it. The top-k + # takes the head/row strides, so these slices need no copy back. + buf = self.topk_indices_buffer + if buf is None: + buf = torch.empty( + (self.num_owned_index_heads, num_tokens, self.topk_blocks), + dtype=torch.int32, + device=iq.device, + ) + + decode_topk: torch.Tensor | None = None + prefill_topk: torch.Tensor | None = None + + if md.num_decodes > 0: + d = md.decode + assert d is not None + assert decode_page16_block_table is not None, ( + "the MSA indexer's top-k emits the attend's page table and " + "needs the page-16 rebase of the attend's decode block table" + ) + decode_topk = buf[:, :nd, :] + sparse_bt, sparse_ctx = self._table_rows(0, nd) + # A decode row's causal length is seq_len - query_len + token + 1, + # which every kernel in the chain derives itself, so speculative + # rows need no extra per-row shape. + self._decode( + iq[:nd], + kv, + d, + decode_topk, + decode_page16_block_table, + sparse_bt, + sparse_ctx, + ) + + if md.num_prefills > 0: + p = md.prefill + assert p is not None + assert prefill_page16_block_table is not None, ( + "the MSA indexer's top-k emits the attend's page table and " + "needs the page-16 rebase of the attend's prefill block table" + ) + assert md.prefill_num_valid_pages is not None + assert md.prefill_row_req_id is not None + assert md.prefill_kv_lens is not None + prefill_topk = buf[:, nd:num_tokens, :] + sparse_bt, sparse_ctx = self._table_rows(nd, num_tokens) + self._prefill( + iq[nd:], + kv, + p, + prefill_topk, + prefill_page16_block_table, + sparse_bt, + sparse_ctx, + md.prefill_num_valid_pages, + md.prefill_row_req_id, + md.prefill_kv_lens, + ) + + return decode_topk, prefill_topk + + +class MiniMaxM3MSAIndexer(nn.Module): + """``MiniMaxM3Indexer``'s surface over the MSA impl. + Fp8 score + top-k for both prefill and decode. + """ + + def __init__( + self, + *, + impl_cls: type[MiniMaxM3IndexerMSAImpl], + sparse_bt_buffer: torch.Tensor | None = None, + sparse_ctx_buffer: torch.Tensor | None = None, + indexer_cp: bool = False, + **impl_kwargs, + ) -> None: + super().__init__() + self.impl = impl_cls(**impl_kwargs) + # Assigned rather than passed: the impl's base ``__init__`` lives in + # common and takes neither the table buffers nor the CP topology. + self.impl.sparse_bt_buffer = sparse_bt_buffer + self.impl.sparse_ctx_buffer = sparse_ctx_buffer + self.impl.indexer_cp = indexer_cp + if indexer_cp: + self.impl.cp_world = get_tensor_model_parallel_world_size() + self.impl.cp_rank = get_tensor_model_parallel_rank() + # Resolves which decode chain runs and allocates its buffers, which + # needs the topology above and so cannot happen in the impl's own + # ``__init__``. + self.impl.init_cp() + + @property + def index_cache(self) -> MiniMaxM3IndexerCache: + return self.impl.index_cache + + @property + def num_index_heads(self) -> int: + return self.impl.num_index_heads + + def forward( + self, + index_query: torch.Tensor, + *, + decode_page16_block_table: torch.Tensor | None = None, + prefill_page16_block_table: torch.Tensor | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + return self.impl( + index_query, + decode_page16_block_table=decode_page16_block_table, + prefill_page16_block_table=prefill_page16_block_table, + ) + + +def msa_indexer_unsupported_reason( + *, + topk_blocks: int, + sparse_block_size: int, + index_head_dim: int, + indexer_kv_dtype: IndexerKVDType | str, + max_model_len: int, + score_type: str = "max", + cp_world: int = 1, + num_index_heads: int = 1, + num_kv_heads: int = 1, + max_query_len: int = 1, + dcp_size: int = 1, +) -> str | None: + """Return why this config cannot use this indexer, or None if it can. + + Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache + dtype, the kernels' shape contract and the max context in blocks. The + default ``cp_world=1`` ignores every argument below it and asks about the + unsharded indexer; above one it adds the checks sharding brings. Callers + fall back on a reason rather than raise. + """ + if not current_platform.is_rocm(): + return ( + "needs ROCm for the fp8 MFMA score/top-k, " + f"got platform={current_platform.device_type!r}" + ) + if not _minimax_m3_aiter_sparse_pa_requested(): + # The top-k emits the attend's page table in page-16 numbering, which + # only addresses the interleaved cache the AITER attend reads. Paired + # with any other attend there is nowhere valid to write it, so this + # indexer is not usable on its own. + return ( + "needs the AITER sparse PA attend, whose page table its top-k " + "emits (rocm_aiter_ops + shuffle KV cache layout)" + ) + if indexer_kv_dtype not in ("fp8", "fp8_e4m3"): + # The score kernels are fp8 MFMA; there is no bf16 instantiation. + return ( + f"needs an fp8 e4m3 index cache, got indexer_kv_dtype={indexer_kv_dtype!r}" + ) + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + return f"needs {' or '.join(SUPPORTED_ARCHS)} for the fp8 MFMA" + if score_type != MSA_SCORE_TYPE: + return f"needs score_type={MSA_SCORE_TYPE!r}, got score_type={score_type!r}" + if topk_blocks != MSA_TOPK_BLOCKS: + return f"needs topk_blocks={MSA_TOPK_BLOCKS}, got topk_blocks={topk_blocks}" + if sparse_block_size != MSA_SPARSE_BLOCK_SIZE: + return ( + f"needs sparse_block_size={MSA_SPARSE_BLOCK_SIZE}, " + f"got sparse_block_size={sparse_block_size}" + ) + if index_head_dim != MSA_INDEX_HEAD_DIM: + return ( + f"needs index_head_dim={MSA_INDEX_HEAD_DIM}, " + f"got index_head_dim={index_head_dim}" + ) + + max_blocks = math.ceil(max_model_len / sparse_block_size) + if max_blocks > MAX_SUPPORTED_BLOCKS: + return ( + f"max_model_len={max_model_len} needs {max_blocks} blocks per row, " + f"more than the {MAX_SUPPORTED_BLOCKS} the top-k is compiled for" + ) + + if cp_world <= 1: + return None + + from aiter.ops.msa_block_select import SCORE_MFMA_COLS, TOPK_CP_MAX_CAND + + if dcp_size > 1: + return ( + f"decode_context_parallel_size={dcp_size} > 1 (KV-cache DCP owns " + "the context axis; indexer CP replaces it, it does not extend it)" + ) + if cp_world > num_index_heads or num_index_heads % cp_world: + return ( + f"needs num_index_heads divisible by cp_world, got " + f"cp_world={cp_world}, num_index_heads={num_index_heads}" + ) + if num_index_heads * max_query_len > SCORE_MFMA_COLS: + return ( + f"num_index_heads={num_index_heads} x query_len={max_query_len} " + f"exceeds the {SCORE_MFMA_COLS} MFMA columns of the score pass" + ) + candidates = cp_world * topk_blocks + if candidates > TOPK_CP_MAX_CAND: + return ( + f"cp_world={cp_world} x topk_blocks={topk_blocks} = {candidates} " + f"exceeds the {TOPK_CP_MAX_CAND} candidates the merge holds" + ) + if num_index_heads // cp_world != num_kv_heads: + return ( + f"this rank would own {num_index_heads // cp_world} index head(s) " + f"but {num_kv_heads} kv head(s); the merge's page table needs " + "them 1:1" + ) + return None + + +def select_msa_indexer_impl_cls( + *, + topk_blocks: int, + sparse_block_size: int, + index_head_dim: int, + indexer_kv_dtype: IndexerKVDType | str, + score_type: str = "max", +) -> type[MiniMaxM3IndexerMSAImpl] | None: + """The MSA indexer impl if this config can use it, else None.""" + reason = msa_indexer_unsupported_reason( + topk_blocks=topk_blocks, + sparse_block_size=sparse_block_size, + index_head_dim=index_head_dim, + indexer_kv_dtype=indexer_kv_dtype, + max_model_len=get_current_vllm_config().model_config.max_model_len, + score_type=score_type, + ) + if reason is not None: + logger.info_once("MiniMax M3 indexer: MSA unavailable (%s)", reason) + return None + logger.info_once( + "MiniMax M3 indexer: selected MSA (Triton fp8 MFMA score + top-k) " + "[topk_blocks=%d, indexer_kv_dtype=%s]", + topk_blocks, + indexer_kv_dtype, + ) + return MiniMaxM3IndexerMSAImpl + + +def msa_indexer_cp_enabled( + vllm_config: VllmConfig, + *, + msa_indexer_selected: bool, +) -> bool: + """Whether to run the selected MSA indexer context-parallel. + + ``select_msa_indexer_impl_cls``'s gate re-asked at the CP world. Call once + per model: the projection layout and head count are decided off one answer. + """ + requested = vllm_config.attention_config.minimax_m3_indexer_cp + sparse_cfg = getattr( + vllm_config.model_config.hf_text_config, "sparse_attention_config", None + ) + if not requested or sparse_cfg is None: + return False + + if not msa_indexer_selected: + # CP has no tensor-parallel fallback to offer: the qkv projection is + # built with index_q replicated, so an indexer that did not take this + # path would read its own heads out of a replicated tensor. The + # selection logged its own reason on the way past. + reason: str | None = "the MSA indexer was not selected, so nothing to shard" + else: + config = vllm_config.model_config.hf_text_config + tp_size = get_tensor_model_parallel_world_size() + spec = vllm_config.speculative_config + n_spec = int(getattr(spec, "num_speculative_tokens", 0) or 0) if spec else 0 + reason = msa_indexer_unsupported_reason( + topk_blocks=int(sparse_cfg["sparse_topk_blocks"]), + sparse_block_size=int(sparse_cfg["sparse_block_size"]), + index_head_dim=int(sparse_cfg["sparse_index_dim"]), + indexer_kv_dtype=vllm_config.attention_config.resolve_indexer_kv_dtype( + "bf16" + ), + max_model_len=vllm_config.model_config.max_model_len, + score_type=sparse_cfg.get("sparse_score_type", "max"), + cp_world=tp_size, + num_index_heads=int(sparse_cfg["sparse_num_index_heads"]), + num_kv_heads=max(1, int(config.num_key_value_heads) // tp_size), + max_query_len=n_spec + 1, + dcp_size=vllm_config.parallel_config.decode_context_parallel_size, + ) + + if reason is not None: + logger.info_once("MiniMax M3 indexer CP: disabled (%s)", reason) + return False + world = get_tensor_model_parallel_world_size() + logger.info_once( + "MiniMax M3 indexer CP: enabled over %d ranks (each scores 1/%d of " + "the blocks for every index head)", + world, + world, + ) + return True diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index f4d83312b185..3222280b65ab 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -89,9 +89,11 @@ spec_decode_needs_target_embed, ) from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model -from vllm.models.minimax_m3.amd.indexer_aiter import ( - MiniMaxM3AiterIndexer, - select_aiter_indexer_impl_cls, +from vllm.models.minimax_m3.amd.indexer_msa import ( + MiniMaxM3IndexerMSAImpl, + MiniMaxM3MSAIndexer, + msa_indexer_cp_enabled, + select_msa_indexer_impl_cls, ) from vllm.models.minimax_m3.amd.ops import ( gemma_fused_add_rmsnorm, @@ -642,7 +644,10 @@ def __init__( prefix: str = "", cache_config: CacheConfig | None = None, topk_indices_buffer: torch.Tensor | None = None, - sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None, + msa_indexer_impl_cls: type[MiniMaxM3IndexerMSAImpl] | None = None, + sparse_bt_buffer: torch.Tensor | None = None, + sparse_ctx_buffer: torch.Tensor | None = None, + indexer_cp: bool = False, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -673,7 +678,8 @@ def __init__( # identically -- including replication when tp_size > num_key_value_heads. sparse_cfg = config.sparse_attention_config self.total_idx_heads = sparse_cfg["sparse_num_index_heads"] - self.num_idx_heads = self.num_kv_heads + self.indexer_cp = indexer_cp + self.num_idx_heads = self.total_idx_heads if indexer_cp else self.num_kv_heads self.idx_head_dim = sparse_cfg["sparse_index_dim"] self.index_q_size = self.num_idx_heads * self.idx_head_dim @@ -688,6 +694,7 @@ def __init__( bias=False, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", + replicate_index_q=indexer_cp, ) # reduce_results=False: the attention all-reduce is fused with the # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer @@ -733,7 +740,7 @@ def __init__( # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main cache # (--attention-config '{"indexer_kv_dtype": ...}'). fp8 e4m3 is what the - # AITER indexer needs; bf16 keeps the Triton indexer. + # MSA indexer needs; bf16 keeps the platform-neutral indexer. self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( "bf16" ) @@ -741,25 +748,18 @@ def __init__( # Shared top-k buffer: the indexer writes the selected blocks into it and # the attend impl reads them back (no Python value crosses the break). self.topk_indices_buffer = topk_indices_buffer - # Indexer and main attention are separate impls. The AITER indexer is - # ROCm-only, so it is selected here rather than inside the neutral - # MiniMaxM3Indexer, which knows nothing about it and would pick Triton. - indexer_impl_cls = select_aiter_indexer_impl_cls( - topk_blocks=sparse_cfg["sparse_topk_blocks"], - sparse_block_size=sparse_cfg["sparse_block_size"], - num_index_heads=self.num_idx_heads, - index_head_dim=self.idx_head_dim, - indexer_kv_dtype=self.indexer_kv_dtype, - score_type=sparse_cfg.get("sparse_score_type", "max"), - ) - # The attend gate below needs this: an emitted table is what lets it - # serve more than one KV head per rank. Buffers to write the table into - # only arrive when the model already resolved the attend to the AITER - # path the table addresses, so that is not rechecked here. - self.indexer_emits_table = ( - indexer_impl_cls is not None and sparse_table_buffers is not None - ) - # The backend names the metadata builder, which for the AITER path is + # Indexer and main attention are separate impls. This indexer is + # ROCm-only, so the model selects it rather than the neutral + # MiniMaxM3Indexer, which knows nothing about it and would pick its own. + # The selection arrives from the model, which has to make it anyway to + # size the buffers its top-k emits into, so there is no second probe + # here to drift out of step with it -- every input to that selection is + # global, so a layer could only ever recompute the same answer. + # + # The attend gate further down needs this too: an emitted table is what + # lets it serve more than one KV head per rank. + self.indexer_emits_table = msa_indexer_impl_cls is not None + # The backend names the metadata builder, which for the MSA path is # the one that rebases the block table the indexer's top-k resolves # through, so both have to come from the same selection. self.attn_backend, main_impl_cls = select_main_backend_and_impl_cls( @@ -800,19 +800,28 @@ def __init__( indexer_kv_dtype=self.indexer_kv_dtype, topk_indices_buffer=topk_indices_buffer, ) - self.sparse_bt_buffer: torch.Tensor | None = None - self.sparse_ctx_buffer: torch.Tensor | None = None - if indexer_impl_cls is None: + # Read back off the layer by the AITER attend, which builds the table + # itself when they are absent. + self.sparse_bt_buffer = sparse_bt_buffer + self.sparse_ctx_buffer = sparse_ctx_buffer + # Same condition the forward dispatches on, so the indexer built here + # and the branch that calls it cannot disagree about which one this is. + if msa_indexer_impl_cls is None: # Self-contained nn.Module: owns its side cache, selects its impl. self.indexer = MiniMaxM3Indexer(**indexer_kwargs) else: - if self.indexer_emits_table: - assert sparse_table_buffers is not None - self.sparse_bt_buffer, self.sparse_ctx_buffer = sparse_table_buffers - self.indexer = MiniMaxM3AiterIndexer( - impl_cls=indexer_impl_cls, - sparse_bt_buffer=self.sparse_bt_buffer, - sparse_ctx_buffer=self.sparse_ctx_buffer, + # The impl has no signature for the Triton fused-table call the + # branch above falls back to, so a selection that reached us + # without its buffers would have nowhere valid to go. + assert sparse_bt_buffer is not None and sparse_ctx_buffer is not None, ( + "the MSA indexer emits the attend's page table and cannot be " + "selected without the buffers it emits into" + ) + self.indexer = MiniMaxM3MSAIndexer( + impl_cls=msa_indexer_impl_cls, + sparse_bt_buffer=sparse_bt_buffer, + sparse_ctx_buffer=sparse_ctx_buffer, + indexer_cp=indexer_cp, **indexer_kwargs, ) @@ -1212,9 +1221,9 @@ def _run_attention( decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None if not self.skip_index_topk: assert index_query is not None - # The AITER indexer emits the attend table into its persistent + # The MSA indexer emits the attend table into its persistent # buffers, resolving its selection through the page-16 rebase of the - # attend's block table. The Triton indexer can instead fuse decode + # attend's block table. The neutral indexer can instead fuse decode # top-k with main's per-forward sparse-table allocation. if self.indexer_emits_table: attn_metadata = get_forward_context().attn_metadata @@ -1284,7 +1293,10 @@ def __init__( force_sparse_attn: bool = False, force_moe: bool = False, topk_indices_buffer: torch.Tensor | None = None, - sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None, + msa_indexer_impl_cls: type[MiniMaxM3IndexerMSAImpl] | None = None, + sparse_bt_buffer: torch.Tensor | None = None, + sparse_ctx_buffer: torch.Tensor | None = None, + indexer_cp: bool = False, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -1305,7 +1317,10 @@ def __init__( prefix=f"{prefix}.self_attn", cache_config=cache_config, topk_indices_buffer=topk_indices_buffer, - sparse_table_buffers=sparse_table_buffers, + msa_indexer_impl_cls=msa_indexer_impl_cls, + sparse_bt_buffer=sparse_bt_buffer, + sparse_ctx_buffer=sparse_ctx_buffer, + indexer_cp=indexer_cp, ) else: self.self_attn = MiniMaxM3Attention( @@ -1398,7 +1413,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # layers (mirrors DeepseekV4); the indexer writes its per-head decode/ # prefill block selection into it, the attend reads it back. sparse_cfg = getattr(config, "sparse_attention_config", None) - self.sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None + self.msa_indexer_impl_cls: type[MiniMaxM3IndexerMSAImpl] | None = None + self.sparse_bt_buffer: torch.Tensor | None = None + self.sparse_ctx_buffer: torch.Tensor | None = None + self.indexer_cp = False if sparse_cfg is not None: tp_size = get_tensor_model_parallel_world_size() num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size) @@ -1414,36 +1432,53 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( "bf16" ) - # Probed with the head count the layer will pass (it derives - # num_idx_heads from num_kv_heads, not from sparse_num_index_heads), - # so this cannot land on the other side of the MFMA column limit - # from the selection the layer makes. - emits_table = ( - select_aiter_indexer_impl_cls( - topk_blocks=sparse_cfg["sparse_topk_blocks"], - sparse_block_size=sparse_cfg["sparse_block_size"], - num_index_heads=num_kv_heads, - index_head_dim=sparse_cfg["sparse_index_dim"], - indexer_kv_dtype=indexer_kv_dtype, - score_type=sparse_cfg.get("sparse_score_type", "max"), - ) - is not None + # The only MSA indexer probe. It has to live here rather than in + # the layer because the buffers it gates are one allocation shared + # by all of them, so the answer is needed before any layer exists, + # and the impl itself is kept rather than collapsed to a bool so + # the layer has nothing left to re-derive. Asked before the CP + # question and not after, because CP is a choice about how to run + # this indexer and there is no such choice until it exists; the + # probe owes nothing to CP in return, since what it checks is the + # platform and the kernels' shape contract, neither of which a + # head count enters. + impl_cls = select_msa_indexer_impl_cls( + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + index_head_dim=sparse_cfg["sparse_index_dim"], + indexer_kv_dtype=indexer_kv_dtype, + score_type=sparse_cfg.get("sparse_score_type", "max"), ) - if emits_table and minimax_m3_use_aiter_sparse_pa( + if impl_cls is not None and minimax_m3_use_aiter_sparse_pa( num_kv_heads, emits_sparse_block_table=True ): # One row per (token, kv head): the table addresses the attend's - # cache, not the index cache. + # cache, not the index cache. Allocated here with the selection + # that needs them, and only here -- one allocation serves every + # layer, which is why the probe above cannot live in the layer. rows = max_tokens * num_kv_heads - self.sparse_table_buffers = ( - torch.empty( - rows, - sparse_cfg["sparse_topk_blocks"] - * (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE), - dtype=torch.int32, - ), - torch.empty(rows, dtype=torch.int32), + self.msa_indexer_impl_cls = impl_cls + self.sparse_bt_buffer = torch.empty( + rows, + sparse_cfg["sparse_topk_blocks"] + * (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE), + dtype=torch.int32, ) + self.sparse_ctx_buffer = torch.empty(rows, dtype=torch.int32) + # Resolved once here and handed to every layer. The gate reads + # only config, topology and the selection above, so re-asking per + # layer would give the same answer -- but the projection's shard + # layout and the indexer's head count have to be decided off one + # answer, and one call site is what guarantees that. The selection + # is handed in rather than looked up: CP has no tensor-parallel + # fallback to offer, since the qkv projection is built with + # index_q replicated and an indexer that did not take this path + # would read its own heads out of a replicated tensor, so the two + # cannot be allowed to disagree. + self.indexer_cp = msa_indexer_cp_enabled( + vllm_config, + msa_indexer_selected=self.msa_indexer_impl_cls is not None, + ) else: self.topk_indices_buffer = None @@ -1455,7 +1490,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): cache_config=cache_config, quant_config=quant_config, topk_indices_buffer=self.topk_indices_buffer, - sparse_table_buffers=self.sparse_table_buffers, + msa_indexer_impl_cls=self.msa_indexer_impl_cls, + sparse_bt_buffer=self.sparse_bt_buffer, + sparse_ctx_buffer=self.sparse_ctx_buffer, + indexer_cp=self.indexer_cp, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/minimax_m3/amd/mtp.py b/vllm/models/minimax_m3/amd/mtp.py index 2f7dd8a8316d..4c63b3ff01f7 100644 --- a/vllm/models/minimax_m3/amd/mtp.py +++ b/vllm/models/minimax_m3/amd/mtp.py @@ -18,13 +18,15 @@ """ from collections.abc import Iterable +from typing import NamedTuple import regex as re import torch import torch.nn as nn +from transformers import PreTrainedConfig from vllm.config import VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -46,6 +48,15 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) +from vllm.models.minimax_m3.amd.indexer_msa import ( + MiniMaxM3IndexerMSAImpl, + msa_indexer_cp_enabled, + select_msa_indexer_impl_cls, +) +from vllm.models.minimax_m3.amd.ops.sparse_pa import ASM_PAGE_SIZE +from vllm.models.minimax_m3.common.sparse_attention import ( + minimax_m3_use_aiter_sparse_pa, +) from vllm.sequence import IntermediateTensors from .model import ( @@ -54,8 +65,81 @@ ) +class _SparseIndexerContext(NamedTuple): + topk_indices_buffer: torch.Tensor | None = None + msa_indexer_impl_cls: type[MiniMaxM3IndexerMSAImpl] | None = None + sparse_bt_buffer: torch.Tensor | None = None + sparse_ctx_buffer: torch.Tensor | None = None + indexer_cp: bool = False + + +def _build_sparse_indexer_context( + config: PreTrainedConfig, vllm_config: VllmConfig +) -> _SparseIndexerContext: + sparse_cfg = getattr(config, "sparse_attention_config", None) + if sparse_cfg is None: + return _SparseIndexerContext() + + # Reserved top-k indices buffer shared by all the sparse-attention indexer + # layers below; the indexer writes its per-head decode/prefill block + # selection into it, the attend reads it back. + tp_size = get_tensor_model_parallel_world_size() + num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size) + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + topk_indices_buffer = torch.empty( + num_index_heads, + max_tokens, + sparse_cfg["sparse_topk_blocks"], + dtype=torch.int32, + ) + + num_kv_heads = max(1, config.num_key_value_heads // tp_size) + indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("bf16") + + impl_cls = select_msa_indexer_impl_cls( + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + index_head_dim=sparse_cfg["sparse_index_dim"], + indexer_kv_dtype=indexer_kv_dtype, + score_type=sparse_cfg.get("sparse_score_type", "max"), + ) + msa_indexer_impl_cls = None + sparse_bt_buffer = None + sparse_ctx_buffer = None + if impl_cls is not None and minimax_m3_use_aiter_sparse_pa( + num_kv_heads, emits_sparse_block_table=True + ): + # One row per (token, kv head): the table addresses the attend's cache, + # not the index cache. + rows = max_tokens * num_kv_heads + msa_indexer_impl_cls = impl_cls + sparse_bt_buffer = torch.empty( + rows, + sparse_cfg["sparse_topk_blocks"] + * (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE), + dtype=torch.int32, + ) + sparse_ctx_buffer = torch.empty(rows, dtype=torch.int32) + + indexer_cp = msa_indexer_cp_enabled( + vllm_config, msa_indexer_selected=msa_indexer_impl_cls is not None + ) + return _SparseIndexerContext( + topk_indices_buffer, + msa_indexer_impl_cls, + sparse_bt_buffer, + sparse_ctx_buffer, + indexer_cp, + ) + + class MiniMaxM3MultiTokenPredictorLayer(nn.Module): - def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + indexer_ctx: _SparseIndexerContext, + ) -> None: super().__init__() assert vllm_config.speculative_config is not None @@ -79,6 +163,11 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: quant_config=quant_config, force_sparse_attn=True, force_moe=True, + topk_indices_buffer=indexer_ctx.topk_indices_buffer, + msa_indexer_impl_cls=indexer_ctx.msa_indexer_impl_cls, + sparse_bt_buffer=indexer_ctx.sparse_bt_buffer, + sparse_ctx_buffer=indexer_ctx.sparse_ctx_buffer, + indexer_cp=indexer_ctx.indexer_cp, ) self.final_layernorm = MiniMAXGemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -124,10 +213,24 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # standalone checkpoint, and the promoted text_config for a bundled one. config = vllm_config.speculative_config.draft_model_config.hf_config self.num_mtp_layers = config.num_mtp_modules + indexer_ctx = _build_sparse_indexer_context(config, vllm_config) + if indexer_ctx.msa_indexer_impl_cls is None: + indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( + "bf16" + ) + if indexer_kv_dtype != "bf16": + raise NotImplementedError( + f"MiniMax M3 MTP could not select the MSA indexer, leaving " + f"its draft layers on the bf16 indexer, so " + f"indexer_kv_dtype={indexer_kv_dtype!r} has nowhere to go. " + f"See the preceding 'MSA unavailable' log line for why. Drop " + f"indexer_kv_dtype from --attention-config, or speculate " + f"with an EAGLE3 drafter instead." + ) self.layers = torch.nn.ModuleDict( { str(idx): MiniMaxM3MultiTokenPredictorLayer( - vllm_config, f"{prefix}.layers.{idx}" + vllm_config, f"{prefix}.layers.{idx}", indexer_ctx ) for idx in range(self.num_mtp_layers) } From 39504080e56e6438207baba9100b0abf92758f85 Mon Sep 17 00:00:00 2001 From: xaguilar-amd Date: Sat, 24 Oct 2026 14:08:44 +0200 Subject: [PATCH 016/631] [ROCm][MLA][DCP] Drop VLLM_ROCM_AITER_MLA_DCP_VERIFY (#61628) Signed-off-by: Xavier Aguilar Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: vllm-agent --- .../attention/test_rocm_aiter_mla_dcp_cprr.py | 85 ++++++------------- .../test_rocm_aiter_mla_dcp_cprr_numerics.py | 5 -- vllm/envs.py | 11 --- .../attention/backends/mla/rocm_aiter_mla.py | 28 +++--- 4 files changed, 34 insertions(+), 95 deletions(-) diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py index 7c7589801520..0ca3ed98a702 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py @@ -30,23 +30,6 @@ MIN_QLEN = rocm_aiter_mla._MIN_CPRR_QLEN -def test_dcp_verify_env_defaults_to_auto(monkeypatch): - monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False) - assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "auto" - - -@pytest.mark.parametrize("route", ["auto", "asm", "segmented"]) -def test_dcp_verify_env_accepts_routes(monkeypatch, route): - monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route) - assert route == rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY - - -def test_dcp_verify_env_rejects_unknown_route(monkeypatch): - monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", "unknown") - with pytest.raises(ValueError, match="Invalid value"): - _ = rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY - - # -------------------------------------------------------------------------- # _asm_dcp_verify_heads -- which head count the cprr kernel runs at # -------------------------------------------------------------------------- @@ -127,13 +110,13 @@ def test_route_needs_gfx950(monkeypatch): # -------------------------------------------------------------------------- -# AiterMLAMetadataBuilder -- the route the env value resolves to +# AiterMLAMetadataBuilder -- the route the builder selects # -------------------------------------------------------------------------- DCP = 8 -def _build(monkeypatch, route, gathered_heads, segmented_supported=True): +def _build(monkeypatch, gathered_heads, segmented_supported=True): """Run the real builder __init__ on CPU at TP8/DCP8 with speculative decoding. Stubs only what needs a GPU, AITER, or a DCP group, so the route decision @@ -184,11 +167,6 @@ def init_common_builder(self, *args, **kwargs): ), ), ) - if route is None: - monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False) - else: - monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route) - config = SimpleNamespace( speculative_config=SimpleNamespace(num_speculative_tokens=4), parallel_config=SimpleNamespace( @@ -209,53 +187,38 @@ def init_common_builder(self, *args, **kwargs): ) -@pytest.mark.parametrize("route", [None, "auto", "AUTO"]) -def test_auto_selects_cprr_for_k3_gathered_heads(on_gfx950, monkeypatch, route): - """K3 at TP8/DCP8 gathers 96 heads; auto must take cprr padded to 128.""" - builder = _build(monkeypatch, route, gathered_heads=96) +def test_selects_cprr_for_k3_gathered_heads(on_gfx950, monkeypatch): + """K3 at TP8/DCP8 gathers 96 heads; the builder must take cprr padded to 128.""" + builder = _build(monkeypatch, gathered_heads=96) assert builder._asm_dcp_verify assert builder._asm_dcp_verify_heads == 128 assert builder._num_attention_heads == 128 -def test_auto_falls_back_to_segmented_above_largest_native(on_gfx950, monkeypatch): - builder = _build(monkeypatch, "auto", gathered_heads=256) - assert not builder._asm_dcp_verify - assert builder._supports_segmented_dcp_verify - - -def test_auto_without_any_route_names_the_missing_segmented_build( - on_gfx950, monkeypatch +@pytest.mark.parametrize( + "gfx950,gathered_heads", + [ + (True, 256), # above the largest native cprr kernel + (False, 96), # cprr kernels exist only on gfx950 + ], +) +def test_falls_back_to_segmented_when_cprr_cannot_serve( + monkeypatch, gfx950, gathered_heads ): - with pytest.raises(ValueError) as exc: - _build(monkeypatch, "auto", gathered_heads=256, segmented_supported=False) - assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=auto" in str(exc.value) - assert "lacks segmented MLA decode" in str(exc.value) - + """Decode must keep the gathered head count, not cprr's padded one.""" + import vllm.platforms.rocm as rocm -@pytest.mark.parametrize("route", ["asm", "ASM"]) -def test_explicit_asm_is_case_insensitive(on_gfx950, monkeypatch, route): - """env_with_choices returns the raw string; ASM used to mean segmented.""" - assert _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify + monkeypatch.setattr(rocm, "on_gfx950", lambda: gfx950) + builder = _build(monkeypatch, gathered_heads=gathered_heads) + assert not builder._asm_dcp_verify + assert builder._num_attention_heads == gathered_heads -def test_explicit_asm_refuses_unservable_heads(on_gfx950, monkeypatch): +def test_without_any_route_names_the_missing_segmented_build(on_gfx950, monkeypatch): with pytest.raises(ValueError) as exc: - _build(monkeypatch, "ASM", gathered_heads=256) - assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm" in str(exc.value) - assert "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented" in str(exc.value) - - -@pytest.mark.parametrize("route", ["segmented", "SEGMENTED"]) -def test_explicit_segmented_never_selects_cprr(on_gfx950, monkeypatch, route): - assert not _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify - - -def test_auto_off_gfx950_stays_segmented(monkeypatch): - import vllm.platforms.rocm as rocm - - monkeypatch.setattr(rocm, "on_gfx950", lambda: False) - assert not _build(monkeypatch, "auto", gathered_heads=96)._asm_dcp_verify + _build(monkeypatch, gathered_heads=256, segmented_supported=False) + assert "256 DCP-gathered heads" in str(exc.value) + assert "lacks segmented MLA decode" in str(exc.value) # -------------------------------------------------------------------------- diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py index 842fb1f0f0ba..21626e9a675e 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py @@ -78,11 +78,6 @@ def _on_gfx950() -> bool: ) -@pytest.fixture(autouse=True) -def select_asm_route(monkeypatch): - monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", "asm") - - # -------------------------------------------------------------------------- # reference maths # -------------------------------------------------------------------------- diff --git a/vllm/envs.py b/vllm/envs.py index 1d6a4aa598c1..0572a1e3735b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -146,7 +146,6 @@ VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" - VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["auto", "asm", "segmented"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False @@ -1319,16 +1318,6 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_USE_AITER_MLA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1") ), - # Kernel for causal multi-token (spec-decode) verify steps under decode - # context parallelism: "asm" uses AITER's round-robin ASM decode (gfx950), - # "segmented" the Triton segmented MLA path. "auto" (default) takes "asm" - # wherever it can serve the shape and "segmented" otherwise. - "VLLM_ROCM_AITER_MLA_DCP_VERIFY": env_with_choices( - "VLLM_ROCM_AITER_MLA_DCP_VERIFY", - "auto", - ["auto", "asm", "segmented"], - case_sensitive=False, - ), # Small-head (<16) AITER MLA decode kernel selection. Small head counts # (e.g. Kimi-K3: 12 heads/rank at TP8, 6 at TP16) can decode either through # the Gluon small-head kernel or through the padded persistent-scheduling diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 6d52d4506fa5..d923c2d87a04 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -791,14 +791,10 @@ def __init__( parallel_config.decode_context_parallel_size, parallel_config.cp_kv_cache_interleave_size, ) - verify_route = envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY.lower() - asm_dcp_verify_config = ( - _asm_dcp_verify_configured( - parallel_config.decode_context_parallel_size, - parallel_config.cp_kv_cache_interleave_size, - multi_token_decode=vllm_config.speculative_config is not None, - ) - and verify_route != "segmented" + asm_dcp_verify_config = _asm_dcp_verify_configured( + parallel_config.decode_context_parallel_size, + parallel_config.cp_kv_cache_interleave_size, + multi_token_decode=vllm_config.speculative_config is not None, ) super().__init__( kv_cache_spec, @@ -822,17 +818,13 @@ def __init__( ) if self._asm_dcp_verify_heads: self._asm_dcp_verify = True - elif verify_route == "asm" or not supports_segmented_dcp_verify: + elif not supports_segmented_dcp_verify: raise ValueError( - f"VLLM_ROCM_AITER_MLA_DCP_VERIFY={verify_route}, but the " - "round-robin asm decode has no kernel for " - f"{self.num_heads * self.dcp_world_size} DCP-gathered heads " - f"(native counts: {_NATIVE_CPRR_HEADS}). " - + ( - "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented." - if supports_segmented_dcp_verify - else "This AITER build also lacks segmented MLA decode." - ) + "ROCM_AITER_MLA DCP multi-token verify has no kernel for " + f"{self.num_heads * self.dcp_world_size} DCP-gathered heads: " + "the round-robin asm decode only has native counts " + f"{_NATIVE_CPRR_HEADS}, and this AITER build lacks segmented " + "MLA decode." ) self._supports_segmented_dcp_verify = ( supports_segmented_dcp_verify From c0ba7e01ffdd446a087df759e8c8c2bda28de2ea Mon Sep 17 00:00:00 2001 From: xiao-vllm_amd Date: Sat, 24 Oct 2026 08:09:43 -0400 Subject: [PATCH 017/631] =?UTF-8?q?[Perf][ROCm][Kernel]=20Add=20W4A8=20int?= =?UTF-8?q?8-dot=20MXFP4=20GEMV=20for=20RDNA3/RDNA3.5=20d=E2=80=A6=20(#604?= =?UTF-8?q?13)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Xiao Yu Co-authored-by: Cursor Co-authored-by: vllm-agent --- CMakeLists.txt | 3 +- csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu | 254 +++++++++++++++++ csrc/rocm/ops.h | 3 + csrc/rocm/torch_bindings.cpp | 6 + docs/features/quantization/online.md | 9 + .../quantization/test_mxfp4_w4a8_rdna3.py | 269 ++++++++++++++++++ .../quantization/test_rdna3_compile_guards.py | 17 +- vllm/_custom_ops.py | 23 ++ vllm/config/quantization.py | 3 + .../model_executor/kernels/linear/__init__.py | 51 +++- .../kernels/linear/mxfp4/base.py | 25 +- .../kernels/linear/mxfp4/emulation.py | 16 +- .../kernels/linear/mxfp4/humming.py | 18 ++ .../kernels/linear/mxfp4/rdna3_hip.py | 152 ++++++++++ .../layers/quantization/utils/quant_utils.py | 2 + 15 files changed, 840 insertions(+), 11 deletions(-) create mode 100644 csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu create mode 100644 tests/kernels/quantization/test_mxfp4_w4a8_rdna3.py create mode 100644 vllm/model_executor/kernels/linear/mxfp4/rdna3_hip.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 92a2b37ecf68..bac0426cd21e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1630,7 +1630,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") list(APPEND VLLM_ROCM_EXT_SRC "csrc/rocm/q_gemm_rdna3.cu" "csrc/rocm/q_gemm_rdna3_wmma.cu" - "csrc/rocm/moe_q_gemm_rdna3.cu") + "csrc/rocm/moe_q_gemm_rdna3.cu" + "csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu") endif() define_extension_target( diff --git a/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu new file mode 100644 index 000000000000..8cc03ea5397b --- /dev/null +++ b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu @@ -0,0 +1,254 @@ +// MXFP4 x int8 (W4A8) GEMV for RDNA3/RDNA3.5 (gfx11) small-batch decode. +// +// Activations are quantized to int8 per 32-element group (symmetric, one fp32 +// scale per group), E2M1 weight codes map exactly to int8 (2 * e2m1), and each +// 32-element group is reduced with eight v_dot4_i32_iu8 instructions. One fp32 +// multiply per group applies the E8M0 weight scale and the activation scale. +// The scheme follows ggml's MMVQ / q8_1 path (llama.cpp, MIT License). +// +// Layouts (Quark OCP MX checkpoints, no repack): +// a [M, K] fp16 / bf16, 1 <= M <= 8 +// b_q [N, K/2] uint8, two E2M1 codes per byte, low nibble = even k +// b_scale [N, K/32] uint8 E8M0 +// out [M, N] same dtype as a + +#include +#include +#include + +#include +#include +#include + +#include + +#include "core/cuda_compat.h" + +#if defined(__GFX11__) + #define __HIP__MXFP4_W4A8_GFX11__ +#endif + +namespace { + +constexpr int kMaxM = 8; +constexpr int kGroup = 32; +constexpr int kWarpsPerBlock = 8; + +template +__device__ __forceinline__ float to_float(T v); +template <> +__device__ __forceinline__ float to_float(half v) { + return __half2float(v); +} +template <> +__device__ __forceinline__ float to_float(__hip_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ half from_float(float v) { + return __float2half(v); +} +template <> +__device__ __forceinline__ __hip_bfloat16 from_float(float v) { + return __float2bfloat16(v); +} + +#if defined(__HIP__MXFP4_W4A8_GFX11__) || !defined(__HIP_DEVICE_COMPILE__) + +// 2 * E2M1 value for each 4-bit code (sign in bit 3). Exact in int8. +__device__ __constant__ int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, + 0, -1, -2, -3, -4, -6, -8, -12}; + +template +__global__ void mxfp4_w4a8_quant_act_kernel(const T* __restrict__ a, + int8_t* __restrict__ a_q, + float* __restrict__ a_s, int M, + int K) { + const int groups = K / kGroup; + const int m = blockIdx.x; + const int g = blockIdx.y * blockDim.x + threadIdx.x; + if (m >= M || g >= groups) return; + + const T* row = a + (int64_t)m * K + g * kGroup; + float v[kGroup]; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < kGroup; ++i) { + v[i] = to_float(row[i]); + amax = fmaxf(amax, fabsf(v[i])); + } + const float scale = fmaxf(amax / 127.0f, 1e-12f); + a_s[(int64_t)m * groups + g] = scale; + + int8_t* out = a_q + (int64_t)m * K + g * kGroup; + #pragma unroll + for (int i = 0; i < kGroup; ++i) { + out[i] = (int8_t)fmaxf(-127.f, fminf(127.f, rintf(v[i] / scale))); + } +} + +// One wave (32 lanes) per output column n. Lanes stride over the K groups; +// each lane decodes its weight group once and reuses it for all M rows. +template +__global__ void mxfp4_w4a8_gemv_rdna3_kernel( + const int* __restrict__ a_q, const float* __restrict__ a_s, + const uint8_t* __restrict__ b_q, const uint8_t* __restrict__ b_scale, + T* __restrict__ out, int N, int K) { + const int groups = K / kGroup; + const int lane = threadIdx.x; + const int n = blockIdx.x * blockDim.y + threadIdx.y; + if (n >= N) return; + + const uint8_t* w_row = b_q + (int64_t)n * (K / 2); + const uint8_t* s_row = b_scale + (int64_t)n * groups; + + float acc[M]; + #pragma unroll + for (int m = 0; m < M; ++m) acc[m] = 0.f; + + for (int g = lane; g < groups; g += 32) { + const uint8_t* gp = w_row + g * (kGroup / 2); + int w4[8]; + #pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t two = gp[2 * j] | (gp[2 * j + 1] << 8); + const int w0 = kE2M1x2[two & 0xF]; + const int w1 = kE2M1x2[(two >> 4) & 0xF]; + const int w2 = kE2M1x2[(two >> 8) & 0xF]; + const int w3 = kE2M1x2[(two >> 12) & 0xF]; + w4[j] = (w0 & 0xFF) | ((w1 & 0xFF) << 8) | ((w2 & 0xFF) << 16) | + ((w3 & 0xFF) << 24); + } + // 2^(e - 127) / 2: undoes the factor 2 folded into kE2M1x2. + const float w_scale = __int_as_float((int)s_row[g] << 23) * 0.5f; + #pragma unroll + for (int m = 0; m < M; ++m) { + const int* ap = a_q + (int64_t)m * (K / 4) + g * (kGroup / 4); + int isum = 0; + #pragma unroll + for (int j = 0; j < 8; ++j) { + isum = __builtin_amdgcn_sudot4(true, w4[j], true, ap[j], isum, false); + } + acc[m] += (float)isum * (w_scale * a_s[(int64_t)m * groups + g]); + } + } + + #pragma unroll + for (int m = 0; m < M; ++m) { + float v = acc[m]; + #pragma unroll + for (int o = 16; o > 0; o >>= 1) v += VLLM_SHFL_XOR_SYNC_WIDTH(v, o, 32); + if (lane == 0) out[(int64_t)m * N + n] = from_float(v); + } +} + +#else // non-gfx11 device pass: empty kernels for symbol parity. + +template +__global__ void mxfp4_w4a8_quant_act_kernel(const T*, int8_t*, float*, int, + int) {} + +template +__global__ void mxfp4_w4a8_gemv_rdna3_kernel(const int*, const float*, + const uint8_t*, const uint8_t*, T*, + int, int) {} + +#endif // __HIP__MXFP4_W4A8_GFX11__ || !__HIP_DEVICE_COMPILE__ + +bool mxfp4_w4a8_on_gfx11() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string arch = dprops->gcnArchName; + return arch.rfind("gfx11", 0) == 0; + }(); + return result; +} + +template +void launch(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale, at::Tensor& out, int M, int N, int K) { + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int groups = K / kGroup; + + auto a_q = torch::empty({M, K}, a.options().dtype(torch::kInt8)); + auto a_s = torch::empty({M, groups}, a.options().dtype(torch::kFloat32)); + + const dim3 q_block(256); + const dim3 q_grid(M, (groups + 255) / 256); + mxfp4_w4a8_quant_act_kernel<<>>( + reinterpret_cast(a.data_ptr()), a_q.data_ptr(), + a_s.data_ptr(), M, K); + + const dim3 block(32, kWarpsPerBlock); + const dim3 grid((N + kWarpsPerBlock - 1) / kWarpsPerBlock); + const int* a_q32 = reinterpret_cast(a_q.data_ptr()); + const float* a_s_ptr = a_s.data_ptr(); + const uint8_t* b_q_ptr = b_q.data_ptr(); + const uint8_t* b_s_ptr = b_scale.data_ptr(); + T* out_ptr = reinterpret_cast(out.data_ptr()); + +#define MXFP4_W4A8_LAUNCH(MM) \ + case MM: \ + mxfp4_w4a8_gemv_rdna3_kernel<<>>( \ + a_q32, a_s_ptr, b_q_ptr, b_s_ptr, out_ptr, N, K); \ + break; + + switch (M) { + MXFP4_W4A8_LAUNCH(1) + MXFP4_W4A8_LAUNCH(2) + MXFP4_W4A8_LAUNCH(3) + MXFP4_W4A8_LAUNCH(4) + MXFP4_W4A8_LAUNCH(5) + MXFP4_W4A8_LAUNCH(6) + MXFP4_W4A8_LAUNCH(7) + MXFP4_W4A8_LAUNCH(8) + } +#undef MXFP4_W4A8_LAUNCH +} + +} // namespace + +torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale) { + TORCH_CHECK(a.is_cuda() && b_q.is_cuda() && b_scale.is_cuda(), + "mxfp4_w4a8_gemv_rdna3: all inputs must be on the GPU"); + TORCH_CHECK(mxfp4_w4a8_on_gfx11(), + "mxfp4_w4a8_gemv_rdna3 requires an RDNA3/RDNA3.5 (gfx11) GPU"); + TORCH_CHECK(a.scalar_type() == at::kHalf || a.scalar_type() == at::kBFloat16, + "mxfp4_w4a8_gemv_rdna3: a must be float16 or bfloat16"); + TORCH_CHECK( + b_q.scalar_type() == at::kByte && b_scale.scalar_type() == at::kByte, + "mxfp4_w4a8_gemv_rdna3: b_q and b_scale must be uint8"); + TORCH_CHECK(a.dim() == 2 && b_q.dim() == 2 && b_scale.dim() == 2, + "mxfp4_w4a8_gemv_rdna3: inputs must be 2D"); + TORCH_CHECK( + a.is_contiguous() && b_q.is_contiguous() && b_scale.is_contiguous(), + "mxfp4_w4a8_gemv_rdna3: inputs must be contiguous"); + + const int M = a.size(0); + const int K = a.size(1); + const int N = b_q.size(0); + TORCH_CHECK(M >= 1 && M <= kMaxM, + "mxfp4_w4a8_gemv_rdna3 supports 1 <= M <= ", kMaxM, + ", got M=", M); + TORCH_CHECK(K % kGroup == 0, + "mxfp4_w4a8_gemv_rdna3: K must be a multiple of 32"); + TORCH_CHECK(b_q.size(1) * 2 == K, + "mxfp4_w4a8_gemv_rdna3: b_q must be [N, K/2]"); + TORCH_CHECK(b_scale.size(0) == N && b_scale.size(1) == K / kGroup, + "mxfp4_w4a8_gemv_rdna3: b_scale must be [N, K/32]"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto out = torch::empty({M, N}, a.options()); + + if (a.scalar_type() == at::kHalf) { + launch(a, b_q, b_scale, out, M, N, K); + } else { + launch<__hip_bfloat16>(a, b_q, b_scale, out, M, N, K); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index de963913ea5d..f36cd3d90184 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -24,6 +24,9 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, const at::Tensor& scale_a, const at::Tensor& scale_b, const int64_t CuCount); +torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale); + torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, bool use_v2_format); diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 077c3f35cef1..5c34d6676b96 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -54,6 +54,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { #endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 + // MXFP4 weights x int8 activations GEMV for RDNA3/RDNA3.5 (gfx11) decode, + // 1 <= M <= 8. Weights [N, K/2] uint8 E2M1, scales [N, K/32] uint8 E8M0. + rocm_ops.def( + "mxfp4_w4a8_gemv_rdna3(Tensor a, Tensor b_q, Tensor b_scale) -> Tensor"); + rocm_ops.impl("mxfp4_w4a8_gemv_rdna3", torch::kCUDA, &mxfp4_w4a8_gemv_rdna3); + // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). rocm_ops.def( "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index e70be35ed911..dfe474f637bb 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -99,6 +99,15 @@ vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8 Combine with `--moe-backend` to pin a specific kernel family. +For MXFP4 linear layers on RDNA3/RDNA3.5 (gfx11), `int8_per_group_32` selects +the W4A8 int8-dot GEMV for decode (batches of up to 8 tokens); larger batches +run weight-only. On W4A4 checkpoints this replaces the MXFP4 activation +quantization: + +```bash +vllm serve --quantization-config.linear.activation int8_per_group_32 +``` + ### Re-quantizing layers to a different precision An online target can re-quantize a checkpoint-quantized (prequantized) layer at load time. diff --git a/tests/kernels/quantization/test_mxfp4_w4a8_rdna3.py b/tests/kernels/quantization/test_mxfp4_w4a8_rdna3.py new file mode 100644 index 000000000000..aaf32356ba1f --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_w4a8_rdna3.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the RDNA3/RDNA3.5 (gfx11) MXFP4 W4A8 GEMV and its linear kernel. + +Run `pytest tests/kernels/quantization/test_mxfp4_w4a8_rdna3.py`. +""" + +from unittest.mock import patch + +import pytest +import torch + +import vllm.model_executor.kernels.linear as linear_kernels +from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.kernels.linear import ( + EmulationMxfp4LinearKernel, + MxFp4LinearLayerConfig, + Rdna3W4A8MxFp4LinearKernel, + init_mxfp4_linear_kernel, +) +from vllm.model_executor.kernels.linear.mxfp4.rdna3_hip import ( + MAX_W4A8_BATCH_SIZE, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt8Dynamic32Sym, + kMxfp4Dynamic, + kMxfp4Static, + kMxfp6E3M2Dynamic, +) +from vllm.platforms import current_platform + +_E2M1_MAG = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +E2M1 = torch.tensor(_E2M1_MAG + [-v for v in _E2M1_MAG], dtype=torch.float64) + + +def _on_gfx11() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx11 + + return on_gfx11() + + +def _has_op() -> bool: + return hasattr(torch.ops, "_rocm_C") and hasattr( + torch.ops._rocm_C, "mxfp4_w4a8_gemv_rdna3" + ) + + +gfx11_only = pytest.mark.skipif( + not (_on_gfx11() and _has_op()), + reason="requires gfx11 and _rocm_C.mxfp4_w4a8_gemv_rdna3", +) + + +def _random_mxfp4(N: int, K: int, device: str = "cuda"): + w_q = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device=device) + # E8M0 scales near 2^0 keep the products well inside fp16/bf16 range. + w_s = torch.randint(122, 131, (N, K // 32), dtype=torch.uint8, device=device) + return w_q, w_s + + +def _dequant_ref(w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor: + """[N, K/2] uint8 E2M1 + [N, K/32] uint8 E8M0 -> [N, K] float64.""" + N, k_half = w_q.shape + codes = torch.stack([w_q & 0xF, w_q >> 4], dim=-1).reshape(N, 2 * k_half) + vals = E2M1.to(w_q.device)[codes.long()] + scale = torch.pow(2.0, w_s.double() - 127.0).repeat_interleave(32, dim=1) + return vals * scale + + +def _w4a8_ref(x: torch.Tensor, w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor: + """Reference for the W4A8 math: int8 per-32 activation quant, exact dot.""" + M, K = x.shape + xg = x.double().view(M, K // 32, 32) + a_s = (xg.abs().amax(dim=-1) / 127.0).clamp_min(1e-12) + a_q = torch.round(xg / a_s[..., None]).clamp(-127, 127) + x_qdq = (a_q * a_s[..., None]).view(M, K) + return x_qdq @ _dequant_ref(w_q, w_s).t() + + +@gfx11_only +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("M", list(range(1, MAX_W4A8_BATCH_SIZE + 1))) +@pytest.mark.parametrize("N,K", [(64, 256), (256, 512), (128, 4096), (40, 5120)]) +def test_w4a8_gemv_matches_reference(dtype, M, N, K): + torch.manual_seed(0) + w_q, w_s = _random_mxfp4(N, K) + x = (torch.randn(M, K, device="cuda") * 0.1).to(dtype) + + out = torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s) + assert out.shape == (M, N) and out.dtype == dtype + + ref = _w4a8_ref(x, w_q, w_s) + # Bound by the output rounding of the accumulated magnitude; a quant code + # that rounds the other way on a tie stays well inside this. + accum = x.double().abs() @ _dequant_ref(w_q, w_s).abs().t() + tol = 2 * torch.finfo(dtype).eps * accum + 1e-3 + err = (out.double() - ref).abs() + assert bool((err <= tol).all()), f"max excess {(err - tol).max().item():.3e}" + + +@gfx11_only +@pytest.mark.parametrize("M", [1, 4, 8]) +def test_w4a8_gemv_close_to_bf16_dequant(M): + """The int8 activation rounding stays a small relative error.""" + torch.manual_seed(0) + N, K = 256, 4096 + w_q, w_s = _random_mxfp4(N, K) + x = (torch.randn(M, K, device="cuda") * 0.1).to(torch.bfloat16) + + out = torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s).double() + ref = x.double() @ _dequant_ref(w_q, w_s).t() + rel = (out - ref).norm() / ref.norm() + assert rel < 1e-2, f"relative error {rel.item():.3e}" + + +@gfx11_only +@pytest.mark.parametrize("M", [0, MAX_W4A8_BATCH_SIZE + 1]) +def test_w4a8_gemv_rejects_out_of_range_m(M): + w_q, w_s = _random_mxfp4(64, 256) + x = torch.randn(M, 256, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="1 <= M <= 8"): + torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s) + + +@gfx11_only +def test_w4a8_gemv_opcheck(): + w_q, w_s = _random_mxfp4(64, 256) + x = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) + torch.library.opcheck(torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3, (x, w_q, w_s)) + + +@gfx11_only +def test_w4a8_gemv_cuda_graph(): + w_q, w_s = _random_mxfp4(256, 512) + x = torch.randn(2, 512, device="cuda", dtype=torch.bfloat16) + eager = torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(x, w_q, w_s) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(captured, eager, rtol=0, atol=0) + + +def _make_layer(N: int, K: int): + w_q, w_s = _random_mxfp4(N, K) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(w_q, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_s, requires_grad=False) + return layer + + +@gfx11_only +@pytest.mark.parametrize("M", [1, 8, 9, 64]) +def test_linear_kernel_dispatch(M): + """M <= 8 takes the W4A8 GEMV, larger M the weight-only dequant + GEMM + fallback.""" + pytest.importorskip("quark") + torch.manual_seed(0) + N, K = 256, 512 + kernel = Rdna3W4A8MxFp4LinearKernel( + MxFp4LinearLayerConfig(activation_quant_key=kInt8Dynamic32Sym) + ) + layer = _make_layer(N, K) + kernel.process_weights_after_loading(layer) + x = (torch.randn(M, K, device="cuda") * 0.1).to(torch.bfloat16) + bias = torch.randn(N, device="cuda", dtype=torch.bfloat16) + + with patch( + "vllm._custom_ops.mxfp4_w4a8_gemv_rdna3", + wraps=torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3, + ) as gemv: + out = kernel.apply_weights(layer, x, bias) + uses_w4a8 = M <= MAX_W4A8_BATCH_SIZE + assert gemv.called == uses_w4a8 + + ref = x.double() @ _dequant_ref(layer.weight, layer.weight_scale).t() + ref = ref + bias.double() + rel = (out.double() - ref).norm() / ref.norm() + assert rel < 2e-2, f"relative error {rel.item():.3e}" + + +@gfx11_only +def test_init_selects_w4a8_only_for_int8_activations(monkeypatch): + pytest.importorskip("quark") + monkeypatch.setattr( + linear_kernels, "_user_linear_activation_override", lambda: None + ) + kernel = init_mxfp4_linear_kernel(activation_quant_key=kInt8Dynamic32Sym) + assert isinstance(kernel, Rdna3W4A8MxFp4LinearKernel) + + for key in (None, kMxfp4Dynamic): + kernel = init_mxfp4_linear_kernel(activation_quant_key=key) + assert not isinstance(kernel, Rdna3W4A8MxFp4LinearKernel) + + +@gfx11_only +def test_init_selects_w4a8_from_user_override(monkeypatch): + """--quantization-config.linear.activation int8_per_group_32 on a W4A4 + checkpoint selects the W4A8 kernel.""" + pytest.importorskip("quark") + monkeypatch.setattr( + linear_kernels, "_user_linear_activation_override", lambda: kInt8Dynamic32Sym + ) + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) + assert isinstance(kernel, Rdna3W4A8MxFp4LinearKernel) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("override", [None, kMxfp4Dynamic, kInt8Dynamic32Sym]) +@pytest.mark.parametrize("checkpoint", [None, kMxfp4Dynamic]) +def test_resolve_linear_activation_key(override, checkpoint, monkeypatch): + monkeypatch.setattr( + linear_kernels, "_user_linear_activation_override", lambda: override + ) + resolved = linear_kernels._resolve_linear_activation_key(checkpoint) + assert resolved == (override if override is not None else checkpoint) + + +@pytest.mark.cpu_test +def test_quantization_config_parses_int8_activation(): + args = QuantizationConfigArgs(linear={"activation": "int8_per_group_32"}) + assert args.linear is not None + assert args.linear.weight is None + assert args.linear.activation == kInt8Dynamic32Sym + + +@pytest.mark.cpu_test +def test_supports_quant_scheme_dispatches_on_activation(): + assert Rdna3W4A8MxFp4LinearKernel._supports_quant_scheme( + kMxfp4Static, kInt8Dynamic32Sym + ) + for key in (None, kMxfp4Dynamic, kMxfp6E3M2Dynamic): + assert not Rdna3W4A8MxFp4LinearKernel._supports_quant_scheme(kMxfp4Static, key) + # Kernels without an int8 path keep their previous activation formats. + assert not EmulationMxfp4LinearKernel._supports_quant_scheme( + kMxfp4Static, kInt8Dynamic32Sym + ) + for key in (None, kMxfp4Dynamic, kMxfp6E3M2Dynamic): + assert EmulationMxfp4LinearKernel._supports_quant_scheme(kMxfp4Static, key) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("key", [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic]) +def test_can_implement_rejects_non_int8_activation(key): + config = MxFp4LinearLayerConfig(activation_quant_key=key) + can_implement, reason = Rdna3W4A8MxFp4LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.cpu_test +def test_is_supported_requires_rocm_gfx11(): + with patch( + "vllm.model_executor.kernels.linear.mxfp4.rdna3_hip.current_platform.is_rocm", + return_value=False, + ): + is_supported, reason = Rdna3W4A8MxFp4LinearKernel.is_supported() + assert not is_supported + assert reason diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py index 194441f0f909..d78b014ed3de 100644 --- a/tests/kernels/quantization/test_rdna3_compile_guards.py +++ b/tests/kernels/quantization/test_rdna3_compile_guards.py @@ -53,11 +53,17 @@ reason="This test verifies non-gfx1100 builds — skip on gfx1100", ) -RDNA3_OPS = ["gptq_gemm_rdna3", "gptq_gemm_rdna3_wmma", "moe_gptq_gemm_rdna3"] +RDNA3_OPS = [ + "gptq_gemm_rdna3", + "gptq_gemm_rdna3_wmma", + "moe_gptq_gemm_rdna3", + "mxfp4_w4a8_gemv_rdna3", +] RDNA3_CU_FILES = [ "q_gemm_rdna3.cu", "q_gemm_rdna3_wmma.cu", "moe_q_gemm_rdna3.cu", + "mxfp4_w4a8_gemv_rdna3.cu", ] @@ -334,11 +340,16 @@ def test_register_fake_guarded_by_hasattr(self): if match is None: continue preceding = src[: match.start()] - last_hasattr = preceding.rfind(f'hasattr(torch.ops._rocm_C, "{op}")') - assert last_hasattr != -1, ( + guards = list( + re.finditer( + rf'hasattr\(\s*torch\.ops\._rocm_C,\s*"{op}"\s*\)', preceding + ) + ) + assert guards, ( f'register_fake("_rocm_C::{op}") is not preceded by a ' f"hasattr check — would crash on CDNA import" ) + last_hasattr = guards[-1].start() gap = preceding[last_hasattr:].count("\n") assert gap <= 5, ( f"hasattr guard for {op} is {gap} lines before " diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 850917f0efa9..e3459466f1ed 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2143,6 +2143,29 @@ def _wvSplitK_int4_g_fake( ) +def mxfp4_w4a8_gemv_rdna3( + a: torch.Tensor, + b_q: torch.Tensor, + b_scale: torch.Tensor, +) -> torch.Tensor: + # a: [M, K] fp16/bf16 with 1 <= M <= 8; b_q: [N, K/2] uint8 E2M1; + # b_scale: [N, K/32] uint8 E8M0. Returns [M, N] in a.dtype. + return torch.ops._rocm_C.mxfp4_w4a8_gemv_rdna3(a, b_q, b_scale) + + +if hasattr(torch.ops, "_rocm_C") and hasattr( + torch.ops._rocm_C, "mxfp4_w4a8_gemv_rdna3" +): + + @register_fake("_rocm_C::mxfp4_w4a8_gemv_rdna3") + def _mxfp4_w4a8_gemv_rdna3_fake( + a: torch.Tensor, + b_q: torch.Tensor, + b_scale: torch.Tensor, + ) -> torch.Tensor: + return torch.empty((a.size(0), b_q.size(0)), dtype=a.dtype, device=a.device) + + def wvSplitKrc( a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None ) -> torch.Tensor: diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index face2c3a7ba7..0991d0f5d002 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -23,6 +23,7 @@ kFp8StaticChannelSym, kFp8StaticTensorSym, kInt4Static32, + kInt8Dynamic32Sym, kInt8StaticChannelSym, kMxfp4Dynamic, kMxfp4Static, @@ -43,6 +44,8 @@ "mxfp4": kMxfp4Dynamic, "nvfp4_per_token": kNvfp4DynamicToken, "int8_per_channel_static": kInt8StaticChannelSym, + # Dynamic per-32-element int8 activations (RDNA3 MXFP4 W4A8 GEMV). + "int8_per_group_32": kInt8Dynamic32Sym, # Load-time MXFP4-to-int4 for Kimi-K3 on gfx942. "int4_per_group_32": kInt4Static32, } diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 26f3cceaa90b..8befe5006966 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -91,6 +91,9 @@ from vllm.model_executor.kernels.linear.mxfp4.pytorch import ( TorchMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.rdna3_hip import ( + Rdna3W4A8MxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.xpu import ( XPUMxFp4LinearKernel, ) @@ -241,7 +244,10 @@ prefers_humming, prioritize_humming, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) from vllm.platforms import PlatformEnum, current_platform logger = init_logger(__name__) @@ -625,6 +631,7 @@ def _resolve_backend_kernels( EmulationMxfp4LinearKernel, ], PlatformEnum.ROCM: [ + Rdna3W4A8MxFp4LinearKernel, AiterMxfp4LinearKernel, EmulationMxfp4LinearKernel, ], @@ -1003,11 +1010,45 @@ def init_mxfp8_linear_kernel( ) +def _user_linear_activation_override() -> QuantKey | None: + """User's linear activation override from quantization_config, or None.""" + from vllm.config import get_current_vllm_config_or_none + from vllm.config.quantization import QuantizationConfigArgs + + config = get_current_vllm_config_or_none() + if config is None or config.model_config is None: + return None + args = config.model_config.quantization_config + if not isinstance(args, QuantizationConfigArgs) or args.linear is None: + return None + return args.linear.activation + + +def _resolve_linear_activation_key( + model_activation_key: QuantKey | None, +) -> QuantKey | None: + """Apply ``--quantization-config.linear.activation`` to the checkpoint's + activation key. The override wins: it may replace a checkpoint's MXFP4 + activations with a format a kernel computes directly (e.g. int8).""" + user_override = _user_linear_activation_override() + if user_override is None or user_override == model_activation_key: + return model_activation_key + if model_activation_key is not None: + logger.warning_once( + "Checkpoint declares linear activation=%s; using " + "quantization_config.linear.activation=%s instead.", + model_activation_key, + user_override, + ) + return user_override + + def init_mxfp4_linear_kernel( activation_quant_key: QuantKey | None = None, ) -> MxFp4LinearKernel: """Select and instantiate the best MXFP4 linear kernel for the current platform.""" + activation_quant_key = _resolve_linear_activation_key(activation_quant_key) config = MxFp4LinearLayerConfig( activation_quant_key=activation_quant_key, ) @@ -1035,6 +1076,13 @@ def init_mxfp4_linear_kernel( failure_reasons.append(f"{kernel_cls.__name__}: {reason}") continue + if not kernel_cls._supports_quant_scheme(kMxfp4Static, activation_quant_key): + failure_reasons.append( + f"{kernel_cls.__name__}: quantization scheme " + f"{kMxfp4Static}x{activation_quant_key} not supported" + ) + continue + can_implement, reason = kernel_cls.can_implement(config) if not can_implement: failure_reasons.append(f"{kernel_cls.__name__}: {reason}") @@ -1373,6 +1421,7 @@ def register_linear_kernel( "EmulationMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", "MarlinMxFp4LinearKernel", + "Rdna3W4A8MxFp4LinearKernel", "FlashInferCutedslMxfp8LinearKernel", "FlashInferCutlassMxfp8LinearKernel", "FlashInferTrtllmMxfp8LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/base.py b/vllm/model_executor/kernels/linear/mxfp4/base.py index de0a2a216327..7fe129824b6b 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/base.py +++ b/vllm/model_executor/kernels/linear/mxfp4/base.py @@ -6,7 +6,11 @@ import torch -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp4Static, +) @dataclass @@ -30,8 +34,8 @@ class MxFp4LinearKernel(ABC): 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. + priority order, calling ``is_supported``, ``_supports_quant_scheme`` and + ``can_implement`` to find the best match for the current hardware. """ def __init__(self, config: MxFp4LinearLayerConfig) -> None: @@ -47,6 +51,21 @@ def is_supported( """Return whether this kernel can run on the current platform.""" raise NotImplementedError + @classmethod + def _supports_quant_scheme( + cls, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Return whether this kernel implements *weight_key* x *activation_key*. + + Selection uses this to dispatch on the requested activation format + (from the checkpoint or ``--quantization-config.linear.activation``). + The default covers MXFP4 weights with MXFP4 or unquantized activations; + kernels for other activation formats override it. + """ + return weight_key == kMxfp4Static and activation_key in (None, kMxfp4Dynamic) + @classmethod @abstractmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: diff --git a/vllm/model_executor/kernels/linear/mxfp4/emulation.py b/vllm/model_executor/kernels/linear/mxfp4/emulation.py index e3e64bf0a3be..38c519ba9f92 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/emulation.py +++ b/vllm/model_executor/kernels/linear/mxfp4/emulation.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Dynamic, + kMxfp4Static, kMxfp6E2M3Dynamic, kMxfp6E3M2Dynamic, ) @@ -46,9 +47,8 @@ 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 - ) + qdq_fn: Callable[[torch.Tensor], torch.Tensor] = lambda x: x + self.quant_dequant_func = qdq_fn else: self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[ config.activation_quant_key @@ -60,6 +60,16 @@ def is_supported( ) -> tuple[bool, str | None]: return True, None + @classmethod + def _supports_quant_scheme( + cls, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return weight_key == kMxfp4Static and ( + activation_key is None or activation_key in _ACTIVATION_QUANT_DEQUANT_FUNCS + ) + @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: if config.activation_quant_key not in ( diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index 0bfb06077036..b38bb9d4d756 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -10,6 +10,10 @@ prepare_humming_linear_layer_config, quant_key_to_input_schema, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) from vllm.platforms import current_platform from vllm.utils.humming import create_humming_locks from vllm.utils.import_utils import has_humming @@ -35,6 +39,20 @@ def is_supported( return True, None + @classmethod + def _supports_quant_scheme( + cls, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + if weight_key != kMxfp4Static: + return False + try: + quant_key_to_input_schema(activation_key) + except ValueError: + return False + return True + @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: try: diff --git a/vllm/model_executor/kernels/linear/mxfp4/rdna3_hip.py b/vllm/model_executor/kernels/linear/mxfp4/rdna3_hip.py new file mode 100644 index 000000000000..407f270d46f0 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/rdna3_hip.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP4 linear kernel for RDNA3/RDNA3.5 (gfx11) with an int8-activation decode +path. + +Routes on the number of activation rows M: + 1 <= M <= MAX_W4A8_BATCH_SIZE (decode, speculative verify): + HIP W4A8 GEMV (``_rocm_C.mxfp4_w4a8_gemv_rdna3``). Activations are quantized + to int8 per 32-element group and reduced with int8 dot instructions. + otherwise (prefill, profiling): + dequantize the weight and run a high-precision GEMM on the unquantized + activations (weight-only). + +Both paths read the checkpoint layout directly ([N, K/2] uint8 E2M1 codes, +[N, K/32] uint8 E8M0 scales), so there is no weight repack. + +Selected when int8 activations are requested, e.g. with +``--quantization-config.linear.activation int8_per_group_32``. For W4A4 +checkpoints this replaces the MXFP4 activation QDQ: decode uses int8 +activations and prefill uses the model dtype, which kept accuracy where mixing +MXFP4-QDQ prefill with int8 decode did not. +""" + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt8Dynamic32Sym, + kMxfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_quark +from vllm.utils.torch_utils import direct_register_custom_op + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +MAX_W4A8_BATCH_SIZE = 8 +_MXFP4_GROUP_SIZE = 32 + + +def _rdna3_mxfp4_w4a8_apply_impl( + x_2d: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + import vllm._custom_ops as ops + + M = x_2d.shape[0] + if 0 < M <= MAX_W4A8_BATCH_SIZE and x_2d.dtype in (torch.float16, torch.bfloat16): + out = ops.mxfp4_w4a8_gemv_rdna3(x_2d, weight, weight_scale) + if bias is not None: + out.add_(bias) + return out + + return F.linear(x_2d, dequant_mxfp4(weight, weight_scale, x_2d.dtype), bias) + + +def _rdna3_mxfp4_w4a8_apply_fake( + x_2d: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + return torch.empty( + (x_2d.shape[0], weight.shape[0]), dtype=x_2d.dtype, device=x_2d.device + ) + + +direct_register_custom_op( + op_name="rdna3_mxfp4_w4a8_apply", + op_func=_rdna3_mxfp4_w4a8_apply_impl, + mutates_args=[], + fake_impl=_rdna3_mxfp4_w4a8_apply_fake, +) + + +class Rdna3W4A8MxFp4LinearKernel(MxFp4LinearKernel): + """MXFP4 GEMM for RDNA3/RDNA3.5: W4A8 int8-dot GEMV for M <= 8, dequant + + high-precision GEMM above.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "only targets ROCm" + from vllm.platforms.rocm import on_gfx11 + + if not on_gfx11(): + return False, "requires RDNA3/RDNA3.5 (gfx11)" + if not ( + hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "mxfp4_w4a8_gemv_rdna3") + ): + return False, "_rocm_C.mxfp4_w4a8_gemv_rdna3 is not built" + return True, None + + @classmethod + def _supports_quant_scheme( + cls, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return weight_key == kMxfp4Static and activation_key == kInt8Dynamic32Sym + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kInt8Dynamic32Sym: + return False, "only supports dynamic per-32-group int8 activations" + if not has_quark(): + return False, "amd-quark package not available (needed for M > 8)" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight.data + weight_scale = layer.weight_scale.data + if weight.dtype != torch.uint8 or weight_scale.dtype != torch.uint8: + raise ValueError( + "Rdna3W4A8MxFp4LinearKernel expects uint8 packed weights and " + f"uint8 E8M0 scales, got {weight.dtype} and {weight_scale.dtype}" + ) + N, k_half = weight.shape + K = 2 * k_half + if K % _MXFP4_GROUP_SIZE != 0 or weight_scale.shape != ( + N, + K // _MXFP4_GROUP_SIZE, + ): + raise ValueError( + "Rdna3W4A8MxFp4LinearKernel expects weight [N, K/2] and " + f"weight_scale [N, K/32] with K % 32 == 0, got " + f"{tuple(weight.shape)} and {tuple(weight_scale.shape)}" + ) + layer.weight = Parameter(weight.contiguous(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale.contiguous(), requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + x_2d = x.reshape(-1, x.shape[-1]) + if not x_2d.is_contiguous(): + x_2d = x_2d.contiguous() + out = torch.ops.vllm.rdna3_mxfp4_w4a8_apply( + x_2d, layer.weight, layer.weight_scale, bias + ) + return out.reshape(*x.shape[:-1], out.shape[-1]) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 03d4df386ab2..8cdbb3e596ae 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -297,6 +297,8 @@ def __str__(self): kInt8DynamicTokenAsym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=False) kInt8StaticTensorAsym = QuantKey(torch.int8, kStaticTensorScale, symmetric=False) kInt8DynamicTensorAsym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=False) +kDynamic32Scale = ScaleDesc(torch.float32, False, GroupShape(1, 32)) +kInt8Dynamic32Sym = QuantKey(torch.int8, kDynamic32Scale, symmetric=True) # INT4 W4A8 quantization keys From 7822663cb48093a949835afc7af52ee13b0946ce Mon Sep 17 00:00:00 2001 From: frida-andersson Date: Sat, 24 Oct 2026 07:12:00 -0500 Subject: [PATCH 018/631] [DeepSeek][ROCm] V4.1-Flash: enable --monokernel on MI325X (#62713) Signed-off-by: Frida Andersson Co-authored-by: vllm-agent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../deepseek-ai/DeepSeek-V4.1-Flash.yaml | 41 +++++++++++++++---- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml index 2181f03f599e..7092c63910e7 100644 --- a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml +++ b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml @@ -42,9 +42,10 @@ model: default_frontend: rust docker_image: # ROCm 10.0 nightly channel. MI355X's --monokernel (vllm-project/vllm#60397, - # #61082), paged MXFP4 indexer (#58671) and Gluon sparse MLA (#53492) are on - # main but not in 0.31.0, so AMD stays on the nightly image the AMD rows were - # verified on; the MI355X pip tab is hidden in the variants below. + # #61082), paged MXFP4 indexer (#58671) and Gluon sparse MLA (#53492), and + # MI325X's gfx942 monokernel (#60731), are on main but not in 0.31.0, so AMD + # stays on the nightly image the AMD rows were verified on; the MI325X and + # MI355X pip tabs are hidden in the variants below. amd: "vllm/vllm-openai-rocm:nightly-rocm100" architecture: moe parameter_count: "552B" @@ -147,8 +148,12 @@ variants: extra_env: VLLM_USE_V2_MODEL_RUNNER: "1" PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" - # --monokernel and the ROCm MXFP4 indexer path are on main, not in the - # 0.31.0 wheel, which rejects the flag. Drop this once a release ships them. + # --monokernel is on main, not in the 0.31.0 wheel, which rejects the + # flag. MI355X also needs the ROCm MXFP4 indexer. Drop these once a + # release ships them. + mi325x: + install: + pip: false mi355x: install: pip: false @@ -278,6 +283,8 @@ strategy_overrides: - '{"cpu_offload":true}' - "--compilation-config" - '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + # gfx942 mono decode (vllm-project/vllm#60731): TP4, up to 42 rows. + - "--monokernel" extra_env: VLLM_USE_BREAKABLE_CUDAGRAPH: "0" mi355x: @@ -695,10 +702,10 @@ guide: | and the ROCm Engram host-memory gate ([vllm-project/vllm#57491](https://github.com/vllm-project/vllm/pull/57491)). The launch-day `deepseekv41-flash-0909` image is no longer needed. - On AMD the Docker tab uses `vllm/vllm-openai-rocm:nightly-rocm100`, the ROCm 10.0 - nightly the AMD rows were verified on. **MI355X needs it**: the TP2 command uses - `--monokernel` and the ROCm MXFP4 indexer, which are on `main` but not in 0.31.0, - so the Install block offers only Docker on MI355X. MI300X, MI325X and MI350X also - run from the ROCm pip wheel. + nightly the AMD rows were verified on. **MI355X and MI325X need it**: their + commands pass `--monokernel`, and MI355X also passes the ROCm MXFP4 indexer. + Both are on `main` but not in 0.31.0, so the Install block offers only Docker + on those two SKUs. MI300X and MI350X also run from the ROCm pip wheel. - The Rust OpenAI frontend is selected by default in the command builder. Switch to Python if you encounter unsupported features or compatibility issues. - Expect a long first load: `VLLM_ENGINE_READY_TIMEOUT_S=3600` is set for that reason. @@ -877,6 +884,11 @@ guide: | (`--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`), with prefill eager. + This recipe does not pass `--monokernel`. With the kernels' weight copies + included, #60731 measured 170.11 GiB of weights a rank. At + `--gpu-memory-utilization 0.9` a 192 GiB MI300X leaves about 2.7 GiB beside + those weights. + The offload resolves on gfx942 only on builds carrying [vllm-project/vllm#57491](https://github.com/vllm-project/vllm/pull/57491), the same gate MI325X and MI355X need. Those builds also default `cpu_offload` @@ -901,6 +913,17 @@ guide: | (`--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`); prefill runs eagerly. + `--monokernel` selects the gfx942 mono decode kernels from + [vllm-project/vllm#60731](https://github.com/vllm-project/vllm/pull/60731). + They run at TP4 only, on decode steps of up to 42 rows (7 requests with + five DSpark drafts). Larger steps, prefill, and mixed batches take the + regular path, and TP2 stays on vLLM's path. With their weight copies + included, #60731 measured 170.11 GiB of weights a rank, up from 98.88 GiB. + At `--gpu-memory-utilization 0.9` the KV pool still held 11.18 million + tokens. Use a ROCm nightly from 2026-10-20 or later; + `nightly-rocm100-df3bd094aeea829f27baa0a5282830464f601d9f` is the first. + Earlier builds have the flag for gfx950 only and refuse to start on MI325X. + The offload resolves on gfx942 only on builds carrying [vllm-project/vllm#57491](https://github.com/vllm-project/vllm/pull/57491), the same gate MI355X needs. Those builds also default `cpu_offload` on through From c0f9fff649e732253c5472d3a4d28389605b4a01 Mon Sep 17 00:00:00 2001 From: CharlesWu Date: Sat, 24 Oct 2026 20:25:43 +0800 Subject: [PATCH 019/631] [UMBP][3/10] Add shared lookup, transfer, hybrid, and observability (#61791) Signed-off-by: CharlesWu Signed-off-by: tjtanaa Co-authored-by: Yichao Zhu Co-authored-by: tjtanaa --- tests/v1/kv_connector/umbp_test_utils.py | 113 +- tests/v1/kv_connector/unit/test_umbp_core.py | 879 +++++++------- .../kv_connector/v1/umbp/__init__.py | 5 +- .../kv_connector/v1/umbp/connector.py | 164 ++- .../kv_transfer/kv_connector/v1/umbp/data.py | 154 ++- .../kv_connector/v1/umbp/runtime/base.py | 20 +- .../kv_connector/v1/umbp/scheduler.py | 1056 +++++++++++++++-- .../kv_transfer/kv_connector/v1/umbp/stats.py | 126 ++ .../kv_connector/v1/umbp/worker.py | 194 ++- 9 files changed, 2152 insertions(+), 559 deletions(-) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/stats.py diff --git a/tests/v1/kv_connector/umbp_test_utils.py b/tests/v1/kv_connector/umbp_test_utils.py index baa0ce9eddaa..4f7ca892e68f 100644 --- a/tests/v1/kv_connector/umbp_test_utils.py +++ b/tests/v1/kv_connector/umbp_test_utils.py @@ -29,6 +29,7 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MambaSpec, ) @@ -178,6 +179,7 @@ def poll(self, job: TransferJobState) -> TransferJobState | None: if job.status in ( TransferJobStatus.COMPLETED, TransferJobStatus.FAILED, + TransferJobStatus.CANCELLED, ): return job return None @@ -187,6 +189,14 @@ def publish(self, job: TransferJobState) -> None: raise RuntimeError("cannot publish an incomplete embedded job") self._store.publish(id(job)) + def cancel(self, job: TransferJobState) -> TransferJobState: + if job.status.value not in ("completed", "failed"): + job.cancel("preempted") + return job + + def take_evicted_keys(self) -> Sequence[str]: + return () + def close(self) -> None: return @@ -254,7 +264,29 @@ def _kv_cache_config() -> KVCacheConfig: ) -def _vllm_config(extra: dict, **parallel_overrides) -> SimpleNamespace: +def _hybrid_kv_cache_config() -> KVCacheConfig: + full = FullAttentionSpec( + block_size=16, num_kv_heads=2, head_size=8, dtype=torch.float16 + ) + mamba = MambaSpec( + block_size=16, + shapes=((4,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + return KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], full), + KVCacheGroupSpec(["mamba"], mamba), + ], + ) + + +def _vllm_config( + extra: dict, prefix_match_unit: int | None = None, **parallel_overrides +) -> SimpleNamespace: parallel = { "tensor_parallel_size": 1, "pipeline_parallel_size": 1, @@ -269,7 +301,7 @@ def _vllm_config(extra: dict, **parallel_overrides) -> SimpleNamespace: cache_config=SimpleNamespace( block_size=16, enable_prefix_caching=True, - prefix_match_unit=None, + prefix_match_unit=prefix_match_unit, ), model_config=SimpleNamespace( model="test-model", @@ -285,6 +317,69 @@ def _vllm_config(extra: dict, **parallel_overrides) -> SimpleNamespace: ) +def make_request(request_id: str = "req", **overrides) -> SimpleNamespace: + """Build the common scheduler-side request shape used by UMBP tests.""" + num_tokens = overrides.pop("num_tokens", 32) + num_prompt_tokens = overrides.pop("num_prompt_tokens", num_tokens) + values = { + "request_id": request_id, + "req_id": request_id, + "num_tokens": num_tokens, + "num_prompt_tokens": num_prompt_tokens, + "num_computed_tokens": 0, + "num_in_flight_tokens": 0, + "block_hashes": [b"a", b"b"], + "block_ids": ([1, 2],), + "prompt_token_ids": list(range(num_prompt_tokens)), + "all_token_ids": list(range(num_tokens)), + } + values.update(overrides) + return SimpleNamespace(**values) + + +def make_scheduler_output( + *, + new_reqs=(), + cached_reqs=None, + scheduled_tokens=None, + finished=(), + preempted=(), + block_state=None, +) -> SimpleNamespace: + """Build a scheduler output with empty defaults for unrelated fields.""" + return SimpleNamespace( + finished_req_ids=set(finished), + preempted_req_ids=set(preempted), + scheduled_new_reqs=list(new_reqs), + scheduled_cached_reqs=( + SimpleNamespace(req_ids=[]) if cached_reqs is None else cached_reqs + ), + num_scheduled_tokens=dict(scheduled_tokens or {}), + kv_connector_block_state=block_state, + ) + + +def make_kv_caches(*, fill: bool = False) -> dict[str, torch.Tensor]: + """Build the standard two-layer CPU KV cache used by connector tests.""" + caches = { + name: torch.empty_strided( + (8, 2, 16, 8), + (512, 256, 8, 1), + dtype=torch.float16, + ) + for name in ("layer1", "layer2") + } + for index, cache in enumerate(caches.values()): + if fill: + cache.copy_( + torch.arange(cache.numel(), dtype=torch.float16).reshape(cache.shape) + + index + ) + else: + cache.zero_() + return caches + + class _SchedulerHandle: def __init__(self, hits): self.hits = hits @@ -329,6 +424,7 @@ def poll(self, job): if job.status in ( TransferJobStatus.COMPLETED, TransferJobStatus.FAILED, + TransferJobStatus.CANCELLED, ): return job return None @@ -392,8 +488,19 @@ def wait(self, job): return job -class _DelayedStoreWorkerHandle(_WorkerHandle): +class _CancellableWorkerHandle(_WorkerHandle): def __init__(self): + self.cancelled = [] + + def cancel(self, job): + self.cancelled.append(job) + job.cancel("preempted") + return job + + +class _DelayedStoreWorkerHandle(_CancellableWorkerHandle): + def __init__(self): + super().__init__() self.jobs = [] self.waited = [] self.publications = [] diff --git a/tests/v1/kv_connector/unit/test_umbp_core.py b/tests/v1/kv_connector/unit/test_umbp_core.py index dfe4c830136d..310208cd76cc 100644 --- a/tests/v1/kv_connector/unit/test_umbp_core.py +++ b/tests/v1/kv_connector/unit/test_umbp_core.py @@ -10,6 +10,7 @@ from tests.v1.kv_connector.umbp_test_utils import ( MemoryRuntime, + _CancellableWorkerHandle, _DelayedLoadWorkerHandle, _DelayedStoreWorkerHandle, _EmbeddedRuntime, @@ -21,6 +22,9 @@ _WorkerHandle, handshake, install_memory_runtime, + make_kv_caches, + make_request, + make_scheduler_output, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.umbp.connector import ( @@ -28,6 +32,8 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.umbp.data import ( BlockIdentityCodec, + BlockLoadBatch, + BlockStoreBatch, BlockTransferPlan, KVLayoutDescriptor, KVLayoutPlanner, @@ -64,6 +70,7 @@ KVQuantMode, UniformTypeKVCacheSpecs, ) +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder mori_available = importlib.util.find_spec("mori") is not None @@ -508,6 +515,10 @@ def test_materialization_preserves_plan_metadata_after_reregistration(): 2, request_id="req", group_id=0, + block_hash=b"hash", + parent_block_hash=b"parent", + token_ids=(1, 2), + block_size=16, ) previous = None for _ in range(2): @@ -561,12 +572,21 @@ def test_load_spec_reports_only_external_tokens(): def test_worker_metadata_aggregates_store_events(): - metadata = UMBPConnectorWorkerMetadata(store_events={7: StoreEventResult(1)}) - other = UMBPConnectorWorkerMetadata(store_events={7: StoreEventResult(1)}) + metadata = UMBPConnectorWorkerMetadata( + store_events={7: StoreEventResult(1, {(0, b"a")})} + ) + other = UMBPConnectorWorkerMetadata( + store_events={7: StoreEventResult(1, {(1, b"b")})} + ) metadata.aggregate(other) - assert metadata.store_events == {7: StoreEventResult(2)} - assert other.store_events == {7: StoreEventResult(1)} + assert metadata.store_events == {7: StoreEventResult(2, {(0, b"a"), (1, b"b")})} + assert other.store_events == {7: StoreEventResult(1, {(1, b"b")})} + + restored = MsgpackDecoder(UMBPConnectorWorkerMetadata).decode( + MsgpackEncoder().encode(metadata) + ) + assert restored == metadata @pytest.mark.parametrize("mode", ["embedded"]) @@ -605,11 +625,7 @@ def test_scheduler_tracks_consecutive_load_plan(): handle, BlockIdentityCodec(UMBPNamespace("test")), ) - request = SimpleNamespace( - request_id="req", - num_tokens=33, - block_hashes=[b"a", b"b"], - ) + request = make_request("req", num_tokens=33) assert scheduler.get_num_new_matched_tokens(request, 0) == (16, False) assert handle.queries[0] == [ @@ -633,9 +649,7 @@ def lookup(self, keys): handle, BlockIdentityCodec(UMBPNamespace("lookup-retry")), ) - request = SimpleNamespace( - request_id="retry", num_tokens=32, block_hashes=[b"a", b"b"] - ) + request = make_request("retry") try: assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) handle.hits = [True] @@ -656,9 +670,7 @@ def test_scheduler_single_group_lookup_drops_last_block_for_eagle(use_eagle, exp _SchedulerHandle([True, True]), BlockIdentityCodec(UMBPNamespace("eagle")), ) - request = SimpleNamespace( - request_id="req", num_tokens=33, block_hashes=[b"a", b"b"] - ) + request = make_request("req", num_tokens=33) assert scheduler.get_num_new_matched_tokens(request, 0) == (expected, False) @@ -671,11 +683,7 @@ def test_scheduler_emits_async_load_without_scheduled_model_tokens(): _SchedulerHandle([True]), BlockIdentityCodec(UMBPNamespace("async-load")), ) - request = SimpleNamespace( - request_id="async-load", - num_tokens=32, - block_hashes=[b"a", b"b"], - ) + request = make_request("async-load") assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True) scheduler.update_state_after_alloc( @@ -683,38 +691,49 @@ def test_scheduler_emits_async_load_without_scheduled_model_tokens(): SimpleNamespace(get_block_ids=lambda group_ids: ([7],)), 16, ) - metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={}, - ) - ) + metadata = scheduler.build_connector_meta(make_scheduler_output()) assert [plan.block_id for plan in metadata.load_requests["async-load"]] == [7] assert scheduler._pending_loads == {} +def test_load_batch_metadata_roundtrip_preserves_destinations_and_identity(): + batch = BlockLoadBatch(["g1", "g0"], [7, 3], [1, 0], "r") + meta = UMBPConnectorMetadata(load_requests={"r": batch}) + restored = ( + MsgpackDecoder(UMBPConnectorMetadata) + .decode(MsgpackEncoder().encode(meta)) + .load_requests["r"] + ) + assert list(restored) == [ + BlockTransferPlan("g1", 7, group_id=1, request_id="r"), + BlockTransferPlan("g0", 3, group_id=0, request_id="r"), + ] + assert restored[:1] == [restored[0]] + with pytest.raises(ValueError, match="equal lengths"): + BlockLoadBatch(["missing-destination"], [], [0], "r") + + def test_scheduler_cached_decode_uses_save_watermark_and_new_blocks(): + vllm_config = _vllm_config( + { + "mode": "embedded", + "load_async": False, + "save_decode_cache": True, + } + ) + vllm_config.kv_events_config = SimpleNamespace(enable_kv_cache_events=True) scheduler = UMBPStoreConnectorScheduler( - _vllm_config( - { - "mode": "embedded", - "load_async": False, - "save_decode_cache": True, - } - ), + vllm_config, _kv_cache_config(), _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("decode")), ) - request = SimpleNamespace( - request_id="decode", - req_id="decode", + request = make_request( + "decode", num_tokens=48, num_prompt_tokens=32, + all_token_ids=list(range(48)), block_hashes=[b"a", b"b", b"c"], block_ids=([1, 2, 3],), num_computed_tokens=32, @@ -726,23 +745,22 @@ def test_scheduler_cached_decode_uses_save_watermark_and_new_blocks(): ) def output(num_computed_tokens, new_block_ids): - return SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + return make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["decode"], new_block_ids=[new_block_ids], num_computed_tokens=[num_computed_tokens], ), - num_scheduled_tokens={"decode": 1}, + scheduled_tokens={"decode": 1}, ) first = scheduler.build_connector_meta(output(32, ([3],))) - assert [plan.block_id for plan in first.store_plans] == [1, 2] + assert [plan.block_id for plan in first.iter_store_plans()] == [1, 2] second = scheduler.build_connector_meta(output(48, ())) - assert [plan.block_id for plan in second.store_plans] == [3] + second_plans = list(second.iter_store_plans()) + assert [plan.block_id for plan in second_plans] == [3] + assert second_plans[0].token_ids == tuple(range(32, 48)) def test_scheduler_decode_store_skips_unverified_speculative_tokens(): @@ -759,9 +777,8 @@ def test_scheduler_decode_store_skips_unverified_speculative_tokens(): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("spec")), ) - request = SimpleNamespace( - request_id="spec", - req_id="spec", + request = make_request( + "spec", num_tokens=47, num_prompt_tokens=32, block_hashes=[b"a", b"b"], @@ -775,19 +792,16 @@ def test_scheduler_decode_store_skips_unverified_speculative_tokens(): ) # One verified token plus three draft tokens: 46 + 4 = 50 > 47 tokens. metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["spec"], new_block_ids=[()], num_computed_tokens=[46], ), - num_scheduled_tokens={"spec": 4}, + scheduled_tokens={"spec": 4}, ) ) - assert [plan.block_id for plan in metadata.store_plans] == [1, 2] + assert [plan.block_id for plan in metadata.iter_store_plans()] == [1, 2] def test_scheduler_skips_decode_blocks_without_save_decode_cache(): @@ -797,9 +811,8 @@ def test_scheduler_skips_decode_blocks_without_save_decode_cache(): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("no-decode")), ) - request = SimpleNamespace( - request_id="r", - req_id="r", + request = make_request( + "r", num_tokens=48, num_prompt_tokens=32, block_hashes=[b"a", b"b", b"c"], @@ -810,14 +823,11 @@ def test_scheduler_skips_decode_blocks_without_save_decode_cache(): scheduler._request_trackers["r"].saved_tokens = 32 # A decode step under synchronous scheduling: num_tokens - 1 are computed. metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["r"], new_block_ids=[None], num_computed_tokens=[47] ), - num_scheduled_tokens={"r": 1}, + scheduled_tokens={"r": 1}, ) ) @@ -834,35 +844,54 @@ def test_kv_consumer_never_stores(kv_role): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("role")), ) - request = SimpleNamespace( - request_id="r", - req_id="r", - num_tokens=32, - block_hashes=[b"a", b"b"], + request = make_request( + "r", block_ids=([3, 4],), - num_computed_tokens=0, ) scheduler.update_state_after_alloc( request, SimpleNamespace(get_block_ids=lambda group_ids: ([3, 4],)), 0, ) + scheduled = SimpleNamespace(req_id="r", num_computed_tokens=0, block_ids=([3, 4],)) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[ - SimpleNamespace(req_id="r", num_computed_tokens=0, block_ids=([3, 4],)) - ], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"r": 32}, - ) + make_scheduler_output(new_reqs=[scheduled], scheduled_tokens={"r": 32}) ) - stored = [plan.block_id for plan in metadata.store_plans] + stored = [plan.block_id for plan in metadata.iter_store_plans()] assert stored == ([] if kv_role == "kv_consumer" else [3, 4]) assert bool(metadata.store_requests) == (kv_role != "kv_consumer") +@pytest.mark.parametrize("resident", [False, True]) +def test_restored_prefix_skips_resident_blocks(resident): + """A restored block that is still in the pool is not stored again.""" + codec = BlockIdentityCodec(UMBPNamespace("restored-residency")) + handle = _SchedulerHandle({}) + scheduler = UMBPStoreConnectorScheduler( + _vllm_config({"mode": "embedded"}), _kv_cache_config(), handle, codec + ) + request = make_request( + "req", + num_tokens=49, + block_hashes=[b"a", b"b", b"c"], + block_ids=([1, 2, 3],), + ) + scheduler.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda group_ids: request.block_ids), 0 + ) + scheduler._request_trackers["req"].load_spec = LoadSpec(0, 32) + handle.hits = {codec.key(h, 0): resident for h in (b"a", b"b")} + meta = scheduler.build_connector_meta( + make_scheduler_output(new_reqs=[request], scheduled_tokens={"req": 48}) + ) + + stored = [b"c"] if resident else [b"a", b"b", b"c"] + assert [plan.key for plan in meta.iter_store_plans()] == [ + codec.key(h, 0) for h in stored + ] + scheduler.close() + + def test_store_event_owns_refs_until_all_ranks_finish(): scheduler = UMBPStoreConnectorScheduler( _vllm_config( @@ -884,11 +913,9 @@ def test_store_event_owns_refs_until_all_ranks_finish(): free_blocks=lambda blocks: freed.extend(blocks), ) ) - request = SimpleNamespace( - request_id="store-refs", - req_id="store-refs", + request = make_request( + "store-refs", num_tokens=16, - num_computed_tokens=0, block_hashes=[b"a"], block_ids=([4],), ) @@ -896,13 +923,7 @@ def test_store_event_owns_refs_until_all_ranks_finish(): request, SimpleNamespace(get_block_ids=lambda group_ids: ([4],)), 0 ) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[request], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={request.req_id: 16}, - ) + make_scheduler_output(new_reqs=[request], scheduled_tokens={request.req_id: 16}) ) one_rank = SimpleNamespace( kv_connector_worker_meta=UMBPConnectorWorkerMetadata( @@ -924,7 +945,7 @@ def test_store_event_owns_refs_until_all_ranks_finish(): assert freed == [block] -def test_store_events_hold_independent_refs_to_the_same_gpu_block(): +def test_store_events_deduplicate_in_flight_shared_prefix_objects(): scheduler = UMBPStoreConnectorScheduler( _vllm_config({"mode": "embedded"}), _kv_cache_config(), @@ -941,11 +962,9 @@ def test_store_events_hold_independent_refs_to_the_same_gpu_block(): ) events = [] for request_id in ("first", "second"): - request = SimpleNamespace( - request_id=request_id, - req_id=request_id, + request = make_request( + request_id, num_tokens=16, - num_computed_tokens=0, block_hashes=[b"a"], block_ids=([4],), ) @@ -955,19 +974,13 @@ def test_store_events_hold_independent_refs_to_the_same_gpu_block(): 0, ) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[request], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={request_id: 16}, - ) + make_scheduler_output(new_reqs=[request], scheduled_tokens={request_id: 16}) ) events.append(metadata.store_event) - assert pinned == [4, 4] - assert events == [0, 1] + assert pinned == [4] + assert events == [0, -1] - for event_id in (1, 1, 0): + for event_id in (-1, 0): scheduler.update_connector_output( SimpleNamespace( kv_connector_worker_meta=UMBPConnectorWorkerMetadata( @@ -975,7 +988,7 @@ def test_store_events_hold_independent_refs_to_the_same_gpu_block(): ), ) ) - assert len(freed) == (2 if event_id == 0 else 1) + assert len(freed) == (1 if event_id == 0 else 0) assert scheduler.has_pending_push_work() == (event_id != 0) @@ -991,13 +1004,8 @@ def test_scheduler_resumed_request_replaces_stale_block_table(): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("resumed")), ) - request = SimpleNamespace( - request_id="resumed", - req_id="resumed", - num_tokens=32, - num_prompt_tokens=32, - num_computed_tokens=0, - block_hashes=[b"a", b"b"], + request = make_request( + "resumed", block_ids=([7, 8],), ) scheduler.update_state_after_alloc( @@ -1006,23 +1014,20 @@ def test_scheduler_resumed_request_replaces_stale_block_table(): 0, ) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["resumed"], new_block_ids=[([7, 8],)], num_computed_tokens=[31], resumed_req_ids={"resumed"}, ), - num_scheduled_tokens={"resumed": 1}, + scheduled_tokens={"resumed": 1}, ) ) tracker = scheduler._request_trackers["resumed"] assert tracker.block_ids == ([7, 8],) - assert [plan.block_id for plan in metadata.store_plans] == [7, 8] + assert [plan.block_id for plan in metadata.iter_store_plans()] == [7, 8] def test_scheduler_stores_from_cores_current_block_table(): @@ -1032,12 +1037,10 @@ def test_scheduler_stores_from_cores_current_block_table(): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("current-table")), ) - request = SimpleNamespace( - request_id="r", - req_id="r", + request = make_request( + "r", num_tokens=64, num_prompt_tokens=64, - num_computed_tokens=0, block_hashes=[b"a", b"b", b"c", b"d"], ) scheduler.update_state_after_alloc( @@ -1046,22 +1049,44 @@ def test_scheduler_stores_from_cores_current_block_table(): # Core has since freed block 3, e.g. a sliding-window block out of the window. current = {"r": ([NULL_BLOCK_ID, 4],)} metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["r"], new_block_ids=[None], num_computed_tokens=[0], resumed_req_ids=set(), ), - num_scheduled_tokens={"r": 32}, - kv_connector_block_state=SimpleNamespace(get_block_ids=current.get), + scheduled_tokens={"r": 32}, + block_state=SimpleNamespace( + get_block_ids=current.get, boundary_state_offloads={} + ), ) ) - assert [plan.block_id for plan in metadata.store_plans] == [4] + assert [plan.block_id for plan in metadata.iter_store_plans()] == [4] + + +def test_eager_store_deduplicates_same_step_shared_prefix_keys(): + scheduler = UMBPStoreConnectorScheduler( + _vllm_config({"mode": "embedded"}), + _kv_cache_config(), + _SchedulerHandle([]), + BlockIdentityCodec(UMBPNamespace("same-step-dedup")), + ) + first = make_request("first", block_ids=([1, 2],)) + second = make_request("second", block_ids=([3, 4],)) + reserved: set[str] = set() + + first_batch = scheduler._store_plans( + first, RequestTracker(), 32, skip_keys=reserved + ) + reserved.update(first_batch.keys) + second_batch = scheduler._store_plans( + second, RequestTracker(), 32, skip_keys=reserved + ) + + assert [plan.block_id for plan in first_batch] == [1, 2] + assert len(second_batch) == 0 def test_scheduler_readmission_after_failed_load_uses_new_blocks(): @@ -1071,12 +1096,10 @@ def test_scheduler_readmission_after_failed_load_uses_new_blocks(): _SchedulerHandle([]), BlockIdentityCodec(UMBPNamespace("readmitted")), ) - request = SimpleNamespace( - request_id="readmitted", - req_id="readmitted", + request = make_request( + "readmitted", num_tokens=64, num_prompt_tokens=64, - num_computed_tokens=0, block_hashes=[b"a", b"b", b"c", b"d"], ) scheduler.update_state_after_alloc( @@ -1087,22 +1110,19 @@ def test_scheduler_readmission_after_failed_load_uses_new_blocks(): request, SimpleNamespace(get_block_ids=lambda group_ids: ([5, 6],)), 0 ) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( + make_scheduler_output( + cached_reqs=SimpleNamespace( req_ids=["readmitted"], new_block_ids=[None], num_computed_tokens=[0], resumed_req_ids=set(), ), - num_scheduled_tokens={"readmitted": 32}, + scheduled_tokens={"readmitted": 32}, ) ) assert scheduler._request_trackers["readmitted"].block_ids == ([5, 6],) - assert [plan.block_id for plan in metadata.store_plans] == [5, 6] + assert [plan.block_id for plan in metadata.iter_store_plans()] == [5, 6] @pytest.mark.parametrize("fail_load", [False, True]) @@ -1182,7 +1202,7 @@ def test_async_stores_keep_independent_events_across_steps(fail_first): assert not handle.waited -def test_preemption_waits_for_pending_stores_from_previous_steps(): +def test_preemption_cancels_pending_stores_from_previous_steps(): handle = _DelayedStoreWorkerHandle() worker = UMBPStoreConnectorWorker(handle) for event_id, request_id in enumerate(("first", "other", "first")): @@ -1196,12 +1216,14 @@ def test_preemption_waits_for_pending_stores_from_previous_steps(): ) worker.wait_for_save() worker.handle_preemptions(UMBPConnectorMetadata(preempted_request_ids={"first"})) - assert handle.waited == [handle.jobs[0], handle.jobs[2]] worker.get_finished(set()) + assert handle.cancelled == [handle.jobs[0], handle.jobs[2]] assert worker.build_connector_worker_meta().store_events == { 0: StoreEventResult(1), 2: StoreEventResult(1), } + assert not handle.publications + assert not handle.waited handle.jobs[1].complete() worker.get_finished(set()) assert worker.build_connector_worker_meta().store_events == {1: StoreEventResult(1)} @@ -1220,25 +1242,63 @@ def test_async_store_drains_before_buffers_can_be_reused(): assert worker.build_connector_worker_meta().store_events == {4: StoreEventResult(1)} -def test_preemption_waits_only_for_the_preempted_requests_loads(): - handle = _WaitRecordingDelayedLoadHandle() +def test_worker_preemption_cancels_only_matching_request(): + handle = _CancellableWorkerHandle() worker = UMBPStoreConnectorWorker(handle) + first = BlockTransferPlan("first", 1, request_id="first") + second = BlockTransferPlan("second", 2, request_id="second") + worker.enqueue_stores( + UMBPConnectorMetadata( + store_event=21, + store_plans=[first, second], + store_requests={"first": [first], "second": [second]}, + ) + ) + + worker.handle_preemptions(UMBPConnectorMetadata(preempted_request_ids={"first"})) + + assert len(handle.cancelled) == 1 + assert handle.cancelled[0].plans == (first,) + assert worker.build_connector_worker_meta().store_events == {} + worker.wait_for_save() + assert handle.published.plans == (second,) + assert worker.build_connector_worker_meta().store_events == { + 21: StoreEventResult(1) + } + + +def test_async_load_cancellation_preserves_other_requests(): + handle = _CancellableWorkerHandle() + worker = UMBPStoreConnectorWorker(handle) + plans = { + name: [ + BlockTransferPlan( + name, + block_id, + request_id=name, + ranges=(KVRange("layer1", 0, block_id, 1000, 16, 16, 0),), + ) + ] + for name, block_id in (("first", 1), ("second", 2)) + } worker.start_load_kv( None, UMBPConnectorMetadata( async_load=True, - load_requests={ - name: [BlockTransferPlan(name, block_id, request_id=name)] - for name, block_id in (("first", 1), ("second", 2)) - }, + load_requests=plans, ), ) worker.handle_preemptions(UMBPConnectorMetadata(preempted_request_ids={"first"})) + worker.wait_for_layer_load("") + # Asynchronous loads settle through get_finished, not the forward pass. + finished = worker.get_finished(set()) - assert [job.plans[0].request_id for job in handle.waited] == ["first"] - assert worker.get_finished(set()) == (None, None) - handle.load_job.complete() - assert worker.get_finished(set()) == (None, {"second"}) + assert len(handle.cancelled) == 1 + assert handle.cancelled[0].plans[0].request_id == "first" + assert worker.get_kv_connector_stats().reduce()["load_completed"] == 1 + assert finished == (None, {"second"}) + worker.wait_for_layer_load("") + assert worker.get_kv_connector_stats() is None def test_worker_reports_load_completion_and_store_event_without_deferring_request(): @@ -1283,6 +1343,7 @@ def store(self, plans): metadata = worker.build_connector_worker_meta() assert metadata.store_events == {9: StoreEventResult(1)} assert not hasattr(worker.runtime, "published") + assert worker.get_kv_connector_stats().reduce()["store_failed"] == 1 @pytest.mark.parametrize("forward", [False, True]) @@ -1325,22 +1386,17 @@ def test_connector_finalizes_each_store_once_when_forward_hook_is_skipped( assert next_step == ({} if forward else expected) assert plan.key in runtime.store + stats = connector.get_kv_connector_stats().reduce() + assert stats["store_submitted"] == 2 + assert stats["store_completed"] == 2 + def test_worker_preserves_scheduler_supplied_ranges(): handle = _WorkerHandle() worker = UMBPStoreConnectorWorker( handle, KVLayoutPlanner.from_kv_cache_config(_kv_cache_config()) ) - worker.register_kv_caches( - { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, - ) - for name in ("layer1", "layer2") - } - ) + worker.register_kv_caches(make_kv_caches()) supplied = BlockTransferPlan( key="custom", block_id=1, @@ -1365,90 +1421,8 @@ def test_worker_preserves_scheduler_supplied_ranges(): assert handle.loaded_plans == [supplied] -def test_embedded_connector_core_flow(monkeypatch): - runtime = _EmbeddedRuntime() - monkeypatch.setitem( - UMBPRuntimeFactory._builders, - "embedded", - lambda config: runtime, - ) - config = _kv_cache_config() - vllm_config = _vllm_config({"mode": "embedded", "load_async": False}) - producer = SimpleNamespace( - request_id="producer", - req_id="producer", - num_tokens=32, - prompt_token_ids=list(range(32)), - block_hashes=[b"a", b"b"], - block_ids=([1, 2],), - num_computed_tokens=0, - ) - scheduler_output = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[producer], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"producer": 32}, - ) - - scheduler_connector = UMBPStoreConnector( - vllm_config, KVConnectorRole.SCHEDULER, config - ) - worker_connector = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, config) - handshake(scheduler_connector, worker_connector) - caches = { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, - ) - for name in ("layer1", "layer2") - } - worker_connector.register_kv_caches(caches) - - assert scheduler_connector.get_num_new_matched_tokens(producer, 0) == (0, False) - scheduler_connector.update_state_after_alloc( - producer, - SimpleNamespace(get_block_ids=lambda group_ids: ([1, 2],)), - 0, - ) - metadata = scheduler_connector.build_connector_meta(scheduler_output) - worker_connector.bind_connector_metadata(metadata) - worker_connector.wait_for_save() - assert runtime.store - - consumer = SimpleNamespace( - request_id="consumer", - req_id="consumer", - num_tokens=33, - block_hashes=[b"a", b"b"], - ) - assert scheduler_connector.get_num_new_matched_tokens(consumer, 0) == ( - 32, - False, - ) - scheduler_connector.update_state_after_alloc( - consumer, - SimpleNamespace(get_block_ids=lambda group_ids: ([7, 8],)), - 32, - ) - consumer_output = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[consumer], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"consumer": 32}, - ) - load_metadata = scheduler_connector.build_connector_meta(consumer_output) - worker_connector.bind_connector_metadata(load_metadata) - worker_connector.start_load_kv(None) - worker_connector.wait_for_layer_load("layer0") - - assert worker_connector.get_block_ids_with_load_errors() == set() - assert not worker_connector.get_transfer_results({"consumer"}).finished_recving - - def test_embedded_runtime_register_store_and_load(monkeypatch): + """Scheduler-built batches restore stored bytes into different blocks.""" install_memory_runtime(monkeypatch) config = _kv_cache_config() vllm_config = _vllm_config( @@ -1458,107 +1432,77 @@ def test_embedded_runtime_register_store_and_load(monkeypatch): "key_namespace": "builtin-embedded-test", } ) - source_caches = { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, - ) - for name in ("layer1", "layer2") - } - for index, cache in enumerate(source_caches.values()): - cache.copy_( - torch.arange(cache.numel(), dtype=torch.float16).reshape(cache.shape) - + index - ) - + source_caches = make_kv_caches(fill=True) + destination_caches = make_kv_caches() scheduler_connector = UMBPStoreConnector( vllm_config, KVConnectorRole.SCHEDULER, config ) - worker_connector = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, config) - worker_connector.register_kv_caches(source_caches) - producer = SimpleNamespace( - request_id="builtin-producer", - req_id="builtin-producer", - num_tokens=32, - block_hashes=[b"builtin-a", b"builtin-b"], - block_ids=([1, 2],), - num_computed_tokens=0, - ) - producer_output = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[producer], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"builtin-producer": 32}, - ) - scheduler_connector.update_state_after_alloc( - producer, - SimpleNamespace(get_block_ids=lambda group_ids: producer.block_ids), - 0, - ) - store_metadata = scheduler_connector.build_connector_meta(producer_output) - worker_connector.bind_connector_metadata(store_metadata) - worker_connector.wait_for_save() - - consumer = SimpleNamespace( - request_id="builtin-consumer", - req_id="builtin-consumer", - num_tokens=33, - block_hashes=[b"builtin-a", b"builtin-b"], - ) - assert scheduler_connector.get_num_new_matched_tokens(consumer, 0) == ( - 32, - False, - ) - scheduler_connector.update_state_after_alloc( - consumer, - SimpleNamespace(get_block_ids=lambda group_ids: ([5, 6],)), - 32, - ) - consumer_output = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[consumer], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"builtin-consumer": 32}, - ) - load_metadata = scheduler_connector.build_connector_meta(consumer_output) - - destination_caches = { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, + store_worker = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, config) + load_worker = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, config) + try: + store_worker.register_kv_caches(source_caches) + load_worker.register_kv_caches(destination_caches) + handshake(scheduler_connector, store_worker) + + producer = make_request( + "builtin-producer", + block_hashes=[b"builtin-a", b"builtin-b"], + block_ids=([1, 2],), ) - for name in source_caches - } - for cache in destination_caches.values(): - cache.zero_() - worker_connector.register_kv_caches(destination_caches) - worker_connector.bind_connector_metadata(load_metadata) - worker_connector.start_load_kv(None) - worker_connector.wait_for_layer_load("layer1") - assert worker_connector.get_finished({"builtin-consumer"}) == (None, None) - worker_connector.wait_for_layer_load("layer2") - assert worker_connector.get_finished({"builtin-consumer"}) == (None, None) - load_errors = worker_connector.get_block_ids_with_load_errors() - if load_errors: - raise AssertionError(f"embedded load errors: {sorted(load_errors)}") - - for name in source_caches: - if not torch.equal(source_caches[name][1], destination_caches[name][5]): - raise AssertionError( - f"{name} block 1 round-trip mismatch: " - f"src={source_caches[name][1, 0, 0, 0].item()} " - f"dst={destination_caches[name][5, 0, 0, 0].item()}" + scheduler_connector.update_state_after_alloc( + producer, + SimpleNamespace(get_block_ids=lambda group_ids: producer.block_ids), + 0, + ) + store_metadata = scheduler_connector.build_connector_meta( + make_scheduler_output( + new_reqs=[producer], scheduled_tokens={producer.req_id: 32} ) - if not torch.equal(source_caches[name][2], destination_caches[name][6]): - raise AssertionError( - f"{name} block 2 round-trip mismatch: " - f"src={source_caches[name][2, 0, 0, 0].item()} " - f"dst={destination_caches[name][6, 0, 0, 0].item()}" + ) + store_worker.bind_connector_metadata(store_metadata) + store_worker.wait_for_save() + store_worker.clear_connector_metadata() + + consumer = make_request( + "builtin-consumer", + num_tokens=33, + block_hashes=producer.block_hashes, + block_ids=([5, 6],), + ) + assert scheduler_connector.get_num_new_matched_tokens(consumer, 0) == ( + 32, + False, + ) + scheduler_connector.update_state_after_alloc( + consumer, + SimpleNamespace(get_block_ids=lambda group_ids: consumer.block_ids), + 32, + ) + load_metadata = scheduler_connector.build_connector_meta( + make_scheduler_output( + new_reqs=[consumer], scheduled_tokens={consumer.req_id: 1} ) + ) + assert set(load_metadata.load_requests) == {consumer.req_id} + + load_worker.bind_connector_metadata(load_metadata) + load_worker.start_load_kv(None) + load_worker.wait_for_layer_load("layer1") + load_worker.wait_for_layer_load("layer2") + assert load_worker.get_finished({consumer.req_id}) == (None, None) + assert load_worker.get_block_ids_with_load_errors() == set() + load_worker.clear_connector_metadata() + + for name, source in source_caches.items(): + destination = destination_caches[name] + assert torch.equal(source[1], destination[5]), f"{name} block 1" + assert torch.equal(source[2], destination[6]), f"{name} block 2" + untouched = [block for block in range(8) if block not in (5, 6)] + assert not destination[untouched].any(), f"{name} stray write" + finally: + scheduler_connector.shutdown() + store_worker.shutdown() + load_worker.shutdown() @pytest.mark.parametrize(("tp_size", "dcp_size"), [(2, 1), (2, 2), (4, 2), (8, 8)]) @@ -1577,16 +1521,14 @@ def test_embedded_tp_dcp_rank_store_completeness(monkeypatch, tp_size, dcp_size) "world_size": tp_size, } token_count = 32 * dcp_size - request = SimpleNamespace( - request_id="ranked-producer", - req_id="ranked-producer", + request = make_request( + "ranked-producer", num_tokens=token_count, block_hashes=[b"ranked-a", b"ranked-b"], block_ids=([1, 2],), - num_computed_tokens=0, ) - consumer = SimpleNamespace( - request_id="ranked-consumer", + consumer = make_request( + "ranked-consumer", num_tokens=token_count + 1, block_hashes=request.block_hashes, ) @@ -1600,18 +1542,11 @@ def test_embedded_tp_dcp_rank_store_completeness(monkeypatch, tp_size, dcp_size) 0, ) metadata = scheduler.build_connector_meta( - SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[request], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={request.req_id: token_count}, + make_scheduler_output( + new_reqs=[request], scheduled_tokens={request.req_id: token_count} ) ) - caches = { - name: torch.empty_strided((8, 2, 16, 8), (512, 256, 8, 1), dtype=torch.float16) - for name in ("layer1", "layer2") - } + caches = make_kv_caches() try: for tp_rank in range(tp_size): assert scheduler.get_num_new_matched_tokens(consumer, 0) == (0, False) @@ -1634,78 +1569,6 @@ def test_embedded_tp_dcp_rank_store_completeness(monkeypatch, tp_size, dcp_size) scheduler.shutdown() -def test_embedded_tp_dcp_rank_local_store_flow(monkeypatch): - runtime = _EmbeddedRuntime() - monkeypatch.setitem( - UMBPRuntimeFactory._builders, - "embedded", - lambda config: runtime, - ) - config = _kv_cache_config() - vllm_config = _vllm_config( - {"mode": "embedded", "load_async": False}, - rank=1, - tensor_parallel_size=2, - decode_context_parallel_size=2, - world_size=2, - ) - request = SimpleNamespace( - request_id="ranked", - req_id="ranked", - num_tokens=32, - block_hashes=[b"a", b"b"], - block_ids=([3, 4],), - num_computed_tokens=0, - ) - scheduler_output = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[request], - scheduled_cached_reqs=SimpleNamespace(req_ids=[]), - num_scheduled_tokens={"ranked": 32}, - ) - - scheduler_connector = UMBPStoreConnector( - vllm_config, KVConnectorRole.SCHEDULER, config - ) - worker_connector = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, config) - handshake(scheduler_connector, worker_connector) - scheduler_connector.update_state_after_alloc( - request, - SimpleNamespace(get_block_ids=lambda group_ids: ([3, 4],)), - 0, - ) - scheduler_output.scheduled_new_reqs = [ - SimpleNamespace( - req_id="ranked", - num_computed_tokens=0, - block_ids=([3, 4],), - ) - ] - worker_connector.register_kv_caches( - { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, - ) - for name in ("layer1", "layer2") - } - ) - - metadata = scheduler_connector.build_connector_meta(scheduler_output) - worker_connector.bind_connector_metadata(metadata) - worker_connector.wait_for_save() - - namespace = UMBPNamespace.from_vllm_config(vllm_config, config).value - assert runtime.store == { - f"umbp:vllm:v1:{namespace}:tp1:pcp0:dcp1:pp0:g0:61", - } - assert runtime.scheduler_args is not None - assert runtime.scheduler_args[1].local_namespace == (1, 0, 1, 0) - assert runtime.scheduler_args[1].rank_count == 2 - - def test_embedded_logical_hit_requires_all_tp_dcp_objects(monkeypatch): runtime = _EmbeddedRuntime() monkeypatch.setitem( @@ -1727,8 +1590,8 @@ def test_embedded_logical_hit_requires_all_tp_dcp_objects(monkeypatch): ) scheduler = connector.connector_scheduler assert scheduler is not None - request = SimpleNamespace( - request_id="logical", + request = make_request( + "logical", num_tokens=65, block_hashes=[b"a", b"b"], ) @@ -1762,21 +1625,13 @@ def fail(plans): monkeypatch.setattr(handle, "store", fail) codec = BlockIdentityCodec(UMBPNamespace("rank-local"), tp_rank=tp_rank) layout = KVLayoutPlanner.from_kv_cache_config(_kv_cache_config()) - layout.register_kv_caches( - { - name: torch.empty_strided( - (8, 2, 16, 8), - (512, 256, 8, 1), - dtype=torch.float16, - ) - for name in ("layer1", "layer2") - } - ) + layout.register_kv_caches(make_kv_caches()) worker = UMBPStoreConnectorWorker(handle, layout, codec=codec) plan = BlockTransferPlan( key=BlockIdentityCodec(UMBPNamespace("rank-local")).key(b"hash", 0), block_id=1, group_id=0, + block_hash=b"hash", ) metadata = UMBPConnectorMetadata(store_plans=[plan], store_event=7) @@ -1785,34 +1640,47 @@ def fail(plans): assert handle.store_calls[0][0].key == codec.key(b"hash", 0) worker_meta = worker.build_connector_worker_meta() - assert worker_meta.store_events == {7: StoreEventResult(1)} + assert worker_meta.store_events == { + 7: StoreEventResult( + 1, + {(0, b"hash")} if fail_store else set(), + ) + } assert hasattr(handle, "published") != fail_store @pytest.mark.parametrize("tp_rank", [0, 3]) -def test_load_localizes_rank_without_changing_request_or_destination(tp_rank): +@pytest.mark.parametrize("fast_path", [False, True]) +def test_batched_load_localizes_rank_without_changing_request_or_destination( + tp_rank, fast_path +): + """Fast and fallback loads target the same rank and GPU blocks.""" codec = BlockIdentityCodec(UMBPNamespace("rank-local"), tp_rank=tp_rank) scheduler_codec = replace(codec, tp_rank=0) calls = [] def load(plans): calls.append(plans) - job = TransferJobState(tuple(plans)) + job = TransferJobState(plans if fast_path else tuple(plans)) job.start() job.complete([codec.key(b"a", 0)]) job.fail([codec.key(b"b", 1)], "injected load failure") return job - worker = UMBPStoreConnectorWorker(SimpleNamespace(load=load), codec=codec) - plans = [ - BlockTransferPlan( - scheduler_codec.key(b"a", 0), 7, request_id="request", group_id=0 - ), - BlockTransferPlan( - scheduler_codec.key(b"b", 1), 9, request_id="request", group_id=1 - ), - ] - worker.start_load_kv(None, UMBPConnectorMetadata(load_requests={"request": plans})) + handle = SimpleNamespace(load=load) + if fast_path: + handle.load_blocks = load + worker = UMBPStoreConnectorWorker(handle, codec=codec) + batch = BlockLoadBatch( + [scheduler_codec.key(b"a", 0), scheduler_codec.key(b"b", 1)], + [7, 9], + [0, 1], + "request", + ) + worker.start_load_kv( + None, + UMBPConnectorMetadata(async_load=True, load_requests={"request": batch}), + ) assert len(calls) == 1 assert [(p.key, p.block_id, p.group_id, p.request_id) for p in calls[0]] == [ @@ -1821,6 +1689,98 @@ def load(plans): ] job = worker._load_jobs["request"][None] assert job.failed_block_ids == {9} + assert batch.keys == [scheduler_codec.key(b"a", 0), scheduler_codec.key(b"b", 1)] + + +@pytest.mark.parametrize("fail_store", [False, True]) +@pytest.mark.parametrize("tp_rank", [0, 3]) +@pytest.mark.parametrize("fast_path", [False, True]) +def test_batched_store_localizes_rank_without_changing_scheduler_batch( + fail_store, tp_rank, fast_path, monkeypatch +): + """Fast and fallback stores target the same rank and publish alike.""" + codec = BlockIdentityCodec(UMBPNamespace("rank-local"), tp_rank=tp_rank) + scheduler_codec = replace(codec, tp_rank=0) + handle = _StoreRecordingWorkerHandle() + calls = [] + + def store(plans): + calls.append(plans) + job = TransferJobState(plans if fast_path else tuple(plans)) + job.start() + if fail_store: + # Only the second object fails; the whole job stays unpublished. + job.complete([codec.key(b"a", 0)]) + job.fail([codec.key(b"b", 0)], "injected store failure") + else: + job.complete() + return job + + monkeypatch.setattr(handle, "store", store) + if fast_path: + monkeypatch.setattr(handle, "store_blocks", store, raising=False) + layout = KVLayoutPlanner.from_kv_cache_config(_kv_cache_config()) + layout.register_kv_caches(make_kv_caches()) + worker = UMBPStoreConnectorWorker(handle, layout, codec=codec) + scheduler_keys = [scheduler_codec.key(b"a", 0), scheduler_codec.key(b"b", 0)] + batch = BlockStoreBatch( + keys=list(scheduler_keys), + block_ids=[1, 2], + group_ids=[0, 0], + block_hashes=[b"a", b"b"], + parent_block_hashes=[None, b"a"], + token_ids=[tuple(range(16)), tuple(range(16, 32))], + block_sizes=[16, 16], + request_id="request", + ) + metadata = UMBPConnectorMetadata(store_event=11, store_requests={"request": batch}) + + worker.enqueue_stores(metadata) + worker.wait_for_save() + + assert len(calls) == 1 + if fast_path: + assert isinstance(calls[0], BlockStoreBatch) + else: + assert all(plan.ranges for plan in calls[0]) + assert [ + ( + plan.key, + plan.block_id, + plan.group_id, + plan.request_id, + plan.block_hash, + plan.parent_block_hash, + plan.token_ids, + plan.block_size, + ) + for plan in calls[0] + ] == [ + (codec.key(b"a", 0), 1, 0, "request", b"a", None, tuple(range(16)), 16), + ( + codec.key(b"b", 0), + 2, + 0, + "request", + b"b", + b"a", + tuple(range(16, 32)), + 16, + ), + ] + # The scheduler's batch is shared metadata and must not be rewritten. + assert batch.keys == scheduler_keys + assert metadata.store_requests == {"request": batch} + assert worker.build_connector_worker_meta().store_events == { + 11: StoreEventResult(1, {(0, b"a"), (0, b"b")} if fail_store else set()) + } + if fail_store: + assert not hasattr(handle, "published") + else: + assert list(handle.published.keys) == [ + codec.key(b"a", 0), + codec.key(b"b", 0), + ] @pytest.mark.parametrize("empty_group", [False, True]) @@ -1877,3 +1837,4 @@ def test_worker_submits_stores_once(grouped): assert worker.get_finished({"req"}) == (None, None) completed = worker.build_connector_worker_meta() assert completed.store_events == {19: StoreEventResult(1)} + assert worker.get_kv_connector_stats().reduce()["store_num_bytes"] == 32 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py index 06eee4bfabe3..e6e4cf3ff1b1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .connector import UMBPStoreConnector +from .connector import UMBPStoreConnector, UMBPStoreKVEvents from .data import ( BlockIdentityCodec, BlockTransferPlan, @@ -19,9 +19,11 @@ UMBPNamespace, ) from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory +from .stats import UMBPStoreConnectorStats __all__ = [ "UMBPStoreConnector", + "UMBPStoreKVEvents", "BlockIdentityCodec", "BlockTransferPlan", "KVLayoutPlanner", @@ -38,4 +40,5 @@ "UMBPNamespace", "UMBPRuntimeConfig", "UMBPRuntimeFactory", + "UMBPStoreConnectorStats", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py index e828ac4ee072..283010cbbf09 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py @@ -4,17 +4,24 @@ from __future__ import annotations +from collections.abc import Iterable from typing import TYPE_CHECKING, Any import torch from vllm.config import VllmConfig +from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorHandshakeMetadata, KVConnectorMetadata, KVConnectorRole, KVConnectorTransferResults, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + PromMetric, + PromMetricT, ) from vllm.forward_context import ForwardContext from vllm.v1.attention.backend import AttentionMetadata @@ -32,6 +39,7 @@ ) from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory from .scheduler import UMBPStoreConnectorScheduler +from .stats import UMBPStoreConnectorStats, UMBPStorePromMetrics from .worker import UMBPStoreConnectorWorker if TYPE_CHECKING: @@ -40,7 +48,55 @@ from vllm.v1.request import Request -class UMBPStoreConnector(KVConnectorBase_V1): +class UMBPStoreKVEvents(KVConnectorKVEvents): + """Ordered, deduplicated UMBP events ready for publication.""" + + def __init__(self, events: list[KVCacheEvent] | None = None) -> None: + self._events = list(dict.fromkeys(events or ())) + self._num_workers = 1 + + def add_events(self, events: list[KVCacheEvent]) -> None: + """Merge duplicate worker reports within one engine step.""" + seen = set(self._events) + for event in events: + if event in seen: + continue + seen.add(event) + self._events.append(event) + + def append_events(self, events: Iterable[KVCacheEvent]) -> None: + """Append events from a later step while preserving repetitions.""" + self._events.extend(events) + + def aggregate(self) -> UMBPStoreKVEvents: + self._num_workers = 1 + return self + + def pop_common_events(self) -> list[KVCacheEvent]: + events = self._events + self._events = [] + return events + + def has_events(self) -> bool: + return bool(self._events) + + def increment_workers(self, count: int = 1) -> None: + if count <= 0: + raise ValueError("count must be positive") + self._num_workers += count + + def get_all_events(self) -> list[KVCacheEvent]: + return list(self._events) + + def get_number_of_workers(self) -> int: + return self._num_workers + + def clear_events(self) -> None: + self._events.clear() + self._num_workers = 1 + + +class UMBPStoreConnector(KVConnectorBase_V1, SupportsHMA): """vLLM connector lifecycle shared by every UMBP runtime mode.""" def __init__( @@ -58,10 +114,24 @@ def __init__( raise NotImplementedError( f"{runtime_config.mode} UMBP does not support pipeline parallelism" ) + # A restored prefix skips the draft model's prefill too, so its KV must + # come from the pool along with the target's. + cacheable = set(kv_cache_config.prefix_cacheable_group_ids) + skipped_drafts = [ + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if getattr(group, "is_eagle_group", False) and group_id not in cacheable + ] + if skipped_drafts: + raise NotImplementedError( + f"{runtime_config.mode} UMBP cannot restore draft-model KV cache " + f"groups {skipped_drafts}: they are not prefix-cacheable" + ) runtime = UMBPRuntimeFactory.build(runtime_config) self._runtime = runtime self._topology = topology self._store_submitted = False + self._kv_cache_events: UMBPStoreKVEvents | None = None self._namespace: UMBPNamespace | None = None self.connector_scheduler: UMBPStoreConnectorScheduler | None = None self.connector_worker: UMBPStoreConnectorWorker | None = None @@ -83,6 +153,11 @@ def __init__( ), layout, codec=self._codec(self._namespace), + report_failed_requests=len(kv_cache_config.kv_cache_groups) > 1, + enable_kv_cache_events=bool( + vllm_config.kv_events_config + and vllm_config.kv_events_config.enable_kv_cache_events + ), ) def _codec(self, namespace: UMBPNamespace) -> BlockIdentityCodec: @@ -154,6 +229,16 @@ def bind_gpu_block_pool(self, gpu_block_pool: Any) -> None: assert self.connector_scheduler is not None self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool) + def bind_kv_cache_manager(self, kv_cache_manager: Any) -> None: + super().bind_kv_cache_manager(kv_cache_manager) + assert self.connector_scheduler is not None + + def local_cache_hit(group_id: int, block_hash: bytes) -> bool: + manager = kv_cache_manager.coordinator.single_type_managers[group_id] + return bool(manager.block_pool.get_cached_block(block_hash, [group_id])) + + self.connector_scheduler.bind_local_cache_lookup(local_cache_hit) + def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: @@ -166,9 +251,51 @@ def request_finished( assert self.connector_scheduler is not None return self.connector_scheduler.request_finished(request, (block_ids,)) + def request_finished_all_groups( + self, + request: Request, + block_ids: tuple[list[int], ...], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, block_ids) + + def register_finished_partial_tail( + self, + request: Request, + block_ids: tuple[list[int], ...], + partial_tail_offloads: list[tuple[int, int, int]], + ) -> bool: + assert self.connector_scheduler is not None + return self.connector_scheduler.register_finished_partial_tail( + request, block_ids, partial_tail_offloads + ) + def update_connector_output(self, connector_output: Any) -> None: assert self.connector_scheduler is not None + events = connector_output.kv_cache_events + if isinstance(events, UMBPStoreKVEvents): + worker_events = events.get_all_events() + invalidate = getattr(self.connector_scheduler, "invalidate_kv_events", None) + if invalidate is not None: + invalidate(worker_events) + if self._kv_cache_events is None: + self._kv_cache_events = events + else: + self._kv_cache_events.append_events(worker_events) self.connector_scheduler.update_connector_output(connector_output) + scheduler_events = self.connector_scheduler.take_events() + if scheduler_events: + if self._kv_cache_events is None: + self._kv_cache_events = UMBPStoreKVEvents(scheduler_events) + else: + self._kv_cache_events.append_events(scheduler_events) + + def take_events(self) -> Iterable[KVCacheEvent]: + if self._kv_cache_events is None: + return + yield from self._kv_cache_events.pop_common_events() + if not self._kv_cache_events.has_events(): + self._kv_cache_events = None def has_pending_push_work(self) -> bool: assert self.connector_scheduler is not None @@ -245,6 +372,7 @@ def get_transfer_results( return KVConnectorTransferResults( finished_sending=set(finished_sending or ()), finished_recving=set(finished_recving or ()), + failed_recving=self.connector_worker.get_failed_recving(), ) def get_block_ids_with_load_errors(self) -> set[int]: @@ -255,6 +383,40 @@ def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata: assert self.connector_worker is not None return self.connector_worker.build_connector_worker_meta() + def get_kv_connector_stats(self) -> UMBPStoreConnectorStats | None: + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> UMBPStoreConnectorStats | None: + return UMBPStoreConnectorStats(data=data or {}) + + @classmethod + def build_prom_metrics( + cls, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> UMBPStorePromMetrics: + return UMBPStorePromMetrics( + vllm_config, + metric_types, + labelnames, + per_engine_labelvalues, + ) + + def get_kv_connector_kv_cache_events(self) -> UMBPStoreKVEvents | None: + assert self.connector_worker is not None + events = self.connector_worker.get_kv_events() + if not self.connector_worker.enable_kv_cache_events: + return None + # An empty container still counts this worker toward the quorum. + return UMBPStoreKVEvents(events) + def shutdown(self) -> None: if self.connector_scheduler is not None: self.connector_scheduler.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py index 9374c9ced802..3535facf0416 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py @@ -15,7 +15,7 @@ from dataclasses import dataclass, field, replace from dataclasses import fields as dataclass_fields from enum import Enum -from typing import Any, NamedTuple +from typing import Any, NamedTuple, overload import torch @@ -521,6 +521,103 @@ class BlockTransferPlan: ranges: tuple[KVRange, ...] = () request_id: str | None = None group_id: int | None = None + block_hash: bytes | None = None + parent_block_hash: bytes | None = None + token_ids: tuple[int, ...] = () + block_size: int = 0 + medium: str = "CPU" + + +@dataclass(frozen=True) +class BlockLoadBatch(Sequence[BlockTransferPlan]): + """Whole-block loads with request metadata shared across the batch.""" + + keys: list[str] + block_ids: list[int] + group_ids: list[int] + request_id: str + + def __post_init__(self) -> None: + if len(self.keys) != len(self.block_ids) or len(self.keys) != len( + self.group_ids + ): + raise ValueError("load batch columns must have equal lengths") + + def __len__(self) -> int: + return len(self.keys) + + @overload + def __getitem__(self, index: int) -> BlockTransferPlan: ... + + @overload + def __getitem__(self, index: slice) -> list[BlockTransferPlan]: ... + + def __getitem__( + self, index: int | slice + ) -> BlockTransferPlan | list[BlockTransferPlan]: + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + return BlockTransferPlan( + key=self.keys[index], + block_id=self.block_ids[index], + group_id=self.group_ids[index], + request_id=self.request_id, + ) + + +@dataclass(frozen=True) +class BlockStoreBatch(Sequence[BlockTransferPlan]): + """Columnar whole-block stores with request metadata shared by the batch.""" + + keys: list[str] + block_ids: list[int] + group_ids: list[int] + block_hashes: list[bytes] + parent_block_hashes: list[bytes | None] + token_ids: list[tuple[int, ...]] + block_sizes: list[int] + request_id: str + + def __post_init__(self) -> None: + length = len(self.keys) + columns = ( + self.block_ids, + self.group_ids, + self.block_hashes, + self.parent_block_hashes, + self.token_ids, + self.block_sizes, + ) + if any(len(column) != length for column in columns): + raise ValueError("store batch columns must have equal lengths") + + def __len__(self) -> int: + return len(self.keys) + + @overload + def __getitem__(self, index: int) -> BlockTransferPlan: ... + + @overload + def __getitem__(self, index: slice) -> list[BlockTransferPlan]: ... + + def __getitem__( + self, index: int | slice + ) -> BlockTransferPlan | list[BlockTransferPlan]: + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + return BlockTransferPlan( + key=self.keys[index], + block_id=self.block_ids[index], + request_id=self.request_id, + group_id=self.group_ids[index], + block_hash=self.block_hashes[index], + parent_block_hash=self.parent_block_hashes[index], + token_ids=self.token_ids[index], + block_size=self.block_sizes[index], + ) + + def localized(self, keys: list[str]) -> BlockStoreBatch: + return replace(self, keys=keys) @dataclass @@ -529,6 +626,8 @@ class LoadSpec: local_tokens: int external_tokens: int + # Group block indices start at the beginning of the request. + block_hashes_by_group: tuple[tuple[bytes | None, ...], ...] = () @property def num_tokens_to_load(self) -> int: @@ -581,22 +680,41 @@ class TransferJobStatus(str, Enum): RUNNING = "running" COMPLETED = "completed" FAILED = "failed" + CANCELLED = "cancelled" @dataclass class TransferJobState: """Per-key completion state; one failed key does not fail its siblings.""" - plans: tuple[BlockTransferPlan, ...] + plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch | BlockStoreBatch status: TransferJobStatus = TransferJobStatus.PENDING completed_keys: set[str] = field(default_factory=set) failed_keys: set[str] = field(default_factory=set) error: str | None = None + plan_bytes: tuple[int, ...] | None = None @property def keys(self) -> Sequence[str]: + if isinstance(self.plans, (BlockLoadBatch, BlockStoreBatch)): + return self.plans.keys return tuple(plan.key for plan in self.plans) + @property + def completed_bytes(self) -> int: + if self.plan_bytes is not None: + return sum( + size + for key, size in zip(self.keys, self.plan_bytes, strict=True) + if key in self.completed_keys + ) + return sum( + item.length + for plan in self.plans + if plan.key in self.completed_keys + for item in plan.ranges + ) + def start(self) -> None: if self.status != TransferJobStatus.PENDING: raise RuntimeError(f"cannot start job in state {self.status}") @@ -611,8 +729,20 @@ def fail(self, keys: Sequence[str], error: str | None = None) -> None: self.error = error self._finish_if_done() + def cancel(self, error: str = "cancelled") -> None: + self.error = error + self.status = TransferJobStatus.CANCELLED + @property def failed_block_ids(self) -> set[int]: + if isinstance(self.plans, (BlockLoadBatch, BlockStoreBatch)): + return { + block_id + for key, block_id in zip( + self.plans.keys, self.plans.block_ids, strict=True + ) + if key in self.failed_keys + } return {plan.block_id for plan in self.plans if plan.key in self.failed_keys} def _finish_if_done(self) -> None: @@ -631,20 +761,36 @@ class UMBPConnectorMetadata(KVConnectorMetadata): async_load: bool = False store_plans: list[BlockTransferPlan] = field(default_factory=list) - load_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) - store_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) + load_requests: dict[str, list[BlockTransferPlan] | BlockLoadBatch] = field( + default_factory=dict + ) + store_requests: dict[str, list[BlockTransferPlan] | BlockStoreBatch] = field( + default_factory=dict + ) preempted_request_ids: set[str] = field(default_factory=set) store_event: int = -1 + def iter_store_plans(self): + yield from self.store_plans + for plans in self.store_requests.values(): + yield from plans + + def store_plan_groups(self): + if self.store_plans: + yield self.store_plans + yield from self.store_requests.values() + @dataclass class StoreEventResult: """Number of workers that finished one store event.""" completed_workers: int = 0 + failed_objects: set[tuple[int, bytes]] = field(default_factory=set) def merge(self, other: StoreEventResult) -> None: self.completed_workers += other.completed_workers + self.failed_objects.update(other.failed_objects) @dataclass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py index db6a32d84672..fb44355e02e3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py @@ -35,11 +35,23 @@ def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None: def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: """Start loading the requested plans.""" + def load_blocks( + self, plans: Sequence[BlockTransferPlan] + ) -> TransferJobState | None: + """Optionally load logical blocks; None requests ordinary materialization.""" + return None + def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: """Start storing the requested plans.""" + def store_blocks( + self, plans: Sequence[BlockTransferPlan] + ) -> TransferJobState | None: + """Optionally store logical blocks; None requests ordinary materialization.""" + return None + def wait(self, job: TransferJobState) -> TransferJobState: - """Block until the job is final and its buffers are safe to reuse.""" + """Return a final state only when buffers are safe to reuse; otherwise raise.""" def poll(self, job: TransferJobState) -> TransferJobState | None: """Return a finished job without blocking, or None if still pending.""" @@ -47,6 +59,12 @@ def poll(self, job: TransferJobState) -> TransferJobState | None: def publish(self, job: TransferJobState) -> None: """Accept a completed store; a runtime may defer visibility until then.""" + def cancel(self, job: TransferJobState) -> TransferJobState: + """Cancel if possible, otherwise wait until buffers are safe to reuse.""" + + def take_evicted_keys(self) -> Sequence[str]: + """Return locally published keys evicted since the last call.""" + def close(self) -> None: """Release worker-side resources.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py index dd7ecadb167a..a7e88e3a4308 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py @@ -5,30 +5,52 @@ from __future__ import annotations import logging -from collections.abc import Sequence -from dataclasses import dataclass, field +import time +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field, replace from functools import lru_cache -from typing import Any +from itertools import islice +from typing import Any, cast from vllm.config import VllmConfig +from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.logger import init_logger from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import ( + get_kv_cache_coordinator, + replay_boundaries, +) from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import ( + BlockHash, + KVCacheBlock, + get_block_hash, + make_block_hash_with_group_id, + maybe_convert_block_hash, resolve_dcp_kv_block_size, resolve_kv_cache_block_sizes, ) from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, KVCacheConfig, + KVCacheSpec, + MambaSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + iter_layer_specs, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from .data import ( BlockIdentityCodec, + BlockLoadBatch, + BlockStoreBatch, BlockTransferPlan, LoadSpec, RankTopology, @@ -53,11 +75,98 @@ def _rank_keys( return codec.keys_for_topology(block_hash, topology, (group_id,)) +def _unwrap_store_spec(spec: KVCacheSpec) -> KVCacheSpec: + if isinstance(spec, UniformTypeKVCacheSpecs): + return next(iter(spec.kv_cache_specs.values())) + return spec + + @dataclass class _StoreEvent: - plans: tuple[BlockTransferPlan, ...] + plan_groups: tuple[Sequence[BlockTransferPlan], ...] result: StoreEventResult = field(default_factory=StoreEventResult) + def iter_plans(self): + for plans in self.plan_groups: + yield from plans + + def iter_keys(self): + for plans in self.plan_groups: + if isinstance(plans, BlockStoreBatch): + yield from plans.keys + else: + yield from (plan.key for plan in plans) + + def block_ids(self): + for plans in self.plan_groups: + if isinstance(plans, BlockStoreBatch): + yield from plans.block_ids + else: + yield from (plan.block_id for plan in plans) + + +@dataclass(frozen=True) +class _LookupContext: + request: Request + local_tokens: int + num_tokens: int + hashes: tuple[bytes, ...] + objects: dict[int, list[tuple[bytes, tuple[str, ...]]]] + keys: list[str] + + +@dataclass(frozen=True) +class _PendingLookup: + context: _LookupContext + future: Future[Sequence[bool]] + + +@dataclass +class _EventObject: + serial: int + plan: BlockTransferPlan | None + published: bool = False + + +@dataclass(frozen=True) +class _EventLookup: + objects: tuple[tuple[tuple[int, bytes], int], ...] + future: Future[Sequence[bool]] + + +class _LookupBlockPool(BlockPool): + """Request-scoped residency view for the core hit-window algorithms.""" + + def __init__(self, hash_block_size: int) -> None: + super().__init__( + num_gpu_blocks=1, + enable_caching=True, + hash_block_size=hash_block_size, + enable_kv_cache_events=False, + ) + self.hits: list[set[bytes]] = [] + self.hit_blocks: dict[bytes, KVCacheBlock] = {} + + def get_cached_block( + self, block_hash: bytes, kv_cache_group_ids: list[int] + ) -> list[KVCacheBlock] | None: + if any(block_hash not in self.hits[group] for group in kv_cache_group_ids): + return None + block = self.hit_blocks.get(block_hash) + if block is None: + block = KVCacheBlock( + block_id=len(self.hit_blocks) + 1, + _block_hash=make_block_hash_with_group_id( + BlockHash(block_hash), kv_cache_group_ids[0] + ), + ) + self.hit_blocks[block_hash] = block + return [block] * len(kv_cache_group_ids) + + def clear_lookup(self) -> None: + self.hits = [] + self.hit_blocks.clear() + class UMBPStoreConnectorScheduler: """Own vLLM prefix matching while the runtime owns lookup transport. @@ -79,44 +188,183 @@ def __init__( ) self.kv_cache_config = kv_cache_config dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self._dcp_size = dcp_size + self._retention_interval = kv_cache_config.prefix_cache_retention_interval self.group_block_sizes = { group_id: resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_size) for group_id, group in enumerate(kv_cache_config.kv_cache_groups) } + # Mamba "align" block tables are not append-only: a position may hold + # the live state the next step writes, or a speculative block core will + # relocate, so positional stores would read or pin the wrong block. + # Core hands off each committed boundary state instead + # (_add_boundary_state_plans). + self._boundary_handoff_groups = frozenset( + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if isinstance(group.kv_cache_spec, MambaSpec) + and getattr(group.kv_cache_spec, "mamba_cache_mode", None) == "align" + ) self.runtime = runtime + self.enable_kv_cache_events = bool( + vllm_config.kv_events_config + and vllm_config.kv_events_config.enable_kv_cache_events + ) self.codec = codec self.topology = topology or RankTopology() transfer_config = vllm_config.kv_transfer_config assert transfer_config is not None extra = transfer_config.kv_connector_extra_config self.load_async = bool(extra.get("load_async", True)) + if not self.load_async and ( + len(kv_cache_config.kv_cache_groups) > 1 + or any( + isinstance(spec, MambaSpec) + or getattr(spec, "cache_role", None) == SparseCacheRole.INDEXER + for group in kv_cache_config.kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ): + # Core accepts only request-level load failures for several cache + # groups, and acts on them only while the request waits for an + # asynchronous load. Mamba-style and sparse-attention indexer + # layers also read restored state before any wait_for_layer_load. + logger.warning( + "UMBP loads asynchronously for models with several KV cache " + "groups, Mamba-style or sparse-attention indexer layers; " + "load_async=false is ignored" + ) + self.load_async = True self.enable_lookup = bool(extra.get("enable_lookup", True)) + self.lookup_async = bool(extra.get("lookup_async", False)) self.save_decode_cache = bool(extra.get("save_decode_cache", False)) # A kv_consumer only restores from the pool, as with other store # connectors; it never offloads its own KV. self.store_enabled = getattr(transfer_config, "kv_role", None) != "kv_consumer" + self.enable_partial_hash_hits = bool( + extra.get("enable_partial_hash_hits", False) + ) # Request.block_hashes are computed at this size; keys index into them. self.hash_block_size = resolved_hash_block_size + lookup_groups = [ + kv_cache_config.kv_cache_groups[group_id] + for group_id in kv_cache_config.prefix_cacheable_group_ids + ] + model_config = getattr(vllm_config, "model_config", None) speculative_config = getattr(vllm_config, "speculative_config", None) - self.use_eagle = bool( + use_eagle = bool( speculative_config is not None and speculative_config.use_eagle_block_drop() ) - self._pending_loads: dict[str, list[BlockTransferPlan]] = {} + self.use_eagle = use_eagle + explicit_eagle_groups = { + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if group.is_eagle_group + } + self._eagle_store_groups = ( + explicit_eagle_groups + if explicit_eagle_groups or not use_eagle + else set(kv_cache_config.prefix_cacheable_group_ids) + ) + max_model_len = getattr( + model_config, + "max_model_len", + kv_cache_config.num_blocks * self.block_size, + ) + lookup_config = replace( + kv_cache_config, + kv_cache_groups=lookup_groups, + num_blocks=1, + ) + self._lookup_coordinator = ( + get_kv_cache_coordinator( + kv_cache_config=lookup_config, + max_model_len=max_model_len, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + use_eagle=use_eagle, + enable_caching=True, + enable_kv_cache_events=False, + dcp_world_size=dcp_size, + pcp_world_size=getattr( + vllm_config.parallel_config, + "prefill_context_parallel_size", + 1, + ), + scheduler_block_size=self.block_size, + hash_block_size=self.hash_block_size, + # As core's scheduler builds its coordinator. + num_prefill_lookahead=vllm_config.num_prefill_lookahead_tokens, + allow_partial_hash_hits=self.enable_partial_hash_hits, + ) + if lookup_groups + and ( + len(lookup_groups) > 1 + or not isinstance( + _unwrap_store_spec(lookup_groups[0].kv_cache_spec), + FullAttentionSpec, + ) + ) + else None + ) + self._lookup_pool = _LookupBlockPool(self.hash_block_size) + self._local_cache_lookup: Callable[[int, bytes], bool] | None = None + if self._lookup_coordinator is not None: + self._lookup_coordinator.block_pool = self._lookup_pool + for manager in self._lookup_coordinator.single_type_managers: + manager.block_pool = self._lookup_pool + manager._null_block = self._lookup_pool.null_block + self._pending_loads: dict[str, list[BlockTransferPlan] | BlockLoadBatch] = {} self._load_specs: dict[str, LoadSpec] = {} self._lookup_units = { - group_id: self.group_block_sizes[group_id] + group_id: ( + self.hash_block_size + if self.enable_partial_hash_hits + and ( + self._lookup_coordinator is None + or self._lookup_coordinator.enable_partial_hash_hits + ) + else self.group_block_sizes[group_id] + ) for group_id in kv_cache_config.prefix_cacheable_group_ids } + self._pending_lookups: dict[str, _PendingLookup] = {} + self._lookup_executor = ThreadPoolExecutor( + max_workers=int(extra.get("lookup_workers", 2)), + thread_name_prefix="umbp-lookup", + ) self._requests: dict[str, Request] = {} self._request_trackers: dict[str, RequestTracker] = {} + self._pending_finished_stores: dict[str, list[BlockTransferPlan]] = {} self._gpu_block_pool: BlockPool | None = None self._store_events: dict[int, _StoreEvent] = {} + self._kv_events: list[KVCacheEvent] = [] + self._event_objects: dict[tuple[int, bytes], _EventObject] = {} + self._pending_event_objects: dict[tuple[int, bytes], int] = {} + self._event_lookup: _EventLookup | None = None + self._event_serial = 0 + self._next_event_reconcile = 0.0 + self._event_max_objects = extra.get("kv_event_max_objects", 65536) + if type(self._event_max_objects) is not int or self._event_max_objects <= 0: + raise ValueError("kv_event_max_objects must be a positive integer") + self._event_check_batch_size = 256 + self._event_publications_suppressed = 0 self._num_workers = getattr(vllm_config.parallel_config, "world_size", 1) self._store_event_counter = 0 def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: self._gpu_block_pool = gpu_block_pool + def bind_local_cache_lookup(self, lookup: Callable[[int, bytes], bool]) -> None: + self._local_cache_lookup = lookup + + def _has_local_block(self, group_id: int, block_hash: bytes) -> bool: + if self._local_cache_lookup is not None: + return self._local_cache_lookup(group_id, block_hash) + pool = self._gpu_block_pool + return pool is not None and bool( + pool.get_cached_block(cast(BlockHash, block_hash), [group_id]) + ) + def bind_namespace( self, runtime: UMBPSchedulerHandle, codec: BlockIdentityCodec ) -> None: @@ -133,26 +381,58 @@ def get_num_new_matched_tokens( "UMBP lookup is disabled until workers report their namespace" ) return 0, False - hashes = request.block_hashes - if not hashes or request.num_tokens < self.block_size: + hashes = tuple(request.block_hashes) + pending = self._pending_lookups.get(request.request_id) + if pending is not None: + context = pending.context + if ( + context.request is not request + or context.local_tokens != num_computed_tokens + or context.num_tokens != request.num_tokens + or context.hashes != hashes + ): + self._pending_lookups.pop(request.request_id) + pending.future.cancel() + pending = None + align = ( + self.hash_block_size if self.enable_partial_hash_hits else self.block_size + ) + if not hashes or request.num_tokens < align: return 0, False if num_computed_tokens % self.block_size != 0: return 0, False group_ids = self.kv_cache_config.prefix_cacheable_group_ids + if not group_ids: + return 0, False # vLLM must execute at least the final prompt token. Returning a hit # for the entire prompt leaves the scheduler with no new token to run. max_external_token = request.num_tokens - 1 if max_external_token <= num_computed_tokens: self._load_specs.pop(request.request_id, None) return 0, False - logical_objects_by_group, keys = self._lookup_objects( - request, num_computed_tokens, hashes - ) + if pending is not None and not pending.future.done(): + return None, False + if pending is None: + context = self._build_lookup_context(request, num_computed_tokens, hashes) + if self.lookup_async: + self._pending_lookups[request.request_id] = _PendingLookup( + context, + self._lookup_executor.submit(self.runtime.lookup, context.keys), + ) + return None, False + else: + context = self._pending_lookups.pop(request.request_id).context try: - hits = list(self.runtime.lookup(keys)) + hits = list( + pending.future.result() + if pending is not None + else self.runtime.lookup(context.keys) + ) except Exception as exc: logger.debug("UMBP lookup failed request=%s: %s", request.request_id, exc) return 0, False + keys = context.keys + logical_objects_by_group = context.objects if len(hits) != len(keys): logger.debug( "UMBP lookup returned an invalid result length request=%s", @@ -174,16 +454,39 @@ def get_num_new_matched_tokens( for index in range(0, len(group_hits), per_rank) ] ) - group_id = group_ids[0] - matched_blocks = 0 - for hit in logical_hits_by_group[group_id]: - if not hit: - break - matched_blocks += 1 - if self.use_eagle: - # As core's coordinator: the drafter recomputes the last block. - matched_blocks = max(0, matched_blocks - 1) - need_to_load = matched_blocks * self._lookup_units[group_id] + if self._lookup_coordinator is None: + group_id = group_ids[0] + group_hits = logical_hits_by_group[group_id] + units_per_block = ( + self.group_block_sizes[group_id] // self._lookup_units[group_id] + ) + matched_units = 0 + for index in range(units_per_block - 1, len(group_hits), units_per_block): + if not group_hits[index]: + break + matched_units = index + 1 + # Interior hashes are sparse: only the saved tail needs to exist. + for index in range( + min(matched_units + units_per_block - 1, len(group_hits)) - 1, + matched_units - 1, + -1, + ): + if group_hits[index]: + matched_units = index + 1 + break + if self.use_eagle: + # As core's coordinator: the drafter recomputes the last block. + matched_units = max(0, matched_units - units_per_block) + need_to_load = matched_units * self._lookup_units[group_id] + block_hashes_by_group: tuple[tuple[bytes | None, ...], ...] = () + else: + need_to_load, block_hashes_by_group = self._coordinate_external_hits( + list(context.hashes), + num_computed_tokens, + max_external_token - num_computed_tokens, + logical_objects_by_group, + logical_hits_by_group, + ) if logger.isEnabledFor(logging.DEBUG): logger.debug( "UMBP lookup result request=%s keys=%s group_hits=%s load_tokens=%d", @@ -198,24 +501,99 @@ def get_num_new_matched_tokens( self._load_specs[request.request_id] = LoadSpec( local_tokens=num_computed_tokens, external_tokens=matched_tokens, + block_hashes_by_group=block_hashes_by_group, ) - return need_to_load, self.load_async + load_async = self.load_async + if block_hashes_by_group and not any( + block_hash is not None + for group_id, group_hashes in zip( + group_ids, block_hashes_by_group, strict=True + ) + for index, block_hash in enumerate(group_hashes) + if index >= num_computed_tokens // self.group_block_sizes[group_id] + ): + load_async = False + return need_to_load, load_async - def _lookup_objects( - self, request: Request, local_tokens: int, hashes: Sequence[bytes] - ) -> tuple[dict[int, list[tuple[bytes, tuple[str, ...]]]], list[str]]: + def _build_lookup_context( + self, request: Request, local_tokens: int, hashes: tuple[bytes, ...] + ) -> _LookupContext: assert self.codec is not None objects: dict[int, list[tuple[bytes, tuple[str, ...]]]] = {} keys: list[str] = [] + lookup_limit = request.num_tokens - 1 + if self.use_eagle and self._lookup_coordinator is not None: + lookup_limit = min(request.num_tokens, len(hashes) * self.hash_block_size) for group_id, unit in self._lookup_units.items(): group_objects = [] - for token_end in range(local_tokens + unit, request.num_tokens, unit): + for token_end in range(local_tokens + unit, lookup_limit + 1, unit): block_hash = self._object_hash_at_token_end(hashes, token_end) rank_keys = _rank_keys(self.codec, self.topology, block_hash, group_id) group_objects.append((block_hash, rank_keys)) keys.extend(rank_keys) objects[group_id] = group_objects - return objects, keys + return _LookupContext( + request, local_tokens, request.num_tokens, hashes, objects, keys + ) + + def _coordinate_external_hits( + self, + request_hashes: list[bytes], + local_tokens: int, + max_external_tokens: int, + logical_objects_by_group: dict[int, list[tuple[bytes, tuple[str, ...]]]], + logical_hits_by_group: dict[int, list[bool]], + ) -> tuple[int, tuple[tuple[bytes | None, ...], ...]]: + """Apply the core KV cache coordinator to runtime lookup results.""" + coordinator = self._lookup_coordinator + assert coordinator is not None + pool = self._lookup_pool + try: + pool.hits = [ + { + block_hash + for (block_hash, _), hit in zip( + logical_objects_by_group[group_id], + logical_hits_by_group[group_id], + strict=True, + ) + if hit + } + for group_id in self.kv_cache_config.prefix_cacheable_group_ids + ] + if local_tokens: + for group_index, group_id in enumerate( + self.kv_cache_config.prefix_cacheable_group_ids + ): + for token_end in range( + self._lookup_units[group_id], + local_tokens + 1, + self._lookup_units[group_id], + ): + block_hash = self._object_hash_at_token_end( + request_hashes, token_end + ) + if self._has_local_block(group_id, block_hash): + pool.hits[group_index].add(block_hash) + hit_blocks, hit_length, _ = coordinator.find_longest_cache_hit( + cast(list[BlockHash], request_hashes), + local_tokens + max_external_tokens, + ) + block_hashes_by_group = tuple( + tuple( + None + if block.is_null or block.block_hash is None + else get_block_hash(block.block_hash) + for block in group_blocks + ) + for group_blocks in hit_blocks + ) + return max(0, hit_length - local_tokens), block_hashes_by_group + finally: + self._reset_lookup_coordinator() + + def _reset_lookup_coordinator(self) -> None: + self._lookup_pool.clear_lookup() def update_state_after_alloc( self, @@ -258,11 +636,17 @@ def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: meta = UMBPConnectorMetadata(async_load=self.load_async) + reserved_store_keys = { + key for event in self._store_events.values() for key in event.iter_keys() + } for request_id in scheduler_output.finished_req_ids: self._pending_loads.pop(request_id, None) self._load_specs.pop(request_id, None) self._request_trackers.pop(request_id, None) self._requests.pop(request_id, None) + pending = self._pending_lookups.pop(request_id, None) + if pending is not None: + pending.future.cancel() for request in scheduler_output.scheduled_new_reqs: load_plans = self._pending_loads.pop(request.req_id, []) self._load_specs.pop(request.req_id, None) @@ -280,25 +664,26 @@ def build_connector_meta( for group_id in (self.kv_cache_config.prefix_cacheable_group_ids) ) store_plans = ( - [] + None if not self.store_enabled or tracked_request is None else self._store_plans( tracked_request, tracker, total_tokens, block_ids_override=selected_block_ids, + skip_keys=reserved_store_keys, ) ) logger.debug( "UMBP store planning request=%s tokens=%d plans=%d keys=%s", request.req_id, total_tokens, - len(store_plans), - tuple(plan.key for plan in store_plans), + len(store_plans) if store_plans is not None else 0, + tuple(store_plans.keys) if store_plans is not None else (), ) - meta.store_plans.extend(store_plans) if store_plans: meta.store_requests[request.req_id] = store_plans + reserved_store_keys.update(store_plans.keys) block_state = getattr(scheduler_output, "kv_connector_block_state", None) cached_reqs = scheduler_output.scheduled_cached_reqs @@ -359,10 +744,11 @@ def build_connector_meta( cached_tracker, total_tokens, block_ids_override=cached_tracker.block_ids, + skip_keys=reserved_store_keys, ) - meta.store_plans.extend(store_plans) if store_plans: meta.store_requests[request_id] = store_plans + reserved_store_keys.update(store_plans.keys) # Async loads are admitted without scheduling model tokens, so those # requests are absent from both scheduled request collections above. @@ -380,23 +766,115 @@ def build_connector_meta( if preempted_tracker := self._request_trackers.get(request_id): preempted_tracker.reset() - event_plans = tuple(meta.store_plans) - if event_plans: + if ( + self.store_enabled + and block_state is not None + and block_state.boundary_state_offloads + ): + self._add_boundary_state_plans( + block_state.boundary_state_offloads, + meta, + set(scheduler_output.finished_req_ids), + set(scheduler_output.preempted_req_ids or ()), + ) + + finished_stores = self._pending_finished_stores + self._pending_finished_stores = {} + for plans in finished_stores.values(): + meta.store_plans.extend(plans) + event_plan_groups = tuple(meta.store_plan_groups()) + if event_plan_groups: event = self._store_event_counter self._store_event_counter += 1 meta.store_event = event - self._store_events[event] = _StoreEvent(event_plans) + self._store_events[event] = _StoreEvent(event_plan_groups) if self._gpu_block_pool is not None: self._gpu_block_pool.touch( [ self._gpu_block_pool.blocks[block_id] for block_id in dict.fromkeys( - plan.block_id for plan in event_plans + block_id + for plans in event_plan_groups + for block_id in ( + plans.block_ids + if isinstance(plans, BlockStoreBatch) + else (plan.block_id for plan in plans) + ) ) ] ) + # Transfer ownership from request-finish pins to the event. + for plans in finished_stores.values(): + self._gpu_block_pool.free_blocks( + self._gpu_block_pool.blocks[block_id] + for block_id in dict.fromkeys(plan.block_id for plan in plans) + ) return meta + def _add_boundary_state_plans( + self, + offloads: dict[str, list[tuple[int, int, int]]], + meta: UMBPConnectorMetadata, + finished: set[str], + preempted: set[str], + ) -> None: + """Store exact Mamba/hybrid boundary blocks handed off by core.""" + if self.codec is None: + return + prefix_groups = set(self.kv_cache_config.prefix_cacheable_group_ids) + for request_id, entries in offloads.items(): + if request_id in finished or request_id in preempted: + logger.debug( + "UMBP boundary offloads dropped request=%s entries=%d: %s", + request_id, + len(entries), + "finished" if request_id in finished else "preempted", + ) + continue + request = self._requests.get(request_id) + tracker = self._request_trackers.get(request_id) + if request is None or tracker is None: + logger.debug( + "UMBP boundary offloads dropped request=%s entries=%d: untracked", + request_id, + len(entries), + ) + continue + plans: list[BlockTransferPlan] = [] + for group_id, block_id, boundary_tokens in entries: + if group_id not in prefix_groups or block_id == NULL_BLOCK_ID: + continue + if boundary_tokens <= 0: + continue + hash_index = boundary_tokens // self.hash_block_size - 1 + if not 0 <= hash_index < len(request.block_hashes): + continue + block_hash = request.block_hashes[hash_index] + plans.append( + BlockTransferPlan( + key=self.codec.key(block_hash, group_id), + block_id=block_id, + request_id=request_id, + group_id=group_id, + block_hash=block_hash, + parent_block_hash=( + request.block_hashes[hash_index - 1] + if self.enable_kv_cache_events and hash_index > 0 + else None + ), + block_size=self.group_block_sizes[group_id], + ) + ) + logger.debug( + "UMBP boundary offloads request=%s offered=%d stored=%d boundaries=%s", + request_id, + len(entries), + len(plans), + sorted({entry[2] for entry in entries}), + ) + if plans: + meta.store_plans.extend(plans) + def _tracker_for_request(self, request_id: str) -> RequestTracker: return self._request_trackers.setdefault(request_id, RequestTracker()) @@ -406,32 +884,64 @@ def _load_plans_for_external_tokens( tracker: RequestTracker, block_groups: tuple[list[int], ...], num_external_tokens: int, - ) -> list[BlockTransferPlan]: - """Plan the logical objects selected by the external hit window.""" + ) -> BlockLoadBatch: + """Batch the logical objects selected by the external hit window.""" assert self.codec is not None group_ids = self.kv_cache_config.prefix_cacheable_group_ids spec = tracker.load_spec local_tokens = spec.local_tokens if spec is not None else 0 end_tokens = local_tokens + num_external_tokens - plans: list[BlockTransferPlan] = [] + keys: list[str] = [] + destinations: list[int] = [] + groups: list[int] = [] + if spec is not None and spec.block_hashes_by_group: + for group_index, group_id in enumerate(group_ids): + group_block_size = self.group_block_sizes[group_id] + start_block = local_tokens // group_block_size + for block_index, block_hash in enumerate( + spec.block_hashes_by_group[group_index] + ): + if block_index < start_block or block_hash is None: + continue + block_id = block_groups[group_index][block_index] + if block_id == NULL_BLOCK_ID: + raise RuntimeError( + "UMBP external hit mapped a runtime object to a null " + f"destination block for group {group_id}" + ) + keys.append(self.codec.key(block_hash, group_id=group_id)) + destinations.append(block_id) + groups.append(group_id) + return BlockLoadBatch(keys, destinations, groups, request.request_id) for group_index, group_id in enumerate(group_ids): group_block_size = self.group_block_sizes[group_id] start_block = local_tokens // group_block_size num_full_blocks = end_tokens // group_block_size for block_index in range(start_block, num_full_blocks): token_end = (block_index + 1) * group_block_size - block_hash = self._object_hash_at_token_end( - request.block_hashes, token_end - ) - plans.append( - BlockTransferPlan( - key=self.codec.key(block_hash, group_id=group_id), - block_id=block_groups[group_index][block_index], - request_id=request.request_id, + keys.append( + self.codec.key( + self._object_hash_at_token_end(request.block_hashes, token_end), group_id=group_id, ) ) - return plans + destinations.append(block_groups[group_index][block_index]) + groups.append(group_id) + partial_tokens = end_tokens % self.block_size + if partial_tokens and self.enable_partial_hash_hits: + hash_index = end_tokens // self.hash_block_size - 1 + for group_index, group_id in enumerate(group_ids): + block_index = end_tokens // self.group_block_sizes[group_id] + group_block_ids = block_groups[group_index] + if ( + block_index >= len(group_block_ids) + or group_block_ids[block_index] == NULL_BLOCK_ID + ): + raise RuntimeError("UMBP partial hit has no destination block") + keys.append(self.codec.key(request.block_hashes[hash_index], group_id)) + destinations.append(group_block_ids[block_index]) + groups.append(group_id) + return BlockLoadBatch(keys, destinations, groups, request.request_id) def _object_hash_at_token_end( self, hashes: Sequence[bytes], token_end: int @@ -449,11 +959,12 @@ def _store_plans( tracker: RequestTracker, token_count: int, block_ids_override: tuple[list[int], ...] | None = None, - ) -> list[BlockTransferPlan]: + skip_keys: set[str] | None = None, + ) -> BlockStoreBatch: """Describe full blocks produced by a scheduled prefill.""" - if self.codec is None: - return [] request_id = getattr(request, "request_id", None) or request.req_id + if self.codec is None: + return BlockStoreBatch([], [], [], [], [], [], [], request_id) group_ids = self.kv_cache_config.prefix_cacheable_group_ids # Scheduled speculative tokens are not yet verified and have no block # hash; a block is storable only once its hash exists. @@ -462,39 +973,186 @@ def _store_plans( # Hash-aligned rather than block-aligned: a group with smaller blocks # than the largest must not wait for that group's block to fill. save_to = tracker.mark_saved(token_count, self.hash_block_size) - plans: list[BlockTransferPlan] = [] + keys: list[str] = [] + source_block_ids: list[int] = [] + plan_group_ids: list[int] = [] + block_hashes: list[bytes] = [] + parent_block_hashes: list[bytes | None] = [] + token_ids_by_block: list[tuple[int, ...]] = [] + block_sizes: list[int] = [] block_groups = ( block_ids_override if block_ids_override is not None else tuple(request.block_ids[group_id] for group_id in group_ids) ) - block_ranges = { + candidate_ranges = { group_id: range( previous_saved_tokens // self.group_block_sizes[group_id], min(save_to // self.group_block_sizes[group_id], len(block_ids)), ) for group_id, block_ids in zip(group_ids, block_groups) + if group_id not in self._boundary_handoff_groups + } + block_ranges = { + group_id: tuple( + index + for relative_index, index in enumerate(indices) + if mask is None or mask[relative_index] + ) + for group_id, indices in candidate_ranges.items() + for mask in (self._store_mask(request, group_id, indices),) } + resident_blocks = ( + self._restored_store_blocks( + request, block_ranges, tracker.load_spec.external_tokens, block_groups + ) + if tracker.load_spec is not None + else {} + ) for group_id, block_ids in zip(group_ids, block_groups): if group_id not in block_ranges: continue group_block_size = self.group_block_sizes[group_id] + resident = resident_blocks.get(group_id, ()) for index in block_ranges[group_id]: - if block_ids[index] == NULL_BLOCK_ID: + if block_ids[index] == NULL_BLOCK_ID or index in resident: continue token_end = (index + 1) * group_block_size block_hash = self._object_hash_at_token_end( request.block_hashes, token_end ) - plans.append( - BlockTransferPlan( - key=self.codec.key(block_hash, group_id=group_id), - block_id=block_ids[index], - request_id=request_id, - group_id=group_id, + key = self.codec.key(block_hash, group_id=group_id) + if skip_keys is not None and key in skip_keys: + continue + keys.append(key) + source_block_ids.append(block_ids[index]) + plan_group_ids.append(group_id) + block_hashes.append(block_hash) + parent_block_hashes.append( + self._object_hash_at_token_end( + request.block_hashes, token_end - group_block_size ) + if self.enable_kv_cache_events and index > 0 + else None + ) + token_ids_by_block.append( + tuple(request.all_token_ids[index * group_block_size : token_end]) + if self.enable_kv_cache_events + else () ) - return plans + block_sizes.append(group_block_size) + return BlockStoreBatch( + keys, + source_block_ids, + plan_group_ids, + block_hashes, + parent_block_hashes, + token_ids_by_block, + block_sizes, + request_id, + ) + + def _restored_store_blocks( + self, + request: Any, + block_ranges: Mapping[int, Sequence[int]], + restored_tokens: int, + block_groups: tuple[list[int], ...], + ) -> dict[int, set[int]]: + """Query residency before constructing plans for restored blocks.""" + assert self.runtime is not None and self.codec is not None + # block_ranges covers only positionally stored groups; block_groups + # follows prefix_cacheable_group_ids. + blocks_of = dict( + zip( + self.kv_cache_config.prefix_cacheable_group_ids, + block_groups, + strict=True, + ) + ) + candidates = { + group_id: [ + index + for index in indices + if index < restored_tokens // self.group_block_sizes[group_id] + if blocks_of[group_id][index] != NULL_BLOCK_ID + ] + for group_id, indices in block_ranges.items() + } + keys = [ + key + for group_id, indices in candidates.items() + for index in indices + for key in _rank_keys( + self.codec, + self.topology, + self._object_hash_at_token_end( + request.block_hashes, (index + 1) * self.group_block_sizes[group_id] + ), + group_id, + ) + ] + if not keys: + return {} + try: + hits = list(self.runtime.lookup(keys)) + except Exception as exc: + logger.debug("UMBP store residency check failed: %s", exc) + return {} + if len(hits) != len(keys): + return {} + ranks = self.topology.rank_count + offset = 0 + resident = {} + for group_id, indices in candidates.items(): + resident[group_id] = { + index + for position, index in enumerate(indices) + if all( + hits[offset + position * ranks : offset + (position + 1) * ranks] + ) + } + offset += len(indices) * ranks + return resident + + def _store_mask( + self, + request: Any, + group_id: int, + indices: range, + ) -> list[bool] | None: + """Keep only blocks that core's cache semantics can later reuse.""" + if not indices: + return [] + group = self.kv_cache_config.kv_cache_groups[group_id] + spec = _unwrap_store_spec(group.kv_cache_spec) + manager_cls = KVCacheSpecRegistry.get_manager_class(spec) + if manager_cls is None: + return None + num_prompt_tokens = getattr(request, "num_prompt_tokens", None) + reachable_boundaries = ( + () + if num_prompt_tokens is None + else replay_boundaries( + num_prompt_tokens, self.block_size, use_eagle=self.use_eagle + ) + ) + mask = manager_cls.reachable_block_mask( + start_block=indices.start, + end_block=indices.stop, + alignment_tokens=self.block_size, + kv_cache_spec=spec, + use_eagle=group_id in self._eagle_store_groups, + retention_interval=self._retention_interval, + reachable_boundaries=reachable_boundaries, + dcp_world_size=self._dcp_size, + ) + if mask is not None and len(mask) != len(indices): + raise RuntimeError( + f"UMBP store mask for group {group_id} returned {len(mask)} " + f"entries for {len(indices)} blocks" + ) + return mask def request_finished( self, request: Request, block_ids: tuple[list[int], ...] @@ -502,8 +1160,89 @@ def request_finished( self._pending_loads.pop(request.request_id, None) self._load_specs.pop(request.request_id, None) self._requests.pop(request.request_id, None) + pending = self._pending_lookups.pop(request.request_id, None) + if pending is not None: + pending.future.cancel() return False, None + def register_finished_partial_tail( + self, + request: Request, + block_ids: tuple[list[int], ...], + partial_tail_offloads: list[tuple[int, int, int]], + ) -> bool: + """Pin finished boundary pages without delaying core request cleanup.""" + pool = self._gpu_block_pool + if ( + not self.store_enabled + or self.codec is None + or not partial_tail_offloads + or pool is None + ): + return False + if request.request_id in self._pending_finished_stores: + return False + prefix_groups = set(self.kv_cache_config.prefix_cacheable_group_ids) + tracker = self._request_trackers.get(request.request_id) + if tracker is None: + return False + boundaries = {entry[2] for entry in partial_tail_offloads} + if len(boundaries) != 1: + logger.warning("UMBP partial tail entries span several boundaries") + return False + boundary = boundaries.pop() + if ( + boundary <= 0 + or boundary % self.hash_block_size != 0 + or boundary != request.num_computed_tokens + or bool(getattr(request, "num_in_flight_tokens", 0)) + ): + return False + hash_index = boundary // self.hash_block_size - 1 + if hash_index >= len(request.block_hashes): + return False + sources = { + group_id: block_id for group_id, block_id, _ in partial_tail_offloads + } + if not sources.keys() <= prefix_groups: + return False + # Core hands off Mamba state pages. Include the matching FA tail page + # so every group can resolve the same fine-grained prefix boundary. + for group_id in prefix_groups: + spec = self.kv_cache_config.kv_cache_groups[group_id].kv_cache_spec + if ( + isinstance(spec, FullAttentionSpec) + and boundary % self.group_block_sizes[group_id] + ): + index = boundary // self.group_block_sizes[group_id] + if index >= len(block_ids[group_id]): + return False + sources.setdefault(group_id, block_ids[group_id][index]) + plans: list[BlockTransferPlan] = [] + for group_id, block_id in sources.items(): + if block_id == NULL_BLOCK_ID or not 0 <= block_id < len(pool.blocks): + return False + if pool.blocks[block_id].is_null: + return False + group_block_size = self.group_block_sizes[group_id] + plans.append( + BlockTransferPlan( + request_id=request.request_id, + block_id=block_id, + group_id=group_id, + key=self.codec.key(request.block_hashes[hash_index], group_id), + block_hash=request.block_hashes[hash_index], + block_size=group_block_size, + ) + ) + # A Mamba page is a state snapshot, not a token-addressable byte array. + # Store full pages for both state and attention, matching restore. + pool.touch( + [pool.blocks[block_id] for block_id in dict.fromkeys(sources.values())] + ) + self._pending_finished_stores[request.request_id] = plans + return False + def update_connector_output(self, output: KVConnectorOutput) -> None: metadata = output.kv_connector_worker_meta if not isinstance(metadata, UMBPConnectorWorkerMetadata): @@ -518,21 +1257,204 @@ def update_connector_output(self, output: KVConnectorOutput) -> None: continue del self._store_events[event_id] if pool is not None: - block_ids = dict.fromkeys(plan.block_id for plan in event.plans) + block_ids = dict.fromkeys(event.block_ids()) pool.free_blocks( pool.blocks[block_id] for block_id in reversed(block_ids) ) + if self.enable_kv_cache_events: + self._queue_event_objects(event) + self._reconcile_kv_events() + + def _queue_event_objects(self, event: _StoreEvent) -> None: + seen: set[tuple[int, bytes]] = set() + for plan in event.iter_plans(): + if plan.group_id is None or plan.block_hash is None: + continue + identity = (plan.group_id, plan.block_hash) + if identity in seen or identity in event.result.failed_objects: + continue + seen.add(identity) + previous = self._event_objects.get(identity) + if previous is None and len(self._event_objects) >= self._event_max_objects: + self._event_publications_suppressed += 1 + logger.warning_once( + "UMBP KV event tracking limit reached; suppressing new " + "BlockStored events until tracked objects are evicted" + ) + continue + self._event_serial += 1 + self._event_objects[identity] = _EventObject( + self._event_serial, plan, previous.published if previous else False + ) + self._pending_event_objects[identity] = self._event_serial + + def invalidate_kv_events(self, events: Sequence[KVCacheEvent]) -> None: + removed = { + (event.group_idx, block_hash) + for event in events + if isinstance(event, BlockRemoved) + for block_hash in event.block_hashes + } + if not removed: + return + for identity in list(self._event_objects): + group_id, block_hash = identity + if ( + group_id, + maybe_convert_block_hash(cast(BlockHash, block_hash)), + ) in removed: + del self._event_objects[identity] + self._pending_event_objects.pop(identity, None) + + def _reconcile_kv_events(self) -> None: + if ( + not self.enable_kv_cache_events + or self.runtime is None + or self.codec is None + ): + return + pending = self._event_lookup + if pending is not None: + if not pending.future.done(): + return + self._event_lookup = None + try: + hits = list(pending.future.result()) + if len(hits) != len(pending.objects) * self.topology.rank_count: + raise ValueError("invalid UMBP event residency result length") + except Exception as exc: + logger.debug("UMBP event residency check failed: %s", exc) + hits = None + for index, (identity, serial) in enumerate(pending.objects): + record = self._event_objects.get(identity) + if record is None or record.serial != serial: + continue + # Rotate completed checks so a bounded batch visits every object. + del self._event_objects[identity] + self._pending_event_objects.pop(identity, None) + if hits is None: + if record.published: + self._event_objects[identity] = record + continue + ranks = self.topology.rank_count + resident = all(hits[index * ranks : (index + 1) * ranks]) + if not resident: + if record.published: + group_id, block_hash = identity + self._kv_events.append( + BlockRemoved( + block_hashes=[ + maybe_convert_block_hash( + cast(BlockHash, block_hash) + ) + ], + medium="CPU", + group_idx=group_id, + locality="LOCAL", + ) + ) + continue + if not record.published: + plan = record.plan + assert plan is not None and plan.block_hash is not None + self._kv_events.append( + BlockStored( + block_hashes=[ + maybe_convert_block_hash( + cast(BlockHash, plan.block_hash) + ) + ], + parent_block_hash=( + maybe_convert_block_hash( + cast(BlockHash, plan.parent_block_hash) + ) + if plan.parent_block_hash is not None + else None + ), + token_ids=list(plan.token_ids), + block_size=plan.block_size, + lora_id=None, + medium=plan.medium, + lora_name=None, + group_idx=plan.group_id, + locality="LOCAL", + ) + ) + record.published = True + record.plan = None + self._event_objects[identity] = record + if not self._event_objects: + return + now = time.monotonic() + candidates = [] + if now >= self._next_event_reconcile: + candidates = list( + islice( + ( + (identity, record.serial) + for identity, record in self._event_objects.items() + if record.published + ), + max(1, self._event_check_batch_size // 2), + ) + ) + self._next_event_reconcile = now + 1.0 + selected = {identity for identity, _ in candidates} + candidates.extend( + islice( + ( + (identity, serial) + for identity, serial in self._pending_event_objects.items() + if identity not in selected + ), + self._event_check_batch_size - len(candidates), + ) + ) + if not candidates: + return + keys = [ + key + for (group_id, block_hash), _ in candidates + for key in _rank_keys(self.codec, self.topology, block_hash, group_id) + ] + self._event_lookup = _EventLookup( + tuple(candidates), + self._lookup_executor.submit(self.runtime.lookup, keys), + ) + + def take_events(self) -> list[KVCacheEvent]: + self._reconcile_kv_events() + events = self._kv_events + self._kv_events = [] + return events def has_pending_push_work(self) -> bool: - return bool(self._store_events) + return ( + bool(self._pending_finished_stores) + or bool(self._store_events) + or self._event_lookup is not None + or bool(self._pending_event_objects) + ) def reset_store(self) -> bool: - if self._store_events: + if self._store_events or self._pending_finished_stores: return False + for pending in self._pending_lookups.values(): + pending.future.cancel() + self._pending_lookups.clear() self._load_specs.clear() self._pending_loads.clear() - return True if self.runtime is None else self.runtime.clear() + cleared = True if self.runtime is None else self.runtime.clear() + if cleared: + self._event_objects.clear() + self._pending_event_objects.clear() + if self._event_lookup is not None: + self._event_lookup.future.cancel() + self._event_lookup = None + return cleared def close(self) -> None: + self._lookup_executor.shutdown(wait=True, cancel_futures=True) + self._pending_lookups.clear() if self.runtime is not None: self.runtime.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/stats.py new file mode 100644 index 000000000000..08bce6676259 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/stats.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared UMBP connector transfer statistics.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) + + +@dataclass +class UMBPStoreConnectorStats(KVConnectorStats): + """Serializable interval counters shared by every UMBP runtime.""" + + data: dict[str, Any] = field(default_factory=dict) + + def record( + self, + operation: str, + *, + submitted: int = 0, + completed: int = 0, + failed: int = 0, + num_bytes: int = 0, + ) -> None: + entry = self.data.setdefault( + operation, + { + "submitted": 0, + "completed": 0, + "failed": 0, + "num_bytes": 0, + }, + ) + entry["submitted"] += submitted + entry["completed"] += completed + entry["failed"] += failed + entry["num_bytes"] += num_bytes + + def reset(self) -> None: + self.data.clear() + + def aggregate(self, other: KVConnectorStats) -> UMBPStoreConnectorStats: + if not isinstance(other, UMBPStoreConnectorStats): + raise TypeError("cannot aggregate incompatible UMBP stats") + result = UMBPStoreConnectorStats() + for operation, values in [*self.data.items(), *other.data.items()]: + result.record( + operation, + submitted=values.get("submitted", 0), + completed=values.get("completed", 0), + failed=values.get("failed", 0), + num_bytes=values.get("num_bytes", 0), + ) + return result + + def reduce(self) -> dict[str, int | float]: + result: dict[str, int | float] = {} + for operation, values in self.data.items(): + for key, value in values.items(): + result[f"{operation}_{key}"] = value + return result + + def is_empty(self) -> bool: + return not self.data + + +class UMBPStorePromMetrics(KVConnectorPromMetrics): + """Prometheus counters for shared UMBP transfer outcomes.""" + + def __init__( + self, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> None: + super().__init__( + vllm_config, + metric_types, + labelnames, + per_engine_labelvalues, + ) + labels = labelnames + ["operation"] + self._submitted = self._counter_cls( + name="vllm:umbp_transfer_submitted_total", + documentation="Number of UMBP objects submitted.", + labelnames=labels, + ) + self._completed = self._counter_cls( + name="vllm:umbp_transfer_completed_total", + documentation="Number of UMBP objects completed.", + labelnames=labels, + ) + self._failed = self._counter_cls( + name="vllm:umbp_transfer_failed_total", + documentation="Number of UMBP objects failed.", + labelnames=labels, + ) + self._bytes = self._counter_cls( + name="vllm:umbp_transfer_bytes_total", + documentation="Bytes in successfully completed UMBP transfer ranges.", + labelnames=labels, + ) + + def observe( + self, + transfer_stats_data: dict[str, Any] | None, + engine_idx: int = 0, + ) -> None: + if not transfer_stats_data: + return + for operation, values in transfer_stats_data.items(): + labels = self.per_engine_labelvalues[engine_idx] + [operation] + self._submitted.labels(*labels).inc(values.get("submitted", 0)) + self._completed.labels(*labels).inc(values.get("completed", 0)) + self._failed.labels(*labels).inc(values.get("failed", 0)) + self._bytes.labels(*labels).inc(values.get("num_bytes", 0)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py index 34101f294c51..ccd2fb0a79ba 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py @@ -4,15 +4,20 @@ from __future__ import annotations -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace +from typing import cast import torch +from vllm.distributed.kv_events import BlockRemoved, KVCacheEvent from vllm.forward_context import ForwardContext from vllm.logger import init_logger +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from .data import ( BlockIdentityCodec, + BlockLoadBatch, + BlockStoreBatch, BlockTransferPlan, KVLayoutPlanner, StoreEventResult, @@ -22,6 +27,7 @@ UMBPConnectorWorkerMetadata, ) from .runtime import UMBPWorkerHandle +from .stats import UMBPStoreConnectorStats logger = init_logger(__name__) @@ -30,6 +36,7 @@ class _StoreBatch: event_id: int jobs: dict[str, TransferJobState] + failed_objects: set[tuple[int, bytes]] = field(default_factory=set) class UMBPStoreConnectorWorker: @@ -41,18 +48,26 @@ def __init__( layout: KVLayoutPlanner | None = None, *, codec: BlockIdentityCodec | None = None, + report_failed_requests: bool = False, + enable_kv_cache_events: bool = False, ) -> None: self.runtime = runtime self.layout = layout self.codec = codec + self.enable_kv_cache_events = enable_kv_cache_events self._load_jobs: dict[str, dict[str | None, TransferJobState]] = {} self._store_jobs: dict[str, TransferJobState] = {} self._pending_stores: list[_StoreBatch] = [] self._worker_meta = UMBPConnectorWorkerMetadata() + self._kv_events: list[KVCacheEvent] = [] self._finished_recving: set[str] = set() + # Core maps block-level load failures to requests only with one group. + self._report_failed_requests = report_failed_requests + self._failed_recving: set[str] = set() self._load_error_block_ids: set[int] = set() self._report_load_completions = False self._active_store_event = -1 + self._stats = UMBPStoreConnectorStats() def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: if self.layout is not None and self.layout.regions: @@ -98,9 +113,35 @@ def start_load_kv( for request_id, plans in metadata.load_requests.items(): if plans: self._submit_loads(request_id, plans) + if not metadata.async_load: + self._drain_load_jobs(wait=True) - def _submit_loads(self, request_id: str, plans: list[BlockTransferPlan]) -> None: - materialized = self._materialize_plans(plans) + def _submit_loads( + self, request_id: str, plans: list[BlockTransferPlan] | BlockLoadBatch + ) -> None: + load_blocks = getattr(self.runtime, "load_blocks", None) + if load_blocks is not None: + localized: list[BlockTransferPlan] | BlockLoadBatch + if isinstance(plans, BlockLoadBatch): + localized = plans + if self.codec is not None: + localized = replace( + plans, + keys=[ + self._localize_key(key, group) + for key, group in zip( + plans.keys, plans.group_ids, strict=True + ) + ], + ) + else: + localized = [self._localize_plan(plan) for plan in plans] + job = load_blocks(localized) + if job is not None: + self._stats.record("load", submitted=len(plans)) + self._load_jobs[request_id] = {None: job} + return + materialized = self._materialize_plans(list(plans)) logger.debug( "UMBP load submitted request=%s plans=%d ranges=%d bytes=%d", request_id, @@ -108,12 +149,19 @@ def _submit_loads(self, request_id: str, plans: list[BlockTransferPlan]) -> None sum(len(plan.ranges) for plan in materialized), sum(item.length for plan in materialized for item in plan.ranges), ) + self._stats.record("load", submitted=len(materialized)) self._load_jobs[request_id] = {None: self.runtime.load(materialized)} def _finish_job( self, request_id: str, result: TransferJobState, *, is_load: bool ) -> bool: operation = "load" if is_load else "store" + self._stats.record( + operation, + completed=len(result.completed_keys), + failed=len(result.failed_keys), + num_bytes=result.completed_bytes, + ) succeeded = result.status == TransferJobStatus.COMPLETED if not succeeded: logger.debug( @@ -124,7 +172,10 @@ def _finish_job( result.error, ) if is_load: - self._load_error_block_ids.update(result.failed_block_ids) + if not self._report_failed_requests: + self._load_error_block_ids.update(result.failed_block_ids) + elif not succeeded: + self._failed_recving.add(request_id) elif succeeded: self.runtime.publish(result) return succeeded @@ -151,54 +202,106 @@ def _drain_store_jobs(self, *, wait: bool) -> None: result = self.runtime.wait(job) if wait else self.runtime.poll(job) if result is None: continue - # A partly failed job is not published, including its stored keys. - self._finish_job(request_id, result, is_load=False) + # A partly failed job is not published. + if not self._finish_job(request_id, result, is_load=False): + batch.failed_objects.update( + (plan.group_id, plan.block_hash) + for plan in result.plans + if plan.group_id is not None and plan.block_hash is not None + ) del batch.jobs[request_id] if not batch.jobs: if batch.event_id >= 0: self._worker_meta.store_events[batch.event_id] = StoreEventResult( - completed_workers=1 + completed_workers=1, + failed_objects=batch.failed_objects, ) self._pending_stores.remove(batch) def _submit_store_plans( - self, plans: list[BlockTransferPlan], request_id: str + self, plans: list[BlockTransferPlan] | BlockStoreBatch, request_id: str ) -> None: if not plans: return - materialized = self._materialize_plans(plans) + store_blocks = getattr(self.runtime, "store_blocks", None) + if store_blocks is not None: + localized: list[BlockTransferPlan] | BlockStoreBatch + if isinstance(plans, BlockStoreBatch): + localized = plans + if self.codec is not None: + localized = plans.localized( + [ + self._localize_key(key, group_id) + for key, group_id in zip( + plans.keys, plans.group_ids, strict=True + ) + ] + ) + else: + localized = [self._localize_plan(plan) for plan in plans] + job = store_blocks(localized) + if job is not None: + self._stats.record("store", submitted=len(plans)) + self._store_jobs[request_id] = job + return + materialized = self._materialize_plans(list(plans)) + self._stats.record( + "store", + submitted=len(materialized), + ) self._store_jobs[request_id] = self.runtime.store(materialized) def enqueue_stores(self, metadata: UMBPConnectorMetadata) -> None: self._active_store_event = metadata.store_event - store_requests = { - request_id: list(plans) - for request_id, plans in metadata.store_requests.items() + store_requests = dict(metadata.store_requests) + grouped_plan_ids = { + id(plan) + for plans in store_requests.values() + if not isinstance(plans, BlockStoreBatch) + for plan in plans } - grouped = {plan for plans in store_requests.values() for plan in plans} for plan in metadata.store_plans: - if plan not in grouped: - store_requests.setdefault( - plan.request_id or "__umbp_batch__", [] - ).append(plan) - grouped.add(plan) + if id(plan) not in grouped_plan_ids: + request_id = plan.request_id or "__umbp_batch__" + if isinstance(store_requests.get(request_id), BlockStoreBatch): + request_id = f"__umbp_flat__:{request_id}" + pending = store_requests.setdefault(request_id, []) + assert isinstance(pending, list) + pending.append(plan) + grouped_plan_ids.add(id(plan)) for request_id, plans in store_requests.items(): self._submit_store_plans(plans, request_id) def handle_preemptions(self, metadata: UMBPConnectorMetadata) -> None: - """Finish request-local jobs before vLLM reuses their GPU blocks.""" + """Cancel request-local jobs before vLLM reuses their GPU blocks.""" preempted = metadata.preempted_request_ids if not preempted: return for request_id in preempted: for job in self._load_jobs.pop(request_id, {}).values(): - self.runtime.wait(job) + self._cancel_job(job) for jobs in [self._store_jobs, *(batch.jobs for batch in self._pending_stores)]: for request_id, job in jobs.items(): - if request_id in preempted or any( - plan.request_id in preempted for plan in job.plans + plan_request_id = ( + job.plans.request_id + if isinstance(job.plans, (BlockLoadBatch, BlockStoreBatch)) + else None + ) + if ( + request_id in preempted + or plan_request_id in preempted + or ( + plan_request_id is None + and any(plan.request_id in preempted for plan in job.plans) + ) ): - jobs[request_id] = self.runtime.wait(job) + jobs[request_id] = self._cancel_job(job) + + def _cancel_job(self, job: TransferJobState) -> TransferJobState: + cancel = getattr(self.runtime, "cancel", None) + if callable(cancel): + return cancel(job) + return self.runtime.wait(job) def get_finished( self, finished_req_ids: set[str] @@ -223,16 +326,61 @@ def _drain_load_jobs(self, *, wait: bool) -> None: if self._report_load_completions: self._finished_recving.add(request_id) + def get_failed_recving(self) -> set[str]: + failed = self._failed_recving - self._load_jobs.keys() + self._failed_recving.difference_update(failed) + return failed + def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata: result = self._worker_meta self._worker_meta = UMBPConnectorWorkerMetadata() return result + def get_kv_events(self) -> list[KVCacheEvent]: + take_evicted = getattr(self.runtime, "take_evicted_keys", None) + if callable(take_evicted): + for key in take_evicted(): + if not self.enable_kv_cache_events: + continue + parsed = self._parse_key(key) + if parsed is None: + continue + group_id, block_hash = parsed + self._kv_events.append( + BlockRemoved( + block_hashes=[ + maybe_convert_block_hash(cast(BlockHash, block_hash)) + ], + medium="CPU", + group_idx=group_id, + locality="LOCAL", + ) + ) + events = self._kv_events + self._kv_events = [] + return events + + @staticmethod + def _parse_key(key: str) -> tuple[int, bytes] | None: + try: + group_part, hash_hex = key.rsplit(":", 1) + group_id = int(group_part.rsplit(":g", 1)[1]) + return group_id, bytes.fromhex(hash_hex) + except (IndexError, ValueError): + return None + def get_block_ids_with_load_errors(self) -> set[int]: failed = self._load_error_block_ids self._load_error_block_ids = set() return failed + def get_kv_connector_stats(self) -> UMBPStoreConnectorStats | None: + if self._stats.is_empty(): + return None + result = UMBPStoreConnectorStats(data=dict(self._stats.data)) + self._stats.reset() + return result + def close(self) -> None: self.wait_for_save() self._drain_store_jobs(wait=True) From 5f63ea461ccd567ece6f50dd3701430c0f321039 Mon Sep 17 00:00:00 2001 From: Sandeep Agrawal Date: Sat, 24 Oct 2026 05:28:36 -0700 Subject: [PATCH 020/631] [Bugfix][Quantization][Spec Decode] Load ModelOpt mixed-precision DFlash and DSpark drafts (#63001) Signed-off-by: Sandeep Agrawal Co-authored-by: Kotaro Yamamoto Co-authored-by: Codex Co-authored-by: vllm-agent --- .buildkite/test_areas/quantization.yaml | 5 + docs/features/quantization/modelopt.md | 9 + tests/config/test_model_arch_config.py | 24 ++ tests/quantization/test_modelopt.py | 319 ++++++++++++++++++ tests/v1/spec_decode/test_dflash2.py | 21 -- .../layers/quantization/modelopt.py | 15 +- vllm/model_executor/models/qwen3_dflash.py | 46 +-- vllm/model_executor/models/utils.py | 56 +++ .../model_arch_config_convertor.py | 2 + 9 files changed, 442 insertions(+), 55 deletions(-) diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 50cfc2ee86fc..2d7d083950b7 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -10,6 +10,9 @@ steps: source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/utils.py + - vllm/transformers_utils/model_arch_config_convertor.py - vllm/models/deepseek_v41/quant_config.py - vllm/utils/humming.py - tests/quantization @@ -33,6 +36,8 @@ steps: source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/utils.py - tests/quantization - vllm/_aiter_ops.py - vllm/platforms/rocm.py diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index a62426152713..603f90a1b1d3 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -20,6 +20,15 @@ following `quantization.quant_algo` values: - `W4A16_NVFP4`: weight-only ModelOpt NVFP4 checkpoints with 16-bit activations (use `quantization="modelopt_fp4"`). - `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`). +- `MIXED_PRECISION`: per-layer formats specified by `quantized_layers` + (use `quantization="modelopt_mixed"`). + +For per-tensor `FP8`, vLLM converts E4M3FN weights and scales to E4M3FNUZ +when required by the GPU, for example on AMD MI300. + +For Qwen3-based DFlash and DSpark drafts, `quantized_layers` and exclusion lists +use the exported checkpoint's module names. vLLM maps these names to the draft's +module prefix and decoder layer indices before constructing quantized layers. !!! note For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index a9409be09c6e..33f1a854a3ad 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -269,6 +269,30 @@ def test_legacy_modelopt_config_without_producer_is_normalized(): assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4" +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("quant_algo", ["MIXED_PRECISION", "mixed_precision"]) +def test_modelopt_mixed_precision_config_is_normalized(legacy, quant_algo): + quantized_layers = {"model.layers.0.mlp.down_proj": {"quant_algo": "FP8"}} + quantization = { + "quant_algo": quant_algo, + "quantized_layers": quantized_layers, + } + quantization_config = {"quantization": quantization} + if legacy: + quantization["modelopt_quant_config"] = {"quant_cfg": {}} + else: + quantization_config["producer"] = {"name": "modelopt"} + hf_config = PreTrainedConfig(quantization_config=quantization_config) + convertor = ModelArchConfigConvertorBase(hf_config, hf_config) + + normalized = convertor.get_quantization_config() + + assert normalized["quant_method"] == "modelopt_mixed" + assert normalized["quantization"]["quantized_layers"] == { + "model.layers.0.mlp.down_proj": {"quant_algo": "FP8"} + } + + def _layer(**overrides) -> ModelArchitectureConfig: fields = dict( architectures=["X"], diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 9a300efa6c5e..15603f6e58e8 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -5,8 +5,10 @@ Run `pytest tests/quantization/test_modelopt.py`. """ +import json import os from contextlib import nullcontext +from copy import deepcopy from types import SimpleNamespace from typing import Any, NoReturn from unittest.mock import MagicMock, Mock, patch @@ -365,6 +367,102 @@ 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 +@pytest.mark.parametrize("prefix", ["", "model"]) +@pytest.mark.parametrize("start_layer_id", [0, 1]) +def test_draft_quant_config_uses_runtime_module_names(prefix, start_layer_id): + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + from vllm.model_executor.models.utils import map_draft_quant_config + + config = _mixed_precision_config( + { + "encoder.fc": {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"}, + "markov_head.markov_w2": {"quant_algo": "FP8"}, + **{ + f"layers.{index}.self_attn.q_proj": {"quant_algo": "FP8"} + for index in range(3) + }, + "midlayer.mlp.gate_proj": {"quant_algo": "FP8"}, + "midlayer.mlp.up_proj": {"quant_algo": "FP8"}, + } + ) + config.exclude_modules = [ + "layers.0.mlp_conv*", + "*layers.2.self_attn.k_proj", + "layers.88.already_global", + "lilicorr.layers.0.mlp.0", + ] + native_layers = config.quantized_layers + native_names = tuple(native_layers) + + map_draft_quant_config( + config, DFlashQwen3Model.hf_to_vllm_mapper, prefix, start_layer_id, 3 + ) + + assert tuple(native_layers) == native_names + root = f"{prefix}." if prefix else "" + assert config.quantized_layers == { + f"{root}fc": {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"}, + f"{root}markov_head.markov_w2": {"quant_algo": "FP8"}, + **{ + f"{root}layers.{start_layer_id + index}.self_attn.q_proj": { + "quant_algo": "FP8" + } + for index in range(3) + }, + f"{root}layers.{start_layer_id}.mlp.gate_proj": {"quant_algo": "FP8"}, + f"{root}layers.{start_layer_id}.mlp.up_proj": {"quant_algo": "FP8"}, + } + assert config.is_layer_excluded(f"{root}layers.{start_layer_id}.mlp_conv.weight") + assert config.is_layer_excluded( + f"{root}layers.{start_layer_id + 2}.self_attn.k_proj" + ) + assert config.exclude_modules.count(f"{root}layers.88.already_global") == 1 + assert f"{root}lilicorr.layers.0.mlp.0" in config.exclude_modules + if start_layer_id: + nested_module = f"{root}lilicorr.layers.{start_layer_id}.mlp.0" + assert nested_module not in config.exclude_modules + + +def test_draft_quant_config_preserves_quark_settings(): + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + from vllm.model_executor.models.utils import map_draft_quant_config + + fp8 = {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False} + config = QuarkConfig( + { + "global_quant_config": {"weight": fp8, "input_tensors": fp8}, + "layer_type_quant_config": {}, + "exclude": ["layers.0.mlp_conv*"], + } + ) + original = deepcopy(config.quant_config) + + map_draft_quant_config(config, DFlashQwen3Model.hf_to_vllm_mapper, "model", 88, 3) + + assert config.quant_config == original + + +@pytest.mark.parametrize("prefix", ["", "model."]) +@pytest.mark.parametrize( + "source,target", + [ + ("encoder.fc", "fc"), + ("encoder.output_norm_enc", "hidden_norm"), + ("midlayer", "layers.0"), + ("encoder.fcx", "encoder.fcx"), + ("encoder.output_norm_enc_extra", "encoder.output_norm_enc_extra"), + ("midlayer_extra", "midlayer_extra"), + ], +) +def test_draft_weight_aliases_match_module_boundaries(prefix, source, target): + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + mapper = DFlashQwen3Model.hf_to_vllm_mapper + assert mapper.map_name(f"{prefix}{source}") == f"{prefix}{target}" + assert mapper.map_name(f"{prefix}{source}.weight") == f"{prefix}{target}.weight" + + def test_modelopt_mixed_precision_composes_gemma4_mappers(): from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.gemma4_mm import ( @@ -521,6 +619,227 @@ def check_model(model): print(f"ModelOpt FP8 output: {output}") +_TINY_MODEL_CONFIG = dict( + vocab_size=128, + hidden_size=128, + intermediate_size=256, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=64, +) + + +def _pack_fp8_weight(weights: dict[str, torch.Tensor], name: str, algorithm: str): + value = weights[f"{name}.weight"].float() + per_channel = algorithm == "FP8_PER_CHANNEL_PER_TOKEN" + amax = value.abs().amax(dim=1) if per_channel else value.abs().max() + scale = amax / torch.finfo(torch.float8_e4m3fn).max + divisor = scale[:, None] if per_channel else scale + weights[f"{name}.weight"] = (value / divisor).to(torch.float8_e4m3fn) + weights[f"{name}.weight_scale"] = scale + if not per_channel: + weights[f"{name}.input_scale"] = torch.tensor(1.0 / 448) + + +def test_modelopt_nested_mixed_checkpoint(tmp_path, monkeypatch, vllm_runner): + """Native and flattened ModelOpt metadata load the same FP8/BF16 weights.""" + from transformers import LlamaConfig, LlamaForCausalLM + + from vllm import SamplingParams + + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + config = LlamaConfig(num_hidden_layers=1, **_TINY_MODEL_CONFIG) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + model = LlamaForCausalLM(config).to(torch.bfloat16) + weights = model.state_dict() + prefix = "model.layers.0.mlp.down_proj" + _pack_fp8_weight(weights, prefix, "FP8") + quantization: dict[str, Any] = { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "quantized_layers": {prefix: {"quant_algo": "FP8"}}, + "exclude_modules": [], + } + metadata = { + "nested": { + "producer": {"name": "modelopt", "version": "0.49.0"}, + "quantization": quantization, + }, + "flat": {"quant_method": "modelopt", **quantization}, + } + + def check_model(loaded): + mlp = loaded.model.layers[0].mlp + quantized = mlp.down_proj + assert isinstance(quantized.quant_method, ModelOptLinearMethod) + assert quantized.weight.dtype == current_platform.fp8_dtype() + assert torch.isfinite(quantized.weight.float()).all() + assert torch.isfinite(quantized.weight_scale).all() + scale = (2 if current_platform.is_fp8_fnuz() else 1) / 448 + torch.testing.assert_close( + quantized.input_scale, + torch.full_like(quantized.input_scale, scale), + rtol=0, + atol=0, + ) + assert isinstance(mlp.gate_up_proj.quant_method, UnquantizedLinearMethod) + assert mlp.gate_up_proj.weight.dtype == torch.bfloat16 + + params = SamplingParams(temperature=0, max_tokens=4, ignore_eos=True) + outputs = [] + for name, quantization_config in metadata.items(): + model_path = tmp_path / name + model.config.quantization_config = quantization_config + model.save_pretrained(model_path, state_dict=dict(weights)) + (model_path / "hf_quant_config.json").write_text( + json.dumps(quantization_config) + ) + with vllm_runner( + str(model_path), + skip_tokenizer_init=True, + trust_remote_code=False, + enforce_eager=True, + enable_flashinfer_autotune=False, + max_model_len=32, + max_num_seqs=1, + kv_cache_memory_bytes=16 << 20, + ) as llm: + llm.apply_model(check_model) + outputs.append(llm.generate([[1, 2, 3]], params)[0][0][0]) + assert len(outputs[0]) == 7 + assert outputs[1] == outputs[0] + + +def test_modelopt_mixed_dspark_checkpoint(tmp_path, monkeypatch, vllm_runner): + """Native draft names select FP8 methods, including offset decoder layers.""" + from safetensors.torch import save_file + from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3Model + + from tests.v1.e2e.spec_decode.utils import get_spec_decode_metric_value + from vllm import SamplingParams + + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + target, draft = tmp_path / "target", tmp_path / "draft" + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + teacher = LlamaForCausalLM( + LlamaConfig(num_hidden_layers=4, **_TINY_MODEL_CONFIG) + ) + teacher.to(torch.bfloat16).save_pretrained(target) + config = Qwen3Config(num_hidden_layers=1, **_TINY_MODEL_CONFIG) + weights = dict(Qwen3Model(config).to(torch.bfloat16).state_dict()) + del weights["embed_tokens.weight"] + for name, shape in ( + ("fc.weight", (128, 384)), + ("markov_head.markov_w1.weight", (128, 128)), + ("markov_head.markov_w2.weight", (128, 128)), + ): + weights[name] = torch.randn(shape, dtype=torch.bfloat16) * 0.01 + weights["hidden_norm.weight"] = torch.ones(128, dtype=torch.bfloat16) + algorithms = { + "fc": "FP8_PER_CHANNEL_PER_TOKEN", + "markov_head.markov_w2": "FP8", + "layers.0.mlp.down_proj": "FP8", + } + for name, algorithm in algorithms.items(): + _pack_fp8_weight(weights, name, algorithm) + metadata = { + "producer": {"name": "modelopt", "version": "0.49.0"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "quantized_layers": { + name: {"quant_algo": algorithm} + for name, algorithm in algorithms.items() + }, + "exclude_modules": [], + }, + } + config.update( + dict( + architectures=["Qwen3DSparkModel"], + block_size=16, + target_layer_ids=[0, 1, 2], + markov_rank=128, + sample_from_anchor=True, + dflash_config={ + "target_layer_ids": [0, 1, 2], + "mask_token_id": 127, + "causal": False, + }, + quantization_config=metadata, + ) + ) + config.save_pretrained(draft) + save_file(weights, str(draft / "model.safetensors")) + (draft / "hf_quant_config.json").write_text(json.dumps(metadata)) + + def check_draft(worker): + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + kFp8StaticTokenSym, + ) + + model = worker.get_draft_model().model + for layer, weight_key, activation_key in ( + (model.fc, kFp8StaticTokenSym, kFp8DynamicTokenSym), + (model.markov_head.markov_w2, kFp8StaticTensorSym, kFp8StaticTensorSym), + (model.layers[0].mlp.down_proj, kFp8StaticTensorSym, kFp8StaticTensorSym), + ): + assert isinstance(layer.quant_method, ModelOptLinearMethod) + assert layer.quant_method.spec.weight == weight_key + assert layer.quant_method.spec.activation == activation_key + assert layer.weight.dtype == current_platform.fp8_dtype() + assert torch.isfinite(layer.weight.float()).all() + assert torch.isfinite(layer.weight_scale).all() + assert layer.weight_scale.numel() == (128 if layer is model.fc else 1) + if activation_key == kFp8StaticTensorSym: + scale = (2 if current_platform.is_fp8_fnuz() else 1) / 448 + torch.testing.assert_close( + layer.input_scale, + torch.full_like(layer.input_scale, scale), + rtol=0, + atol=0, + ) + assert model.markov_head.markov_w1.weight.dtype == torch.bfloat16 + assert isinstance( + model.layers[0].self_attn.qkv_proj.quant_method, UnquantizedLinearMethod + ) + + params = SamplingParams(temperature=0, max_tokens=8, ignore_eos=True) + outputs = [] + for speculation in ( + None, + dict(method="dspark", model=str(draft), num_speculative_tokens=2), + ): + with vllm_runner( + str(target), + speculative_config=speculation, + skip_tokenizer_init=True, + trust_remote_code=False, + dtype="bfloat16", + enforce_eager=True, + enable_flashinfer_autotune=False, + disable_log_stats=False, + max_model_len=32, + max_num_seqs=1, + kv_cache_memory_bytes=16 << 20, + ) as llm: + if speculation: + llm.collective_rpc(check_draft) + outputs.append(llm.generate([[1, 5, 7, 9]], params)[0][0][0]) + if speculation: + after = get_spec_decode_metric_value( + llm.llm.get_metrics(), "vllm:spec_decode_num_draft_tokens" + ) + assert after > 0 + assert len(outputs[0]) == 12 + assert outputs[1] == outputs[0] + + @pytest.mark.skipif( not is_quant_method_supported("modelopt"), reason="ModelOpt FP8 is not supported on this GPU type.", diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index e94ae3591fc3..dc35802f2f60 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -6,9 +6,6 @@ import pytest import torch -from vllm.model_executor.models.qwen3_dflash import ( - _add_global_draft_layer_exclusions, -) from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges from vllm.platforms import current_platform from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator @@ -132,24 +129,6 @@ def test_grouped_conv_triton_matches_reference( ) -def test_draft_quant_exclusions_include_global_layer_indices(): - quant_config = SimpleNamespace( - exclude_modules=[ - "layers.0.mlp_conv*", - "*layers.4.self_attn.q_proj", - "layers.88.already_global", - "lilicorr.layers.0.mlp.0", - ] - ) - - _add_global_draft_layer_exclusions(quant_config, 88, 5) - - assert "layers.88.mlp_conv*" in quant_config.exclude_modules - assert "*layers.92.self_attn.q_proj" in quant_config.exclude_modules - assert quant_config.exclude_modules.count("layers.88.already_global") == 1 - assert "lilicorr.layers.88.mlp.0" not in quant_config.exclude_modules - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("rows,heads", [(8, 1), (64, 2), (256, 1)]) def test_beam_fused_context_kv(rows, heads): diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 309ead57cf37..7121e13f6c1e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -99,6 +99,7 @@ kNvfp4Static, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( + normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -116,6 +117,7 @@ replace_parameter, set_weight_attrs, ) +from vllm.platforms import current_platform from vllm.utils.b12x import B12X_BLOCK_CODECS from vllm.utils.math_utils import cdiv @@ -2164,10 +2166,17 @@ def create_weights(self, layer, role, ctx, shapes, wl) -> None: def process(self, layer, role) -> None: if role is WEIGHT: weight = layer.weight - max_w_scale = layer.weight_scale.max() - if not (layer.weight_scale == layer.weight_scale[0]).all(): + weight_scale = layer.weight_scale + if current_platform.is_fp8_fnuz() and weight.dtype == torch.float8_e4m3fn: + weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight, weight_scale, getattr(layer, "input_scale", None) + ) + if input_scale is not None: + replace_parameter(layer, "input_scale", input_scale) + max_w_scale = weight_scale.max() + if not (weight_scale == weight_scale[0]).all(): max_w_scale, weight = requantize_with_max_scale( - layer.weight, layer.weight_scale, layer.logical_widths + weight, weight_scale, layer.logical_widths ) # Transposed here rather than in the kernel: ModelOpt maps each # fp8 key to exactly one kernel, so the layout is key-determined. diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 485ea952fccb..b0d23b119e68 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -48,6 +48,7 @@ AutoWeightsLoader, WeightsMapper, get_draft_quant_config, + map_draft_quant_config, maybe_prefix, process_eagle_weight, ) @@ -56,31 +57,6 @@ _SLIDING_ATTENTION = "sliding_attention" -_DRAFT_LAYER_PATTERN = re.compile(r"(? None: - """Add runtime layer aliases for checkpoint-local quant exclusions.""" - if quant_config is None or start_layer_id == 0: - return - exclusions = getattr(quant_config, "exclude_modules", None) - if not isinstance(exclusions, list): - return - - def offset_local_layer(match: re.Match[str]) -> str: - layer_idx = int(match.group(1)) - if layer_idx >= num_hidden_layers: - return match.group(0) - return f"layers.{layer_idx + start_layer_id}" - - for exclusion in tuple(exclusions): - global_exclusion = _DRAFT_LAYER_PATTERN.sub(offset_local_layer, exclusion) - if global_exclusion != exclusion and global_exclusion not in exclusions: - exclusions.append(global_exclusion) def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: @@ -412,13 +388,17 @@ class DFlashQwen3Model(nn.Module): decoder_layer_cls = DFlashQwen3DecoderLayer hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={ - "midlayer.": "layers.0.", + orig_to_new_regex={ + # Match a complete module name, including bare names and trailing + # wildcards in quantization configs, but not e.g. midlayer_extra. + re.compile(r"(? "QuantizationConfig | Non ) +def map_draft_quant_config( + quant_config: "QuantizationConfig | None", + weight_mapper: WeightsMapper, + prefix: str, + start_layer_id: int, + num_hidden_layers: int, +) -> None: + """Map draft quantization settings to runtime module names.""" + if quant_config is None: + return + from vllm.model_executor.layers.quantization.modelopt import ModelOptQuantConfigBase + + if isinstance(quant_config, ModelOptQuantConfigBase): + quant_config.apply_vllm_mapper(weight_mapper.get_rename_mapper()) + root = rf"(?:{re.escape(prefix)}\.)?" if prefix else "" + # Match root-level layers.N, optionally prefixed or wildcarded. Offset + # higher IDs first so an offset of 1 cannot cascade from 0 -> 1 -> 2. + runtime_names = { + re.compile(rf"^({root}\*?)layers\.{index}(?=\.|\*|$)"): ( + rf"\g<1>layers.{start_layer_id + index}" + ) + for index in reversed(range(num_hidden_layers)) + if start_layer_id + } + if prefix: + # Preserve leading wildcards and names already rooted at prefix. + runtime_names[re.compile(rf"^(?!\*|{re.escape(prefix)}(?:\.|\*|$))")] = ( + f"{prefix}." + ) + quant_config.apply_vllm_mapper(WeightsMapper(orig_to_new_regex=runtime_names)) + return + + # Other backends may use semantic keys or suffix rules rather than module + # names. Preserve their settings and only add the existing layer aliases. + if start_layer_id == 0: + return + exclusions = getattr(quant_config, "exclude_modules", None) + if not isinstance(exclusions, list): + return + + def offset_local_layer(match: re.Match[str]) -> str: + layer_idx = int(match.group(1)) + if layer_idx >= num_hidden_layers: + return match.group(0) + return f"layers.{layer_idx + start_layer_id}" + + for exclusion in tuple(exclusions): + # Accept layers.N or *layers.N, excluding nested names such as + # lilicorr.layers.N, which are not decoder-layer indices. + global_exclusion = re.sub( + r"(? int: """Extract the layer index from the module name. diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 4f370d8f28ac..b206035322fe 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -264,6 +264,8 @@ def _normalize_quantization_config(self, config: PreTrainedConfig): quant_cfg["quant_method"] = "modelopt" elif quant_algo_upper == "NVFP4": quant_cfg["quant_method"] = "modelopt_fp4" + elif quant_algo_upper == "MIXED_PRECISION": + quant_cfg["quant_method"] = "modelopt_mixed" else: raise ValueError(f"Unknown ModelOpt quant algo: {quant_algo}") From be2de711989bb273e108c41f95d868122b154f9d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 24 Oct 2026 08:36:02 -0700 Subject: [PATCH 021/631] [CI][Bench] Keep captured request headers case-insensitive in Responses test (#63462) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- tests/benchmarks/test_responses_request_func.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/benchmarks/test_responses_request_func.py b/tests/benchmarks/test_responses_request_func.py index 1ecf6115fff0..e7d3f60c66f5 100644 --- a/tests/benchmarks/test_responses_request_func.py +++ b/tests/benchmarks/test_responses_request_func.py @@ -123,7 +123,7 @@ async def _serve( async def handler(request: web.Request) -> web.StreamResponse: if seen is not None: seen["payload"] = await request.json() - seen["headers"] = dict(request.headers) + seen["headers"] = request.headers.copy() if status != 200: return web.Response(status=status, reason="Bad Request") From e7473bf6862289f0647d1f4b01498288b181e063 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sun, 25 Oct 2026 00:22:16 +0800 Subject: [PATCH 022/631] [Kernel] Bump up FlashInfer to 0.7.2.post1 (#63466) Signed-off-by: Jee Jee Li --- requirements/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 972152a22b89..12545f8047c7 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,7 +10,7 @@ torchaudio==2.11.0 torchvision==0.29.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-python==0.7.2 +flashinfer-python==0.7.2.post1 apache-tvm-ffi==0.1.11 tilelang==0.1.12 # >=1.31.0 keeps execution plans alive for the CUDA graphs recorded from them From dd5e754494b48900748044dbd5eaba49d5637841 Mon Sep 17 00:00:00 2001 From: Rushabh Lalwani <59822428+rushabh-46@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:52:11 +0530 Subject: [PATCH 023/631] [Profiling] Add support for profile_prefix payload in HTTP /start_profile endpoint (#51974) Signed-off-by: Rushabh Lalwani Co-authored-by: Rushabh Lalwani Co-authored-by: vllm-agent --- docs/contributing/profiling.md | 6 ++++- docs/serving/offline_inference.md | 2 +- docs/serving/online_serving/README.md | 2 +- .../serve/test_profile_api_router.py | 3 +++ tests/v1/engine/test_async_llm.py | 25 ++++++++++++------- tests/v1/engine/test_llm_engine.py | 15 ++++++++++- tests/v1/worker/test_gpu_profiler.py | 13 +++++++++- vllm/engine/protocol.py | 1 + vllm/entrypoints/llm.py | 4 +++ vllm/entrypoints/serve/profile/api_router.py | 3 +++ vllm/profiler/wrapper.py | 24 +++++++++++++++++- vllm/v1/engine/async_llm.py | 2 ++ vllm/v1/engine/core.py | 2 ++ vllm/v1/engine/core_client.py | 8 ++++++ vllm/v1/engine/llm_engine.py | 2 ++ vllm/v1/executor/abstract.py | 7 ++++-- vllm/v1/worker/gpu_worker.py | 3 +++ 17 files changed, 105 insertions(+), 17 deletions(-) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 5e5e27026372..035529b20298 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -109,7 +109,11 @@ curl -X POST http://localhost:8000/start_profile \ -d '{ "profile_prefix": "sharegpt_run_1", "delay_iterations": 5000, - "max_iterations": 20 + "max_iterations": 20, + "profiler_kwargs": { + "record_shapes": true, + "with_stack": true + } }' ``` diff --git a/docs/serving/offline_inference.md b/docs/serving/offline_inference.md index 9a71612f262b..bc0316221f7e 100644 --- a/docs/serving/offline_inference.md +++ b/docs/serving/offline_inference.md @@ -36,7 +36,7 @@ For further details on pooling models, please refer to [this page](../models/poo For further details on profiling, please refer to [this page](../contributing/profiling.md). -- `LLM.start_profile` - Starts profiling with an optional custom trace prefix. +- `LLM.start_profile` - Starts profiling with an optional custom trace prefix (`profile_prefix`) and optional profiler keyword arguments (`profiler_kwargs`). - `LLM.stop_profile` - Stops the ongoing profiling session. ## Sleep Mode APIs diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index e70138655d33..c19f17d3adb7 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -116,7 +116,7 @@ LoRA dynamic loading & unloading is enabled in the API server. This should ONLY For further details on profiling vLLM, please refer to [this page](../../contributing/profiling.md). -- `/start_profile` - Start PyTorch profiler +- `/start_profile` - Start PyTorch profiler (accepts an optional JSON body with `profiler_kwargs` to pass keyword arguments to the underlying profiler) - `/stop_profile` - Stop PyTorch profiler ### SageMaker APIs diff --git a/tests/entrypoints/serve/test_profile_api_router.py b/tests/entrypoints/serve/test_profile_api_router.py index c5ceddd1bc44..2baad2bf02b6 100644 --- a/tests/entrypoints/serve/test_profile_api_router.py +++ b/tests/entrypoints/serve/test_profile_api_router.py @@ -58,6 +58,7 @@ def test_start_profile_forwards_session_overrides( "profile_prefix": "sharegpt_run-1", "delay_iterations": 5000, "max_iterations": 20, + "profiler_kwargs": {"record_shapes": True}, }, ) @@ -66,6 +67,7 @@ def test_start_profile_forwards_session_overrides( profile_prefix="sharegpt_run-1", delay_iterations=5000, max_iterations=20, + profiler_kwargs={"record_shapes": True}, ) @@ -75,6 +77,7 @@ def test_start_profile_forwards_session_overrides( {"profile_prefix": 1}, {"delay_iterations": "1"}, {"max_iterations": "20"}, + {"profiler_kwargs": "invalid"}, ], ) def test_start_profile_rejects_wrongly_typed_overrides( diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 0fc78bd104dc..2b8c48b8f935 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -157,7 +157,7 @@ async def profile(): assert engine.profiler is None engine_core.profile_async.assert_has_awaits( - [call(True, None, None, None), call(False)] + [call(True, None, None, None, profiler_kwargs=None), call(False)] ) @@ -200,6 +200,7 @@ async def profile(): "session", delay_iterations=5, max_iterations=2, + profiler_kwargs={"record_shapes": True}, ) with pytest.raises(ProfilerAlreadyActiveError) as exc_info: await engine.start_profile("duplicate") @@ -212,9 +213,15 @@ async def profile(): engine_core.profile_async.assert_has_awaits( [ - call(True, "session", 5, 2), + call( + True, + "session", + 5, + 2, + profiler_kwargs={"record_shapes": True}, + ), call(False), - call(True, "next-session", None, None), + call(True, "next-session", None, None, profiler_kwargs=None), call(False), ] ) @@ -239,7 +246,7 @@ async def profile(): start_rpc_entered = asyncio.Event() release_start_rpc = asyncio.Event() - async def profile_async(is_start, *args): + async def profile_async(is_start, *args, **kwargs): if is_start: start_rpc_entered.set() await release_start_rpc.wait() @@ -284,7 +291,7 @@ def test_interrupted_engine_profile_start_skips_frontend_profiler( async def profile(): start_rpc_entered = asyncio.Event() - async def profile_async(is_start, *args): + async def profile_async(is_start, *args, **kwargs): if not is_start: return start_rpc_entered.set() @@ -368,7 +375,7 @@ async def profile(): assert frontend_profiler._running is False assert engine._profile_session_active is False assert engine_core.profile_async.await_args_list == [ - call(True, None, None, None), + call(True, None, None, None, profiler_kwargs=None), call(False), ] @@ -409,10 +416,10 @@ async def profile(): engine_core.profile_async.assert_has_awaits( [ - call(True, None, None, None), - call(True, None, None, None), + call(True, None, None, None, profiler_kwargs=None), + call(True, None, None, None, profiler_kwargs=None), call(False), - call(True, None, None, None), + call(True, None, None, None, profiler_kwargs=None), ] ) diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index 69241889dc44..9c831cac4eb5 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -63,7 +63,20 @@ def test_duplicate_profile_start_rejected_before_engine_core_dispatch(): with pytest.raises(ProfilerAlreadyActiveError): engine.start_profile("duplicate") - engine.engine_core.profile.assert_called_once_with(True, "first", None, None) + engine.engine_core.profile.assert_called_once_with( + True, "first", None, None, profiler_kwargs=None + ) + + +def test_profiler_kwargs_passed_to_engine_core(): + engine = object.__new__(LLMEngine) + engine.engine_core = MagicMock() + engine._profile_session_active = False + + engine.start_profile("prefix", profiler_kwargs={"record_shapes": True}) + engine.engine_core.profile.assert_called_once_with( + True, "prefix", None, None, profiler_kwargs={"record_shapes": True} + ) @pytest.mark.parametrize( diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index 9576bddd9e61..6e73fe7815ac 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -300,6 +300,7 @@ def test_worker_creates_platform_torch_profiler(device_type, activities, expecte worker_name="rank0", local_rank=0, activities=expected, + profiler_kwargs=None, ) @@ -321,6 +322,7 @@ def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): profile_prefix="first", delay_iterations=5, max_iterations=2, + profiler_kwargs={"record_shapes": True}, ) worker.profile(is_start=False) worker.profile(profile_prefix="second") @@ -331,6 +333,10 @@ def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): call("first_rank0"), call("second_rank0"), ] + assert wrapper.return_value.set_profiler_kwargs.call_args_list == [ + call({"record_shapes": True}), + call(None), + ] assert wrapper.return_value.start.call_args_list == [ call(delay_iterations=5, max_iterations=2), call(delay_iterations=None, max_iterations=None), @@ -349,6 +355,7 @@ def test_executor_profile_only_forwards_supplied_override_kwargs(): "bounded", delay_iterations=5, max_iterations=2, + profiler_kwargs={"record_shapes": True}, ) assert executor.collective_rpc.call_args_list == [ @@ -356,7 +363,11 @@ def test_executor_profile_only_forwards_supplied_override_kwargs(): call( "profile", args=(True, "bounded"), - kwargs={"delay_iterations": 5, "max_iterations": 2}, + kwargs={ + "delay_iterations": 5, + "max_iterations": 2, + "profiler_kwargs": {"record_shapes": True}, + }, ), ] diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 149a5a54a113..2351a8a9a491 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -172,6 +172,7 @@ async def start_profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict[str, Any] | None = None, ) -> None: """Start profiling with optional per-session overrides.""" ... diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 8c605d84c667..c8b9d2325d8b 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -803,6 +803,7 @@ def start_profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict[str, Any] | None = None, ) -> None: """Start profiling with optional per-session overrides. @@ -814,12 +815,15 @@ def start_profile( profiling starts. max_iterations: Optional maximum number of worker iterations to profile. Zero means no limit. + profiler_kwargs: Optional dictionary of keyword arguments to pass to + the underlying profiler. """ self.llm_engine.start_profile( profile_prefix, delay_iterations=delay_iterations, max_iterations=max_iterations, + profiler_kwargs=profiler_kwargs, ) def stop_profile(self) -> None: diff --git a/vllm/entrypoints/serve/profile/api_router.py b/vllm/entrypoints/serve/profile/api_router.py index 2051da2b1d2a..12e052c84809 100644 --- a/vllm/entrypoints/serve/profile/api_router.py +++ b/vllm/entrypoints/serve/profile/api_router.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + from fastapi import APIRouter, FastAPI, Request from fastapi.responses import Response from pydantic import BaseModel, ConfigDict @@ -26,6 +28,7 @@ class StartProfileRequest(BaseModel): profile_prefix: str | None = None delay_iterations: int | None = None max_iterations: int | None = None + profiler_kwargs: dict[str, Any] | None = None def engine_client(request: Request) -> EngineClient: diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index 8143fddf396a..57d6bc35530d 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -208,6 +208,10 @@ def set_output_name(self, worker_name: str) -> None: """Set the output name for the next profiling session.""" return None + def set_profiler_kwargs(self, profiler_kwargs: dict | None) -> None: + """Set profiler keyword argument overrides for the next profiling session.""" + return None + def annotate_context_manager(self, name: str): """Return a context manager to annotate profiler traces.""" return nullcontext() @@ -229,6 +233,7 @@ def __init__( local_rank: int, activities: Sequence[TorchProfilerActivity], on_trace_ready: Callable[[torch.profiler.profile], None] | None = None, + profiler_kwargs: dict | None = None, ) -> None: super().__init__(profiler_config) @@ -289,7 +294,7 @@ def __init__( profiler_config.active_iterations, ) - self._profiler_kwargs = dict( + self._base_profiler_kwargs = dict( activities=[TorchProfilerActivityMap[activity] for activity in activities], schedule=profiler_schedule, record_shapes=profiler_config.torch_profiler_record_shapes, @@ -298,6 +303,9 @@ def __init__( with_flops=profiler_config.torch_profiler_with_flops, on_trace_ready=trace_handler, ) + self._profiler_kwargs = dict(self._base_profiler_kwargs) + if profiler_kwargs: + self._profiler_kwargs.update(profiler_kwargs) self.profiler: torch.profiler.profile # Track if we're using a schedule (need to call step()) @@ -328,6 +336,18 @@ def set_output_name(self, worker_name: str) -> None: ) ) + @override + def set_profiler_kwargs(self, profiler_kwargs: dict | None) -> None: + if self._active: + return + on_trace_ready = self._profiler_kwargs["on_trace_ready"] + self._profiler_kwargs = { + **self._base_profiler_kwargs, + "on_trace_ready": on_trace_ready, + } + if profiler_kwargs: + self._profiler_kwargs.update(profiler_kwargs) + def _build_profiler_table( self, sort_key: str, @@ -740,6 +760,7 @@ def create_worker_profiler( *, worker_name: str, local_rank: int, + profiler_kwargs: dict | None = None, ) -> WorkerProfiler: """Create a profiler using a validated config and platform defaults.""" profiler_type = profiler_config.profiler @@ -754,6 +775,7 @@ def create_worker_profiler( worker_name=worker_name, local_rank=local_rank, activities=default_activities if configured is None else tuple(configured), + profiler_kwargs=profiler_kwargs, ) if profiler_type == "cuda": logger.debug("Starting CUDA profiler") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 0ca5a45c3352..9fdf687d1345 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1139,6 +1139,7 @@ async def start_profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict[str, Any] | None = None, ) -> None: async with self._profile_lock: if self._profile_session_guard_enabled and self._profile_session_active: @@ -1156,6 +1157,7 @@ async def start_profile( profile_prefix, delay_iterations, max_iterations, + profiler_kwargs=profiler_kwargs, ) except Exception: self._profile_session_active = False diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index a667db454b00..684a033f6444 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -845,12 +845,14 @@ def profile( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ): self.model_executor.profile( is_start, profile_prefix, delay_iterations=delay_iterations, max_iterations=max_iterations, + profiler_kwargs=profiler_kwargs, ) def reset_mm_cache(self): diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index d9210feeebf3..a33aca737540 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -190,6 +190,7 @@ def profile( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ) -> None: raise NotImplementedError @@ -297,6 +298,7 @@ async def profile_async( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ) -> None: raise NotImplementedError @@ -419,12 +421,14 @@ def profile( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ) -> None: self.engine_core.profile( is_start, profile_prefix, delay_iterations=delay_iterations, max_iterations=max_iterations, + profiler_kwargs=profiler_kwargs, ) def reset_mm_cache(self) -> None: @@ -1070,6 +1074,7 @@ def profile( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ) -> None: self.call_utility( "profile", @@ -1077,6 +1082,7 @@ def profile( profile_prefix, delay_iterations, max_iterations, + profiler_kwargs, ) def reset_mm_cache(self) -> None: @@ -1361,6 +1367,7 @@ async def profile_async( profile_prefix: str | None = None, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ) -> None: await self.call_utility_async( "profile", @@ -1368,6 +1375,7 @@ async def profile_async( profile_prefix, delay_iterations, max_iterations, + profiler_kwargs, ) async def reset_mm_cache_async(self) -> None: diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 330c114a86dd..1ec14aaa4d0b 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -371,6 +371,7 @@ def start_profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ): if self._profile_session_active: raise ProfilerAlreadyActiveError() @@ -384,6 +385,7 @@ def start_profile( profile_prefix, delay_iterations, max_iterations, + profiler_kwargs=profiler_kwargs, ) except BaseException: self._profile_session_active = False diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index e99d1dedf4ab..bdd25a8b4608 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -6,7 +6,7 @@ from collections.abc import Callable from concurrent.futures import Future from functools import cached_property -from typing import TYPE_CHECKING, Literal, TypeVar, overload +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload from vllm.config import VllmConfig from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator @@ -300,12 +300,15 @@ def profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ): - kwargs: dict[str, int] = {} + kwargs: dict[str, Any] = {} if delay_iterations is not None: kwargs["delay_iterations"] = delay_iterations if max_iterations is not None: kwargs["max_iterations"] = max_iterations + if profiler_kwargs is not None: + kwargs["profiler_kwargs"] = profiler_kwargs if kwargs: # Reject before dispatch: a failing collective_rpc may leave # unread worker replies queued. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 04e1e7730b16..5378ddf459e6 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1395,6 +1395,7 @@ def profile( *, delay_iterations: int | None = None, max_iterations: int | None = None, + profiler_kwargs: dict | None = None, ): # Check if profiling is enabled if self.profiler_config is None or self.profiler_config.profiler is None: @@ -1425,9 +1426,11 @@ def profile( self.profiler_config, worker_name=trace_name, local_rank=self.local_rank, + profiler_kwargs=profiler_kwargs, ) self.profiler.set_output_name(trace_name) + self.profiler.set_profiler_kwargs(profiler_kwargs) self.profiler.start( delay_iterations=delay_iterations, max_iterations=max_iterations, From 56e4e42868b4652c66cb009928ddcee859b0d3c9 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 24 Oct 2026 16:26:19 -0400 Subject: [PATCH 024/631] [GLM 5.3 Perf] Skip scale calculation for `concat_and_cache_mla`, 1.2% ~ 3.0% E2E TPOT improvement (#62102) --- csrc/cache_kernels.cu | 23 +++++++++++++++---- csrc/ops.h | 2 +- csrc/torch_bindings.cpp | 3 ++- .../test_fused_deepseek_v32_norm_rope.py | 8 ++++++- vllm/_custom_ops.py | 3 ++- .../layers/attention/mla_attention.py | 2 ++ vllm/models/deepseek_v32/amd/rocm.py | 1 + vllm/models/deepseek_v32/attention.py | 3 +++ vllm/models/deepseek_v32/common/kernels.py | 16 ++++++++----- vllm/models/dots3_note/nvidia/attention.py | 2 ++ vllm/v1/attention/backend.py | 2 ++ vllm/v1/attention/backends/mla/amx_mla.py | 1 + vllm/v1/attention/backends/mla/cpu_mla.py | 1 + 13 files changed, 53 insertions(+), 14 deletions(-) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 351bcf24f9ff..dbcb6ebaef90 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -420,7 +420,8 @@ __global__ void reshape_and_cache_flash_kernel( } } -template +template __global__ void concat_and_cache_mla_kernel( const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] @@ -453,6 +454,9 @@ __global__ void concat_and_cache_mla_kernel( block_idx * block_stride + block_offset * entry_stride + i + offset; if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { dst[dst_idx] = src[src_idx]; + } else if constexpr (skip_scale) { + dst[dst_idx] = + fp8::scaled_convert(src[src_idx], 1.0f); } else { dst[dst_idx] = fp8::scaled_convert(src[src_idx], *scale); @@ -910,8 +914,8 @@ void reshape_and_cache_flash( // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_mla_kernel \ +#define LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, SKIP_SCALE) \ + vllm::concat_and_cache_mla_kernel \ <<>>( \ reinterpret_cast(kv_c.data_ptr()), \ reinterpret_cast(k_pe.data_ptr()), \ @@ -920,6 +924,13 @@ void reshape_and_cache_flash( kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); +#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ + if (skip_scale) { \ + LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, true) \ + } else { \ + LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, false) \ + } + // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. #define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ @@ -938,7 +949,8 @@ void concat_and_cache_mla( torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank // + pe_dim)] torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { + const std::string& kv_cache_dtype, torch::stable::Tensor& scale, + bool skip_scale) { // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -1030,6 +1042,9 @@ void concat_and_cache_mla( } } +#undef CALL_CONCAT_AND_CACHE_MLA +#undef LAUNCH_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] diff --git a/csrc/ops.h b/csrc/ops.h index 7c4463ac3f1a..b137d7db0418 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -712,7 +712,7 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, torch::stable::Tensor& kv_cache, torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); + torch::stable::Tensor& scale, bool skip_scale); void concat_and_cache_mla_grouped( torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a64d844e4d77..4e9530cf3b64 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1021,7 +1021,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor! kv_cache," " Tensor slot_mapping," " str kv_cache_dtype," - " Tensor scale) -> ()"); + " Tensor scale," + " bool skip_scale=False) -> ()"); // Grouped concat_and_cache_mla across all layers. Each layer's cache base // pointer and optional plain-FP8 scale are read from device tensors. diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 5e625c837c74..80f4b5e49c74 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -196,9 +196,14 @@ def test_fused_norm_rope( bs = max_pos # single block covering all tokens mla_dim = cfg.kv_lora + cfg.rope_dim mla_fp8 = mla_dtype == "fp8" + unit_scale = ( + mla_fp8 and cfg == MODEL_CONFIGS[0] and num_tokens == 4 and index_interleave + ) if mla_fp8: mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8) - mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32) + mla_k_scale = torch.tensor( + [1.0 if unit_scale else 0.3], device=dev, dtype=torch.float32 + ) else: mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16) mla_k_scale = None @@ -228,6 +233,7 @@ def test_fused_norm_rope( mla_kv_cache=mla_cache, mla_kv_cache_dtype=mla_dtype, mla_k_scale=mla_k_scale, + skip_scale=unit_scale, has_indexer=True, index_rope_interleave=index_interleave, ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index e3459466f1ed..0341422cd5ea 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2915,9 +2915,10 @@ def concat_and_cache_mla( slot_mapping: torch.Tensor, kv_cache_dtype: str, scale: torch.Tensor, + skip_scale: bool = False, ) -> None: torch.ops._C_cache_ops.concat_and_cache_mla( - kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale + kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale, skip_scale ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 674c50477cad..c8b9498af165 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -815,6 +815,8 @@ def update_kv_cache( slot_mapping, kv_cache_dtype, k_scale, + kv_cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2") + and self._k_scale_float == 1.0, ) def prepare_kv_cache_update( diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py index fd7283f4eafa..5d6d6deb4824 100644 --- a/vllm/models/deepseek_v32/amd/rocm.py +++ b/vllm/models/deepseek_v32/amd/rocm.py @@ -244,6 +244,7 @@ def _fused_attention( mla_kv_cache=mla_kv_cache, mla_kv_cache_dtype=self.kv_cache_dtype, mla_k_scale=mla_k_scale, + skip_scale=self._k_scale_float == 1.0, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, kv_c_out=kv_c_out, diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 508a80c65284..ed079204d283 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -398,6 +398,7 @@ def forward( # type: ignore[override] mla_kv_cache=mla_kv_cache, mla_kv_cache_dtype=self.kv_cache_dtype, mla_k_scale=mla_k_scale, + skip_scale=self._k_scale_float == 1.0, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, kv_c_out=kv_c_out, @@ -546,6 +547,8 @@ def _sparse_indexer_and_attn( cache_slot_mapping, self.kv_cache_dtype, self._k_scale, + self._k_scale_float == 1.0 + and self.kv_cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"), ) num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 6311748a8f00..bdc1f1ddc88a 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -179,6 +179,7 @@ def _fused_norm_rope_kernel( mla_cache_entry_stride, MLA_CACHE_BLOCK_SIZE: tl.constexpr, MLA_CACHE_FP8: tl.constexpr, + SKIP_MLA_SCALE: tl.constexpr, mla_cache_scale_ptr, # fp8_ds_mla cache views (block-scaled fp8 NoPE + unquantized bf16 RoPE). # mla_cache_ptr is the fp8 (1-byte) view, so the block/entry strides above @@ -389,18 +390,19 @@ def _fused_norm_rope_kernel( ) # kv_c_normed (KV_DIM elements) if MLA_CACHE_FP8: - scale = tl.load(mla_cache_scale_ptr) - kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype) + scale = 1.0 + if not SKIP_MLA_SCALE: + scale = tl.load(mla_cache_scale_ptr) + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype) + else: + kv_c_fp8 = kv_c.to(tl.float32).to(fp8_dtype) tl.store(dst + kv_block, kv_c_fp8) else: tl.store(dst + kv_block, kv_c) # k_pe_roped (from registers, interleaved layout) if MLA_CACHE_FP8: tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(fp8_dtype)) - tl.store( - dst + KV_DIM + dim_off * 2 + 1, - (r2 / scale).to(fp8_dtype), - ) + tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(fp8_dtype)) else: tl.store(dst + KV_DIM + dim_off * 2, r1) tl.store(dst + KV_DIM + dim_off * 2 + 1, r2) @@ -543,6 +545,7 @@ def fused_norm_rope( kv_c_out: torch.Tensor | None = None, k_pe_out: torch.Tensor | None = None, index_k_out: torch.Tensor | None = None, + skip_scale: bool = False, ) -> torch.Tensor: assert positions.ndim == 1 assert q_c.ndim == 2 @@ -719,6 +722,7 @@ def fused_norm_rope( mla_entry_stride, mla_block_size, mla_cache_fp8, + mla_cache_fp8 and skip_scale, mla_k_scale, mla_ds_scale_view, mla_ds_rope_view, diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index ce69f6b1824c..ecf5b2e57b0e 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -723,6 +723,7 @@ def do_kv_cache_update( slot_mapping: torch.Tensor, kv_cache_dtype: str, k_scale: torch.Tensor, + skip_scale: bool = False, ) -> None: super().do_kv_cache_update( kv_c_normed, @@ -731,6 +732,7 @@ def do_kv_cache_update( slot_mapping, kv_cache_dtype, k_scale, + skip_scale=skip_scale, ) def forward_mha( # type: ignore[override] diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ae9ae3987e66..d1bc18472ee3 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -1113,6 +1113,7 @@ def do_kv_cache_update( slot_mapping: torch.Tensor, kv_cache_dtype: str, k_scale: torch.Tensor, + skip_scale: bool = False, ) -> None: if kv_cache.numel() == 0: return @@ -1125,6 +1126,7 @@ def do_kv_cache_update( slot_mapping.flatten(), kv_cache_dtype=kv_cache_dtype, scale=k_scale, + skip_scale=skip_scale, ) diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 16689887af0a..a89bef5a83bc 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -311,6 +311,7 @@ def do_kv_cache_update( slot_mapping: torch.Tensor, kv_cache_dtype: str, k_scale: torch.Tensor, + skip_scale: bool = False, ) -> None: # Overrides the default (CUDA-only concat_and_cache_mla). if kv_cache.numel() == 0: diff --git a/vllm/v1/attention/backends/mla/cpu_mla.py b/vllm/v1/attention/backends/mla/cpu_mla.py index 431f6c318533..c89dbf4f58a9 100644 --- a/vllm/v1/attention/backends/mla/cpu_mla.py +++ b/vllm/v1/attention/backends/mla/cpu_mla.py @@ -200,6 +200,7 @@ def do_kv_cache_update( slot_mapping: torch.Tensor, kv_cache_dtype: str, k_scale: torch.Tensor, + skip_scale: bool = False, ) -> None: if kv_cache.numel() == 0: return From 7050ef881ff0835f6890a47db1859a24c65481d6 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:35:17 -0700 Subject: [PATCH 025/631] [Sampling] Add adaptive verification support to batch-sharded sampling (#58466) --- tests/v1/sample/test_dry.py | 5 +- tests/v1/worker/test_gpu_batch_shard.py | 56 ++++++++++++++++++- tests/v1/worker/test_gpu_model_runner_v2.py | 1 + vllm/config/vllm.py | 15 ----- vllm/v1/worker/gpu/input_batch.py | 7 +++ vllm/v1/worker/gpu/model_runner.py | 16 +++++- vllm/v1/worker/gpu/sample/batch_shard.py | 40 +++++++++++-- .../gpu/spec_decode/adaptive_verification.py | 45 ++++++++++++++- 8 files changed, 160 insertions(+), 25 deletions(-) diff --git a/tests/v1/sample/test_dry.py b/tests/v1/sample/test_dry.py index 50a5a5e0dbf3..3e2384b66277 100644 --- a/tests/v1/sample/test_dry.py +++ b/tests/v1/sample/test_dry.py @@ -608,7 +608,10 @@ def test_v2_batch_sharded_sampling_matches_unsharded(monkeypatch): group = SimpleNamespace(rank_in_group=rank, world_size=tp_size) with mock.patch.object(batch_shard, "get_tp_group", return_value=group): sharder = batch_shard.BatchSharder( - max_num_reqs=8, max_num_logits_per_req=1, device=DEVICE + max_num_reqs=8, + max_num_logits_per_req=1, + device=DEVICE, + adaptive_verification_enabled=False, ) local, _, _, _ = sharder.shard_sampler_inputs(batch, None) owned = [int(req_id.split("-")[1]) for req_id in local.req_ids] diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index bb918c206a0e..4ee76818df4f 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -7,6 +7,7 @@ compose to a bit-exact round trip of the replicated computation. """ +from dataclasses import replace from types import SimpleNamespace from unittest import mock @@ -110,7 +111,12 @@ def dev(a: np.ndarray) -> torch.Tensor: ) -def _shard_all_ranks(batch: InputBatch, max_num_reqs: int, tp_size: int): +def _shard_all_ranks( + batch: InputBatch, + max_num_reqs: int, + tp_size: int, + adaptive_verification_enabled: bool = False, +): """Simulate every TP rank by faking get_tp_group() per BatchSharder.""" results = [] for rank in range(tp_size): @@ -120,6 +126,7 @@ def _shard_all_ranks(batch: InputBatch, max_num_reqs: int, tp_size: int): max_num_reqs=max_num_reqs, max_num_logits_per_req=8, device=torch.device(DEVICE), + adaptive_verification_enabled=adaptive_verification_enabled, ) local_batch, sorted_logits_indices, _, metadata = sharder.shard_sampler_inputs( batch, None @@ -128,6 +135,53 @@ def _shard_all_ranks(batch: InputBatch, max_num_reqs: int, tp_size: int): return results +@requires_gpu +@pytest.mark.parametrize("tp_size", [2, 4]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_admitted_counts_give_unpadded_plan(tp_size: int, seed: int): + """With the admitted draft counts on the host, an adaptive verification + batch shards exactly like a batch scheduled with only the admitted drafts.""" + rng = np.random.default_rng(seed) + max_num_reqs = 64 + admitted = _make_batch(rng, num_reqs=23, max_num_reqs=max_num_reqs, max_spec=3) + assert admitted.num_draft_tokens_per_req is not None + admitted_drafts = admitted.num_draft_tokens_per_req + # The host only sees the scheduled counts, an upper bound on the admitted ones. + scheduled_drafts = admitted_drafts + rng.integers(0, 3, size=23).astype(np.int32) + scheduled_cu = np.zeros(24, dtype=np.int32) + np.cumsum(1 + scheduled_drafts, out=scheduled_cu[1:]) + av_batch = replace( + admitted, + cu_num_logits_np=scheduled_cu, + num_draft_tokens_per_req=scheduled_drafts, + num_admitted_draft_tokens=SimpleNamespace(wait=lambda: admitted_drafts), + ) + + expected = _shard_all_ranks(admitted, max_num_reqs, tp_size) + actual = _shard_all_ranks( + av_batch, max_num_reqs, tp_size, adaptive_verification_enabled=True + ) + for (exp_local, exp_sorted, exp_meta), (local, sorted_idx, meta) in zip( + expected, actual + ): + assert meta.num_logits_per_rank == exp_meta.num_logits_per_rank + assert meta.num_local_logits == exp_meta.num_local_logits + assert _np(sorted_idx).tolist() == _np(exp_sorted).tolist() + assert (local.cu_num_logits_np == exp_local.cu_num_logits_np).all() + assert ( + _np(local.cu_num_logits).tolist() == _np(exp_local.cu_num_logits).tolist() + ) + assert ( + _np(local.logits_indices).tolist() == _np(exp_local.logits_indices).tolist() + ) + assert ( + _np(local.expanded_idx_mapping).tolist() + == _np(exp_local.expanded_idx_mapping).tolist() + ) + assert local.num_draft_tokens == exp_local.num_draft_tokens + assert meta.max_num_logits_per_req == exp_meta.max_num_logits_per_req + + def _owned_batch_indices(local_batch: InputBatch) -> np.ndarray: return np.array([int(r.split("-")[1]) for r in local_batch.req_ids], dtype=np.int64) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 87e6c5b77ed1..537d4788abc5 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -94,6 +94,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): num_new_sampled_tokens_per_step=1, ) runner.speculator = None + runner.batch_sharder = None runner.req_states = [] runner.input_buffers = SimpleNamespace(query_start_loc=None) runner.vocab_size = 1 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 248fd9b359df..779b62c19587 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3193,21 +3193,6 @@ def _validate_batch_sharded_sampling(self) -> None: "not forward sampling masks" ) - if ( - self.speculative_config is not None - and self.speculative_config.enable_adaptive_verification - ): - # Adaptive verification picks the per-request draft split on the GPU, - # so cu_num_logits_np is only an upper bound, while the shard plan is - # built from that CPU array. The two disagree once the budget binds. - # TODO(TheEpicDolphin): Support adaptive verification with batch-sharded - # sampling. - blockers.append( - "it does not yet work with adaptive verification, which decides " - "the per-request logits counts on the GPU, where the CPU-side " - "shard plan cannot see them" - ) - if blockers: raise ValueError( "Batch-sharded sampling was explicitly enabled via " diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index ec9b4f73c32a..6d9732b67ab1 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -15,6 +15,9 @@ from vllm.v1.worker.gpu.attn_utils import FastPrefillBatchMetadata from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.pcp_manager import PCPBatchMetadata + from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( + AdmittedDraftCounts, + ) class InputBuffers: @@ -124,6 +127,10 @@ class InputBatch: # None if there are no prefills. prefill_runs_as_decode_np: np.ndarray | None = None + # [num_reqs] admitted draft counts in flight to the host. Set only under + # adaptive verification with batch-sharded sampling, when the budget binds. + num_admitted_draft_tokens: "AdmittedDraftCounts | None" = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f7b114779e90..c539520eb5e1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -450,6 +450,10 @@ def load_model( max_num_reqs=self.max_num_reqs, max_num_logits_per_req=self.decode_query_len, device=self.device, + adaptive_verification_enabled=( + self.speculative_config is not None + and self.speculative_config.enable_adaptive_verification + ), ) logger.info("Batch-sharded sampling enabled.") else: @@ -711,6 +715,7 @@ def initialize_kv_cache( vllm_config=self.vllm_config, target_layer_names=target_attn_layer_names, additional_attn_cg_support=additional_attn_cg_support, + copy_admitted_drafts_to_host=self.batch_sharder is not None, ) self.block_tables = BlockTables( @@ -1434,6 +1439,7 @@ def prepare_inputs( # compact_batch's budget is only primed when drafts are scheduled. adaptive_verification = self.adaptive_verification if draft_tokens else None num_scheduled_tokens_upper_bound = num_scheduled_tokens_np + num_admitted_draft_tokens = None if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across # all verification requests, `reallocate_drafts` will unevenly assign the @@ -1458,9 +1464,12 @@ def prepare_inputs( query_start_loc = self.input_buffers.query_start_loc async_tensor_h2d(query_start_loc_np, out=query_start_loc) if adaptive_verification is not None: - cu_num_logits, query_start_loc, total_num_draft_tokens = ( - adaptive_verification.reallocate_drafts(req_ids, idx_mapping) - ) + ( + cu_num_logits, + query_start_loc, + total_num_draft_tokens, + num_admitted_draft_tokens, + ) = adaptive_verification.reallocate_drafts(req_ids, idx_mapping) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens if num_draft_tokens_per_req is not None: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( @@ -1545,6 +1554,7 @@ def prepare_inputs( num_tokens_after_padding=num_tokens_after_padding, num_draft_tokens=total_num_draft_tokens, num_draft_tokens_per_req=num_draft_tokens_per_req, + num_admitted_draft_tokens=num_admitted_draft_tokens, query_start_loc=query_start_loc, query_start_loc_np=query_start_loc_np, seq_lens=seq_lens, diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index d8cecb7f87b7..9dc0c09688b8 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -45,6 +45,7 @@ class BatchShardMetadata: # tensor. Derived from the global cu_num_logits, so each rank agrees # ont he gather shapes, including empty shards. max_num_logits_per_req: int + adaptive_verification_enabled: bool @triton.jit( @@ -150,11 +151,13 @@ def __init__( max_num_reqs: int, max_num_logits_per_req: int, device: torch.device, + adaptive_verification_enabled: bool, ): tp_group = get_tp_group() self.tp_rank = tp_group.rank_in_group self.tp_size = tp_group.world_size self.device = device + self.adaptive_verification_enabled = adaptive_verification_enabled self._padded_num_reqs = triton.next_power_of_2(max_num_reqs) self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req) self._num_warps = max(1, min(8, self._padded_num_reqs // 128)) @@ -175,7 +178,6 @@ def shard_sampler_inputs( tp_rank = self.tp_rank tp_size = self.tp_size num_reqs = input_batch.idx_mapping_np.shape[0] - num_logits = int(input_batch.cu_num_logits_np[-1]) # Deterministically assign requests to ranks, round-robin over slot # indices so ownership stays balanced when the slot allocator fills @@ -187,9 +189,23 @@ def shard_sampler_inputs( num_local_reqs = local_req_indices_np.shape[0] num_reqs_per_rank_np = np.bincount(req_owner_np, minlength=tp_size) max_num_reqs_per_rank = int(num_reqs_per_rank_np.max()) if num_reqs else 1 + # Get the number of logits per request on the host. + num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np) + num_admitted_draft_tokens_np = None + if input_batch.num_admitted_draft_tokens is not None: + # Under adaptive verification, cu_num_logits_np holds the scheduled counts. + # When the budget binds, the admitted counts are copied to the host so the + # plan matches the device layout. + assert input_batch.num_draft_tokens_per_req is not None + num_admitted_draft_tokens_np = input_batch.num_admitted_draft_tokens.wait() + num_logits_per_req_np = ( + num_logits_per_req_np + - input_batch.num_draft_tokens_per_req + + num_admitted_draft_tokens_np + ) + num_logits = int(num_logits_per_req_np.sum()) # Derive the number of logits owned by each rank, as well as the # local rank. - num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np) num_logits_per_rank_np = np.bincount( req_owner_np, weights=num_logits_per_req_np, minlength=tp_size ).astype(np.int64) @@ -254,7 +270,15 @@ def shard_sampler_inputs( num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req[ local_req_indices_np ] - num_draft_tokens = int(num_draft_tokens_per_req.sum()) + if num_admitted_draft_tokens_np is not None: + # Under adaptive verification num_draft_tokens_per_req holds the + # scheduled counts. Follow the admitted counts so every rank picks + # the same sampler as the replicated path, including a budget of 0. + num_draft_tokens = int( + num_admitted_draft_tokens_np[local_req_indices_np].sum() + ) + elif input_batch.num_draft_tokens > 0: + num_draft_tokens = int(num_draft_tokens_per_req.sum()) local_req_ids = [input_batch.req_ids[i] for i in local_req_indices_np.tolist()] local_batch = replace( @@ -284,6 +308,7 @@ def shard_sampler_inputs( max_num_reqs_per_rank=max_num_reqs_per_rank, gathered_src_indices=gathered_src_indices, max_num_logits_per_req=max_num_logits_per_req, + adaptive_verification_enabled=self.adaptive_verification_enabled, ) return local_batch, sorted_logits_indices, local_grammar_output, metadata @@ -580,7 +605,14 @@ def _gather_logprobs_tensors( logprobs=logprobs, selected_token_ranks=selected_token_ranks, cu_num_generated_tokens=( - global_batch.cu_num_logits_np.tolist() if num_logits != num_reqs else None + None + if metadata.adaptive_verification_enabled or num_logits == num_reqs + else global_batch.cu_num_logits_np.tolist() + ), + cu_num_generated_tokens_tensor=( + global_batch.cu_num_logits.clone() + if metadata.adaptive_verification_enabled + else None ), ) diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 5f91af59e592..fee83fcd351e 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -128,6 +128,19 @@ def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray: return draft_table, verify_table +class AdmittedDraftCounts: + """Per-request admitted draft counts, in flight to the host.""" + + def __init__(self, cpu: torch.Tensor, event: torch.cuda.Event): + self._cpu = cpu + self._event = event + + def wait(self) -> np.ndarray: + with gpu_sync_allowed(): + self._event.synchronize() + return self._cpu.numpy() + + class AdaptiveVerificationManager: def __init__( self, @@ -135,6 +148,7 @@ def __init__( query_start_loc: torch.Tensor, num_bonus_tokens: int, max_total_logits: int, + copy_admitted_drafts_to_host: bool = False, ): self.req_states = req_states self.num_speculative_steps = req_states.num_speculative_steps @@ -164,6 +178,15 @@ def __init__( self._num_non_draft_tokens = torch.empty_like(query_start_loc[:-1]) self._cu_num_logits = torch.empty_like(query_start_loc) + # Batch-sharded sampling sizes its all-to-all from host counts, so it + # needs the admitted split on the host. + self._copy_admitted_drafts_to_host = copy_admitted_drafts_to_host + if copy_admitted_drafts_to_host: + self._admitted_drafts_cpu = torch.empty( + max_num_reqs, dtype=torch.int32, pin_memory=True + ) + self._admitted_drafts_event = torch.cuda.Event(blocking=True) + # Two D2H slots preserve stale inputs for budget selection. self._stale_confidences = [ CpuGpuBuffer( @@ -391,7 +414,10 @@ def compact_batch( def reallocate_drafts( self, req_ids: list[str], idx_mapping: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, int]: + ) -> tuple[torch.Tensor, torch.Tensor, int, AdmittedDraftCounts | None]: + """The admitted draft counts are None unless copy_admitted_drafts_to_host + is set and the budget binds; otherwise the host counts are already exact. + """ batch_budget, self._batch_budget = self._batch_budget, None assert batch_budget is not None num_drafts_per_req, num_non_draft_tokens_per_req, draft_budget = batch_budget @@ -440,10 +466,25 @@ def reallocate_drafts( out=self.query_start_loc[1 : num_reqs + 1], ) self.query_start_loc[num_reqs + 1 :].fill_(num_tokens) + + admitted_drafts = None + if self._copy_admitted_drafts_to_host and ( + 0 < draft_budget < int(scheduled_drafts.sum()) + ): + admitted_cpu = self._admitted_drafts_cpu[:num_reqs] + current_stream = torch.cuda.current_stream(self.req_states.device) + self._copy_stream.wait_stream(current_stream) + with stream(self._copy_stream, current_stream): + admitted_cpu.copy_(capacities, non_blocking=True) + self._admitted_drafts_event.record() + admitted_drafts = AdmittedDraftCounts( + admitted_cpu, self._admitted_drafts_event + ) return ( self._cu_num_logits[: num_reqs + 1], self.query_start_loc, draft_budget, + admitted_drafts, ) @@ -458,6 +499,7 @@ def maybe_create_adaptive_verification_manager( vllm_config: "VllmConfig", target_layer_names: set[str] | None = None, additional_attn_cg_support: tuple[AttentionCGSupport, str | None] | None = None, + copy_admitted_drafts_to_host: bool = False, ) -> AdaptiveVerificationManager | None: if not enable_adaptive_verification: return None @@ -506,4 +548,5 @@ def maybe_create_adaptive_verification_manager( query_start_loc, num_bonus_tokens, max_total_logits=max_total_logits, + copy_admitted_drafts_to_host=copy_admitted_drafts_to_host, ) From 4e591657157a021cab997f911d67f50c0d18a71e Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 24 Oct 2026 15:19:14 -0700 Subject: [PATCH 026/631] [Deprecation] Deprecate LM Format Enforcer structured outputs backend (#63423) Signed-off-by: Woosuk Kwon Co-authored-by: Codex Co-authored-by: vllm-agent --- docs/features/structured_outputs.md | 21 +++++++++++++++++++++ vllm/config/structured_outputs.py | 15 ++++++++++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/docs/features/structured_outputs.md b/docs/features/structured_outputs.md index eb2a7500cb9e..36b341ddf468 100644 --- a/docs/features/structured_outputs.md +++ b/docs/features/structured_outputs.md @@ -17,6 +17,27 @@ available to generate structured outputs. - `structural_tag` -> `{"structured_outputs": {"structural_tag": ...}}` or `StructuredOutputsParams(structural_tag=...)` - `guided_decoding_backend` -> Remove this field from your request +## Migrating from LM Format Enforcer + +!!! warning + The `lm-format-enforcer` structured outputs backend is deprecated starting + in v0.33.0. It remains usable with a warning in v0.33.0, will require an + explicit opt-in in v0.34.0, and will be removed in v0.35.0. + +For online serving, replace +`--structured-outputs-config.backend lm-format-enforcer` with +`--structured-outputs-config.backend auto`, or remove the option to use the +default. For offline inference, replace +`structured_outputs_config={"backend": "lm-format-enforcer"}` with +`structured_outputs_config={"backend": "auto"}` when constructing `LLM`. +You can also explicitly select `xgrammar` or `guidance`. + +Validate your existing JSON schemas, regex patterns, and generated outputs with +the replacement backend before migrating. Backend support and formatting can +differ: in particular, LM Format Enforcer uses Python-style regex syntax, while +XGrammar and Guidance use Rust-style regex syntax. Changing backends can change +which tokens are allowed and therefore the generated output. + ## Online Serving (OpenAI API) You can generate structured outputs using the OpenAI's [Completions](https://platform.openai.com/docs/api-reference/completions) and [Chat](https://platform.openai.com/docs/api-reference/chat) API. diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py index 4a13900b5622..1c48c3e467ca 100644 --- a/vllm/config/structured_outputs.py +++ b/vllm/config/structured_outputs.py @@ -6,8 +6,11 @@ from pydantic import model_validator from vllm.config.utils import config +from vllm.logger import init_logger from vllm.utils.hashing import safe_hash +logger = init_logger(__name__) + StructuredOutputsBackend = Literal[ "auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer" ] @@ -21,7 +24,9 @@ class StructuredOutputsConfig: """Which engine will be used for structured outputs (e.g. JSON schema, regex, etc) by default. With "auto", we will make opinionated choices based on request contents and what the backend libraries currently support, - so the behavior is subject to change in each release.""" + so the behavior is subject to change in each release. + The "lm-format-enforcer" backend is deprecated and will be removed in + v0.35.0. Use "auto", "xgrammar", or "guidance" instead.""" disable_any_whitespace: bool = False """If `True`, json output will always be compact without any whitespace. If `False`, the model may generate whitespace between JSON fields, @@ -69,4 +74,12 @@ def _validate_structured_output_config(self) -> Self: "disable_additional_properties is only supported " "for the guidance backend." ) + if self.backend == "lm-format-enforcer": + logger.warning_once( + "The lm-format-enforcer structured outputs backend is deprecated " + "and will be removed in v0.35.0. Use auto, xgrammar, or guidance " + "instead. See https://docs.vllm.ai/en/latest/features/" + "structured_outputs/#migrating-from-lm-format-enforcer " + "for migration guidance." + ) return self From 71a82756040a31849623e720e039a83857b4b562 Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Sun, 25 Oct 2026 10:18:12 +0800 Subject: [PATCH 027/631] [CI] Trigger recipe previews and the recipes.vllm.ai sync on push (#63259) Signed-off-by: esmeetu Co-authored-by: Claude Opus 5.5 --- .github/workflows/recipes-preview.yml | 69 +++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 .github/workflows/recipes-preview.yml diff --git a/.github/workflows/recipes-preview.yml b/.github/workflows/recipes-preview.yml new file mode 100644 index 000000000000..42127e99b79e --- /dev/null +++ b/.github/workflows/recipes-preview.yml @@ -0,0 +1,69 @@ +name: Recipes Preview + +# Tells vllm-project/recipes when recipes/ changes, instead of waiting for its +# scheduled polling, which GitHub runs hours late: +# - a recipe PR is opened, pushed or reopened: rebuild its preview branch +# `preview/vllm-pr-`, which Vercel deploys at the environment URL below +# (recipes debounces bursts of pushes); +# - a recipe PR is closed: delete that branch; +# - recipe changes reach main: sync recipes.vllm.ai (`recipes-updated`). +# +# Runs on pull_request_target so fork PRs can reach the dispatch token. This +# is safe only because no step checks out or runs PR code, and only the PR +# number leaves this workflow. Keep it that way. The token comes from a GitHub +# App installed only on vllm-project/recipes with contents: write; its key +# lives in the `recipes-preview` environment, which only main may use. + +on: + pull_request_target: + types: [opened, synchronize, reopened, closed] + paths: + - 'recipes/**' + push: + branches: [main] + paths: + - 'recipes/**' + +permissions: + contents: read + +concurrency: + group: recipes-preview-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + dispatch: + if: github.repository == 'vllm-project/vllm' + runs-on: ubuntu-latest + timeout-minutes: 5 + environment: + name: recipes-preview + url: >- + ${{ github.event_name == 'push' && 'https://recipes.vllm.ai' || + (github.event.action != 'closed' && + format('https://vllm-recipes-git-preview-vllm-pr-{0}-inferact-inc.vercel.app', github.event.pull_request.number) || + '') }} + steps: + - name: Create a token for vllm-project/recipes + id: token + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.RECIPES_APP_ID }} + private-key: ${{ secrets.RECIPES_APP_PRIVATE_KEY }} + owner: vllm-project + repositories: recipes + permission-contents: write + + - name: Dispatch to vllm-project/recipes + env: + GH_TOKEN: ${{ steps.token.outputs.token }} + EVENT: ${{ github.event_name }} + PR: ${{ github.event.pull_request.number }} + run: | + if [ "$EVENT" = push ]; then + gh api repos/vllm-project/recipes/dispatches -f event_type=recipes-updated + exit 0 + fi + [[ "$PR" =~ ^[0-9]+$ ]] || { echo "Invalid PR number: $PR"; exit 1; } + gh api repos/vllm-project/recipes/dispatches \ + -f event_type=preview-pr -F "client_payload[pr]=${PR}" From fc2b6ab45ed3e498b52d74fd880a8b40a65ba3a7 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Sat, 24 Oct 2026 19:57:07 -0700 Subject: [PATCH 028/631] [Perf][MRV2] Parallelize staged writes across 1024-element chunks (#62698) Signed-off-by: Juntian Liu Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 Co-authored-by: Nick Hill --- tests/v1/worker/test_gpu_block_table.py | 20 ++++++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 35 ++++++++++++++++++------- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 3e7b88f76178..c77e5ffd4e3c 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -79,6 +79,7 @@ def fail_if_apply_write_called(): assert not block_table._staged_write_starts assert not block_table._staged_write_contents assert not block_table._staged_write_cu_lens + assert block_table._staged_write_max_len == 0 block_tables.append_block_ids( req_index=0, @@ -105,6 +106,25 @@ def fail_if_apply_write_called(): assert block_tables.num_blocks.np[2, 0] == 1 +def test_block_tables_fused_write_spans_multiple_chunks(): + """A fused write longer than one 1024-element chunk is fully copied.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16, 32], + max_num_reqs=2, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4, 4096], + device=device, + ) + long_ids = list(range(1, 2501)) + block_tables.append_block_ids(0, ([7], long_ids), overwrite=True) + block_tables.apply_staged_writes() + torch.accelerator.synchronize() + + assert block_tables.block_tables[0].gpu[0, :1].tolist() == [7] + assert block_tables.block_tables[1].gpu[0, :2500].tolist() == long_ids + + def test_block_tables_apply_staged_writes_single_group(): device = torch.device("cuda") block_tables = BlockTables( diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 7364156848bf..85d74a0ba4a2 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -20,6 +20,9 @@ # concurrent in-flight steps (engine batch_queue_size). _DEFAULT_MAX_CONCURRENCY = 2 +# Elements copied per program in `_apply_write_kernel`. +_WRITE_BLOCK_SIZE = 1024 + def set_default_max_concurrency(n: int) -> None: global _DEFAULT_MAX_CONCURRENCY @@ -166,6 +169,7 @@ def __init__( self._staged_write_starts: list[int] = [] self._staged_write_contents: list[int | float] = [] self._staged_write_cu_lens: list[int] = [] + self._staged_write_max_len = 0 new_buffer = partial(UvaBufferPool, max_concurrency=max_concurrency) @@ -183,8 +187,11 @@ def stage_write( return self._staged_write_indices.append(index) self._staged_write_starts.append(start) + prev_len = len(self._staged_write_contents) self._staged_write_contents.extend(x) - self._staged_write_cu_lens.append(len(self._staged_write_contents)) + cu_len = len(self._staged_write_contents) + self._staged_write_cu_lens.append(cu_len) + self._staged_write_max_len = max(self._staged_write_max_len, cu_len - prev_len) def stage_write_elem(self, index: int, x: int) -> None: assert index >= 0 @@ -192,6 +199,7 @@ def stage_write_elem(self, index: int, x: int) -> None: self._staged_write_starts.append(0) self._staged_write_contents.append(x) self._staged_write_cu_lens.append(len(self._staged_write_contents)) + self._staged_write_max_len = max(self._staged_write_max_len, 1) def apply_write(self) -> None: n = len(self._staged_write_indices) @@ -212,7 +220,7 @@ def apply_write(self) -> None: ) # Write diffs to the GPU buffer - _apply_write_kernel[(n,)]( + _apply_write_kernel[(n, _num_write_chunks(self._staged_write_max_len))]( self.gpu, self.gpu.stride(0), indices_uva, @@ -220,7 +228,7 @@ def apply_write(self) -> None: write_contents, cu_lens_uva, None, - BLOCK_SIZE=1024, + BLOCK_SIZE=_WRITE_BLOCK_SIZE, MULTI_GROUP=False, ) # Clear the staged writes @@ -231,6 +239,7 @@ def clear_staged_writes(self) -> None: self._staged_write_starts.clear() self._staged_write_contents.clear() self._staged_write_cu_lens.clear() + self._staged_write_max_len = 0 class FusedStagedWriter: @@ -260,6 +269,7 @@ def apply( starts: list[int] = [] contents: list[int | float] = [] cu_lens: list[int] = [] + max_len = 0 for group_id, t in enumerate(tensors): n = len(t._staged_write_indices) @@ -272,6 +282,7 @@ def apply( content_base = len(contents) contents.extend(t._staged_write_contents) cu_lens.extend(content_base + cu_len for cu_len in t._staged_write_cu_lens) + max_len = max(max_len, t._staged_write_max_len) if not group_ids: return @@ -282,7 +293,7 @@ def apply( cu_lens_uva = self.cu_lens.copy_to_uva(cu_lens) contents_gpu = async_tensor_h2d(contents, device=self.device, dtype=torch.int32) - _apply_write_kernel[(len(group_ids),)]( + _apply_write_kernel[(len(group_ids), _num_write_chunks(max_len))]( output_ptrs, output_strides, indices_uva, @@ -290,13 +301,18 @@ def apply( contents_gpu, cu_lens_uva, group_ids_uva, - BLOCK_SIZE=1024, + BLOCK_SIZE=_WRITE_BLOCK_SIZE, MULTI_GROUP=True, ) for t in tensors: t.clear_staged_writes() +def _num_write_chunks(max_len: int) -> int: + # One program per _WRITE_BLOCK_SIZE-element chunk. + return max(triton.cdiv(max_len, _WRITE_BLOCK_SIZE), 1) + + @triton.jit def _apply_write_kernel( output_ptr, # MULTI_GROUP: ptr-to-ptrs [num_groups]; else: data ptr @@ -328,11 +344,10 @@ def _apply_write_kernel( row_stride = output_stride row_ptr += row_idx * row_stride + start_idx - for i in range(0, content_len, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < content_len - content = tl.load(write_contents_ptr + cu_start + block, mask=mask) - tl.store(row_ptr + block, content, mask=mask) + block = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = block < content_len + content = tl.load(write_contents_ptr + cu_start + block, mask=mask) + tl.store(row_ptr + block, content, mask=mask) @triton.jit From 8e11a8a49e11b2f0825be7b3c2434bcf23d7f87b Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sun, 25 Oct 2026 10:57:33 +0800 Subject: [PATCH 029/631] [LoRA] Reject incompatible MegaMoE backends during config validation (#63496) Signed-off-by: Jee Jee Li --- vllm/config/lora.py | 10 ++++++++++ vllm/config/vllm.py | 1 + vllm/models/step5/nvidia/model.py | 7 ------- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 9443efd2339c..86aab1ea9df6 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -8,6 +8,7 @@ from pydantic import ConfigDict, Field, model_validator from vllm import envs +from vllm.config.kernel import MEGA_MOE_BACKENDS, KernelConfig from vllm.config.utils import config from vllm.logger import init_logger from vllm.platforms import current_platform @@ -134,6 +135,15 @@ def _validate_lora_config(self) -> Self: envs.VLLM_LORA_ENABLE_DUAL_STREAM = False return self + def verify_with_kernel_config(self, kernel_config: KernelConfig) -> None: + """Reject MegaMoE backends, which do not support LoRA.""" + moe_backend = kernel_config.moe_backend + if moe_backend in MEGA_MOE_BACKENDS or moe_backend == "aiter_mega_moe": + raise ValueError( + f"LoRA is not supported with moe_backend={moe_backend!r}. " + "Disable LoRA or select a MoE backend that supports LoRA." + ) + def verify_with_model_config(self, model_config: ModelConfig): # Speculative decoding drafters don't run LoRA adapters. if model_config.runner_type != "draft" and not model_config.is_lora_supported: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 779b62c19587..e44c4ebbade5 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1572,6 +1572,7 @@ def __post_init__(self): "supported for MLA models." ) if self.lora_config is not None: + self.lora_config.verify_with_kernel_config(self.kernel_config) self.lora_config.verify_with_model_config(self.model_config) if ( diff --git a/vllm/models/step5/nvidia/model.py b/vllm/models/step5/nvidia/model.py index 70d7d728d8fc..ea6b1b7273bf 100644 --- a/vllm/models/step5/nvidia/model.py +++ b/vllm/models/step5/nvidia/model.py @@ -2342,13 +2342,6 @@ def __init__( prefix: str = "", ): super().__init__() - if ( - vllm_config.lora_config is not None - and vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" - ): - raise ValueError( - "Step5 LoRA is not supported with moe_backend='deep_gemm_mega_moe'" - ) self.vllm_config = vllm_config model_config = vllm_config.model_config config = model_config.hf_config From d4d647b2b52d5e15ee1427de0effecd474430972 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sun, 25 Oct 2026 10:58:43 +0800 Subject: [PATCH 030/631] [Kernel] Update DeepGEMM commit (#63253) Signed-off-by: Jee Jee Li Co-authored-by: vllm-agent --- cmake/external_projects/deepgemm.cmake | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index cde756127af7..2fbfa157de88 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,9 +32,12 @@ else() # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 # Mega MoE, plus swiglustep for BF16/FP8/NVFP4 Mega MoE and F8_128x4 - # SFs for the FP8 einsum output (vllm-project/DeepGEMM#40). + # SFs for the FP8 einsum output (vllm-project/DeepGEMM#40), plus + # fp4_fp4_mega_moe and fused post-expert RMSNorm for Muse Spark Open + # (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix + # (vllm-project/DeepGEMM#43). # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. - set(_DEEPGEMM_UPSTREAM_TAG "727f847082d2e7946f70dc02b94949db206310d3") + set(_DEEPGEMM_UPSTREAM_TAG "0820030756945b2daec208e8e8255268360f337d") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) From 57ab2f35ac6e7207bd71cf6cad17b31705f03cd3 Mon Sep 17 00:00:00 2001 From: Ronald Date: Sun, 25 Oct 2026 11:08:23 +0800 Subject: [PATCH 031/631] [Hardware] Use accelerator APIs in weight transfer (#63457) Signed-off-by: Ronald1995 --- vllm/distributed/weight_transfer/ipc_engine.py | 2 +- vllm/distributed/weight_transfer/nccl_engine.py | 4 ++-- .../distributed/weight_transfer/packed_tensor.py | 6 +++--- .../weight_transfer/sharded_rdt_engine.py | 16 ++++++++-------- .../weight_transfer/sharded_rdt_trainer.py | 2 +- .../weight_transfer/sparse_nccl_engine.py | 4 ++-- 6 files changed, 17 insertions(+), 17 deletions(-) diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 13da117b99e2..1848ea10de5c 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -349,7 +349,7 @@ def _all_gather_and_merge_handles( gathered: list[list[dict[str, tuple]] | None] = [None] * world_size torch.distributed.all_gather_object(gathered, handles) torch.distributed.barrier() - torch.cuda.synchronize() + torch.accelerator.synchronize() if self.is_sender: merged: list[dict[str, tuple]] = [] diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 921f74a87e78..1691419fbeb4 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -469,8 +469,8 @@ def _post_send_sync(self) -> None: only add a dependency on the default process group that this backend otherwise does not have. """ - if torch.cuda.is_available(): - torch.cuda.current_stream().synchronize() + if torch.accelerator.is_available(): + torch.accelerator.current_stream().synchronize() 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 8c0813b00cbe..95a681f87d25 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -227,7 +227,7 @@ def packed_nccl_broadcast_consumer( [] for _ in range(num_buffers) ] packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)] - caller_stream = torch.cuda.current_stream() + caller_stream = torch.accelerator.current_stream() prev_received: torch.Event | None = None def receive_and_load(buffer_idx: int) -> None: @@ -351,7 +351,7 @@ def packed_ipc_producer( if total_bytes and total_bytes + flat.numel() > buffer_size_bytes: # Drain queued copies so the consumer sees a fully-written buffer. - torch.cuda.current_stream().synchronize() + torch.accelerator.current_stream().synchronize() yield PackedIpcChunk( names=names, shapes=shapes, @@ -372,7 +372,7 @@ def packed_ipc_producer( total_bytes += flat.numel() if total_bytes: - torch.cuda.current_stream().synchronize() + torch.accelerator.current_stream().synchronize() yield PackedIpcChunk( names=names, shapes=shapes, diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index 2293be3ccda6..195b88808020 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -348,7 +348,7 @@ def __init__( # measurably slows the sync. self._ring_depth = 1 self._dest_buffers: list[dict[torch.dtype, torch.Tensor]] = [{}] - self._slot_read_done: list[Any] = [] # one torch.cuda.Event per slot + self._slot_read_done: list[Any] = [] # one torch.Event per slot # Generation counters guarding the events (see _ensure_proc_worker). self._slot_queued: list[int] = [] self._slot_done: list[int] = [] @@ -982,7 +982,7 @@ def _ensure_proc_worker(self) -> None: import queue import threading - self._slot_read_done = [torch.cuda.Event() for _ in range(self._ring_depth)] + self._slot_read_done = [torch.Event() for _ in range(self._ring_depth)] # Generation handshake for the events above: the RPC thread counts items # queued per slot, this thread counts records, and a pull may synchronize() # only once done[slot] has caught up with queued[slot]. Without it the @@ -1014,7 +1014,7 @@ def _proc_worker_loop(self) -> None: raises is recorded in ``_proc_error`` and re-raised on the RPC thread / at drain, so a failed sync fails loudly rather than corrupting silently. """ - torch.cuda.set_device(self.device) + torch.accelerator.set_device_index(self.device.index) q = self._proc_queue assert q is not None while True: @@ -1311,7 +1311,7 @@ def _process_item(self, item: "_ProcItem") -> None: results = item.results chunk = item.chunk with ( - torch.cuda.device(self.device), + torch.accelerator.device_index(self.device.index), torch.cuda.stream(self._proc_stream), torch.device(self.device), ): @@ -1356,14 +1356,14 @@ def _process_item(self, item: "_ProcItem") -> None: # Running it in order here stalls the RPC thread's slot handshake by # ~0.5-0.75s/iter, since every group's quant delays the next publish. if chunk.quant: - ready = torch.cuda.Event() + ready = torch.Event() ready.record(self._proc_stream) if self._quant_queue is None: self._run_quant(chunk.quant, ready) else: self._quant_queue.put((chunk.quant, ready)) - def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: + def _run_quant(self, layers: "list[Any]", ready: "torch.Event") -> None: """Quant/kernel-copy/reset the given COMPLETED leaf modules, exactly as _layerwise_process. Runs on the quant thread's own stream, ordered after the modules' scatters via ``ready``; touches only the scattered params @@ -1381,7 +1381,7 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: stream = self._quant_stream or self._proc_stream assert stream is not None # created in _ensure_proc_worker before use with ( - torch.cuda.device(self.device), + torch.accelerator.device_index(self.device.index), torch.cuda.stream(stream), torch.device(self.device), ): @@ -1403,7 +1403,7 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: def _quant_worker_loop(self) -> None: """Dedicated quant thread: drains (completed_modules, scatter-done event) batches. Errors surface via _proc_error like the scatter thread's.""" - torch.cuda.set_device(self.device) + torch.accelerator.set_device_index(self.device.index) q = self._quant_queue assert q is not None while True: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 51ea401ef92d..55b51928c152 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -926,7 +926,7 @@ def _pack_group_for_export(self, held: list, slot_idx: int) -> tuple: # nothing orders the two. Without this sync it packs bytes the copy has # not written yet -- silent corruption, and only for ring-packed tensors. if slot.is_cuda: - torch.cuda.current_stream(slot.device).synchronize() + torch.accelerator.current_stream(slot.device).synchronize() return {sid: reduce_args}, views, refs def _publish_async(self, group_idx: int, entries): diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index f9233516a56e..c7e72b78961e 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -368,8 +368,8 @@ def _post_send_sync(self) -> None: than relying on same-stream ordering. See `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no cross-rank barrier.""" - if torch.cuda.is_available(): - torch.cuda.current_stream().synchronize() + if torch.accelerator.is_available(): + torch.accelerator.current_stream().synchronize() def shutdown(self) -> None: if self.model_update_group is not None: From c9dbc92630bc837ec666f68dab02853fa095c8ab Mon Sep 17 00:00:00 2001 From: Ronald Date: Sun, 25 Oct 2026 11:11:08 +0800 Subject: [PATCH 032/631] [Refactor][RL] Share NCCL trainer rendezvous initialization (#63460) Signed-off-by: Ronald1995 --- tests/distributed/test_weight_transfer.py | 112 ++++++++++++++++-- .../weight_transfer/nccl_common.py | 35 +++++- .../weight_transfer/nccl_engine.py | 39 +++--- .../weight_transfer/sparse_nccl_engine.py | 34 ++---- 4 files changed, 165 insertions(+), 55 deletions(-) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 08bc35c65a6a..e629d65da4c8 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -49,6 +49,7 @@ IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo, ) +from vllm.distributed.weight_transfer.nccl_common import run_trainer_rendezvous from vllm.distributed.weight_transfer.nccl_engine import ( NCCLTrainerInitInfo, NCCLTrainerWeightTransferEngine, @@ -1669,16 +1670,63 @@ def test_ipc_trainer_init_ships_packed_to_worker(): assert client.last_init_info == {"packed": True} -def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): +def test_nccl_trainer_rendezvous_runs_concurrently_and_waits_for_worker(): + worker_started = threading.Event() + trainer_started = threading.Event() + worker_finished = threading.Event() + trainer_thread = threading.get_ident() + + def init_worker(): + assert threading.get_ident() != trainer_thread + worker_started.set() + assert trainer_started.wait(timeout=5) + worker_finished.set() + + def init_trainer(): + try: + assert threading.get_ident() == trainer_thread + assert worker_started.wait(timeout=5) + assert not worker_finished.is_set() + finally: + trainer_started.set() + + run_trainer_rendezvous(init_worker=init_worker, init_trainer=init_trainer) + assert worker_finished.is_set() + + +@pytest.mark.parametrize("failure_side", ["worker", "trainer"]) +def test_nccl_trainer_rendezvous_propagates_errors_after_worker_exits(failure_side): + trainer_started = threading.Event() + worker_finished = threading.Event() + error = RuntimeError(f"{failure_side} initialization failed") + + def init_worker(): + assert trainer_started.wait(timeout=5) + worker_finished.set() + if failure_side == "worker": + raise error + + def init_trainer(): + trainer_started.set() + if failure_side == "trainer": + raise error + + with pytest.raises(RuntimeError) as exc_info: + run_trainer_rendezvous(init_worker=init_worker, init_trainer=init_trainer) + assert exc_info.value is error + assert worker_finished.is_set() + + +@pytest.mark.parametrize("packed", [False, True]) +def test_nccl_trainer_init_ships_worker_init_info(monkeypatch, packed): """The sender's trainer_init drives the inference-side init handshake with the worker-shaped init info (rank_offset=1) while opening its own endpoint, and propagates the must-agree wire params to the worker.""" import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod # Bypass the real NCCL rendezvous. - monkeypatch.setattr( - nccl_engine_mod, "open_trainer_endpoint", lambda info: MagicMock() - ) + group = object() + monkeypatch.setattr(nccl_engine_mod, "open_trainer_endpoint", lambda info: group) client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( @@ -1687,7 +1735,7 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): master_port=29500, world_size=3, rank=0, - packed=True, + packed=packed, packed_buffer_size_bytes=1024, packed_num_buffers=3, ), @@ -1697,19 +1745,60 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): assert isinstance(engine, NCCLTrainerWeightTransferEngine) assert engine.is_sender is True - assert engine.packed is True + assert engine.packed is packed + assert engine.model_update_group is group assert client.order == ["init"] assert client.last_init_info == { "master_address": "127.0.0.1", "master_port": 29500, "rank_offset": 1, "world_size": 3, - "packed": True, + "packed": packed, "packed_buffer_size_bytes": 1024, "packed_num_buffers": 3, } +@pytest.mark.parametrize("rank", [0, 1]) +def test_nccl_trainer_init_dispatches_subclass_hooks_only_on_sender(rank): + group = object() + calls = [] + + class Trainer(NCCLTrainerWeightTransferEngine): + @classmethod + def _build_worker_init_payload(cls, init_info): + calls.append(("payload", cls, init_info)) + payload = super()._build_worker_init_payload(init_info) + payload["master_port"] += 1 + return payload + + def _open_trainer_endpoint(self, init_info): + calls.append(("endpoint", self, init_info)) + self.model_update_group = group + + init_info = NCCLTrainerInitInfo( + master_address="127.0.0.1", master_port=29500, world_size=3, rank=rank + ) + client = RecordingClient() + engine = Trainer.trainer_init( + init_info, client=client, source=ModuleSource(_module_with()) + ) + + assert type(engine) is Trainer + if rank == 0: + assert calls == [ + ("payload", Trainer, init_info), + ("endpoint", engine, init_info), + ] + assert client.last_init_info is not None + assert client.last_init_info["master_port"] == 29501 + assert engine.model_update_group is group + else: + assert calls == [] + assert client.order == [] + assert engine.model_update_group is None + + def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): """The worker engine reads packed + buffer geometry from the trainer-supplied init info at the handshake, not from the config or the @@ -2011,12 +2100,12 @@ def _sparse_patch(device: str = "cpu") -> SparseWeightPatch: def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): - """The sender's trainer_init drives the init handshake with the - worker-shaped init info; sparse ships no packed wire params, so the worker - keeps its unpacked defaults. Sparse takes no `source`.""" + """The sender's trainer_init sends worker-shaped init info with + `packed=False` and default buffer settings. Sparse takes no `source`.""" import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod - monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock()) + group = object() + monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: group) client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( @@ -2030,6 +2119,7 @@ def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): ) assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine) + assert engine.model_update_group is group assert client.order == ["init"] assert client.last_init_info == { "master_address": "127.0.0.1", diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index b0ac285cf1ef..b7de9f8becfc 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -9,6 +9,8 @@ """ import socket +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Protocol @@ -125,8 +127,39 @@ def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict: return {key: value for key, value in asdict(init_info).items() if value is not None} +def build_worker_init_payload( + init_info: "NCCLRendezvous", + *, + packed: bool = False, + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, +) -> dict: + # Workers sit at rank_offset 1, after the single trainer sender rank 0. + worker_init_info = NCCLWeightTransferInitInfo( + master_address=init_info.master_address, + master_port=init_info.master_port, + rank_offset=1, + world_size=init_info.world_size, + packed=packed, + packed_buffer_size_bytes=packed_buffer_size_bytes, + packed_num_buffers=packed_num_buffers, + ) + return worker_init_payload(worker_init_info) + + +def run_trainer_rendezvous( + *, init_worker: Callable[[], None], init_trainer: Callable[[], None] +) -> None: + # Workers and trainer must enter the rendezvous concurrently, so run the + # worker initialization RPC in a thread while opening the trainer endpoint. + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(init_worker) + init_trainer() + future.result() + + class NCCLRendezvous(Protocol): - """The TCP rendezvous fields `trainer_init` needs. + """The TCP rendezvous fields used by trainer initialization helpers. Structural so each backend can keep its own trainer init info next to its engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 1691419fbeb4..bc6180d45b5c 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -25,7 +25,8 @@ ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, - worker_init_payload, + build_worker_init_payload, + run_trainer_rendezvous, worker_init_process_group, ) from vllm.distributed.weight_transfer.nccl_common import ( @@ -280,6 +281,18 @@ def __init__( # Metadata the worker already holds; an unchanged round sends none. self._sent_meta: list[ParamMeta] | None = None + @classmethod + def _build_worker_init_payload(cls, init_info: NCCLTrainerInitInfo) -> dict: + return build_worker_init_payload( + init_info, + packed=init_info.packed, + packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, + packed_num_buffers=init_info.packed_num_buffers, + ) + + def _open_trainer_endpoint(self, init_info: NCCLTrainerInitInfo) -> None: + self.model_update_group = open_trainer_endpoint(init_info) + @classmethod def trainer_init( cls, @@ -304,28 +317,12 @@ def trainer_init( # trainer-side gather during send_weights. return engine - # Workers sit at rank_offset 1, after the single trainer sender rank 0. - worker_init_info = NCCLWeightTransferInitInfo( - master_address=init_info.master_address, - master_port=init_info.master_port, - rank_offset=1, - world_size=init_info.world_size, - packed=init_info.packed, - packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, - packed_num_buffers=init_info.packed_num_buffers, + payload = cls._build_worker_init_payload(init_info) + run_trainer_rendezvous( + init_worker=lambda: engine.client.init_weight_transfer_engine(payload), + init_trainer=lambda: engine._open_trainer_endpoint(init_info), ) - # The inference workers block inside init_weight_transfer_engine waiting - # for the NCCL rendezvous, so we kick that off on a side thread while we - # open the trainer endpoint (rank 0); both sides must rendezvous together. - with ThreadPoolExecutor(max_workers=1) as exe: - future = exe.submit( - engine.client.init_weight_transfer_engine, - worker_init_payload(worker_init_info), - ) - engine.model_update_group = open_trainer_endpoint(init_info) - future.result() # surface any inference-side init error - return engine def send_weights(self) -> None: diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index c7e72b78961e..b50a0be2c655 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -30,7 +30,8 @@ ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, - worker_init_payload, + build_worker_init_payload, + run_trainer_rendezvous, worker_init_process_group, ) from vllm.distributed.weight_transfer.nccl_common import ( @@ -229,6 +230,10 @@ def __init__( super().__init__(client=client, source=source, is_sender=is_sender) self.model_update_group: PyNcclCommunicator | None = None + @classmethod + def _build_worker_init_payload(cls, init_info: SparseNCCLTrainerInitInfo) -> dict: + return build_worker_init_payload(init_info) + @classmethod def trainer_init( cls, @@ -241,29 +246,14 @@ def trainer_init( if not engine.is_sender: return engine - # Workers sit at rank_offset 1, after the single trainer sender rank 0. - # Sparse transfers are never packed, so the worker keeps the unpacked - # defaults on its init info. - worker_init_info = NCCLWeightTransferInitInfo( - master_address=init_info.master_address, - master_port=init_info.master_port, - rank_offset=1, - world_size=init_info.world_size, + payload = cls._build_worker_init_payload(init_info) + run_trainer_rendezvous( + init_worker=lambda: engine.client.init_weight_transfer_engine(payload), + init_trainer=lambda: setattr( + engine, "model_update_group", open_trainer_endpoint(init_info) + ), ) - # The inference workers block inside init_weight_transfer_engine waiting - # for the NCCL rendezvous, so we kick that off on a side thread while we - # open the trainer endpoint (rank 0); both sides must rendezvous together. - with ThreadPoolExecutor(max_workers=1) as exe: - future = exe.submit( - engine.client.init_weight_transfer_engine, - worker_init_payload(worker_init_info), - ) - # Open the trainer endpoint as NCCL rank 0 on the current device - # (the init info satisfies the helper's rendezvous protocol). - engine.model_update_group = open_trainer_endpoint(init_info) - future.result() # surface any inference-side init error - return engine def send_weights(self, patches: Iterable[SparseWeightPatch] | None = None) -> None: From 65bc698356b4b0106630788536b6af6a916308e9 Mon Sep 17 00:00:00 2001 From: CharlesWu Date: Sun, 25 Oct 2026 12:18:49 +0800 Subject: [PATCH 033/631] [Bugfix][KV Offload] Reject misaligned chunked-local CPU suffix hits (#63266) Signed-off-by: CharlesWu <296364977+charxwu@users.noreply.github.com> Co-authored-by: CharlesWu <296364977+charxwu@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Shanshan Shen <467638484@qq.com> --- docs/features/kv_offloading_usage.md | 7 ++ .../v1/simple_kv_offload/test_integration.py | 89 ++++++++++++++++++- tests/v1/simple_kv_offload/test_scheduler.py | 35 ++++++++ vllm/v1/simple_kv_offload/manager.py | 10 +++ 4 files changed, 140 insertions(+), 1 deletion(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 60d5fa72f99c..804044b802de 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -478,6 +478,13 @@ selects disk slots and is checked against `disk_capacity_bytes`. Omitting `num_cpu_blocks` preserves automatic per-worker sizing. This option applies to `SimpleCPUOffloadConnector`, separately from `OffloadingConnector` above. +## Chunked-local attention with `SimpleCPUOffloadConnector` + +For chunked-local attention, CPU cache reuse requires the GPU-local prefix to +end at an attention-chunk boundary for every chunked-local cache group. Otherwise, +the connector keeps the GPU prefix and recomputes the remaining tokens. Requests +with no GPU-local prefix, or with a chunk-aligned prefix, can still reuse CPU KV. + ## Further Reading - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput). diff --git a/tests/v1/simple_kv_offload/test_integration.py b/tests/v1/simple_kv_offload/test_integration.py index ea6caa421565..c9119b7bff23 100644 --- a/tests/v1/simple_kv_offload/test_integration.py +++ b/tests/v1/simple_kv_offload/test_integration.py @@ -5,8 +5,11 @@ import time import pytest +import torch +from transformers import Llama4ForCausalLM, Llama4TextConfig -from vllm import LLM, SamplingParams, TokensPrompt +from tests.models.utils import check_logprobs_close +from vllm import LLM, SamplingParams, TokensPrompt, envs from vllm.config import KVTransferConfig from vllm.platforms import current_platform @@ -26,6 +29,90 @@ ] +def test_chunked_local_partial_gpu_hit(vllm_runner, tmp_path, monkeypatch): + """A CPU suffix after a short GPU hit must preserve the cold prediction.""" + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + envs.disable_envs_cache() + config = Llama4TextConfig( + vocab_size=512, + hidden_size=128, + intermediate_size=256, + intermediate_size_mlp=256, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=64, + max_position_embeddings=256, + attention_chunk_size=64, + no_rope_layers=[1, 0], + interleave_moe_layer_step=100, + num_local_experts=2, + use_qk_norm=False, + initializer_range=0.08, + ) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(1234) + model = Llama4ForCausalLM(config) + model.save_pretrained(tmp_path) + del model + + with vllm_runner( + str(tmp_path), + skip_tokenizer_init=True, + dtype="float16", + enforce_eager=True, + compilation_config={"mode": 0}, + attention_config={"backend": "TRITON_ATTN"}, + max_model_len=256, + max_num_batched_tokens=256, + max_num_seqs=1, + block_size=16, + num_gpu_blocks_override=64, + kv_cache_memory_bytes=16 << 20, + enable_prefix_caching=True, + disable_hybrid_kv_cache_manager=False, + kv_transfer_config=KVTransferConfig( + kv_connector="SimpleCPUOffloadConnector", + kv_role="kv_both", + kv_connector_extra_config={"cpu_bytes_to_use": 16 << 20}, + ), + ) as runner: + llm = runner.get_llm() + tokens = [3 + (i * 37) % 500 for i in range(81)] + prompt = TokensPrompt(prompt_token_ids=tokens) + params = SamplingParams( + temperature=0, max_tokens=1, logprobs=5, ignore_eos=True, detokenize=False + ) + cold = llm.generate(prompt, params, use_tqdm=False)[0].outputs[0] + + def reset_gpu_cache(): + deadline = time.monotonic() + 30 + while not llm.reset_prefix_cache(): + assert time.monotonic() < deadline, "CPU stores did not finish" + llm.llm_engine.step() + + reset_gpu_cache() + cpu_hit = llm.generate(prompt, params, use_tqdm=False)[0] + assert cpu_hit.num_cached_tokens == 80 + reset_gpu_cache() + short_params = params.clone() + short_params.skip_reading_prefix_cache = True + llm.generate( + TokensPrompt(prompt_token_ids=tokens[:17]), short_params, use_tqdm=False + ) + mixed_hit = llm.generate(prompt, params, use_tqdm=False)[0] + for name, output in (("CPU hit", cpu_hit), ("mixed hit", mixed_hit)): + cached = output.outputs[0] + check_logprobs_close( + outputs_0_lst=[(list(cold.token_ids), "", cold.logprobs)], + outputs_1_lst=[(list(cached.token_ids), "", cached.logprobs)], + name_0="cold", + name_1=name, + always_check_logprobs=True, + ) + assert mixed_hit.num_cached_tokens == 16 + + def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM: kv_transfer_config = KVTransferConfig( kv_connector="SimpleCPUOffloadConnector", diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 4eff332ce056..d56ab77797f9 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -45,6 +45,7 @@ register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, KVCacheConfig, @@ -1789,6 +1790,40 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None: ) +@pytest.mark.parametrize( + "local_tokens,cached_chunk_blocks,expected_external_tokens", + [(0, 4, 64), (0, 5, 80), (16, 4, 0), (16, 5, 0), (64, 4, 0), (64, 5, 16)], +) +def test_chunked_local_cpu_lookup_requires_chunk_aligned_prefix( + local_tokens: int, cached_chunk_blocks: int, expected_external_tokens: int +) -> None: + """Recompute suffixes whose relative and absolute chunk boundaries differ.""" + config = _make_kv_cache_config(num_blocks=32, num_groups=2) + config.kv_cache_groups[1].kv_cache_spec = ChunkedLocalAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + attention_chunk_size=64, + ) + fix = make_scheduler(num_cpu_blocks=32, kv_cache_config=config) + scheduler = fix.scheduler + request = make_request(num_blocks=5) + pool = scheduler.cpu_block_pool + cached = _allocate_gpu_blocks(pool, request, 5, group_id=0) + cached += _allocate_gpu_blocks(pool, request, cached_chunk_blocks, group_id=1) + pool.free_blocks(cached) + free_before_lookup = pool.get_num_free_blocks() + + assert scheduler.get_num_new_matched_tokens(request, local_tokens) == ( + expected_external_tokens, + expected_external_tokens > 0, + ) + + scheduler.request_finished(request, []) + assert pool.get_num_free_blocks() == free_before_lookup + + # --------------------------------------------------------------------------- # Test 11: TOCTOU between Phase A and Phase B (regression for #39702) # --------------------------------------------------------------------------- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 5d23d6cab1ff..62a022d9e0f2 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -38,6 +38,7 @@ ) from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, FullAttentionSpec, MambaSpec, SlidingWindowSpec, @@ -385,6 +386,15 @@ def get_num_new_matched_tokens( ) return 0, False + # Suffix-relative chunk boundaries must match absolute token positions. + for group_id in self.prefix_cacheable_group_ids: + spec = self.cpu_kv_cache_config.kv_cache_groups[group_id].kv_cache_spec + if ( + isinstance(spec, ChunkedLocalAttentionSpec) + and num_computed_tokens % spec.attention_chunk_size + ): + return 0, False + num_skipped_hashes = num_computed_tokens // self.hash_block_size remaining_hashes = request.block_hashes[num_skipped_hashes:] From 2c19f6793e98384ee38ba9b5df425b7378f3ae95 Mon Sep 17 00:00:00 2001 From: Chenxi Ma <114001934+algebra-MCX@users.noreply.github.com> Date: Sun, 25 Oct 2026 12:45:25 +0800 Subject: [PATCH 034/631] [Bugfix][MRV2] Keep routed-expert rows of every DBO microbatch (#63372) Signed-off-by: Chenxi Ma <114001934+algebra-MCX@users.noreply.github.com> Signed-off-by: Ronald1995 Signed-off-by: Ao Shen Co-authored-by: Claude Opus 5.5 Co-authored-by: Ronald1995 Co-authored-by: Ao Shen --- tests/v1/worker/test_gpu_ubatch_slicing.py | 32 +++++++++++++++++-- .../fused_moe/routed_experts_capturer.py | 9 ++++-- vllm/v1/worker/gpu/ubatch_utils.py | 1 + 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index a02ed6c08577..9a53357047f0 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -29,6 +29,9 @@ VllmConfig, ) from vllm.forward_context import create_forward_context +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker import gpu_ubatch_wrapper as legacy_gpu_ubatch_wrapper @@ -592,19 +595,36 @@ def test_runner_allocates_one_buffer_pair_per_microbatch(): class _YieldingModel(torch.nn.Module): """Toy model that hands off to the other microbatch mid-forward.""" - def __init__(self, trace: list[tuple[str, int]]): + def __init__( + self, + trace: list[tuple[str, int]], + capturer: RoutedExpertsCapturer | None = None, + ): super().__init__() self.trace = trace + self.capturer = capturer def forward(self, input_ids, positions, intermediate_tensors=None, **kwargs): self.trace.append(("enter", dbo_current_ubatch_id())) out = input_ids.float().unsqueeze(-1) * 2 + positions.float().unsqueeze(-1) # Stands in for the expert all-to-all handoff point. dbo_yield() + if self.capturer is not None: + self.capturer.capture(0, input_ids.view(-1, 1).int()) self.trace.append(("exit", dbo_current_ubatch_id())) return out +def _make_capturer(num_tokens: int, device: torch.device) -> RoutedExpertsCapturer: + capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) + capturer.dp_rank = 0 + capturer.tp_size = 1 + capturer.device_buffer = torch.full( + (num_tokens, 1, 1), -1, dtype=torch.int32, device=device + ) + return capturer + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") def test_ubatch_runner_overlaps_and_matches_single_batch(): """Microbatches interleave at the yield point and produce the same output.""" @@ -622,7 +642,8 @@ def test_ubatch_runner_overlaps_and_matches_single_batch(): ) trace: list[tuple[str, int]] = [] - model = _YieldingModel(trace) + capturer = _make_capturer(16, device) + model = _YieldingModel(trace, capturer) output = runner.run(model, model_inputs, ubatch_state) expected = model_inputs["input_ids"].float().unsqueeze(-1) * 2 + model_inputs[ @@ -632,6 +653,9 @@ def test_ubatch_runner_overlaps_and_matches_single_batch(): # Both microbatches reach the handoff before either finishes. assert trace == [("enter", 0), ("enter", 1), ("exit", 0), ("exit", 1)] + torch.testing.assert_close( + capturer.device_buffer[:, 0, 0], model_inputs["input_ids"].int() + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") @@ -850,7 +874,8 @@ def test_capturable_run_replays_as_a_cudagraph(): ], ) - model = _YieldingModel([]) + capturer = _make_capturer(num_tokens, device) + model = _YieldingModel([], capturer) graph = torch.cuda.CUDAGraph() finish = runner.begin_capturable_run( model, model_inputs, ubatch_state, for_capture=True @@ -867,6 +892,7 @@ def test_capturable_run_replays_as_a_cudagraph(): (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device ) torch.testing.assert_close(captured, expected) + assert (capturer.device_buffer == 3).all() def _request_slices(input_batch: InputBatch) -> list[slice]: 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 8ece4da679ef..85ccc888d565 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -17,6 +17,7 @@ from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.platforms import current_platform +from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset logger = logging.getLogger(__name__) @@ -91,8 +92,12 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: f"layer count {self.device_buffer.shape[1]}" ) - local_topk_ids = self._restore_local_topk_ids(topk_ids, get_forward_context()) - self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids + ctx = get_forward_context() + local_topk_ids = self._restore_local_topk_ids(topk_ids, ctx) + offset = dbo_current_ubatch_token_offset() + self.device_buffer[offset : offset + len(local_topk_ids), layer_id] = ( + local_topk_ids + ) def _restore_local_topk_ids( self, topk_ids: torch.Tensor, ctx: ForwardContext diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 3fa9783bf50c..5177ee8e374b 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -462,6 +462,7 @@ def begin_capturable_run( compute_stream=compute_stream, forward_contexts=ubatch_state.forward_contexts, ready_barrier=self.ready_barrier, + token_offsets=[s.token_slice.start for s in ubatch_slices], ) outputs: dict[int, Any] = {} From e4a5b5537344b103ed10de7764378ef09604ef3d Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Sun, 25 Oct 2026 12:47:31 +0800 Subject: [PATCH 035/631] [Doc] MiniMax-M3: InferenceX AgentX lanes as tuned presets (#63534) Signed-off-by: esmeetu Co-authored-by: Claude Opus 5.5 --- recipes/models/MiniMaxAI/MiniMax-M3.yaml | 407 ++++++++++++++--------- 1 file changed, 251 insertions(+), 156 deletions(-) diff --git a/recipes/models/MiniMaxAI/MiniMax-M3.yaml b/recipes/models/MiniMaxAI/MiniMax-M3.yaml index 0f430d515496..4518948489e8 100644 --- a/recipes/models/MiniMaxAI/MiniMax-M3.yaml +++ b/recipes/models/MiniMaxAI/MiniMax-M3.yaml @@ -4,7 +4,7 @@ meta: provider: "MiniMax" description: "MiniMax M3 vision-language MoE (427B total / 26B active) for frontier coding, agent toolchains, and 1M-token reasoning via MSA sparse attention — native multimodal (image + video + computer use); BF16 plus MXFP8, NVIDIA Blackwell NVFP4, and AMD MI355X MXFP4 variants. Runs on NVIDIA (Hopper/Blackwell) and AMD CDNA4/CDNA3." date_added: 2026-06-12 - date_updated: 2026-10-23 + date_updated: 2026-10-25 tasks: - text - multimodal @@ -365,6 +365,256 @@ strategy_overrides: - "--attention_config.indexer_kv_dtype" - "fp8" +# Tuned presets (SCHEMA.md): the InferenceX AgentX lanes, each a delta on the +# baseline command for one GPU and variant. All are opt-in: the builder's +# defaults stay the recipe's own; pick a preset (and a point) to render a lane. +presets: + h200_agentx: + label: "AgentX TP8" + hardware: h200 + strategy: single_node_tp + variants: [mxfp8] + workload: "Agentic (InferenceX AgentX)" + description: "One 8-GPU H200 node, TP8, with the GQA EAGLE3 draft (three tokens) and thinking always on; the scheduler batch is twice the client count." + source: { label: "InferenceX #4429", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4429", key: "minimaxm3-fp8-h200-vllm-agentic-mtp" } + image: "vllm/vllm-openai:nightly-83c078d955cd6e0c63d3b12abe669c8de936a370@sha256:cb12c3016938a4ee98ed176ac833984089460689a0982cff4b5a60bb45b3cded" + select: + features: { spec_decoding: eagle3_gqa, text_only: true, thinking_always_on: true } + # The lane passes these; the builder's MXFP8 command does not. + args: ["--trust-remote-code", "--block-size", "128"] + # --max-num-seqs is twice the client count; the CUDA graph capture size covers + # it at four tokens per sequence (the sequence plus three draft tokens). + points: + c1: { label: "1 client", concurrency: 1, args: ["--max-num-seqs", "2", "--max-cudagraph-capture-size", "8"] } + c2: { label: "2 clients", concurrency: 2, args: ["--max-num-seqs", "4", "--max-cudagraph-capture-size", "16"] } + c4: { label: "4 clients", concurrency: 4, args: ["--max-num-seqs", "8", "--max-cudagraph-capture-size", "32"] } + c6: { label: "6 clients", concurrency: 6, args: ["--max-num-seqs", "12", "--max-cudagraph-capture-size", "48"] } + c8: { label: "8 clients", concurrency: 8, args: ["--max-num-seqs", "16", "--max-cudagraph-capture-size", "64"] } + c10: { label: "10 clients", concurrency: 10, args: ["--max-num-seqs", "20", "--max-cudagraph-capture-size", "80"] } + c12: { label: "12 clients", concurrency: 12, args: ["--max-num-seqs", "24", "--max-cudagraph-capture-size", "96"] } + c14: { label: "14 clients", concurrency: 14, args: ["--max-num-seqs", "28", "--max-cudagraph-capture-size", "112"] } + runbook: | + The H200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp8-h200-vllm-agentic-mtp`, [SemiAnalysisAI/InferenceX#4429](https://github.com/SemiAnalysisAI/InferenceX/pull/4429)) serves the MXFP8 checkpoint with the GQA EAGLE3 draft (three speculative tokens) as one aggregated vLLM worker on a single 8-GPU H200 node: TP8, no expert parallelism, no host KV offload. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length (2.78 for MiniMax-M3 with three draft tokens, thinking on; [`minimaxm3_eagle3_gqa.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/main/inferencex-e2e/infx/golden_al_distribution/minimaxm3_eagle3_gqa.yaml)) so every point is measured at the same acceptance rate. The command above does not set it. + + b200_agentx: + label: "AgentX TP4" + hardware: b200 + strategy: single_node_tp + variants: [nvfp4] + workload: "Agentic (InferenceX AgentX)" + description: "TP4 on four of the node's eight B200 GPUs with the GQA EAGLE3 draft (three tokens, local argmax reduction); the larger points add host KV offload." + source: { label: "InferenceX #4382", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4382", key: "minimaxm3-fp4-b200-vllm-agentic-mtp" } + image: "vllm/vllm-openai:nightly-af7f9488c2210d67e1033ecdc845b087ee7fe92b@sha256:8e2ffbb9f866de678b22e7446903b1ee5a5766c51a948013277477c462a68a77" + tp: 4 + select: + features: { spec_decoding: eagle3_gqa, text_only: true, thinking_always_on: true } + # The variant's B200 block already carries the lane's shared flags (prefix + # caching, 16384 batched tokens, 0.95 memory, autotune off, stream interval). + points: + # Concurrency 15 and 32 leave the MSA decode backend at the image default. + c15: + label: "15 clients" + concurrency: 15 + args: ["--max-num-seqs", "30", "--max-cudagraph-capture-size", "120"] + remove_args: ["--attention_config.minimax_m3_msa_decode_backend"] + select: { kv_offload: "off" } + c30: + label: "30 clients · KV offload" + concurrency: 30 + args: ["--max-num-seqs", "60", "--max-cudagraph-capture-size", "4096"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 256, 320, 384, 448, 512, 640, 768, 896, 1024, 1280, 1536, 1792, 2048, 2560, 3072, 3584, 4096] } + c32: + label: "32 clients · KV offload" + concurrency: 32 + args: ["--max-num-seqs", "64", "--max-cudagraph-capture-size", "256"] + remove_args: ["--attention_config.minimax_m3_msa_decode_backend"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + c38: + label: "38 clients · KV offload" + concurrency: 38 + args: ["--max-num-seqs", "76", "--max-cudagraph-capture-size", "4096"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 384, 448, 512, 640, 768, 896, 1024, 1280, 1536, 1792, 2048, 2560, 3072, 3584, 4096] } + runbook: | + The B200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-b200-vllm-agentic-mtp` and `minimaxm3-fp4-b200-vllm-agentic-mtp-c10`, [SemiAnalysisAI/InferenceX#4382](https://github.com/SemiAnalysisAI/InferenceX/pull/4382)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft as one aggregated vLLM worker on a single 8-GPU B200 node. Most points run TP4; the concurrency-10 arm also has a TP8 point. + + The host-KV-offload points pin 1,073 GB of host memory for KV blocks across the node (lazy offload). + + The **Verified** badge for this GPU is recipe-wide. This preset covers the NVFP4 variant only; the GQA draft override in `features.spec_decoding` also applies to the other variants on this GPU, which were not qualified. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler, as in the command above. + + b200_agentx_c10: + label: "AgentX active-seq caps" + hardware: b200 + strategy: single_node_tp + variants: [nvfp4] + workload: "Agentic (InferenceX AgentX)" + description: "Caps the sequences the scheduler keeps active (--max-num-active-seqs) well below the client count, with a five- or three-token GQA draft without local argmax reduction." + source: { label: "InferenceX #4382", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4382", key: "minimaxm3-fp4-b200-vllm-agentic-mtp-c10" } + image: "vllm/vllm-openai:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9@sha256:98adb2311118263a453ec0acd5d4b72bda3db26cecc0fe8ffc94b9ab03c951f9" + tp: 4 + select: + features: { spec_decoding: eagle3_gqa, text_only: true, thinking_always_on: true } + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "mnnvl" } + json_args: + "--speculative-config": { use_local_argmax_reduction: null } + # --max-num-seqs stays at its default; every point sets --max-cudagraph-capture-size + # to the last of its capture sizes. + points: + c10_tp4: + label: "10 clients · TP4" + concurrency: 10 + args: ["--max-num-active-seqs", "1", "--gpu-memory-utilization", "0.925", "--max-cudagraph-capture-size", "4096"] + select: { kv_offload: "off" } + json_args: + "--speculative-config": { num_speculative_tokens: 5 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096] } + c10_tp8: + label: "10 clients · TP8" + concurrency: 10 + tp: 8 + args: ["--max-num-active-seqs", "1", "--gpu-memory-utilization", "0.9", "--max-cudagraph-capture-size", "4096"] + remove_args: ["--disable-flashinfer-autotune"] # FlashInfer autotune on at this point + select: { kv_offload: "off" } + json_args: + "--speculative-config": { num_speculative_tokens: 5 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096] } + c20_a4: + label: "20 clients · 4 active" + concurrency: 20 + args: ["--max-num-active-seqs", "4", "--gpu-memory-utilization", "0.9", "--max-cudagraph-capture-size", "8192"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--speculative-config": { num_speculative_tokens: 3 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096, 6144, 8192] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + c20_a3: + label: "20 clients · 3 active" + concurrency: 20 + args: ["--max-num-active-seqs", "3", "--gpu-memory-utilization", "0.9", "--max-cudagraph-capture-size", "8192"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--speculative-config": { num_speculative_tokens: 3 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096, 6144, 8192] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + c32_a6: + label: "32 clients · 6 active" + concurrency: 32 + args: ["--max-num-active-seqs", "6", "--gpu-memory-utilization", "0.925", "--max-cudagraph-capture-size", "4096"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--speculative-config": { num_speculative_tokens: 3 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + c32_a16: + label: "32 clients · 16 active" + concurrency: 32 + args: ["--max-num-active-seqs", "16", "--gpu-memory-utilization", "0.925", "--max-cudagraph-capture-size", "4096"] + select: { kv_offload: simple } + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + json_args: + "--speculative-config": { num_speculative_tokens: 3 } + "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512, 640, 768, 1024, 1536, 2048, 3072, 4096] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1073000000000, lazy_offload: true } } + runbook: | + The B200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-b200-vllm-agentic-mtp` and `minimaxm3-fp4-b200-vllm-agentic-mtp-c10`, [SemiAnalysisAI/InferenceX#4382](https://github.com/SemiAnalysisAI/InferenceX/pull/4382)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft as one aggregated vLLM worker on a single 8-GPU B200 node. Most points run TP4; the concurrency-10 arm also has a TP8 point. + + The **Verified** badge for this GPU is recipe-wide. This preset covers the NVFP4 variant only; the GQA draft override in `features.spec_decoding` also applies to the other variants on this GPU, which were not qualified. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler, as in the command above. + + b300_agentx: + label: "AgentX" + hardware: b300 + strategy: single_node_tp + variants: [nvfp4] + workload: "Agentic (InferenceX AgentX)" + description: "One B300 node: TP8 at one client, TP4 at 20 and 24, TP2 with host KV offload at 24 and 28; GQA EAGLE3 draft (three tokens, local argmax reduction)." + source: { label: "InferenceX #4357", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4357", key: "minimaxm3-fp4-b300-vllm-agentic-mtp" } + image: "vllm/vllm-openai:nightly-49cb2bbf71cfd36ee5aa33d1e184b1cddff55294@sha256:8c1434bfee7d7626eef78ef70600434ffba4a1b8813cb1d1c8707ad25f47035e" + select: + features: { spec_decoding: eagle3_gqa, text_only: true, thinking_always_on: true } + json_args: + "--compilation-config": { cudagraph_mode: "FULL_AND_PIECEWISE" } + # The variant's B300 block carries the shared flags. Only the TP2 points keep + # the MSA cutlass decode backend. + points: + c1_tp8: + label: "1 client · TP8" + concurrency: 1 + tp: 8 + args: ["--max-num-seqs", "2", "--gpu-memory-utilization", "0.95"] + remove_args: ["--attention_config.minimax_m3_msa_decode_backend"] + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "trtllm" } + select: { kv_offload: "off" } + json_args: + "--compilation-config": { cudagraph_capture_sizes: [4, 8, 256, 512, 1024, 2048] } + c20_tp4: + label: "20 clients · TP4" + concurrency: 20 + tp: 4 + args: ["--max-num-seqs", "30", "--max-num-active-seqs", "6", "--gpu-memory-utilization", "0.95"] + remove_args: ["--attention_config.minimax_m3_msa_decode_backend"] + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "trtllm" } + select: { kv_offload: "off" } + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 256, 512, 1024, 2048, 3072, 4096, 5120, 6144, 7168, 8192, 10240, 12288, 14336, 16384] } + c24_tp4: + label: "24 clients · TP4" + concurrency: 24 + tp: 4 + args: ["--max-num-seqs", "7", "--gpu-memory-utilization", "0.95"] + remove_args: ["--attention_config.minimax_m3_msa_decode_backend"] + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "mnnvl" } + select: { kv_offload: "off" } + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 256, 512, 1024, 2048, 2560, 3072, 3584, 4096, 5120, 6144, 7168, 8192] } + c24_tp2: + label: "24 clients · TP2 · KV offload" + concurrency: 24 + tp: 2 + args: ["--max-num-seqs", "36", "--gpu-memory-utilization", "0.92"] + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "trtllm", VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + select: { kv_offload: simple } + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 256, 512, 1024, 2048] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 749000000000, lazy_offload: true } } + c28_tp2: + label: "28 clients · TP2 · KV offload" + concurrency: 28 + tp: 2 + args: ["--max-num-seqs", "36", "--gpu-memory-utilization", "0.92"] + env: { VLLM_FLASHINFER_ALLREDUCE_BACKEND: "trtllm", VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + select: { kv_offload: simple } + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 256, 512, 1024, 2048] } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 749000000000, lazy_offload: true } } + runbook: | + The B300 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-b300-vllm-agentic-mtp`, [SemiAnalysisAI/InferenceX#4357](https://github.com/SemiAnalysisAI/InferenceX/pull/4357)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft (three speculative tokens) as one aggregated vLLM worker on a single B300 node: TP8 at one client, TP4 at 20 and 24 clients, and TP2 at 24 and 28 clients with host KV offload. + + The TP2 points pin 749 GB of host memory for KV blocks (lazy offload). + + The **Verified** badge for this GPU is recipe-wide. This preset covers the NVFP4 variant only; the GQA draft override in `features.spec_decoding` also applies to the other variants on this GPU, which were not qualified. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length (2.78 for MiniMax-M3 with three draft tokens, thinking on; [`minimaxm3_eagle3_gqa.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/main/inferencex-e2e/infx/golden_al_distribution/minimaxm3_eagle3_gqa.yaml)) so every point is measured at the same acceptance rate. The command above does not set it. + + guide: | ## Overview @@ -818,161 +1068,6 @@ guide: | --speculative-config '{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3-GQA","num_speculative_tokens":3,"attention_backend":"FLASHINFER","use_local_argmax_reduction":true}' ``` - ## H200: InferenceX AgentX aggregated (vLLM) - - The H200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp8-h200-vllm-agentic-mtp`, [SemiAnalysisAI/InferenceX#4429](https://github.com/SemiAnalysisAI/InferenceX/pull/4429)) serves the MXFP8 checkpoint with the GQA EAGLE3 draft (three speculative tokens) as one aggregated vLLM worker on a single 8-GPU H200 node: TP8, no expert parallelism, no host KV offload. It ran on `vllm/vllm-openai:nightly-83c078d955cd6e0c63d3b12abe669c8de936a370@sha256:cb12c3016938a4ee98ed176ac833984089460689a0982cff4b5a60bb45b3cded` (vLLM `0.32.1rc1.dev361`). The command builder keeps its own pinned H200 MXFP8 image, an earlier nightly; use the image above to reproduce these points. - - In the command builder, pick **H200** and the **MXFP8** variant (the builder keeps the current variant when you switch hardware, so select MXFP8 explicitly), turn on **Spec Decoding → GQA**, **Text-only** and **Thinking Always On**, and use **Single-node TP**. On an 8-GPU H200 node the builder sizes TP to 8. That renders every serving flag and environment variable below except the per-point `--max-num-seqs` and `--max-cudagraph-capture-size`, which stay in this guide. The lane does not set `--max-model-len`. - - Every point uses `--gpu-memory-utilization 0.95`, which the builder already sets for H200. The scheduler batch is twice the client count, and the CUDA graph capture size covers it at four tokens per sequence (the sequence plus three draft tokens): - - | Clients | `--max-num-seqs` | `--max-cudagraph-capture-size` | - |---:|---:|---:| - | 1 | 2 | 8 | - | 2 | 4 | 16 | - | 4 | 8 | 32 | - | 6 | 12 | 48 | - | 8 | 16 | 64 | - | 10 | 20 | 80 | - | 12 | 24 | 96 | - | 14 | 28 | 112 | - - ### Command - - ```bash - VLLM_MARLIN_USE_ATOMIC_ADD=1 \ - vllm serve MiniMaxAI/MiniMax-M3-MXFP8 --tensor-parallel-size 8 --trust-remote-code \ - --attention-backend FLASH_ATTN --linear-backend humming --kv-cache-dtype fp8 \ - --attention-config '{"indexer_kv_dtype":"fp8"}' --block-size 128 --enable-prefix-caching \ - --speculative-config '{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3-GQA","num_speculative_tokens":3,"attention_backend":"FLASH_ATTN"}' \ - --max-num-seqs --max-cudagraph-capture-size --gpu-memory-utilization 0.95 --language-model-only \ - --reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice \ - --default-chat-template-kwargs '{"thinking_mode":"enabled"}' - ``` - - The lane passes `--trust-remote-code`; the builder's MXFP8 command does not. - - The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length (2.78 for MiniMax-M3 with three draft tokens, thinking on; [`minimaxm3_eagle3_gqa.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/main/inferencex-e2e/infx/golden_al_distribution/minimaxm3_eagle3_gqa.yaml)) so every point is measured at the same acceptance rate. The command above does not set it. - - ## B200: InferenceX AgentX aggregated (vLLM) - - The B200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-b200-vllm-agentic-mtp` and `minimaxm3-fp4-b200-vllm-agentic-mtp-c10`, [SemiAnalysisAI/InferenceX#4382](https://github.com/SemiAnalysisAI/InferenceX/pull/4382)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft as one aggregated vLLM worker on a single 8-GPU B200 node. Most points run TP4; the concurrency-10 arm also has a TP8 point. The two keys use different images: `vllm/vllm-openai:nightly-af7f9488c2210d67e1033ecdc845b087ee7fe92b@sha256:8e2ffbb9f866de678b22e7446903b1ee5a5766c51a948013277477c462a68a77` for the first and `vllm/vllm-openai:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9@sha256:98adb2311118263a453ec0acd5d4b72bda3db26cecc0fe8ffc94b9ab03c951f9` for the second. - - In the command builder, pick **B200** and the **NVFP4** variant (the builder keeps the current variant when you switch hardware, so select NVFP4 explicitly), turn on **Spec Decoding → GQA** and **Text-only**, and use **Single-node TP**. On an 8-GPU B200 node the builder sizes TP to 8, so change `--tensor-parallel-size` to `4` for the TP4 points. That renders the flags and environment below as a `vllm serve` command; the per-point settings in the tables stay in this guide. The B200 **Verified** badge is recipe-wide; this profile covers the NVFP4 variant only, and the GQA draft override also applies to the other variants on B200, which were not qualified. - - For the host-KV-offload points, don't use the **KV Offload → Simple** pill as is: replace its `--kv-transfer-config` with the qualified one below and set `VLLM_USE_SIMPLE_KV_OFFLOAD=1`. - - ### `-vllm-agentic-mtp` (three-token draft, `nightly-af7f9488`) - - Every point uses the GQA draft with FlashInfer draft attention and local argmax reduction, `--gpu-memory-utilization 0.95`, and `VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm`. - - | Clients | TP | `--max-num-seqs` | Host KV offload | `--max-cudagraph-capture-size` | - |---:|---:|---:|---|---:| - | 15 | 4 | 30 | none | 120 | - | 30 | 4 | 60 | SimpleCPUOffload | 4096 (explicit list below) | - | 32 | 4 | 64 | SimpleCPUOffload | 256 | - | 38 | 4 | 76 | SimpleCPUOffload | 4096 (explicit list below) | - - The concurrency-30 and -38 points also set `"minimax_m3_msa_decode_backend":"cutlass"` in `--attention-config` (the builder's NVFP4 default) and an explicit `--cudagraph-capture-sizes`: - - - 30: `1,2,4,8`, every 8 from `16` to `232`, then `240,256,320,384,448,512,640,768,896,1024,1280,1536,1792,2048,2560,3072,3584,4096` - - 38: `1,2,4,8`, every 8 from `16` to `248`, then `256,272,288,304,320,384,448,512,640,768,896,1024,1280,1536,1792,2048,2560,3072,3584,4096` - - The other points omit `minimax_m3_msa_decode_backend`, so the pinned image's default applies. - - ### `-vllm-agentic-mtp-c10` (active-sequence caps, `nightly-ac7509e2`) - - These arms cap the number of sequences the scheduler keeps active with `--max-num-active-seqs`, well below the client count, and set `VLLM_FLASHINFER_ALLREDUCE_BACKEND=mnnvl`. They use the GQA draft with FlashInfer draft attention but without local argmax reduction, and leave `--max-num-seqs` at its default. - - | Clients | TP | `--max-num-active-seqs` | Draft tokens | Host KV offload | `--gpu-memory-utilization` | - |---:|---:|---:|---:|---|---:| - | 10 | 4 | 1 | 5 | none | 0.925 | - | 10 | 8 | 1 | 5 | none | 0.9 (FlashInfer autotune on) | - | 20 | 4 | 4 | 3 | SimpleCPUOffload | 0.9 | - | 20 | 4 | 3 | 3 | SimpleCPUOffload | 0.9 | - | 32 | 4 | 6 | 3 | SimpleCPUOffload | 0.925 | - | 32 | 4 | 16 | 3 | SimpleCPUOffload | 0.925 | - - The draft is `{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3-GQA","num_speculative_tokens":,"attention_backend":"FLASHINFER"}`. Every point sets `--max-cudagraph-capture-size` to the last entry of its `--cudagraph-capture-sizes`: `1,2,4,8`, every 4 from `12` to `256`, every 16 from `272` to `512`, then `640,768,1024,1536,2048,3072,4096` (max `4096`). The concurrency-20 points append `6144,8192` (max `8192`). - - ### Host KV offload - - ```bash - export VLLM_USE_SIMPLE_KV_OFFLOAD=1 - --kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":1073000000000,"lazy_offload":true}}' - ``` - - That is 1,073 GB of pinned host memory for KV blocks across the node. - - ### Command - - ```bash - VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm \ - vllm serve nvidia/MiniMax-M3-NVFP4 --tensor-parallel-size 4 --trust-remote-code \ - --attention-config '{"backend":"FLASHINFER","use_trtllm_attention":true,"indexer_kv_dtype":"fp8"}' \ - --kv-cache-dtype fp8 --block-size 128 --enable-prefix-caching --enable-chunked-prefill \ - --speculative-config '{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3-GQA","num_speculative_tokens":3,"attention_backend":"FLASHINFER","use_local_argmax_reduction":true}' \ - --max-num-batched-tokens 16384 --max-num-seqs --max-cudagraph-capture-size \ - --gpu-memory-utilization 0.95 --disable-flashinfer-autotune --stream-interval 20 --language-model-only \ - --reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice \ - --default-chat-template-kwargs '{"thinking_mode":"enabled"}' - ``` - - The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler, as in the command above. - - ## B300: InferenceX AgentX aggregated (vLLM) - - The B300 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-b300-vllm-agentic-mtp`, [SemiAnalysisAI/InferenceX#4357](https://github.com/SemiAnalysisAI/InferenceX/pull/4357)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft (three speculative tokens) as one aggregated vLLM worker on a single B300 node: TP8 at one client, TP4 at 20 and 24 clients, and TP2 at 24 and 28 clients with host KV offload. It ran on `vllm/vllm-openai:nightly-49cb2bbf71cfd36ee5aa33d1e184b1cddff55294@sha256:8c1434bfee7d7626eef78ef70600434ffba4a1b8813cb1d1c8707ad25f47035e`. - - In the command builder, pick **B300** and the **NVFP4** variant (the builder keeps the current variant when you switch hardware, so select NVFP4 explicitly), turn on **Spec Decoding → GQA**, **Text-only** and **Thinking Always On**, and use **Single-node TP**. The builder sizes TP to 1 on B300, so set `--tensor-parallel-size` to `2`, `4` or `8` as in the table. That renders the static flags and environment below; the per-point settings stay in this guide. The TP8 and TP4 points do not set `minimax_m3_msa_decode_backend`, so drop the builder's `--attention_config.minimax_m3_msa_decode_backend cutlass` for them and keep it for the TP2 points. The B300 **Verified** badge is recipe-wide; this profile covers the NVFP4 variant only, and the GQA draft override also applies to the other variants on B300, which were not qualified. - - For the host-KV-offload points, don't use the **KV Offload → Simple** pill as is: replace its `--kv-transfer-config` with the qualified one below and set `VLLM_USE_SIMPLE_KV_OFFLOAD=1`. - - Every point uses the GQA draft with FlashInfer draft attention and local argmax reduction and `"cudagraph_mode":"FULL_AND_PIECEWISE"`. `--max-num-active-seqs` caps how many requests the scheduler admits into RUNNING; the remaining requests wait in the queue. - - | Clients | TP | `--max-num-seqs` | `--max-num-active-seqs` | `--gpu-memory-utilization` | Host KV offload | `VLLM_FLASHINFER_ALLREDUCE_BACKEND` | - |---:|---:|---:|---:|---:|---|---| - | 1 | 8 | 2 | not set | 0.95 | none | `trtllm` | - | 20 | 4 | 30 | 6 | 0.95 | none | `trtllm` | - | 24 | 4 | 7 | not set | 0.95 | none | `mnnvl` | - | 24 | 2 | 36 | not set | 0.92 | SimpleCPUOffload | `trtllm` | - | 28 | 2 | 36 | not set | 0.92 | SimpleCPUOffload | `trtllm` | - - CUDA graph capture sizes (`cudagraph_capture_sizes` in `--compilation-config`): - - - 1 client, TP8: `4,8,256,512,1024,2048` - - 20 clients, TP4: every 4 from `8` to `120`, then `256,512,1024,2048,3072,4096,5120,6144,7168,8192,10240,12288,14336,16384` - - 24 clients, TP4: every 4 from `8` to `144`, then `256,512,1024,2048,2560,3072,3584,4096,5120,6144,7168,8192` - - 24 and 28 clients, TP2: every 4 from `8` to `144`, then `256,512,1024,2048` - - ### Host KV offload - - ```bash - export VLLM_USE_SIMPLE_KV_OFFLOAD=1 - --kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":749000000000,"lazy_offload":true}}' - ``` - - That is 749 GB of pinned host memory for KV blocks. - - ### Command - - ```bash - VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_FLASHINFER_ALLREDUCE_BACKEND= \ - vllm serve nvidia/MiniMax-M3-NVFP4 --tensor-parallel-size --trust-remote-code \ - --attention-config '{"backend":"FLASHINFER","use_trtllm_attention":true,"indexer_kv_dtype":"fp8"}' \ - --kv-cache-dtype fp8 --block-size 128 --enable-prefix-caching --enable-chunked-prefill \ - --speculative-config '{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3-GQA","num_speculative_tokens":3,"attention_backend":"FLASHINFER","use_local_argmax_reduction":true}' \ - --max-num-batched-tokens 16384 --max-num-seqs [--max-num-active-seqs ] \ - --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[]}' \ - --gpu-memory-utilization --no-enable-flashinfer-autotune --stream-interval 20 --language-model-only \ - --reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice \ - --default-chat-template-kwargs '{"thinking_mode":"enabled"}' - ``` - - For the TP2 points, add `"minimax_m3_msa_decode_backend":"cutlass"` to `--attention-config`. - - The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden EAGLE3-GQA acceptance length (2.78 for MiniMax-M3 with three draft tokens, thinking on; [`minimaxm3_eagle3_gqa.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/main/inferencex-e2e/infx/golden_al_distribution/minimaxm3_eagle3_gqa.yaml)) so every point is measured at the same acceptance rate. The command above does not set it. - ## GB200: InferenceX AgentX aggregated (Dynamo + vLLM) The GB200 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm3-fp4-gb200-dynamo-vllm-agentic-agg-mtp`, [SemiAnalysisAI/InferenceX#4281](https://github.com/SemiAnalysisAI/InferenceX/pull/4281)) serves the NVFP4 checkpoint with the GQA EAGLE3 draft (three speculative tokens) as one aggregated vLLM worker behind the Dynamo frontend's KV-aware router. TP4 runs on one 4-GPU GB200 tray; TP8 spans two trays. Image `vllm/vllm-openai:nightly-5555a2d8f7116b2b4877f66516f4128176cf722a@sha256:8896fd7b29fc4ee7ce839f146cb4757984c476ec79d0234ffa8c995dee6e339a`, Dynamo `1.6.0.dev20261016` (`pip install ai-dynamo==1.6.0.dev20261016`). From 444d69110e5edb1990a412001c7be9530d9a25b2 Mon Sep 17 00:00:00 2001 From: Chun Fang Date: Sun, 25 Oct 2026 07:48:04 +0200 Subject: [PATCH 036/631] [Doc][ROCm] Align Kimi-K3 MI355X AgentX recipe with InferenceX #4393 (#63520) Signed-off-by: Chun Fang Co-authored-by: Codex Co-authored-by: Roy Wang Co-authored-by: vllm-agent --- recipes/models/moonshotai/Kimi-K3.yaml | 87 +++++++++++++------------- 1 file changed, 45 insertions(+), 42 deletions(-) diff --git a/recipes/models/moonshotai/Kimi-K3.yaml b/recipes/models/moonshotai/Kimi-K3.yaml index 3033f16a866d..b99886f8d875 100644 --- a/recipes/models/moonshotai/Kimi-K3.yaml +++ b/recipes/models/moonshotai/Kimi-K3.yaml @@ -4,7 +4,7 @@ meta: provider: "Moonshot AI" description: "Pre-release 2.8T-parameter native multimodal MoE with Kimi Delta Attention, Gated MLA, Attention Residuals, and a 1M-token context window" date_added: 2026-07-27 - date_updated: 2026-10-22 + date_updated: 2026-10-24 tasks: - multimodal - text @@ -67,7 +67,7 @@ features: # MI355X serves the InferenceX AgentX lane's draft (see the guide's MI355X # section). The stock Inferact checkpoint needs ROCM_AITER_MLA from # vllm-project/vllm#55966, and the draft's FP8 KV assumes an FP8 target KV. - # 7 draft tokens is the lane's dcp1 value (concurrency 1 and 4), not a universal value. + # 7 draft tokens is the lane's GPU-resident c1/c4 value, not a universal value. mi355x: args: - "--kv-cache-dtype" @@ -1092,34 +1092,36 @@ guide: | ## AMD MI355X: InferenceX AgentX alignment - The MI355X single-node TP8 command follows the [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) agentic-coding (AgentX) lane `kimik3-fp4-mi355x-vllm-agentic-mtp` as of [InferenceX PR #4215](https://github.com/SemiAnalysisAI/InferenceX/pull/4215) at branch revision [`6e72cbc`](https://github.com/SemiAnalysisAI/InferenceX/tree/6e72cbc11234e624d66b8fc5eaddac50068db15c): the entry in [`configs/amd-master.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/6e72cbc11234e624d66b8fc5eaddac50068db15c/inferencex-e2e/configs/amd-master.yaml) and its [lane recipe](https://github.com/SemiAnalysisAI/InferenceX/blob/6e72cbc11234e624d66b8fc5eaddac50068db15c/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4-mtp/agentic.yaml). The lane is a benchmark configuration: text-only, TP8 on one MI355X node, with the tool-call and reasoning parsers on. + The single-node MI355X TP8 AgentX configuration follows [InferenceX #4393](https://github.com/SemiAnalysisAI/InferenceX/pull/4393): its [master entry](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/configs/amd-master.yaml#L766) and [lane recipe](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4-mtp/agentic.yaml) at `dee7824`. It serves text with tool-call and reasoning parsers, using DCP8/A2A at every point. There is no concurrency-2 point. **In the command builder.** Settings the lane uses at every concurrency are part of the MI355X TP8 command: `--gpu-memory-utilization 0.9`, `--max-model-len 1048576`, `--enable-prefix-caching`, `--kv-cache-dtype fp8`, `--attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}'`, `--stream-interval 10`, `VLLM_ROCM_AITER_MLA_ASM_PADDING=asm`, `VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm` and `AITER_FLYDSL_STAGE2_FP8=1`. Weights load with vLLM's default loader (the AMD profile's `--load-format auto`): with `fastsafetensors`, the lane's servers ran out of GPU memory while loading weights on its image. The lane's `--moe-backend auto` is vLLM's default selection, so the builder leaves that flag unset. Other AMD GPUs keep the general AMD profile. - **Spec decoding** serves the lane's draft, the published [`Inferact/Kimi-K3-DSpark`](https://huggingface.co/Inferact/Kimi-K3-DSpark), with its attention on `ROCM_AITER_MLA` ([vllm-project/vllm#55966](https://github.com/vllm-project/vllm/pull/55966)), FP8 KV, probabilistic drafting and block verification. Draft length is tuned per cell in the matrix below. - - **Target quantization** is online FP8 per-channel and differs by cell, so it stays in the guide: concurrency 1 applies `linear: fp8_per_channel` with the ignored-layer list in the command below, and every other cell quantizes only the KDA `in_proj_qkvgfab` projections. The routed experts remain MXFP4. The DSpark draft keeps its BF16 weights and activations with FP8 KV: vLLM builds the named draft from its own unquantized config, so `--quantization-config` reaches only the target. + - **Target quantization** stays in the guide: resident c1 passes no online quantization config, preserving the BF16 target weights required by specialized decode. Every other cell quantizes only the KDA `in_proj_qkvgfab` projections to FP8 per-channel. Routed experts retain their native MXFP4 weights. The DSpark draft remains BF16 with FP8 KV: vLLM builds the named draft from its own unquantized config, so removing c1's target quantization does not change the draft's dtype. - **Decode context parallelism** emits the lane's DCP8 pairing: `--attention-backend ROCM_AITER_MLA`, `--dcp-comm-backend a2a`, `ROCM_AITER_FA` prefill and the direct DCP ops off. - The **Simple** KV-offload option emits the connector the lane uses (`SimpleCPUOffloadConnector`, `kv_both`, `lazy_offload: false`). - Select **Python** in the frontend selector to match the lane, which runs vLLM's default frontend. **Per-concurrency tuning (guide only).** The lane runs the following ten cells. The builder keeps general values (`--max-num-seqs 128`, `--max-num-batched-tokens 4096`, and effective `FULL_DECODE_ONLY` CUDA graphs); use the reproduction command below for exact benchmark cells. Every cell requests `FULL_AND_PIECEWISE` with breakable CUDA graphs disabled (`VLLM_USE_BREAKABLE_CUDAGRAPH=0`); because Kimi-K3 has no `torch.compile` submodule, vLLM resolves that mode to `FULL_DECODE_ONLY` at load. - | Concurrency | DCP | KV offload | Draft tokens | `--max-num-seqs` | `--max-num-batched-tokens` | Largest CUDA graph | Target quantization | + Every cell adds `kda_prefill_backend: fused` in `--additional-config`. Only **resident c1** also sets `decode_backend: latency_specialized` and `VLLM_USE_V2_MODEL_RUNNER=1`. On the pinned image, this [MI355X specialized path](https://github.com/vllm-project/vllm/pull/60683) requires TP8, V2, gfx950, no KV connector and BF16 target projections. It handles pure speculative-verification batches of at most eight rows (c1 with K7); other steps use the normal layers. These lane-specific settings stay out of generic MI355X defaults. + + | Concurrency | DCP | Profile | Draft tokens | `--max-num-seqs` | `--max-num-batched-tokens` | Largest CUDA graph | Target online quantization | | ---: | ---: | --- | ---: | ---: | ---: | ---: | --- | - | 1 | 1 | none | 7 | 1 | 16384 | 16 | `linear` | - | 4 | 1 | none | 7 | 8 | 8192 | 64 | `in_proj_qkvgfab` | - | 4 | 8 | Simple | 5 | 8 | 8192 | 48 | `in_proj_qkvgfab` | - | 10 | 8 | Simple | 4 | 20 | 8192 | 100 | `in_proj_qkvgfab` | - | 16 | 8 | Simple | 4 | 32 | 8192 | 160 | `in_proj_qkvgfab` | - | 24 | 8 | Simple | 4 | 48 | 8192 | 240 | `in_proj_qkvgfab` | - | 32 | 8 | Simple | 4 | 64 | 8192 | 320 | `in_proj_qkvgfab` | - | 48 | 8 | Simple | 4 | 96 | 8192 | 480 | `in_proj_qkvgfab` | - | 52 | 8 | Simple | 4 | 104 | 8192 | 520 | `in_proj_qkvgfab` | - | 60 | 8 | Simple | 4 | 120 | 8192 | 600 | `in_proj_qkvgfab` | + | 1 | 8 | resident | 7 | 1 | 16384 | 16 | none | + | 4 | 8 | resident | 7 | 8 | 8192 | 64 | `in_proj_qkvgfab` | + | 4 | 8 | offload | 5 | 8 | 8192 | 48 | `in_proj_qkvgfab` | + | 10 | 8 | offload | 4 | 20 | 8192 | 100 | `in_proj_qkvgfab` | + | 16 | 8 | offload | 4 | 32 | 8192 | 160 | `in_proj_qkvgfab` | + | 24 | 8 | offload | 4 | 48 | 8192 | 240 | `in_proj_qkvgfab` | + | 32 | 8 | offload | 4 | 64 | 8192 | 320 | `in_proj_qkvgfab` | + | 48 | 8 | offload | 4 | 96 | 8192 | 480 | `in_proj_qkvgfab` | + | 52 | 8 | offload | 4 | 104 | 8192 | 520 | `in_proj_qkvgfab` | + | 60 | 8 | offload | 4 | 120 | 8192 | 600 | `in_proj_qkvgfab` | - DCP8 shards the decode KV read across all eight GPUs and uses the A2A backend. Simple offload cells use 224.875 GB of host DRAM per rank. `--max-num-seqs` is twice the concurrency, except that concurrency 1 admits one sequence. Every cell captures every even CUDA-graph size from 2 through `2 × concurrency × (1 + draft tokens)`. Agentic benchmark runs replace block rejection with the golden acceptance length for every draft length; synthetic acceptance is not suitable for serving or accuracy evaluation. + The `resident` profile keeps KV on GPU; `offload` uses `SimpleCPUOffloadConnector` with 224.875 GB of host DRAM per rank. `--max-num-seqs` is twice the concurrency, except that concurrency 1 admits one sequence. Every cell captures every even CUDA-graph size from 2 through `2 × concurrency × (1 + draft tokens)`. Agentic benchmark runs replace block rejection with the golden acceptance length; synthetic acceptance is not suitable for serving or accuracy evaluation. - **Image.** The MI355X default variant uses `vllm/vllm-openai-rocm:nightly-rocm100`, the rolling ROCm 10.0 nightly. InferenceX PR #4215 pins `vllm/vllm-openai-rocm:nightly-rocm100-49cb2bbf71cfd36ee5aa33d1e184b1cddff55294` (vLLM `49cb2bbf`). This path needs [vllm-project/vllm#53940](https://github.com/vllm-project/vllm/pull/53940) (SiTUv2 a4w4 MoE routing) and #55966. + **Image.** The builder keeps the rolling `vllm/vllm-openai-rocm:nightly-rocm100` image. Run the exact commands below inside `vllm/vllm-openai-rocm:nightly-rocm100-df6603f0583375f244d1ac6d36265462b7490392` (vLLM `df6603f0`), the image qualified by #4393 and carrying #60683. The compatibility limits above refer to this pinned runtime. **MoE environment.** The lane sets `VLLM_ROCM_USE_AITER_MOE_SITUV2=1`, as the AMD profile does, plus `AITER_BF16_FP8_MOE_BOUND=0` and `AITER_FLYDSL_STAGE2_FP8=1`. The SiTUv2 path does not read `AITER_BF16_FP8_MOE_BOUND`. `AITER_FLYDSL_STAGE2_FP8=1` makes the FlyDSL MoE down projection write its per-route outputs as MXFP8 (FP8 values with E8M0 scales) instead of BF16 before they are summed. It affects only the target's MoE layers; the DSpark draft's MLP is dense. @@ -1127,23 +1129,23 @@ guide: | ### Reproducing one lane cell - This is the real-verification form, the one InferenceX's eval jobs use and the one to use for any serving: + Run each cell in a fresh shell inside the pinned image. This real-verification form is for serving and accuracy evaluation; `PROFILE` distinguishes the two concurrency-4/DCP8 cells: ```bash - # One lane cell per server. For CONC=4, choose DCP=1 or DCP=8. - CONC=1 - DCP=1 - case "$CONC:$DCP" in - 1:1) K=7; BATCHED=16384; OFFLOAD=0 ;; - 4:1) K=7; BATCHED=8192; OFFLOAD=0 ;; - 4:8) K=5; BATCHED=8192; OFFLOAD=1 ;; - 10:8|16:8|24:8|32:8|48:8|52:8|60:8) K=4; BATCHED=8192; OFFLOAD=1 ;; - *) echo "CONC=$CONC DCP=$DCP is not a lane cell" >&2; exit 1 ;; + PROFILE=resident + CLIENTS=1 + case "$PROFILE:$CLIENTS" in + resident:1) K=7; BATCHED=16384; OFFLOAD=0 ;; + resident:4) K=7; BATCHED=8192; OFFLOAD=0 ;; + offload:4) K=5; BATCHED=8192; OFFLOAD=1 ;; + offload:10|offload:16|offload:24|offload:32|offload:48|offload:52|offload:60) K=4; BATCHED=8192; OFFLOAD=1 ;; + *) echo "PROFILE=$PROFILE CLIENTS=$CLIENTS is not a lane cell" >&2; exit 1 ;; esac export VLLM_ROCM_USE_AITER=1 export SAFETENSORS_FAST_GPU=1 export VLLM_ROCM_USE_AITER_MOE_SITUV2=1 + export AITER_BF16_FP8_MOE_BOUND=0 export AITER_FLYDSL_STAGE2_FP8=1 export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm export VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm @@ -1156,17 +1158,20 @@ guide: | # 1.56 TB of weights from shared storage, and long agentic prefills. export VLLM_ENGINE_READY_TIMEOUT_S=7200 export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200 - # On MEC firmware older than 177 (rocm-smi --showfw), also export HSA_NO_SCRATCH_RECLAIM=1. + export PYTHONNOUSERSITE=1 + # Required on the qualification nodes' MEC firmware 38; fixed in firmware 177. + export HSA_NO_SCRATCH_RECLAIM=1 GRAPH_MODE=FULL_AND_PIECEWISE export VLLM_USE_BREAKABLE_CUDAGRAPH=0 - MAX_SEQS=$((CONC == 1 ? 1 : 2 * CONC)) - MAX_CAPTURE=$((2 * CONC * (1 + K))) + MAX_SEQS=$((CLIENTS == 1 ? 1 : 2 * CLIENTS)) + MAX_CAPTURE=$((2 * CLIENTS * (1 + K))) COMPILATION="{\"mode\":3,\"cudagraph_mode\":\"$GRAPH_MODE\",\"max_cudagraph_capture_size\":$MAX_CAPTURE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$(seq -s, 2 2 "$MAX_CAPTURE")]}" - EXTRA=() - if [ "$CONC" -eq 1 ]; then - EXTRA+=(--quantization-config '{"linear":"fp8_per_channel","ignore":["*kv_b_proj","*vision_tower*","*mm_projector*","*f_b_proj","*fused_qkv_a_proj","*self_attn.g_proj","*.layers.0.mlp.down_proj"]}') + EXTRA=(--additional-config '{"kda_prefill_backend":"fused"}') + if [ "$PROFILE:$CLIENTS" = resident:1 ]; then + export VLLM_USE_V2_MODEL_RUNNER=1 + EXTRA=(--additional-config '{"decode_backend":"latency_specialized","kda_prefill_backend":"fused"}') else EXTRA+=(--quantization-config '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}') fi @@ -1180,14 +1185,12 @@ guide: | export PYTHONHASHSEED=42 EXTRA+=(--kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":224875000000,"lazy_offload":false}}') fi - if [ "$DCP" -gt 1 ]; then - EXTRA+=(--attention-backend ROCM_AITER_MLA --decode-context-parallel-size "$DCP" --dcp-comm-backend a2a) - fi - vllm serve moonshotai/Kimi-K3 \ --trust-remote-code \ --moe-backend auto \ --tensor-parallel-size 8 \ + --decode-context-parallel-size 8 \ + --dcp-comm-backend a2a \ --gpu-memory-utilization 0.9 \ --language-model-only \ --max-num-seqs "$MAX_SEQS" \ @@ -1216,17 +1219,17 @@ guide: | | Draft tokens | Golden acceptance length | Lane cells | | ---: | ---: | --- | - | 7 | 3.84 | dcp1 concurrency 1 and 4 | - | 5 | 3.62 | dcp8 concurrency 4 | + | 7 | 3.84 | resident concurrency 1 and 4, DCP8 | + | 5 | 3.62 | offload concurrency 4, DCP8 | | 4 | 3.36 | dcp8 concurrency 10, 16, 24, 32, 48, 52 and 60 | - These lengths come from [`kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/6e72cbc11234e624d66b8fc5eaddac50068db15c/inferencex-e2e/infx/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml): SPEED-Bench coding prompts with thinking on, measured on B300 for `Inferact/Kimi-K3-DSpark` with probabilistic drafting and block verification. They hold only for that draft at those draft lengths, not for `RedHatAI/Kimi-K3-speculator.dspark`. Synthetic acceptance skips real verification: never use it to serve users or to evaluate accuracy. + These lengths come from [`kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml): SPEED-Bench coding prompts with thinking on, measured on B300 for `Inferact/Kimi-K3-DSpark` with probabilistic drafting and block verification. They hold only for that draft at those draft lengths, not for `RedHatAI/Kimi-K3-speculator.dspark`. Synthetic acceptance skips real verification: never use it to serve users or to evaluate accuracy. ### Evidence and limits - - InferenceX [run 39002920671](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/39002920671) ran all ten cells on InferenceX PR #4215's commit `8a2aec3`, whose lane recipe is unchanged at `6e72cbc`. Its evals use the real-verification form above: the Kimi tool-call suite (`kimi_tool_call_schema_full`, 0.978-0.998) at every cell except dcp8 concurrency 4, and GSM8K (0.969) at dcp8 concurrency 60. The lane's current change is [SemiAnalysisAI/InferenceX#4215](https://github.com/SemiAnalysisAI/InferenceX/pull/4215). + - InferenceX [run 40037005947, attempt 2](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40037005947/attempts/2) passed all ten performance cells and its accuracy jobs at source `0df7e892b73c53d66a486c2a085e994b67653042`. The Kimi serving recipe, model, image and selected matrix are unchanged at audited head `dee7824`; later changes are unrelated upstream work and changelog text. Accuracy jobs use real verification and are deduplicated by [evaluation identity](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/evals/EVALS.md#L33), so this is not a claim of a separate accuracy job for every matrix cell. - Performance figures on the InferenceX dashboard are reported by SemiAnalysis and have not been reproduced for this recipe. - - Model policy (vLLM as the native engine, `Inferact/Kimi-K3-DSpark` as the agreed draft): [`MODELS.md`](https://github.com/SemiAnalysisAI/InferenceX/blob/6e72cbc11234e624d66b8fc5eaddac50068db15c/inferencex-e2e/docs/MODELS.md). + - Model policy (vLLM as the native engine, `Inferact/Kimi-K3-DSpark` as the agreed draft): [`MODELS.md`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/docs/MODELS.md). ## AMD MI355X: disaggregated prefill and decode From 66bf931544202f159cf83c66a812826b7d5c84fb Mon Sep 17 00:00:00 2001 From: Chenxi Ma <114001934+algebra-MCX@users.noreply.github.com> Date: Sun, 25 Oct 2026 13:56:21 +0800 Subject: [PATCH 037/631] [Bugfix][MRV2] Report the first failed DBO microbatch (#63476) Signed-off-by: Chenxi Ma <114001934+algebra-MCX@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 --- tests/v1/worker/test_gpu_ubatch_slicing.py | 26 ++++++++++++++++------ vllm/v1/worker/gpu/ubatch_utils.py | 2 +- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index 9a53357047f0..406b21f7f2a4 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -55,7 +55,11 @@ maybe_create_ubatch_slices, split_attn_metadata, ) -from vllm.v1.worker.ubatching import dbo_current_ubatch_id, dbo_yield +from vllm.v1.worker.ubatching import ( + dbo_current_ubatch_id, + dbo_register_recv_hook, + dbo_yield, +) MAX_NUM_REQS = 32 MAX_NUM_TOKENS = 128 @@ -659,11 +663,12 @@ def test_ubatch_runner_overlaps_and_matches_single_batch(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") -def test_ubatch_runner_names_the_microbatch_that_failed(): +@pytest.mark.parametrize("failed_ubatch", [0, 1]) +def test_ubatch_runner_names_the_microbatch_that_failed(failed_ubatch): """A failing microbatch surfaces as a named error, not a downstream KeyError. - The model here never yields, so the microbatches run back to back and no - handoff is left outstanding when one of them dies. A microbatch that fails + The first recorded failure is reported, not microbatch 0's later error from + registering a recv hook on the exited microbatch 1. A microbatch that fails *while its sibling is parked at a yield* hangs the step instead -- the shared handoff protocol in `ubatching.py` has no way to unwind a parked microbatch, and the V1 runner has the same gap. Fixing that means changing @@ -672,12 +677,19 @@ def test_ubatch_runner_names_the_microbatch_that_failed(): vllm_config = _make_dbo_config() device = torch.device("cuda:0") runner = _make_execution_runner(vllm_config) + failed_threads: list[threading.Thread] = [] class _FailingModel(torch.nn.Module): def forward(self, input_ids, positions, **kwargs): - if dbo_current_ubatch_id() == 0: + if dbo_current_ubatch_id() == failed_ubatch: + failed_threads.append(threading.current_thread()) raise ValueError("boom") - return input_ids.float().unsqueeze(-1) + if failed_ubatch == 0: + return input_ids.float().unsqueeze(-1) + dbo_yield() + failed_threads[0].join(timeout=5.0) + assert not failed_threads[0].is_alive() + dbo_register_recv_hook(lambda: None) model_inputs = _make_model_inputs(8, device) ubatch_state = _make_ubatch_state( @@ -704,7 +716,7 @@ def _call(): assert not caller.is_alive(), "UBatchRunner.run hung on a failing microbatch" assert isinstance(result["error"], RuntimeError) - assert "Microbatch 0" in str(result["error"]) + assert f"Microbatch {failed_ubatch}" in str(result["error"]) assert isinstance(result["error"].__cause__, ValueError) diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 5177ee8e374b..608231a3ed42 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -521,7 +521,7 @@ def finish() -> Any: stack.close() if errors: - failed = min(errors) + failed = next(iter(errors)) raise RuntimeError( f"Microbatch {failed} of {self.num_ubatches} failed" ) from errors[failed] From fc0cfcdc05b1717cadb17a0e24a6b1f611aec456 Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:03:52 +0800 Subject: [PATCH 038/631] [Bugfix][KV Connector] Serve and store OffloadingConnector partial tails with EAGLE/MTP groups (#61019) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: kimseunghyun Co-authored-by: vllm-agent --- .../offloading_connector/test_scheduler.py | 401 +++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 193 +++++++-- 2 files changed, 561 insertions(+), 33 deletions(-) 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 4b52f45e1c29..2f0a7d2bb396 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -207,6 +207,8 @@ def test_swa_offload_window_covers_unaligned_hit( def _make_partial_tail_scheduler( recurrent_only: bool = False, + eagle: bool = False, + cache_hit_alignment_tokens: int | None = 4, ) -> OffloadingConnectorScheduler: vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True}) vllm_config.cache_config.prefix_match_unit = None if recurrent_only else 4 @@ -215,6 +217,10 @@ def _make_partial_tail_scheduler( enable_kv_cache_events=True, publisher="null" ) kv_cache_config = _make_mamba_hybrid_kv_cache_config() + # Core resolves fine-grained hits for this layout at prefix_match_unit. + kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens + # The drafter shares the full-attention group, as in Qwen3-Next MTP. + kv_cache_config.kv_cache_groups[0].is_eagle_group = eagle if recurrent_only: kv_cache_config.kv_cache_groups = kv_cache_config.kv_cache_groups[1:] spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) @@ -224,9 +230,10 @@ def _make_partial_tail_scheduler( def _make_partial_tail_request( scheduler: OffloadingConnectorScheduler, kv_transfer_params: dict[str, Any] | None = None, + request_id: str = "req", ) -> MagicMock: request = MagicMock() - request.request_id = "req" + request.request_id = request_id request.kv_transfer_params = kv_transfer_params request.kv_hints = None request.num_prompt_tokens = 30 @@ -890,6 +897,398 @@ def lookup(key, req_context): assert req_status.partial_tail_boundary is None +def _key_tokens(scheduler: OffloadingConnectorScheduler, keys) -> dict[int, int]: + req_context = scheduler._req_status["req"].req_context + return { + get_offload_group_idx(key): req_context.get_offload_key_position(key) + for key in keys + } + + +@pytest.mark.parametrize( + ("cache_hit_alignment_tokens", "proof_tokens"), [(4, 28), (16, 40), (None, 40)] +) +def test_eagle_proof_margin_follows_core_hit_alignment( + cache_hit_alignment_tokens, proof_tokens +): + """Core drops one hash unit past a fine-grained hit, else one block.""" + scheduler = _make_partial_tail_scheduler( + eagle=True, cache_hit_alignment_tokens=cache_hit_alignment_tokens + ) + attention, mamba = scheduler.config.kv_group_configs + assert scheduler._partial_tail_key_tokens(attention, 24) == proof_tokens + assert scheduler._partial_tail_key_tokens(mamba, 24) == 24 + + +def test_eagle_partial_tail_store_keys_attention_at_proof(): + """The EAGLE group stores its tail block under the hash one unit later. + + With EAGLE block drop, core puts the Mamba checkpoint of a 30-token prompt + at 24. Core matches the attention group up to 28 and then drops 4 tokens, + so the attention block is stored under the hash that ends at 28. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + request = _make_partial_tail_request(scheduler) + request.num_computed_tokens = 30 + req_status = scheduler._req_status["req"] + req_status.loaded_from_offload = True + req_status.group_states[0].block_ids[:] = [11, 12] + req_status.group_states[1].block_ids[:] = [0, 21] + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + output = SimpleNamespace( + num_scheduled_tokens={}, + kv_connector_block_state=KVConnectorBlockState( + req_ids=set(), + resolve_block_ids={}.__getitem__, + boundary_state_offloads={"req": [(1, 99, 24)]}, + ), + ) + jobs = scheduler._build_partial_tail_store_jobs(output) + + [job_id] = jobs + offered_keys = scheduler.manager.prepare_store.call_args.args[0] + assert _key_tokens(scheduler, offered_keys) == {0: 28, 1: 24} + assert [get_offload_block_hash(key) for key in offered_keys] == [b"h6", b"h5"] + src_spec = jobs[job_id].src_spec + assert isinstance(src_spec, GPULoadStoreSpec) + assert src_spec.block_ids.tolist() == [12, 99] + assert src_spec.block_indices == [1, 1] + + events = list( + scheduler._events_tracker.take_events( + [ + OffloadingEvent( + keys=list(scheduler._jobs[job_id].keys), + medium=Medium.CPU, + removed=False, + ) + ] + ) + ) + events_by_group = {event.group_idx: event for event in events} + assert events_by_group[0].token_ids == list(range(16, 28)) + + +@pytest.mark.parametrize( + ("num_prompt_tokens", "num_computed_tokens", "num_scheduled_tokens", "stored"), + [ + (30, 24, 0, False), + (30, 27, 0, False), + (30, 24, 3, False), + # The step computes the proof token 28 exactly. + (30, 24, 4, True), + # Tokens past the prompt end at 26 do not compute the proof. + (26, 24, 6, False), + ], +) +def test_eagle_partial_tail_store_needs_proof_in_prompt( + num_prompt_tokens, num_computed_tokens, num_scheduled_tokens, stored +): + """Store the tail only if this step computes the proof at 28 in the prompt.""" + scheduler = _make_partial_tail_scheduler(eagle=True) + request = _make_partial_tail_request(scheduler) + request.num_prompt_tokens = num_prompt_tokens + request.num_computed_tokens = num_computed_tokens + req_status = scheduler._req_status["req"] + req_status.loaded_from_offload = True + req_status.group_states[0].block_ids[:] = [11, 12] + req_status.group_states[1].block_ids[:] = [0, 21] + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + output = SimpleNamespace( + num_scheduled_tokens={"req": num_scheduled_tokens}, + kv_connector_block_state=KVConnectorBlockState( + req_ids=set(), + resolve_block_ids={}.__getitem__, + boundary_state_offloads={"req": [(1, 99, 24)]}, + ), + ) + + assert bool(scheduler._build_partial_tail_store_jobs(output)) == stored + assert scheduler.manager.prepare_store.called == stored + + +@pytest.mark.parametrize("missing", [None, b"h6", b"h3"]) +def test_eagle_partial_lookup_loads_attention_proof(missing): + """An EAGLE partial hit needs the proof key and the popped chunk. + + The complete-chunk lookup pops the only EAGLE chunk (0-16), so the hit at + 24 lies in the second block after the complete boundary. It needs the + attention chunk 0-16, the attention proof at 28 and the Mamba state at 24. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + request = _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.num_locally_computed_tokens = 0 + req_status.update_offload_keys() + + def lookup(key, req_context): + block_hash = get_offload_block_hash(key) + if get_offload_group_idx(key) == 1: + ready = block_hash == b"h5" + else: + ready = block_hash != missing + return LookupResult.HIT if ready else LookupResult.MISS + + scheduler.manager.lookup.side_effect = lookup + if missing is not None: + assert scheduler._lookup(req_status) == 0 + assert req_status.partial_tail_boundary is None + return + assert scheduler._lookup(req_status) == 24 + assert req_status.partial_tail_boundary == 24 + + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32)], + [KVCacheBlock(0, is_null=True), KVCacheBlock(41)], + ) + ), + num_external_tokens=24, + ) + [load_job_id] = scheduler._current_batch_load_jobs + load_job = scheduler._current_batch_load_jobs[load_job_id] + assert { + (get_offload_group_idx(key), get_offload_block_hash(key)) + for key in scheduler._jobs[load_job_id].keys + } == {(0, b"h3"), (0, b"h6"), (1, b"h5")} + dst_spec = load_job.dst_spec + assert isinstance(dst_spec, GPULoadStoreSpec) + assert dst_spec.block_ids.tolist() == [31, 32, 41] + assert dst_spec.group_sizes == [2, 1] + + +@pytest.mark.parametrize( + ("attention_missing", "expected_hit", "expected_lookups"), + [ + # The complete-chunk lookup of h3 also serves the boundaries 24 and 20. + (None, 20, ["0h3", "0h6", "1h5", "0h5", "1h4"]), + # A missing chunk h3 is looked up once more and skips the boundaries + # 24, 20 and 16. + (b"h3", 0, ["0h3", "0h3", "0h3", "0h2", "1h1", "0h1", "1h0"]), + ], +) +def test_eagle_partial_lookup_probes_popped_chunk_once( + attention_missing, expected_hit, expected_lookups +): + """Boundaries past the popped EAGLE chunk reuse its complete-chunk lookup. + + The Mamba state is stored only at 20 (h4). Before this lookup was shared, + the same hits took 7 and 9 lookups. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.num_locally_computed_tokens = 0 + req_status.update_offload_keys() + lookups = [] + + def lookup(key, req_context): + group_idx = get_offload_group_idx(key) + block_hash = get_offload_block_hash(key) + lookups.append(f"{group_idx}{block_hash.decode()}") + ready = ( + block_hash == b"h4" if group_idx == 1 else block_hash != attention_missing + ) + return LookupResult.HIT if ready else LookupResult.MISS + + scheduler.manager.lookup.side_effect = lookup + assert scheduler._lookup(req_status) == expected_hit + assert lookups == expected_lookups + + +def test_eagle_partial_lookup_after_recurrent_group_tightens_hit(): + """A recurrent group can lower the hit after the EAGLE group pops a chunk. + + The 62-token prompt has attention chunks h3, h7 and h11. The EAGLE group + hits all three and pops h11. The Mamba group has a complete state only at + 16 (h3), so the complete hit drops to 16. The EAGLE lookup still confirmed + h7, so the partial hit at 36 loads h7 without another lookup of it. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + request = _make_partial_tail_request(scheduler) + request.num_prompt_tokens = request.num_tokens = 62 + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(15)] + request.all_token_ids = list(range(62)) + req_status = scheduler._req_status["req"] + req_status.num_locally_computed_tokens = 0 + req_status.update_offload_keys() + lookups = [] + + def lookup(key, req_context): + group_idx = get_offload_group_idx(key) + block_hash = get_offload_block_hash(key) + lookups.append(f"{group_idx}{block_hash.decode()}") + ready = group_idx == 0 or block_hash in (b"h3", b"h8") + return LookupResult.HIT if ready else LookupResult.MISS + + scheduler.manager.lookup.side_effect = lookup + assert scheduler._lookup(req_status) == 36 + assert req_status.partial_tail_boundary == 36 + assert lookups.count("0h7") == 1 + + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32), KVCacheBlock(33)], + [ + KVCacheBlock(0, is_null=True), + KVCacheBlock(0, is_null=True), + KVCacheBlock(41), + ], + ) + ), + num_external_tokens=36, + ) + [load_job_id] = scheduler._current_batch_load_jobs + assert { + (get_offload_group_idx(key), get_offload_block_hash(key)) + for key in scheduler._jobs[load_job_id].keys + } == {(0, b"h3"), (0, b"h7"), (0, b"h9"), (1, b"h8")} + + +@pytest.mark.parametrize("cow_on_proof_step", [False, True]) +def test_eagle_partial_tail_round_trip_serves_mamba_checkpoint(cow_on_proof_step): + """Store, finish and resend a 30-token prompt with an EAGLE group. + + The resend hits the Mamba checkpoint at 24 only through the EAGLE + group's proof key at 28. With ``cow_on_proof_step``, the steps follow core: + the prefill stops at 24, and the CoW hand-off arrives on the step that + computes tokens 24-30. The producer models a turn that loaded its prefix + from the offload tier, so that its tail is stored. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + stored: set = set() + scheduler.manager.lookup.side_effect = lambda key, req_context: ( + LookupResult.HIT if key in stored else LookupResult.MISS + ) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + def run_step(num_scheduled_tokens, handoffs=None, finished=()): + output = SchedulerOutput.make_empty() + output.num_scheduled_tokens = num_scheduled_tokens + output.finished_req_ids = set(finished) + if handoffs is not None: + output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids={}.__getitem__, + boundary_state_offloads=handoffs, + ) + meta = scheduler.build_connector_meta(output) + for job_id in meta.store_jobs: + stored.update(scheduler._jobs[job_id].keys) + scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1 for job_id in meta.store_jobs} + ) + ) + ) + + producer = _make_partial_tail_request(scheduler) + producer.num_computed_tokens = 0 + producer.status = RequestStatus.RUNNING + assert scheduler.get_num_new_matched_tokens(producer, 0) == (0, False) + producer_status = scheduler._req_status["req"] + producer_status.loaded_from_offload = True + producer_status.group_states[0].block_ids[:] = [11, 12] + producer_status.group_states[1].block_ids[:] = [0, 21] + if cow_on_proof_step: + run_step({"req": 24}) + producer.num_computed_tokens = 24 + run_step({"req": 6}, handoffs={"req": [(1, 99, 24)]}) + producer.num_computed_tokens = 30 + else: + run_step({"req": 30}) + producer.num_computed_tokens = 30 + run_step({"req": 1}, handoffs={"req": [(1, 99, 24)]}) + producer.is_finished.return_value = True + scheduler.request_finished(producer) + run_step({}, finished={"req"}) + + stored_positions = { + (get_offload_group_idx(key), get_offload_block_hash(key)) for key in stored + } + assert stored_positions == {(0, b"h3"), (0, b"h6"), (1, b"h5")} + + consumer = _make_partial_tail_request(scheduler, request_id="resend") + assert scheduler.get_num_new_matched_tokens(consumer, 0) == (24, True) + scheduler.update_state_after_alloc( + consumer, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32)], + [KVCacheBlock(0, is_null=True), KVCacheBlock(41)], + ) + ), + num_external_tokens=24, + ) + [load_job_id] = scheduler._current_batch_load_jobs + assert scheduler._jobs[load_job_id].keys == stored + + +@pytest.mark.parametrize("host_hit", [False, True]) +def test_eagle_partial_tail_store_needs_offload_load(host_hit): + """Store an EAGLE partial tail only after a load from the offload tier. + + Without a host hit, the GPU prefix cache keeps the tail, and the hand-off + that computes the proof at 28 stores nothing. After the request loads 24 + tokens from the offload tier, the same hand-off stores the tail. + """ + scheduler = _make_partial_tail_scheduler(eagle=True) + scheduler.manager.lookup.side_effect = lambda key, req_context: ( + LookupResult.HIT + if host_hit + and (get_offload_group_idx(key) == 0 or get_offload_block_hash(key) == b"h5") + else LookupResult.MISS + ) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + request = _make_partial_tail_request(scheduler) + num_external_tokens = 24 if host_hit else 0 + assert scheduler.get_num_new_matched_tokens(request, 0) == ( + num_external_tokens, + host_hit, + ) + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(11), KVCacheBlock(12)], + [KVCacheBlock(0, is_null=True), KVCacheBlock(21)], + ) + ), + num_external_tokens=num_external_tokens, + ) + req_status = scheduler._req_status["req"] + req_status.group_states[0].block_ids[:] = [11, 12] + req_status.group_states[1].block_ids[:] = [0, 21] + request.num_computed_tokens = 24 + + output = SimpleNamespace( + num_scheduled_tokens={"req": 6}, + kv_connector_block_state=KVConnectorBlockState( + req_ids=set(), + resolve_block_ids={}.__getitem__, + boundary_state_offloads={"req": [(1, 99, 24)]}, + ), + ) + + assert bool(scheduler._build_partial_tail_store_jobs(output)) == host_hit + assert scheduler.manager.prepare_store.called == host_hit + + def test_scheduler_reports_allocation_failure(request_runner): runner = request_runner( block_size=4, 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 284538f8631e..43b938e2849d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -31,6 +31,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.kv_cache_utils import eagle_proof_margin from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager from vllm.v1.kv_cache_interface import ( @@ -300,7 +301,8 @@ def from_spec( ) # Partial tails currently require one physical block per offload chunk # and uniform, non-windowed groups so one boundary identifies every - # group's source. EAGLE and DCP need additional hand-off semantics. + # group's source. EAGLE groups key their partial tail at a proof + # position past the boundary. DCP needs additional hand-off semantics. supports_partial_tail = ( spec.blocks_per_chunk == 1 and len(group_block_sizes) == 1 @@ -310,7 +312,6 @@ def from_spec( or config.requires_cow_source for config in kv_group_configs ) - and not any(config.is_eagle_group for config in kv_group_configs) and vllm_config.parallel_config.decode_context_parallel_size == 1 ) @@ -353,6 +354,9 @@ class RequestGroupState: # Number of offloaded chunks hit (including GPU prefix cache) # when the request first started num_hit_chunks: int = 0 + # End (exclusive) of the chunks that the current complete-chunk lookup + # confirmed as ready hits, including a popped EAGLE chunk. + lookup_hit_end_chunk_idx: int = 0 @dataclass(slots=True) @@ -377,6 +381,8 @@ class RequestOffloadState: # Fine-grained token boundary selected beyond the last complete offload # chunk. It is consumed when the corresponding load is scheduled. partial_tail_boundary: int | None = None + # True once a load from the offload tier was scheduled for this request. + loaded_from_offload: bool = False # Loaded attention keys cover their token positions. Recurrent state # supports the whole reused external prefix. load_key_ranges: list[tuple[int, int, OffloadKey]] = field(default_factory=list) @@ -607,6 +613,29 @@ def _sliding_window_sort_key(i: int) -> int: if self.config.supports_partial_tail else 0 ) + # Core's prefix cache matches an EAGLE group this many tokens past a + # hit and then drops them. A partial tail stores and looks up EAGLE + # groups at that proof position, with the margin rule of the core + # coordinator. + alignment = kv_cache_config.cache_hit_alignment_tokens + partial_hash_hits = ( + alignment is not None and alignment < self._partial_tail_block_size + ) + self._eagle_proof_margins: dict[int, int] = { + config.group_idx: eagle_proof_margin( + self._partial_tail_block_size, + self.config.tokens_per_hash, + partial_hash_hits + and config.manager_cls.supports_fine_grained_hash_lookup, + ) + for config in self.config.kv_group_configs + if self.config.supports_partial_tail + and config.is_eagle_group + and not config.requires_cow_source + } + self._max_eagle_proof_margin = max( + self._eagle_proof_margins.values(), default=0 + ) self._cow_source_groups = frozenset( config.group_idx for config in self.config.kv_group_configs @@ -859,6 +888,8 @@ def _lookup_complete_chunks( num_hit_tokens: int = 0 defer_lookup = False lookup_groups = self._lookup_groups + for state in req_status.group_states: + state.lookup_hit_end_chunk_idx = 0 # Tracks which eagle groups have already popped their volatile trailing chunk # in the current convergence iteration. Reset when a non-eagle group @@ -928,6 +959,9 @@ def _lookup_complete_chunks( group_config, start_chunk_idx, ) + group_state.lookup_hit_end_chunk_idx = start_chunk_idx + ( + num_hit_chunks or 0 + ) else: required_window = sliding_window_size_in_chunks if is_eagle_unverified: @@ -1049,8 +1083,20 @@ def _lookup( local_tokens = req_status.num_locally_computed_tokens complete_boundary = local_tokens + complete_hit tokens_per_hash = self.config.tokens_per_hash - block_end = complete_boundary + self._partial_tail_block_size + # A partial tail carries its own EAGLE proof, so it may also lie in + # the block after the chunk that the complete-chunk lookup pops. + num_search_blocks = 2 if self._max_eagle_proof_margin else 1 + block_end = complete_boundary + ( + num_search_blocks * self._partial_tail_block_size + ) max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) + if self._max_eagle_proof_margin: + # The proof key needs a hash of this request past the boundary. + max_boundary = min( + max_boundary, + len(req_status.req.block_hashes) * tokens_per_hash + - self._max_eagle_proof_margin, + ) if max_num_new_tokens is not None: max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) if req_status.max_load_tokens is not None: @@ -1062,41 +1108,97 @@ def _lookup( if max_boundary <= complete_boundary: return complete_hit + # A boundary past the block of complete_boundary also loads that + # block's complete chunk. Look it up once, and only if it is needed. + # Partial tails use one block per chunk, so a block index is a chunk + # index. + block_size = self._partial_tail_block_size + bridge_chunk_idx = complete_boundary // block_size + bridge_boundary = (bridge_chunk_idx + 1) * block_size + # Bridge keys that the complete-chunk lookup did not confirm. + unconfirmed_bridge_keys: list[tuple[GroupOffloadConfig, OffloadKey]] = [] pending = False - for boundary in range(max_boundary, complete_boundary, -tokens_per_hash): - boundary_pending = False - boundary_missed = False - boundary_keys = [] - for group_config in self.config.kv_group_configs: - key = self._make_boundary_key( - req_status.req, - group_config.group_idx, - boundary, - req_status.req_context, - ) - boundary_keys.append(key) + if max_boundary >= bridge_boundary: + for group_config, group_state in zip( + self.config.kv_group_configs, req_status.group_states + ): + if group_config.requires_cow_source: + continue + key = group_state.offload_keys[bridge_chunk_idx] + pending |= key in (self._chunks_being_loaded or ()) + # An EAGLE group always confirmed the chunk that it popped. + if group_state.lookup_hit_end_chunk_idx <= bridge_chunk_idx: + unconfirmed_bridge_keys.append((group_config, key)) + for _, key in unconfirmed_bridge_keys: result = self.manager.lookup(key, req_status.req_context) if result is LookupResult.MISS: - boundary_missed = True + max_boundary = bridge_boundary - tokens_per_hash break - if result in (LookupResult.HIT_PENDING, LookupResult.RETRY): - boundary_pending = True - - pending |= boundary_pending - if not boundary_missed and not boundary_pending: - for group_config, key in zip( - self.config.kv_group_configs, boundary_keys - ): - self._events_tracker.record_partial_lookup( - req_status.req, group_config, boundary, key - ) + pending |= result is not LookupResult.HIT + bridge_pending = pending + for boundary in range(max_boundary, complete_boundary, -tokens_per_hash): + loads_bridge = boundary >= bridge_boundary + hit, boundary_pending = self._lookup_partial_tail( + req_status, boundary, loads_bridge and bridge_pending + ) + if hit: + # The complete-chunk lookup recorded the confirmed bridge keys. + if loads_bridge: + for group_config, key in unconfirmed_bridge_keys: + self._events_tracker.record_lookup( + req_status.req, group_config, bridge_chunk_idx, key + ) req_status.partial_tail_boundary = boundary return boundary - local_tokens + pending |= boundary_pending if pending and complete_hit == 0: return None return complete_hit + def _partial_tail_key_tokens( + self, group_config: GroupOffloadConfig, boundary: int + ) -> int: + """Token position whose hash keys a group's partial-tail block.""" + return boundary + self._eagle_proof_margins.get(group_config.group_idx, 0) + + def _lookup_partial_tail( + self, + req_status: RequestOffloadState, + boundary: int, + bridge_pending: bool, + ) -> tuple[bool, bool]: + """Look up every key a partial-tail hit at ``boundary`` loads. + + ``bridge_pending`` is True if a complete chunk that the hit also loads + is not ready yet. + Returns ``(hit, pending)``. ``hit`` is True only if every key is ready. + ``pending`` is True if a key looked up before the first miss is not + ready yet. + """ + req = req_status.req + req_context = req_status.req_context + pending = bridge_pending + boundary_keys: list[tuple[GroupOffloadConfig, int, OffloadKey]] = [] + for group_config in self.config.kv_group_configs: + key_tokens = self._partial_tail_key_tokens(group_config, boundary) + key = self._make_boundary_key( + req, group_config.group_idx, key_tokens, req_context + ) + result = self.manager.lookup(key, req_context) + if result is LookupResult.MISS: + return False, pending + pending |= result is not LookupResult.HIT + boundary_keys.append((group_config, key_tokens, key)) + + if pending: + return False, True + for group_config, key_tokens, key in boundary_keys: + self._events_tracker.record_partial_lookup( + req, group_config, key_tokens, key + ) + return True, False + def on_new_request(self, request: Request) -> None: """Called when a new request is added to the scheduler.""" req_context = _create_req_context(request) @@ -1240,6 +1342,7 @@ def update_state_after_alloc( self._acquire_step() req_status = self._req_status[request.request_id] + req_status.loaded_from_offload = True num_locally_computed_tokens = req_status.num_locally_computed_tokens num_cached_tokens = num_locally_computed_tokens + num_external_tokens @@ -1303,7 +1406,9 @@ def update_state_after_alloc( self._make_boundary_key( request, group_config.group_idx, - partial_tail_boundary, + self._partial_tail_key_tokens( + group_config, partial_tail_boundary + ), req_status.req_context, ) ) @@ -1510,6 +1615,10 @@ def _build_partial_tail_store_jobs( continue req_status = self._req_status.get(req_id) assert req_status is not None + if self._max_eagle_proof_margin and not req_status.loaded_from_offload: + # Requests that did not load from the offload tier keep their + # EAGLE tail in the GPU prefix cache; skip the store. + continue boundaries = {boundary for _, _, boundary in entries} assert len(boundaries) == 1 boundary = boundaries.pop() @@ -1528,6 +1637,20 @@ def _build_partial_tail_store_jobs( assert boundary % self.config.tokens_per_hash == 0 assert boundary <= max_boundary + if self._max_eagle_proof_margin: + # The CoW hand-off arrives on the step that computes the prompt + # tokens past the boundary, and it is not offered again. Workers + # submit a store at the next step's start, after this step's + # forward, so count the prompt tokens this step computes. + prompt_tokens_after_step = min( + req.num_computed_tokens + + scheduler_output.num_scheduled_tokens.get(req_id, 0), + req.num_prompt_tokens, + ) + if boundary + self._max_eagle_proof_margin > prompt_tokens_after_step: + # The EAGLE proof is not computed by the end of this step. + continue + cow_blocks = {group_idx: block_id for group_idx, block_id, _ in entries} assert self._cow_source_groups.issubset(cow_blocks) @@ -1539,11 +1662,15 @@ def _build_partial_tail_store_jobs( for group in self.config.kv_group_configs ): continue + key_tokens = [ + self._partial_tail_key_tokens(group, boundary) + for group in self.config.kv_group_configs + ] keys = [ self._make_boundary_key( - req, group.group_idx, boundary, req_status.req_context + req, group.group_idx, tokens, req_status.req_context ) - for group in self.config.kv_group_configs + for group, tokens in zip(self.config.kv_group_configs, key_tokens) ] block_ids = [ cow_blocks[group.group_idx] @@ -1562,10 +1689,12 @@ def _build_partial_tail_store_jobs( if not store_output.keys_to_store: continue - for group_config, key in zip(self.config.kv_group_configs, keys): + for group_config, tokens, key in zip( + self.config.kv_group_configs, key_tokens, keys + ): if key in store_output.keys_to_store: self._events_tracker.record_partial_store( - req, group_config, boundary, key + req, group_config, tokens, key ) group_by_key = {key: idx for idx, key in enumerate(keys)} From a55937f255287bc638a839cc7a441fc7e9c55aae Mon Sep 17 00:00:00 2001 From: Uros Markovic Date: Sun, 25 Oct 2026 07:59:13 +0100 Subject: [PATCH 039/631] [ROCm][Perf] Restore 32-bit buffer addressing for TRITON_ATTN KV cache views (#60084) Signed-off-by: Uros Markovic Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> --- vllm/v1/attention/backends/triton_attn.py | 6 ++++++ vllm/v1/attention/backends/utils.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 42504c71ff8d..232943dc38ba 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -33,6 +33,7 @@ from vllm.v1.attention.backends.utils import ( compute_mm_prefix_range_tensor, get_num_attention_heads_from_layers, + set_triton_ptr_range, ) from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( @@ -573,6 +574,7 @@ def __init__( self.use_td = current_platform.is_xpu() else: self.use_td = td_override + self._set_kv_ptr_range = current_platform.is_rocm() def forward( self, @@ -659,6 +661,8 @@ def forward( ): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + if self._set_kv_ptr_range: + set_triton_ptr_range(kv_cache, key_cache, value_cache) descale_shape = ( attn_metadata.query_start_loc.shape[0] - 1, key_cache.shape[2], @@ -825,6 +829,8 @@ def do_kv_cache_update( if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + if self._set_kv_ptr_range: + set_triton_ptr_range(kv_cache, key_cache, value_cache) triton_reshape_and_cache_flash( key, value, diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 6ad995a305a3..40b283a3f1c5 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools import math from collections import defaultdict from collections.abc import Callable, Iterable, Sequence @@ -1015,6 +1016,23 @@ def reorder_batch_to_split_decodes_and_prefills( return True +def set_triton_ptr_range(region: torch.Tensor, *views: torch.Tensor) -> None: + """Bound the bytes Triton kernels may access through ``views``. + + Triton's AMD backend uses 32-bit buffer addressing only for pointer + arguments within 2 GiB. It checks ``ptr_range()`` first and otherwise the + size of the whole storage, so views into one large allocation (such as a + layer of the KV cache) lose it. Each view gets the distance from its + ``data_ptr()`` to the end of ``region``; kernels must not access past it. + """ + if region.numel() == 0: + return + last = sum((n - 1) * stride for n, stride in zip(region.shape, region.stride())) + region_end = region.data_ptr() + (last + 1) * region.element_size() + for view in views: + view.ptr_range = functools.partial(int, region_end - view.data_ptr()) + + def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor: """Reshapes the query tensor for the specified batch size, so that it has shape (batch_size, seq_len, num_heads, head_dim). From 921ed7fa11c41b89ce36e81f50fb3b1c293dc8f7 Mon Sep 17 00:00:00 2001 From: alberto Date: Sun, 25 Oct 2026 07:55:59 +0000 Subject: [PATCH 040/631] [Bugfix][KV Offload] Reset a resumed streaming session's block table in OffloadingConnector and SimpleCPUOffloadConnector (#60810) Signed-off-by: Alberto Perdomo Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: vllm-agent --- .../offloading_connector/test_scheduler.py | 18 +++++++++ .../unit/offloading_connector/utils.py | 9 ++++- tests/v1/simple_kv_offload/test_scheduler.py | 37 +++++++++++++++++++ .../kv_transfer/kv_connector/utils.py | 6 ++- .../kv_connector/v1/offloading/scheduler.py | 10 ++--- vllm/v1/simple_kv_offload/manager.py | 13 ++++--- 6 files changed, 79 insertions(+), 14 deletions(-) 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 2f0a7d2bb396..f1b483af408d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1727,6 +1727,24 @@ def take_raw_events(): assert events[0].token_ids == token_ids +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_streaming_session_stores_new_chunks(request_runner, async_scheduling: bool): + """A streaming session resumed for its next input chunk arrives with its + full block table, so its new chunks must be stored from their own blocks.""" + runner = request_runner( + block_size=4, num_gpu_blocks=100, async_scheduling=async_scheduling + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.new_request(token_ids=list(range(1, 9)), req_id="s", resumable=True) + runner.run(decoded_tokens=[100, EOS_TOKEN_ID], expected_stored=(0, 1)) + + runner.new_request(token_ids=list(range(200, 212)), req_id="s", resumable=True) + runner.run(decoded_tokens=[300, EOS_TOKEN_ID], expected_stored=(2, 3, 4)) + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_offloading_connector(request_runner, async_scheduling: bool): block_size = 4 diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index e5152fb201a5..303c42df6d2c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -373,8 +373,12 @@ def new_request( token_ids: list[int], kv_transfer_params: dict | None = None, skip_reading_prefix_cache: bool = False, + req_id: str | None = None, + resumable: bool = False, ): - self.req_id += 1 + if req_id is None: + self.req_id += 1 + req_id = str(self.req_id) sampling_params = SamplingParams( max_tokens=1000, @@ -383,11 +387,12 @@ def new_request( sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) req = Request( - request_id=str(self.req_id), + request_id=req_id, prompt_token_ids=token_ids, sampling_params=sampling_params, pooling_params=None, block_hasher=self._block_hasher, + resumable=resumable, ) if kv_transfer_params is not None: req.kv_transfer_params = kv_transfer_params diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index d56ab77797f9..3de021396f17 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1686,6 +1686,43 @@ def test_eager_store_resets_block_cursor_after_preemption( assert resumed_meta.store_gpu_blocks == [resumed_blocks.blocks[0][2].block_id] +def test_eager_store_streaming_session_stores_new_chunks() -> None: + """A resumed streaming session is rescheduled as a new request carrying + its full block table without being preempted. The accumulated block IDs + must be replaced, not extended, so the new chunk's blocks get stored.""" + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=False) + sched = fix.scheduler + req = make_request(num_blocks=4) + + gpu_blocks = fix.gpu_block_pool.get_new_blocks(2) + fix.gpu_block_pool.cache_full_blocks(req, gpu_blocks, 0, 2, BLOCK_SIZE, 0) + req.num_computed_tokens = 2 * BLOCK_SIZE + first_turn = KVCacheBlocks(blocks=(gpu_blocks,)) + sched.update_state_after_alloc(req, first_turn, num_external_tokens=0) + first_meta = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 2 * BLOCK_SIZE}, + new_reqs={req.request_id: first_turn.get_block_ids()}, + ) + ) + assert len(first_meta.store_gpu_blocks) == 2 + simulate_store_completion(sched, first_meta.store_event) + + # The paused session keeps its blocks and grows by a new chunk. + gpu_blocks = gpu_blocks + fix.gpu_block_pool.get_new_blocks(2) + fix.gpu_block_pool.cache_full_blocks(req, gpu_blocks, 2, 4, BLOCK_SIZE, 0) + req.num_computed_tokens = 4 * BLOCK_SIZE + second_turn = KVCacheBlocks(blocks=(gpu_blocks,)) + sched.update_state_after_alloc(req, second_turn, num_external_tokens=0) + second_meta = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 2 * BLOCK_SIZE}, + new_reqs={req.request_id: second_turn.get_block_ids()}, + ) + ) + assert second_meta.store_gpu_blocks == [b.block_id for b in gpu_blocks[2:]] + + # --------------------------------------------------------------------------- # Test 10: Partial GPU prefix hit + CPU load + new compute blocks # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index a9c2da36188e..436666f3dc74 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -321,12 +321,14 @@ def yield_req_data( scheduler_output, ) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]: """Yields: - (req_id, new_block_id_groups, preempted) + (req_id, block_id_groups, is_full_table) + ``is_full_table`` is True when ``block_id_groups`` is the request's whole + block table (new and resumed requests) rather than this step's additions. """ # new requests for req_data in scheduler_output.scheduled_new_reqs: - yield req_data.req_id, req_data.block_ids, False + yield req_data.req_id, req_data.block_ids, True # cached requests cached_reqs = scheduler_output.scheduled_cached_reqs 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 43b938e2849d..6bed634430a8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1476,23 +1476,23 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: # Used to detect sliding window blocks that got re-allocated. new_block_ids_end: dict[str, tuple[int, ...]] = {} - for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output): + for req_id, block_id_groups, is_full_table in yield_req_data(scheduler_output): req_status = self._req_status[req_id] req_status.update_offload_keys() - if preempted: + if is_full_table: for group_state in req_status.group_states: group_state.block_ids.clear() - if new_block_id_groups: + if block_id_groups: if self._sliding_window_groups: new_block_ids_end[req_id] = tuple( len(req_status.group_states[grp_idx].block_ids) for grp_idx in self._sliding_window_groups ) - req_status.update_block_id_groups(new_block_id_groups) + req_status.update_block_id_groups(block_id_groups) for group_config in self.config.kv_group_configs: - new_blocks = new_block_id_groups[group_config.group_idx] + new_blocks = block_id_groups[group_config.group_idx] for bid in new_blocks: if bid != 0: self._current_batch_allocated_block_ids.add(bid) diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 62a022d9e0f2..022ecd7757ab 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -779,16 +779,19 @@ def _prepare_eager_store_specs( state.block_ids = tuple([] for _ in range(num_groups)) state.num_stored_blocks = [0] * num_groups - for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output): + for req_id, block_id_groups, is_full_table in yield_req_data(scheduler_output): state = self._reqs_to_store.get(req_id) if state is None or state.finished: continue - if new_block_id_groups: - for g in range(min(num_groups, len(new_block_id_groups))): - if new_block_id_groups[g] is not None: + if is_full_table: + state.block_ids = tuple([] for _ in range(num_groups)) + + if block_id_groups: + for g in range(min(num_groups, len(block_id_groups))): + if block_id_groups[g] is not None: # Accumulate new block IDs. - state.block_ids[g].extend(new_block_id_groups[g]) + state.block_ids[g].extend(block_id_groups[g]) num_new_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) if num_new_tokens == 0: From 2f1041e8dd05415d9984d0a83ad957d54a049eb4 Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Sun, 25 Oct 2026 16:30:41 +0800 Subject: [PATCH 041/631] [Doc] Kimi-K3: InferenceX AgentX lanes as tuned presets; DCP only in presets (#63552) Signed-off-by: esmeetu Co-authored-by: Claude Opus 5.5 --- recipes/models/moonshotai/Kimi-K3.yaml | 1809 ++++++++++++++---------- 1 file changed, 1029 insertions(+), 780 deletions(-) diff --git a/recipes/models/moonshotai/Kimi-K3.yaml b/recipes/models/moonshotai/Kimi-K3.yaml index b99886f8d875..b9661b0b9fac 100644 --- a/recipes/models/moonshotai/Kimi-K3.yaml +++ b/recipes/models/moonshotai/Kimi-K3.yaml @@ -102,54 +102,10 @@ features: description: "Skip the vision encoder for text-only workloads. Mutually exclusive with encoder_parallel." args: - "--language-model-only" - decode_context_parallelism: - description: "Decode context parallelism (DCP): shard the decode KV cache across the tensor-parallel ranks for decode-heavy long-context serving. The DCP size must divide the tensor-parallel size, so the args below assume the TP8 single-node layout. Pairs the TOKENSPEED_MLA decode backend with TRTLLM_RAGGED MLA prefill under FP8 KV; on MI355X it pairs ROCM_AITER_MLA decode with ROCM_AITER_FA prefill instead." - # Scoped to single_node_tp: --decode-context-parallel-size is pinned to 8 - # here, which only divides a TP of 8 or more. - strategies: - - single_node_tp - args: - - "--decode-context-parallel-size" - - "8" - - "--dcp-comm-backend" - - "a2a" - - "--attention-backend" - - "TOKENSPEED_MLA" - - "--attention-config" - - '{"use_prefill_query_quantization":true,"mla_prefill_backend":"TRTLLM_RAGGED"}' - hardware_overrides: - # TOKENSPEED_MLA and TRTLLM_RAGGED are not ROCm backends. MI355X uses the - # InferenceX AgentX lane's DCP8 pairing with the direct DCP ops pinned off. - mi355x: - args: - - "--decode-context-parallel-size" - - "8" - - "--dcp-comm-backend" - - "a2a" - - "--attention-backend" - - "ROCM_AITER_MLA" - - "--attention-config" - - '{"mla_prefill_backend":"ROCM_AITER_FA"}' - env: - VLLM_USE_DIRECT_DCP_A2A: "0" - VLLM_USE_DIRECT_DCP_Q_GATHER: "0" - VLLM_USE_DIRECT_DCP_KV_GATHER: "0" - -# DCP size selector, qualified on GB200/GB300 only. A size is offered when it -# equals the tensor-parallel size of every role: TP8/PP2 aggregated -> DCP8 -# (InferenceX #4322), TP16 per disaggregated role -> DCP16 (#4322), TP8 GB300 -# pools -> DCP8 (#4047) and TP16 aggregated -> DCP16 (low-latency profile). -decode_context_parallel: - sizes: - gb200: [8, 16] - gb300: [8, 16] - comm_backend: a2a - match_tp: true opt_in_features: - spec_decoding - text_only - - decode_context_parallelism variants: default: @@ -276,7 +232,6 @@ strategy_min_gpus: multi_node_tp: 16 ascend_910c: 64 - hardware_overrides: hopper: extra_args: @@ -349,32 +304,6 @@ hardware_overrides: - '{"cudagraph_mode":"FULL_DECODE_ONLY","custom_ops":["+fused_rms_norm_gated"]}' strategy_overrides: - # MI355X TP8: the InferenceX AgentX lane's settings that hold at every - # concurrency it sweeps; its per-concurrency tuning stays in the guide. An - # exact-GPU block layers after the recipe `amd` baseline and wins last-wins - # dedupe there, so other AMD GPUs and every other strategy keep the baseline. - single_node_tp: - hardware_overrides: - mi355x: - # No --load-format override: fastsafetensors ran out of GPU memory while - # loading on the lane's image, so the `amd` baseline's `auto` applies. - extra_args: - - "--gpu-memory-utilization" - - "0.9" - - "--max-model-len" - - "1048576" - - "--enable-prefix-caching" - - "--kv-cache-dtype" - - "fp8" - - "--attention-config" - - '{"mla_prefill_backend":"ROCM_AITER_FA"}' - - "--stream-interval" - - "10" - extra_env: - VLLM_ROCM_AITER_MLA_ASM_PADDING: "asm" - VLLM_ROCM_AITER_MLA_DCP_VERIFY: "asm" - AITER_FLYDSL_STAGE2_FP8: "1" - # Use MegaMoE for DEP on GB200/GB300 NVL trays. B200/B300 leave # --moe-backend unset and use vLLM's default automatic selection. multi_node_dep: @@ -420,51 +349,6 @@ strategy_overrides: # can hold a PP stage past vLLM's 300 s model-execution default. VLLM_ENGINE_READY_TIMEOUT_S: "3600" VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" - # GB200: the InferenceX AgentX aggregated profile (SemiAnalysisAI/InferenceX#3748). - # A GB200 tray has 4 GPUs, so the strategy's TP-per-node x PP-across-nodes would - # give TP4 x PP4 on the 16-GPU floor; the lane runs TP8 x PP2 + DCP8 across the - # four trays instead (same --nnodes 4 rendezvous). Exact-GPU block: layers after - # the recipe Blackwell baseline, so its fastsafetensors, 1M context, fp8 KV and - # prefix caching stay. - hardware_overrides: - gb200: - extra_args: - - "--tensor-parallel-size" - - "8" - - "--pipeline-parallel-size" - - "2" - - "--decode-context-parallel-size" - - "8" - - "--dcp-comm-backend" - - "a2a" - - "--attention-config" - - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' - - "--mamba-ssm-cache-dtype" - - "bfloat16" - - "--gpu-memory-utilization" - - "0.88" - - "--moe-backend" - - "auto" - - "--language-model-only" - # The K4 draft matches the lane's concurrency 24-56 rows; 64 is the C=32 row. - - "--max-num-seqs" - - "64" - - "--speculative-config" - - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' - extra_env: - VLLM_USE_DIRECT_DCP_A2A: "1" - VLLM_USE_DIRECT_DCP_Q_GATHER: "1" - VLLM_USE_DIRECT_DCP_KV_GATHER: "1" - VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" - VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" - VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0" - VLLM_RPC_TIMEOUT: "600000" - MC_SLICE_SIZE: "1048576" - MC_STORE_MEMCPY: "1" - MC_WORKERS_PER_CTX: "4" - MC_ENABLE_DEST_DEVICE_AFFINITY: "1" - UCX_MEMTYPE_CACHE: "n" - UCX_MEMTYPE_REG_WHOLE: "n" pd_cluster: router: @@ -476,172 +360,9 @@ strategy_overrides: blackwell: extra_env: UCX_TLS: "rc,cuda_copy" - # GB200: the InferenceX AgentX 1P1D profile (SemiAnalysisAI/InferenceX#3748), - # one TP16 + DCP16 pool per role across four trays with the DSpark K4 draft. - # Shared serving flags restated here as for GB300 below; role blocks add the - # per-role memory and batch settings. - gb200: - extra_args: - - "--decode-context-parallel-size" - - "16" - - "--dcp-comm-backend" - - "a2a" - - "--attention-backend" - - "TOKENSPEED_MLA" - - "--attention-config" - - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' - - "--kv-cache-dtype" - - "fp8" - - "--mamba-ssm-cache-dtype" - - "bfloat16" - - "--enable-prefix-caching" - - "--prefix-match-unit" - - "128" - - "--enable-cumem-allocator" - - "--load-format" - - "safetensors" - - "--safetensors-load-strategy" - - "lazy" - - "--stream-interval" - - "10" - - "--language-model-only" - - "--speculative-config" - - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' - extra_env: - VLLM_USE_DIRECT_DCP_A2A: "1" - VLLM_USE_DIRECT_DCP_Q_GATHER: "1" - VLLM_USE_DIRECT_DCP_KV_GATHER: "1" - VLLM_ALLREDUCE_USE_FLASHINFER: "1" - VLLM_USE_NCCL_SYMM_MEM: "0" - VLLM_USE_V2_MODEL_RUNNER: "1" - VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0" - VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" - NCCL_CUMEM_ENABLE: "1" - NCCL_MNNVL_ENABLE: "1" - NCCL_NVLS_ENABLE: "1" - NCCL_P2P_LEVEL: "NVL" - NCCL_NET_PLUGIN: "none" - VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" - MC_SLICE_SIZE: "1048576" - MC_ENABLE_DEST_DEVICE_AFFINITY: "1" - UCX_MEMTYPE_CACHE: "n" - UCX_MEMTYPE_REG_WHOLE: "n" - UCX_RCACHE_MAX_UNRELEASED: "1024" - UCX_TCP_AF_PRIO: "inet" - # The lane left UCX on its default transports and raised the startup - # timeouts (this block replaces the recipe Blackwell baseline). - UCX_TLS: "all" - # The prefill MoE workspace (about 3.3 GiB at 16384 tokens) failed to - # allocate from reserved but fragmented memory; blocks above 512 MiB - # are not split. - PYTORCH_CUDA_ALLOC_CONF: "max_split_size_mb:512" - VLLM_ENGINE_READY_TIMEOUT_S: "3600" - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" - # GB300: the InferenceX AgentX xPyD profile (SemiAnalysisAI/InferenceX#4047), - # TP8 + DCP8 pools across two trays with the DSpark K4 draft. The pd_cluster - # Blackwell block above replaces the recipe Blackwell baseline, so the - # shared serving flags are restated here; role blocks below add the rest. - # It renders the low-concurrency specialized-decode profile (#4373: c16-c32), - # so both roles keep FP32 KV-cache SSM state, pin the DCP KV interleave to 1 - # and leave FlashInfer autotune off. NIXL P/D would otherwise raise the - # interleave to the KV block size after the specialized decode has read it, - # and its spec-verify causal mask would use the stale value. The uncapped - # c48+ points keep BF16 state, NIXL's interleave and autotune on (see the guide). - gb300: - extra_args: - - "--decode-context-parallel-size" - - "8" - - "--dcp-comm-backend" - - "a2a" - - "--attention-backend" - - "TOKENSPEED_MLA" - - "--kv-cache-dtype" - - "fp8" - - "--mamba-ssm-cache-dtype" - - "float32" - - "--cp-kv-cache-interleave-size" - - "1" - - "--enable-prefix-caching" - - "--prefix-match-unit" - - "128" - - "--gpu-memory-utilization" - - "0.92" - - "--max-cudagraph-capture-size" - - "512" - - "--enable-cumem-allocator" - - "--no-enable-flashinfer-autotune" - - "--load-format" - - "fastsafetensors" - - "--stream-interval" - - "10" - - "--language-model-only" - - "--speculative-config" - - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' - extra_env: - VLLM_USE_DIRECT_DCP_A2A: "1" - VLLM_USE_DIRECT_DCP_Q_GATHER: "1" - VLLM_USE_DIRECT_DCP_KV_GATHER: "1" - VLLM_ALLREDUCE_USE_FLASHINFER: "1" - VLLM_USE_NCCL_SYMM_MEM: "0" - VLLM_USE_V2_MODEL_RUNNER: "1" - VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" - NCCL_CUMEM_ENABLE: "1" - NCCL_MNNVL_ENABLE: "1" - NCCL_NVLS_ENABLE: "1" - NCCL_P2P_LEVEL: "NVL" - VLLM_CONNECTOR_PREFETCH_DEPTH: "8" - VLLM_CONNECTOR_PREFETCH_KV_CAP: "0.65" - VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" - MC_SLICE_SIZE: "1048576" - MC_MAX_MR_SIZE: "4294967296" - MC_ENABLE_DEST_DEVICE_AFFINITY: "1" - UCX_MEMTYPE_CACHE: "n" - UCX_MEMTYPE_REG_WHOLE: "n" - UCX_RCACHE_MAX_UNRELEASED: "1024" - UCX_TCP_AF_PRIO: "inet" - # The qualified workers left UCX on its default transports; this - # replaces the Blackwell PD block's rc,cuda_copy restriction. - UCX_TLS: "all" - # The pd_cluster Blackwell block replaces the recipe baseline, so the - # startup timeouts are restated: loading and autotune take ~25 min. - VLLM_ENGINE_READY_TIMEOUT_S: "3600" - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" - VLLM_RPC_TIMEOUT: "600000" prefill: - # GB200 prefill is one TP16 + DCP16 pool across four trays (#3748); every - # other GPU keeps the TEP8 prefill pool. - parallelism: - default: tep - gb200: tp + parallelism: tep tp: 8 - nodes: - gb200: 4 - gb300: 2 - hardware_overrides: - gb200: - # The lane leaves FlashInfer autotune at vLLM's default (on); this undoes - # the shared prefill flag below. - extra_args: - - "--enable-flashinfer-autotune" - - "--gpu-memory-utilization" - - "0.88" - - "--max-num-seqs" - - "32" - - "--max-cudagraph-capture-size" - - "512" - extra_env: - VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" - gb300: - extra_args: - - "--moe-backend" - - "deep_gemm_mega_moe" - - "--attention-config" - - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' - - "--max-num-seqs" - - "64" - extra_env: - VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" - VLLM_ENABLE_GEMM_RS: "1" vllm_args: - "--max-num-batched-tokens" - "16384" @@ -650,65 +371,9 @@ strategy_overrides: env: VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "1" decode: - # GB300 decode is a TP8 + DCP8 pool without expert parallelism (#4047); - # every other GPU keeps the DEP decode pool. - parallelism: - default: dep - gb200: tp - gb300: tp + parallelism: dep tp: 1 - nodes: - gb200: 4 - gb300: 2 env: {} - hardware_overrides: - # GB200 decode is one TP16 + DCP16 pool without expert parallelism (#3748). - # Same undo of the DEP-only decode flags as GB300 below; 32 sequences and a - # 160-token decode graph cap are the lane's concurrency-16 values. Decode - # batches at most 480 tokens, so 1024 bounds the FlashInfer MoE autotune; - # the 16384 default bucket needs a 2.9 GiB workspace that runs out of - # memory at 0.92 (InferenceX #4322). - gb200: - extra_args: - - "--no-data-parallel-hybrid-lb" - - "--max-num-batched-tokens" - - "1024" - - "--compilation-config" - - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' - - "--gpu-memory-utilization" - - "0.92" - - "--max-num-seqs" - - "32" - - "--max-cudagraph-capture-size" - - "160" - extra_env: - VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" - gb300: - # Undo the DEP-only decode flags below for the TP pool. Hybrid LB is - # invalid on the --headless follower; 16384 is vLLM's API-server default - # on >=160 GB GPUs, which the qualified decode ran with. - # FULL_AND_PIECEWISE is vLLM's V1 default; it replaces the strategy's - # FULL_DECODE_ONLY, which the qualified decode did not set. - # Specialized decode (InferenceX #4373, c16-c32), decode only: it needs no - # expert parallelism and FULL decode CUDA graphs. --max-num-seqs is - # 64 // (K + 1) = 12 for the K4 draft, and graphs are captured for every - # batch size up to 64. The lane also drops tokenspeed_mla_min_split_kv. - extra_args: - - "--no-data-parallel-hybrid-lb" - - "--max-num-batched-tokens" - - "16384" - - "--additional-config" - - '{"decode_backend":"latency_specialized"}' - - "--max-num-seqs" - - "12" - - "--max-cudagraph-capture-size" - - "64" - - "--compilation-config" - - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64]}' - - "--attention-config" - - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' - extra_env: - VLLM_ENABLE_GEMM_RS: "0" vllm_args: - "--data-parallel-hybrid-lb" - "--no-enable-prefix-caching" @@ -718,6 +383,1031 @@ strategy_overrides: - "32" - "--no-disable-hybrid-kv-cache-manager" +# Tuned presets: the SemiAnalysis InferenceX AgentX lanes, each a delta on +# the baseline command for one GPU + strategy, on by default there (Baseline +# in the builder drops it). Their runbooks render on the preset's card. +presets: + gb200_agentx_agg: + label: "AgentX aggregated" + hardware: gb200 + strategy: multi_node_tp_pp + workload: "Agentic (InferenceX AgentX)" + default: true + description: "TP8 x PP2 + DCP8 across four GB200 trays with the DSpark K4 draft and an embedded Mooncake store." + source: { label: "InferenceX #3748", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/3748", key: "kimik3-fp4-gb200-dynamo-vllm-agentic-dspark-mooncake-agg", revision: b7eee97 } + image: "vllm/vllm-openai:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f" + # vLLM's own Python frontend; the K4 draft is in args, so the generic + # Spec decoding toggle (a K8 draft) stays off. + select: + frontend: python + features: { spec_decoding: false } + # GB200: the InferenceX AgentX aggregated lane (SemiAnalysisAI/InferenceX#3748). + # A GB200 tray has 4 GPUs, so the strategy's TP-per-node x PP-across-nodes would + # give TP4 x PP4 on the 16-GPU floor; the lane runs TP8 x PP2 + DCP8 across the + # four trays instead (same --nnodes 4 rendezvous). Exact-GPU block: layers after + # the recipe Blackwell baseline, so its fastsafetensors, 1M context, fp8 KV and + # prefix caching stay. + args: + - "--tensor-parallel-size" + - "8" + - "--pipeline-parallel-size" + - "2" + - "--decode-context-parallel-size" + - "8" + - "--dcp-comm-backend" + - "a2a" + - "--attention-config" + - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' + - "--mamba-ssm-cache-dtype" + - "bfloat16" + - "--gpu-memory-utilization" + - "0.88" + - "--moe-backend" + - "auto" + - "--language-model-only" + # The K4 draft matches the lane's concurrency 24-56 rows; 64 is the C=32 row. + - "--speculative-config" + - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' + env: + VLLM_USE_DIRECT_DCP_A2A: "1" + VLLM_USE_DIRECT_DCP_Q_GATHER: "1" + VLLM_USE_DIRECT_DCP_KV_GATHER: "1" + VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" + VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" + VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0" + VLLM_RPC_TIMEOUT: "600000" + MC_SLICE_SIZE: "1048576" + MC_STORE_MEMCPY: "1" + MC_WORKERS_PER_CTX: "4" + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + UCX_MEMTYPE_CACHE: "n" + UCX_MEMTYPE_REG_WHOLE: "n" + json_args: + "--compilation-config": { cudagraph_mode: "FULL_AND_PIECEWISE" } + # Client concurrency points: K is num_speculative_tokens (K4 in args), S is + # --max-num-seqs; every graph list ends with the shared 256..8192 tail. + points: + c1: + label: "C1" + concurrency: 1 + args: ["--max-num-seqs", "2"] + json_args: + "--speculative-config": { num_speculative_tokens: 7 } + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192] } + c4: + label: "C4" + concurrency: 4 + args: ["--max-num-seqs", "8"] + json_args: + "--speculative-config": { num_speculative_tokens: 7 } + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 32, 40, 48, 56, 64, 128, 256, 512, 1024, 2048, 4096, 8192] } + c8: + label: "C8" + concurrency: 8 + args: ["--max-num-seqs", "16"] + json_args: + "--speculative-config": { num_speculative_tokens: 7 } + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 40, 56, 64, 80, 104, 128, 256, 512, 1024, 2048, 4096, 8192] } + c24: + label: "C24" + concurrency: 24 + args: ["--max-num-seqs", "48"] + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 40, 56, 64, 80, 104, 128, 256, 512, 1024, 2048, 4096, 8192] } + c32: + label: "C32" + concurrency: 32 + default: true + args: ["--max-num-seqs", "64"] + json_args: + "--compilation-config": { cudagraph_capture_sizes: [1, 7, 18, 34, 53, 64, 75, 100, 128, 256, 512, 1024, 2048, 4096, 8192] } + c48: + label: "C48" + concurrency: 48 + args: ["--max-num-seqs", "96"] + json_args: + "--compilation-config": { cudagraph_capture_sizes: [1, 7, 18, 34, 53, 64, 75, 100, 128, 256, 512, 1024, 2048, 4096, 8192] } + c56: + label: "C56" + concurrency: 56 + args: ["--max-num-seqs", "112"] + json_args: + "--compilation-config": { cudagraph_capture_sizes: [1, 7, 18, 34, 53, 64, 75, 100, 128, 256, 512, 1024, 2048, 4096, 8192] } + runbook: | + The GB200 AgentX lanes in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb200-dynamo-vllm-agentic-dspark-mooncake-*`, [SemiAnalysisAI/InferenceX#4322](https://github.com/SemiAnalysisAI/InferenceX/pull/4322), pinned to [`b7eee97`](https://github.com/SemiAnalysisAI/InferenceX/tree/b7eee972451f9297941d4cbe6d42d489d7ec8363/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/gb200-fp4/agentx), which builds on #3748) serve the default `moonshotai/Kimi-K3` checkpoint (MXFP4 weights) with DSpark speculative decoding and a Mooncake KV store on 4-GPU NVL trays. They are benchmark configurations; the settings below are the production-useful parts. The builder's **Mooncake** option generates the `mooncake_master` and store config; the lane-specific connector settings are below. Throughput runs use InferenceX's golden acceptance lengths for DSpark. + + Served by vLLM's own frontend (no Dynamo) with an embedded RDMA Mooncake store (`global_segment_size` 200GB, `mlx5_0..mlx5_3`). The preset renders the TP8 x PP2 + DCP8 layout across four trays with the lane's K4/K7 draft, `--max-num-seqs` and CUDA graph sizes at each concurrency point. The builder's **Mooncake** KV Offload option renders a plain `MooncakeStoreConnector` with a 100GB segment; use the lane's connector and 200GB store from the command below. + + ```bash + export VLLM_USE_DIRECT_DCP_A2A=1 VLLM_USE_DIRECT_DCP_Q_GATHER=1 VLLM_USE_DIRECT_DCP_KV_GATHER=1 + export VLLM_ALLREDUCE_USE_FLASHINFER=1 NCCL_CUMEM_ENABLE=1 VLLM_USE_V2_MODEL_RUNNER=1 + export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0 + export MC_STORE_MEMCPY=1 MC_ENABLE_DEST_DEVICE_AFFINITY=1 MC_WORKERS_PER_CTX=4 + export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 + export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000 + vllm serve moonshotai/Kimi-K3 --trust-remote-code \ + --tensor-parallel-size 8 --pipeline-parallel-size 2 \ + --decode-context-parallel-size 8 --dcp-comm-backend a2a \ + --attention-backend TOKENSPEED_MLA \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ + --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 \ + --max-model-len 1048576 --gpu-memory-utilization 0.88 --max-num-seqs --max-num-batched-tokens 8192 \ + --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","method":"dspark","num_speculative_tokens":,"attention_backend":"TOKENSPEED_MLA","draft_sample_method":"probabilistic"}' \ + --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[,256,512,1024,2048,4096,8192]}' \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}' \ + --load-format fastsafetensors --moe-backend auto --no-enable-flashinfer-autotune \ + --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 --enable-auto-tool-choice \ + --language-model-only + ``` + + gb200_agentx_pd: + label: "AgentX 1P1D" + hardware: gb200 + strategy: pd_cluster + workload: "Agentic (InferenceX AgentX)" + default: true + description: "One TP16 + DCP16 prefill pool and one decode pool across four GB200 trays each, served by Dynamo with the DSpark K4 draft." + source: { label: "InferenceX #4322", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4322", key: "kimik3-fp4-gb200-dynamo-vllm-agentic-dspark-mooncake-disagg", revision: b7eee97 } + image: + vllm: "vllm/vllm-openai:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f" + dynamo: "ffbbaaa6c44370e798687881317535ac17d66a76" + # Served by Dynamo with the Rust frontend. Dynamo workers take engine + # arguments only, so the API-server tool-call flags are off (see the + # runbook for the --dyn-* parsers); the K4 draft is in args below. + select: + pd_router: dynamo + frontend: rust + features: { tool_calling: false, spec_decoding: false } + # GB200: the InferenceX AgentX 1P1D lane (SemiAnalysisAI/InferenceX#3748), + # one TP16 + DCP16 pool per role across four trays with the DSpark K4 draft. + # Shared serving flags restated here as for GB300 below; role blocks add the + # per-role memory and batch settings. + args: + - "--decode-context-parallel-size" + - "16" + - "--dcp-comm-backend" + - "a2a" + - "--attention-backend" + - "TOKENSPEED_MLA" + - "--attention-config" + - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' + - "--kv-cache-dtype" + - "fp8" + - "--mamba-ssm-cache-dtype" + - "bfloat16" + - "--enable-prefix-caching" + - "--prefix-match-unit" + - "128" + - "--enable-cumem-allocator" + - "--load-format" + - "safetensors" + - "--safetensors-load-strategy" + - "lazy" + - "--stream-interval" + - "10" + - "--language-model-only" + - "--speculative-config" + - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' + env: + VLLM_USE_DIRECT_DCP_A2A: "1" + VLLM_USE_DIRECT_DCP_Q_GATHER: "1" + VLLM_USE_DIRECT_DCP_KV_GATHER: "1" + VLLM_ALLREDUCE_USE_FLASHINFER: "1" + VLLM_USE_NCCL_SYMM_MEM: "0" + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0" + VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" + NCCL_CUMEM_ENABLE: "1" + NCCL_MNNVL_ENABLE: "1" + NCCL_NVLS_ENABLE: "1" + NCCL_P2P_LEVEL: "NVL" + NCCL_NET_PLUGIN: "none" + VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" + MC_SLICE_SIZE: "1048576" + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + UCX_MEMTYPE_CACHE: "n" + UCX_MEMTYPE_REG_WHOLE: "n" + UCX_RCACHE_MAX_UNRELEASED: "1024" + UCX_TCP_AF_PRIO: "inet" + # The lane left UCX on its default transports and raised the startup + # timeouts (this block replaces the recipe Blackwell baseline). + UCX_TLS: "all" + # The prefill MoE workspace (about 3.3 GiB at 16384 tokens) failed to + # allocate from reserved but fragmented memory; blocks above 512 MiB + # are not split. + PYTORCH_CUDA_ALLOC_CONF: "max_split_size_mb:512" + VLLM_ENGINE_READY_TIMEOUT_S: "3600" + VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" + roles: + prefill: + # GB200 prefill is one TP16 + DCP16 pool across four trays (#3748). + parallelism: tp + nodes: 4 + # The lane leaves FlashInfer autotune at vLLM's default (on); this undoes + # the shared prefill flag below. + args: + - "--enable-flashinfer-autotune" + - "--gpu-memory-utilization" + - "0.88" + - "--max-cudagraph-capture-size" + - "512" + env: + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" + decode: + parallelism: tp + nodes: 4 + # GB200 decode is one TP16 + DCP16 pool without expert parallelism (#3748). + # Same undo of the DEP-only decode flags as GB300 below; 32 sequences and a + # 160-token decode graph cap are the lane's concurrency-16 values. Decode + # batches at most 480 tokens, so 1024 bounds the FlashInfer MoE autotune; + # the 16384 default bucket needs a 2.9 GiB workspace that runs out of + # memory at 0.92 (InferenceX #4322). + args: + - "--no-data-parallel-hybrid-lb" + - "--max-num-batched-tokens" + - "1024" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' + env: + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" + # --max-num-seqs is 2x concurrency on both roles; the decode graph cap is + # max-num-seqs x (4 + 1) tokens, and decode memory utilization is 0.92 up + # to a 320-token cap and 0.90 above it. + points: + c16: + label: "C16" + concurrency: 16 + default: true + roles: + prefill: { args: ["--max-num-seqs", "32"] } + decode: { args: ["--gpu-memory-utilization", "0.92", "--max-num-seqs", "32", "--max-cudagraph-capture-size", "160"] } + c32: + label: "C32" + concurrency: 32 + roles: + prefill: { args: ["--max-num-seqs", "64"] } + decode: { args: ["--gpu-memory-utilization", "0.92", "--max-num-seqs", "64", "--max-cudagraph-capture-size", "320"] } + c48: + label: "C48" + concurrency: 48 + roles: + prefill: { args: ["--max-num-seqs", "96"] } + decode: { args: ["--gpu-memory-utilization", "0.90", "--max-num-seqs", "96", "--max-cudagraph-capture-size", "480"] } + runbook: | + The GB200 AgentX lanes in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb200-dynamo-vllm-agentic-dspark-mooncake-*`, [SemiAnalysisAI/InferenceX#4322](https://github.com/SemiAnalysisAI/InferenceX/pull/4322), pinned to [`b7eee97`](https://github.com/SemiAnalysisAI/InferenceX/tree/b7eee972451f9297941d4cbe6d42d489d7ec8363/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/gb200-fp4/agentx), which builds on #3748) serve the default `moonshotai/Kimi-K3` checkpoint (MXFP4 weights) with DSpark speculative decoding and a Mooncake KV store on 4-GPU NVL trays. They are benchmark configurations; the settings below are the production-useful parts. The builder's **Mooncake** option generates the `mooncake_master` and store config; the lane-specific connector settings are below. Throughput runs use InferenceX's golden acceptance lengths for DSpark. + + Image `vllm/vllm-openai:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f` with Dynamo `ffbbaaa6c44370e798687881317535ac17d66a76`, etcd and an embedded RDMA Mooncake store. The lane's request plane is TCP (`DYN_REQUEST_PLANE=tcp`) and it starts no NATS service. Each role is one worker across four trays. + + Run each command on all four trays of the role with the same `--master-addr` (the first tray's address), `--nnodes 4`, `--node-rank 0` to `3`, and `--headless` on ranks 1-3. Prefill (`P_GPU`/`P_CAP` and `S` are the selected point's values): + + ```bash + python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ + --disaggregation-mode prefill --request-plane tcp \ + --nnodes 4 --node-rank <0-3> --master-addr --device-ids 0,1,2,3 \ + --tensor-parallel-size 16 --decode-context-parallel-size 16 --dcp-comm-backend a2a \ + --attention-backend TOKENSPEED_MLA \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ + --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' \ + --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 --language-model-only \ + --gpu-memory-utilization --max-num-seqs --max-cudagraph-capture-size \ + --enable-cumem-allocator --stream-interval 10 --load-format safetensors --safetensors-load-strategy lazy \ + --trust-remote-code \ + --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \ + --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' + ``` + + Decode (`D_GPU`/`D_CAP` and `S` are the selected point's values). `--max-num-batched-tokens 1024` bounds the FlashInfer MoE autotune to the token counts decode can reach (at most 480); at vLLM's 16384-token default the autotune bucket needs a 2.9 GiB workspace that runs out of memory at 0.92, and the other ranks hang: + + ```bash + python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ + --disaggregation-mode decode --request-plane tcp \ + --nnodes 4 --node-rank <0-3> --master-addr --device-ids 0,1,2,3 \ + --tensor-parallel-size 16 --decode-context-parallel-size 16 --dcp-comm-backend a2a \ + --attention-backend TOKENSPEED_MLA \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ + --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' \ + --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 --language-model-only \ + --gpu-memory-utilization --max-num-seqs --max-num-batched-tokens 1024 --max-cudagraph-capture-size \ + --enable-cumem-allocator --stream-interval 10 --load-format safetensors --safetensors-load-strategy lazy \ + --trust-remote-code \ + --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' + ``` + + For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there. + + - **Environment (both roles):** the three direct-DCP variables, `VLLM_ALLREDUCE_USE_FLASHINFER=1`, `VLLM_USE_V2_MODEL_RUNNER=1`, `VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0`, `VLLM_USE_NCCL_SYMM_MEM=0`, `VLLM_SSM_CONV_STATE_LAYOUT=DS`, `VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0`, `NCCL_MNNVL_ENABLE=1`, `NCCL_NVLS_ENABLE=1`, `NCCL_CUMEM_ENABLE=1`, `NCCL_P2P_LEVEL=NVL`, `NCCL_NET_PLUGIN=none`, `VLLM_MOONCAKE_LOAD_RECV_THREADS=4`, `MC_SLICE_SIZE=1048576`, `MC_ENABLE_DEST_DEVICE_AFFINITY=1`, `UCX_MEMTYPE_CACHE=n`, `UCX_MEMTYPE_REG_WHOLE=n`, `UCX_RCACHE_MAX_UNRELEASED=1024`, `UCX_TCP_AF_PRIO=inet`, `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0`, `VLLM_ENGINE_READY_TIMEOUT_S=3600`, `VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800`, `ETCD_LEASE_TTL=600`, `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512` and `DYN_REQUEST_PLANE=tcp`. + - **Frontend:** the Dynamo frontend runs with `--router-mode least-loaded --router-session-affinity-ttl-secs 900`, so each conversation sticks to one worker and reuses its prefix cache. The Rust frontend is selected with `VLLM_USE_RUST_FRONTEND=1`, with `DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600` and `DYN_TOKENIZER_CACHE_BYTES=8589934592`. AgentX sends the conversation ID in the `X-Dynamo-Session-Id` request header. + - **Mooncake master:** `--default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10`, 128GB global segment, 4GB local buffer, RDMA over `mlx5_0..mlx5_3`. + + gb300_agentx_pd: + label: "AgentX xPyD" + hardware: gb300 + strategy: pd_cluster + workload: "Agentic (InferenceX AgentX)" + default: true + description: "TP8 + DCP8 prefill and decode workers across two GB300 trays each, served by Dynamo with the DSpark K4 draft." + source: { label: "InferenceX #4047", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4047", key: "kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg" } + image: + vllm: "vllm/vllm-openai:nightly-1b2a666391f2a24228a1436b64d547c34cc6b6ce" + dynamo: "ffbbaaa6c44370e798687881317535ac17d66a76" + # Served by Dynamo with the Rust frontend; tool-call flags off under Dynamo, + # and the K4 draft is in args (specialized decode rejects the toggle's K8). + select: + pd_router: dynamo + frontend: rust + features: { tool_calling: false, spec_decoding: false } + # GB300: the InferenceX AgentX xPyD lane (SemiAnalysisAI/InferenceX#4047), + # TP8 + DCP8 pools across two trays with the DSpark K4 draft. The pd_cluster + # Blackwell block above replaces the recipe Blackwell baseline, so the + # shared serving flags are restated here; role blocks below add the rest. + # It renders the low-concurrency specialized-decode profile (#4373: c16-c32), + # so both roles keep FP32 KV-cache SSM state, pin the DCP KV interleave to 1 + # and leave FlashInfer autotune off. NIXL P/D would otherwise raise the + # interleave to the KV block size after the specialized decode has read it, + # and its spec-verify causal mask would use the stale value. The uncapped + # c48+ points keep BF16 state, NIXL's interleave and autotune on (see the guide). + args: + - "--decode-context-parallel-size" + - "8" + - "--dcp-comm-backend" + - "a2a" + - "--attention-backend" + - "TOKENSPEED_MLA" + - "--kv-cache-dtype" + - "fp8" + - "--mamba-ssm-cache-dtype" + - "float32" + - "--cp-kv-cache-interleave-size" + - "1" + - "--enable-prefix-caching" + - "--prefix-match-unit" + - "128" + - "--gpu-memory-utilization" + - "0.92" + - "--max-cudagraph-capture-size" + - "512" + - "--enable-cumem-allocator" + - "--no-enable-flashinfer-autotune" + - "--load-format" + - "fastsafetensors" + - "--stream-interval" + - "10" + - "--language-model-only" + - "--speculative-config" + - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' + env: + VLLM_USE_DIRECT_DCP_A2A: "1" + VLLM_USE_DIRECT_DCP_Q_GATHER: "1" + VLLM_USE_DIRECT_DCP_KV_GATHER: "1" + VLLM_ALLREDUCE_USE_FLASHINFER: "1" + VLLM_USE_NCCL_SYMM_MEM: "0" + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" + NCCL_CUMEM_ENABLE: "1" + NCCL_MNNVL_ENABLE: "1" + NCCL_NVLS_ENABLE: "1" + NCCL_P2P_LEVEL: "NVL" + VLLM_CONNECTOR_PREFETCH_DEPTH: "8" + VLLM_CONNECTOR_PREFETCH_KV_CAP: "0.65" + VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" + MC_SLICE_SIZE: "1048576" + MC_MAX_MR_SIZE: "4294967296" + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + UCX_MEMTYPE_CACHE: "n" + UCX_MEMTYPE_REG_WHOLE: "n" + UCX_RCACHE_MAX_UNRELEASED: "1024" + UCX_TCP_AF_PRIO: "inet" + # The qualified workers left UCX on its default transports; this + # replaces the Blackwell PD block's rc,cuda_copy restriction. + UCX_TLS: "all" + # The pd_cluster Blackwell block replaces the recipe baseline, so the + # startup timeouts are restated: loading and autotune take ~25 min. + VLLM_ENGINE_READY_TIMEOUT_S: "3600" + VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" + VLLM_RPC_TIMEOUT: "600000" + roles: + prefill: + nodes: 2 + args: + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--attention-config" + - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' + - "--max-num-seqs" + - "64" + env: + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0" + VLLM_ENABLE_GEMM_RS: "1" + decode: + # GB300 decode is a TP8 + DCP8 pool without expert parallelism (#4047). + parallelism: tp + nodes: 2 + # Undo the DEP-only decode flags below for the TP pool. Hybrid LB is + # invalid on the --headless follower; 16384 is vLLM's API-server default + # on >=160 GB GPUs, which the qualified decode ran with. + # FULL_AND_PIECEWISE is vLLM's V1 default; it replaces the strategy's + # FULL_DECODE_ONLY, which the qualified decode did not set. + # Specialized decode (InferenceX #4373, c16-c32), decode only: it needs no + # expert parallelism and FULL decode CUDA graphs. --max-num-seqs is + # 64 // (K + 1) = 12 for the K4 draft, and graphs are captured for every + # batch size up to 64. The lane also drops tokenspeed_mla_min_split_kv. + args: + - "--no-data-parallel-hybrid-lb" + - "--max-num-batched-tokens" + - "16384" + - "--additional-config" + - '{"decode_backend":"latency_specialized"}' + - "--max-num-seqs" + - "12" + - "--max-cudagraph-capture-size" + - "64" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64]}' + - "--attention-config" + - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' + env: + VLLM_ENABLE_GEMM_RS: "0" + runbook: | + The GB300 AgentX disaggregated lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg`, [SemiAnalysisAI/InferenceX#4047](https://github.com/SemiAnalysisAI/InferenceX/pull/4047), points refined in [#4373](https://github.com/SemiAnalysisAI/InferenceX/pull/4373)) runs Dynamo + vLLM with several prefill and decode workers. Prefill and decode workers are both TP8 + DCP8 across two 4-GPU GB300 trays; prefill workers also use EP8 with MegaMoE, and decode workers use no expert parallelism. KV moves over NIXL, and all workers share an RDMA Mooncake store. Image `vllm/vllm-openai:nightly-1b2a666391f2a24228a1436b64d547c34cc6b6ce`, Dynamo `ffbbaaa6c44370e798687881317535ac17d66a76`. + + | Clients | Prefill workers | Decode workers | GPUs | Decode profile | Decode `--max-num-seqs` | + |---:|---:|---:|---:|---|---:| + | 16 | 1 | 3 | 32 | [Specialized](#specialized-decode-low-concurrency) | 12 | + | 24-32 | 1 | 2 | 24 | [Specialized](#specialized-decode-low-concurrency) | 12 | + | 48-64 | 1 | 2 | 24 | Regular | 128 | + | 240 | 2 | 3 | 40 | Regular | 128 | + | 240 | 3 | 4 | 56 | Regular | 128 | + + Prefill workers always use `--max-num-seqs 64`. The regular profile is the Workers commands below; the specialized one changes both roles as described under [Specialized decode](#specialized-decode-low-concurrency). + + The preset renders one prefill pool and one decode pool of the specialized profile (the 1P2D c24-c32 row); its decode uses `--max-num-seqs 12`. For the regular points use the Workers commands below. It turns **Spec decoding** off, because the toggle's generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft was not qualified and specialized decode allows at most 7 draft tokens, and **Tool calling** off (see the parser note under Workers). Other pool modes are not supported on GB300: the DEP decode pool renders TP1 with `--decode-context-parallel-size 8`, which vLLM rejects. + + Documented here only: worker counts beyond one pool per role, the Mooncake store connector (the builder's Mooncake option renders a plain `MooncakeStoreConnector` without the qualified NIXL MultiConnector, recompute policy and load options), the Mooncake services, prefill `--kv-events-config`, and the frontend's router mode and session affinity (the builder renders round-robin). + + ### Services + + Start etcd, NATS and a Mooncake master once: + + ```bash + etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 + printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf + mooncake_master --enable_http_metadata_server=true \ + --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10 + ``` + + The workers' Mooncake store uses P2P handshake metadata, RDMA over the node's NICs, a 160GB global segment and a 4GB local buffer, with offload disabled. + + ### Environment (every worker) + + ```bash + export VLLM_USE_DIRECT_DCP_A2A=1 VLLM_USE_DIRECT_DCP_Q_GATHER=1 VLLM_USE_DIRECT_DCP_KV_GATHER=1 + export VLLM_ALLREDUCE_USE_FLASHINFER=1 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0 VLLM_USE_NCCL_SYMM_MEM=0 + export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_P2P_LEVEL=NVL + export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=1 VLLM_SSM_CONV_STATE_LAYOUT=DS + export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0 + export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 MC_MAX_MR_SIZE=4294967296 MC_ENABLE_DEST_DEVICE_AFFINITY=1 + export VLLM_CONNECTOR_PREFETCH_DEPTH=8 VLLM_CONNECTOR_PREFETCH_KV_CAP=0.65 + export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n UCX_RCACHE_MAX_UNRELEASED=1024 UCX_TCP_AF_PRIO=inet + export DYN_REQUEST_PLANE=tcp DYN_TCP_REQUEST_TIMEOUT=120 ETCD_LEASE_TTL=600 + export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000 + ``` + + Set `VLLM_ENABLE_GEMM_RS=1` on prefill workers and `VLLM_ENABLE_GEMM_RS=0` on decode workers. Loading the 2.8T checkpoint across trays and autotuning takes a while, which is why the engine-ready, execute-model and RPC timeouts are raised. + + ### Workers + + Each worker spans two trays. Run the command on both trays with the same `--master-addr` (the first tray's address), `--nnodes 2` and `--node-rank 0` or `1`. + + Prefill: + + ```bash + python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ + --disaggregation-mode prefill --request-plane tcp \ + --nnodes 2 --node-rank <0|1> --master-addr --device-ids 0,1,2,3 \ + --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \ + --enable-expert-parallel --moe-backend deep_gemm_mega_moe \ + --attention-backend TOKENSPEED_MLA \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' \ + --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ + --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \ + --enable-prefix-caching --prefix-match-unit 128 \ + --gpu-memory-utilization 0.92 --max-num-seqs 64 --max-cudagraph-capture-size 512 \ + --enable-cumem-allocator --enable-flashinfer-autotune \ + --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \ + --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \ + --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' + ``` + + Decode (regular profile, `--max-num-seqs 128`): + + ```bash + python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ + --disaggregation-mode decode --request-plane tcp \ + --nnodes 2 --node-rank <0|1> --master-addr --device-ids 0,1,2,3 \ + --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \ + --attention-backend TOKENSPEED_MLA \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"tokenspeed_mla_min_split_kv":8}' \ + --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ + --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \ + --enable-prefix-caching --prefix-match-unit 128 \ + --gpu-memory-utilization 0.92 --max-num-seqs 128 --max-cudagraph-capture-size 512 \ + --enable-cumem-allocator --enable-flashinfer-autotune \ + --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \ + --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' + ``` + + For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there. + + ### Specialized decode (low concurrency) + + The c16-c32 points use the `latency_specialized` decode backend from [GB300 low-latency preset](#preset-gb300_low_latency) on the decode workers (1P3D for c16, 1P2D for c24-c32), with the Workers commands above changed as follows. The backend's requirements are listed there: V2 model runner (set in Environment), FULL decode CUDA graphs, BF16 convolution state with FP32 recurrent state, DSpark with at most 7 draft tokens, no pipeline or expert parallelism, and DCP of 1 or TP. The decode workers already meet the last two. + + Both roles: + + ```bash + --mamba-ssm-cache-dtype float32 --cp-kv-cache-interleave-size 1 --no-enable-flashinfer-autotune + # replacing --mamba-ssm-cache-dtype bfloat16 and --enable-flashinfer-autotune + ``` + + Pin `--cp-kv-cache-interleave-size 1` on both roles. Left unset, NIXL P/D raises the DCP KV interleave to the KV block size after the specialized decode has read it, so its spec-verify causal mask uses the stale value. + + Decode only: + + ```bash + --additional-config '{"decode_backend":"latency_specialized"}' \ + --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ + --max-num-seqs 12 --max-cudagraph-capture-size 64 \ + --compilation-config "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[$(seq -s, 1 64)]}" + # replacing the regular decode's --attention-config (it also sets tokenspeed_mla_min_split_kv), + # --max-num-seqs 128 and --max-cudagraph-capture-size 512 + ``` + + `--max-num-seqs` is `64 // (K + 1)` with the K4 draft, so a full verify batch (12 x 5 = 60 tokens) fits the 64-token graph, and CUDA graphs are captured for every batch size up to 64. The prefill workers do not set `--additional-config`: the backend rejects the expert parallelism they use. The uncapped c48+ points keep the regular profile, so TTFT is not traded for ITL by a standing decode queue. + + ### Frontend + + ```bash + export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TOKENIZER_CACHE_BYTES=8589934592 DYN_TCP_REQUEST_TIMEOUT=120 + python3 -m dynamo.frontend --router-mode least-loaded --router-session-affinity-ttl-secs 900 + ``` + + Session affinity keeps each conversation on one worker so it reuses that worker's prefix cache; send the conversation ID in the `X-Dynamo-Session-Id` request header. + + gb300_low_latency: + label: "Low latency" + hardware: gb300 + strategy: multi_node_tp + workload: "Low latency (client concurrency 1-4)" + description: "One aggregated engine on the latency_specialized decode backend with the DSpark draft: TP16/DCP16 over four trays at concurrency 1, TP8/DCP8 over two at 2-4." + source: { label: "recipes #1134", href: "https://github.com/vllm-project/recipes/pull/1134" } + image: "vllm/vllm-openai:nightly" + select: + frontend: python + features: { spec_decoding: false } + env: + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_ALLREDUCE_USE_FLASHINFER: "1" + VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: "1" + VLLM_USE_DIRECT_DCP_A2A: "1" + VLLM_USE_DIRECT_DCP_Q_GATHER: "1" + VLLM_USE_DIRECT_DCP_KV_GATHER: "1" + VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" + VLLM_RPC_TIMEOUT: "600000" + VLLM_LOG_STATS_INTERVAL: "1" + args: + - "--dcp-comm-backend" + - "a2a" + - "--additional-config" + - '{"decode_backend":"latency_specialized"}' + - "--attention-backend" + - "TOKENSPEED_MLA" + - "--attention-config" + - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' + - "--kv-cache-dtype" + - "fp8" + - "--mamba-ssm-cache-dtype" + - "float32" + - "--gpu-memory-utilization" + - "0.92" + - "--max-num-seqs" + - "8" + - "--max-num-batched-tokens" + - "16384" + - "--speculative-config" + - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":7,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' + - "--enable-prefix-caching" + - "--prefix-match-unit" + - "128" + - "--disable-flashinfer-autotune" + - "--moe-backend" + - "auto" + - "--load-format" + - "fastsafetensors" + - "--stream-interval" + - "10" + - "--language-model-only" + - "--trust-remote-code" + points: + tp16: + label: "C1 · TP16 (4 trays)" + concurrency: 1 + default: true + select: { nodes: 4 } + args: ["--decode-context-parallel-size", "16", "--max-cudagraph-capture-size", "48"] + json_args: + "--speculative-config": { num_speculative_tokens: 5 } + "--compilation-config": { cudagraph_mode: "FULL_AND_PIECEWISE", cudagraph_capture_sizes: [5, 6, 10, 12, 20, 24, 40, 48] } + tp8: + label: "C2-C4 · TP8 (2 trays)" + concurrency: 4 + select: { nodes: 2 } + args: ["--decode-context-parallel-size", "8", "--max-cudagraph-capture-size", "64"] + json_args: + "--compilation-config": { cudagraph_mode: "FULL_AND_PIECEWISE", cudagraph_capture_sizes: [7, 8, 14, 16, 28, 32, 56, 64] } + runbook: | + The `latency_specialized` decode backend provides a CuTe-based implementation for low-concurrency Kimi-K3 serving on Blackwell. The configurations below serve the default [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) MXFP4 checkpoint with [Inferact/Kimi-K3-DSpark](https://huggingface.co/Inferact/Kimi-K3-DSpark), using one aggregated engine across 4-GPU GB300 NVLink compute nodes and the native Python vLLM frontend. + + The explicit opt-in is `--additional-config '{"decode_backend":"latency_specialized"}'`; `auto` keeps the stock decode path in this image. The specialized path requires the V2 model runner, FULL decode CUDA graphs, and BF16 convolution state with FP32 recurrent state (`--mamba-ssm-cache-dtype float32`). Keep `FULL_AND_PIECEWISE` and the capture sizes together: the lists cover DSpark draft and verification shapes. Prefill and eager/piecewise steps still use the stock implementation. This recipe uses PP1, no expert parallelism, and at most seven draft tokens; do not combine it with `--enforce-eager`, ReplaySSM/RecoverSSM, or adaptive verification. See the pinned [backend requirements](https://github.com/vllm-project/vllm/blob/49cb2bbf71cfd36ee5aa33d1e184b1cddff55294/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py). + + Make the same target checkpoint available at the same local path on every node, and make `Inferact/Kimi-K3-DSpark` available through the mounted Hugging Face cache. Run the preset's commands inside the nightly image on every allocated node, with four visible GPUs, host networking, sufficient shared memory, and the devices and network access required by your NVLink/RDMA deployment. Use the reachable address of rank 0 for `MASTER_ADDR` on every node and a distinct `NODE_RANK` from `0` through `NODES - 1`. The lane served the checkpoint from a local path with `--served-model-name moonshotai/Kimi-K3` and `--device-ids 0,1,2,3`; the builder serves the Hugging Face id. + + mi355x_agentx: + label: "AgentX TP8" + hardware: mi355x + strategy: single_node_tp + workload: "Agentic (InferenceX AgentX)" + default: true + description: "The lane's every-concurrency settings on one MI355X node, for any concurrency; the per-concurrency cells are the AgentX cells preset." + source: { label: "InferenceX #4393", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4393", key: "kimik3-fp4-mi355x-vllm-agentic-mtp" } + # MI355X TP8: the InferenceX AgentX lane's settings that hold at every + # concurrency it sweeps; its per-concurrency tuning stays in the runbook. An + # exact-GPU block layers after the recipe `amd` baseline and wins last-wins + # dedupe there, so other AMD GPUs and every other strategy keep the baseline. + # No --load-format override: fastsafetensors ran out of GPU memory while + # loading on the lane's image, so the `amd` baseline's `auto` applies. + args: + - "--gpu-memory-utilization" + - "0.9" + - "--max-model-len" + - "1048576" + - "--enable-prefix-caching" + - "--kv-cache-dtype" + - "fp8" + - "--attention-config" + - '{"mla_prefill_backend":"ROCM_AITER_FA"}' + - "--stream-interval" + - "10" + env: + VLLM_ROCM_AITER_MLA_ASM_PADDING: "asm" + VLLM_ROCM_AITER_MLA_DCP_VERIFY: "asm" + AITER_FLYDSL_STAGE2_FP8: "1" + runbook: | + The single-node MI355X TP8 AgentX configuration follows [InferenceX #4393](https://github.com/SemiAnalysisAI/InferenceX/pull/4393): its [master entry](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/configs/amd-master.yaml#L766) and [lane recipe](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4-mtp/agentic.yaml) at `dee7824`. It serves text with tool-call and reasoning parsers, using DCP8/A2A at every point. There is no concurrency-2 point. + + **In the command builder.** Settings the lane uses at every concurrency are part of the MI355X TP8 command: `--gpu-memory-utilization 0.9`, `--max-model-len 1048576`, `--enable-prefix-caching`, `--kv-cache-dtype fp8`, `--attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}'`, `--stream-interval 10`, `VLLM_ROCM_AITER_MLA_ASM_PADDING=asm`, `VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm` and `AITER_FLYDSL_STAGE2_FP8=1`. Weights load with vLLM's default loader (the AMD profile's `--load-format auto`): with `fastsafetensors`, the lane's servers ran out of GPU memory while loading weights on its image. The lane's `--moe-backend auto` is vLLM's default selection, so the builder leaves that flag unset. Other AMD GPUs keep the general AMD profile. + + - **Spec decoding** serves the lane's draft, the published [`Inferact/Kimi-K3-DSpark`](https://huggingface.co/Inferact/Kimi-K3-DSpark), with its attention on `ROCM_AITER_MLA` ([vllm-project/vllm#55966](https://github.com/vllm-project/vllm/pull/55966)), FP8 KV, probabilistic drafting and block verification. Draft length is tuned per cell in the matrix below. + - **Target quantization** is set per cell in the **AgentX cells** preset: resident c1 passes no online quantization config, preserving the BF16 target weights required by specialized decode. Every other cell quantizes only the KDA `in_proj_qkvgfab` projections to FP8 per-channel. Routed experts retain their native MXFP4 weights. The DSpark draft remains BF16 with FP8 KV: vLLM builds the named draft from its own unquantized config, so removing c1's target quantization does not change the draft's dtype. + - **DCP8** (`--decode-context-parallel-size 8`, `--attention-backend ROCM_AITER_MLA`, `--dcp-comm-backend a2a`, `ROCM_AITER_FA` prefill, and the direct DCP ops off) is part of the **AgentX cells** preset; this preset serves without DCP. + - The **Simple** KV-offload option emits the connector the lane uses (`SimpleCPUOffloadConnector`, `kv_both`, `lazy_offload: false`). + + **MoE environment.** The lane sets `VLLM_ROCM_USE_AITER_MOE_SITUV2=1`, as the AMD profile does, plus `AITER_BF16_FP8_MOE_BOUND=0` and `AITER_FLYDSL_STAGE2_FP8=1`. The SiTUv2 path does not read `AITER_BF16_FP8_MOE_BOUND`. `AITER_FLYDSL_STAGE2_FP8=1` makes the FlyDSL MoE down projection write its per-route outputs as MXFP8 (FP8 values with E8M0 scales) instead of BF16 before they are summed. It affects only the target's MoE layers; the DSpark draft's MLP is dense. + + **QuickReduce.** The lane also enables INT4 QuickReduce (`VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4`) with `VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB=16`, so only tensor-parallel all-reduces of 16 MiB or more use the INT4 codec. INT4 all-reduce is lossy, so it is not in the builder. The DSpark draft's all-reduces carry only its draft-step tokens (at most 600 tokens × 7,168 BF16 values, about 8.2 MiB, at concurrency 60), and its context projection has no all-reduce, so the 16 MiB floor keeps the draft out of the INT4 path in every lane cell. + + The lane's ten per-concurrency cells, with DCP8, the lane's draft and the qualified image, are the [AgentX cells](#preset-mi355x_agentx_cells) preset. + + mi355x_agentx_cells: + label: "AgentX cells" + hardware: mi355x + strategy: single_node_tp + workload: "Agentic (InferenceX AgentX)" + description: "The lane's ten per-concurrency cells on one MI355X node: TP8 with DCP8 and the DSpark draft, GPU-resident or with host KV offload." + source: { label: "InferenceX #4393", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4393", key: "kimik3-fp4-mi355x-vllm-agentic-mtp", revision: dee7824e } + image: "vllm/vllm-openai-rocm:nightly-rocm100-df6603f0583375f244d1ac6d36265462b7490392" + # The lane runs vLLM's default (Python) frontend; Spec decoding renders its + # Inferact/Kimi-K3-DSpark draft on MI355X (K7, edited per cell below). + select: + frontend: python + features: { spec_decoding: true, text_only: true } + # The AgentX TP8 preset's every-concurrency settings plus the lane's DCP8 + # pairing (ROCM_AITER_MLA decode, a2a, direct DCP ops off). + args: + - "--gpu-memory-utilization" + - "0.9" + - "--max-model-len" + - "1048576" + - "--enable-prefix-caching" + - "--kv-cache-dtype" + - "fp8" + - "--attention-config" + - '{"mla_prefill_backend":"ROCM_AITER_FA"}' + - "--stream-interval" + - "10" + - "--decode-context-parallel-size" + - "8" + - "--dcp-comm-backend" + - "a2a" + - "--attention-backend" + - "ROCM_AITER_MLA" + env: + VLLM_ROCM_AITER_MLA_ASM_PADDING: "asm" + VLLM_ROCM_AITER_MLA_DCP_VERIFY: "asm" + AITER_FLYDSL_STAGE2_FP8: "1" + VLLM_USE_DIRECT_DCP_A2A: "0" + VLLM_USE_DIRECT_DCP_Q_GATHER: "0" + VLLM_USE_DIRECT_DCP_KV_GATHER: "0" + # The lane cells (InferenceX #4393). Resident cells keep KV on GPU; offload + # cells use Simple with 224.875 GB of host DRAM per rank and a fixed hash seed. + points: + r1: + label: "C1 · resident" + concurrency: 1 + select: { kv_offload: "off" } + args: + - "--max-num-seqs" + - "1" + - "--max-num-batched-tokens" + - "16384" + - "--additional-config" + - '{"decode_backend":"latency_specialized","kda_prefill_backend":"fused"}' + env: { VLLM_USE_V2_MODEL_RUNNER: "1" } + json_args: + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 16, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16]} + r4: + label: "C4 · resident" + concurrency: 4 + select: { kv_offload: "off" } + args: + - "--max-num-seqs" + - "8" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + json_args: + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 64, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64]} + o4: + label: "C4 · offload" + concurrency: 4 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "8" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 5 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 48, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o10: + label: "C10 · offload" + concurrency: 10 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "20" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 100, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o16: + label: "C16 · offload" + concurrency: 16 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "32" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 160, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o24: + label: "C24 · offload" + concurrency: 24 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "48" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 240, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206, 208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o32: + label: "C32 · offload" + concurrency: 32 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "64" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 320, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206, 208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240, 242, 244, 246, 248, 250, 252, 254, 256, 258, 260, 262, 264, 266, 268, 270, 272, 274, 276, 278, 280, 282, 284, 286, 288, 290, 292, 294, 296, 298, 300, 302, 304, 306, 308, 310, 312, 314, 316, 318, 320]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o48: + label: "C48 · offload" + concurrency: 48 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "96" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 480, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206, 208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240, 242, 244, 246, 248, 250, 252, 254, 256, 258, 260, 262, 264, 266, 268, 270, 272, 274, 276, 278, 280, 282, 284, 286, 288, 290, 292, 294, 296, 298, 300, 302, 304, 306, 308, 310, 312, 314, 316, 318, 320, 322, 324, 326, 328, 330, 332, 334, 336, 338, 340, 342, 344, 346, 348, 350, 352, 354, 356, 358, 360, 362, 364, 366, 368, 370, 372, 374, 376, 378, 380, 382, 384, 386, 388, 390, 392, 394, 396, 398, 400, 402, 404, 406, 408, 410, 412, 414, 416, 418, 420, 422, 424, 426, 428, 430, 432, 434, 436, 438, 440, 442, 444, 446, 448, 450, 452, 454, 456, 458, 460, 462, 464, 466, 468, 470, 472, 474, 476, 478, 480]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o52: + label: "C52 · offload" + concurrency: 52 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "104" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 520, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206, 208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240, 242, 244, 246, 248, 250, 252, 254, 256, 258, 260, 262, 264, 266, 268, 270, 272, 274, 276, 278, 280, 282, 284, 286, 288, 290, 292, 294, 296, 298, 300, 302, 304, 306, 308, 310, 312, 314, 316, 318, 320, 322, 324, 326, 328, 330, 332, 334, 336, 338, 340, 342, 344, 346, 348, 350, 352, 354, 356, 358, 360, 362, 364, 366, 368, 370, 372, 374, 376, 378, 380, 382, 384, 386, 388, 390, 392, 394, 396, 398, 400, 402, 404, 406, 408, 410, 412, 414, 416, 418, 420, 422, 424, 426, 428, 430, 432, 434, 436, 438, 440, 442, 444, 446, 448, 450, 452, 454, 456, 458, 460, 462, 464, 466, 468, 470, 472, 474, 476, 478, 480, 482, 484, 486, 488, 490, 492, 494, 496, 498, 500, 502, 504, 506, 508, 510, 512, 514, 516, 518, 520]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + o60: + label: "C60 · offload" + concurrency: 60 + select: { kv_offload: simple } + args: + - "--max-num-seqs" + - "120" + - "--max-num-batched-tokens" + - "8192" + - "--additional-config" + - '{"kda_prefill_backend":"fused"}' + - "--quantization-config" + - '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}' + env: { PYTHONHASHSEED: "42" } + json_args: + "--speculative-config": { num_speculative_tokens: 4 } + "--compilation-config": {mode: 3, cudagraph_mode: "FULL_AND_PIECEWISE", max_cudagraph_capture_size: 600, cudagraph_capture_sizes: [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160, 162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206, 208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240, 242, 244, 246, 248, 250, 252, 254, 256, 258, 260, 262, 264, 266, 268, 270, 272, 274, 276, 278, 280, 282, 284, 286, 288, 290, 292, 294, 296, 298, 300, 302, 304, 306, 308, 310, 312, 314, 316, 318, 320, 322, 324, 326, 328, 330, 332, 334, 336, 338, 340, 342, 344, 346, 348, 350, 352, 354, 356, 358, 360, 362, 364, 366, 368, 370, 372, 374, 376, 378, 380, 382, 384, 386, 388, 390, 392, 394, 396, 398, 400, 402, 404, 406, 408, 410, 412, 414, 416, 418, 420, 422, 424, 426, 428, 430, 432, 434, 436, 438, 440, 442, 444, 446, 448, 450, 452, 454, 456, 458, 460, 462, 464, 466, 468, 470, 472, 474, 476, 478, 480, 482, 484, 486, 488, 490, 492, 494, 496, 498, 500, 502, 504, 506, 508, 510, 512, 514, 516, 518, 520, 522, 524, 526, 528, 530, 532, 534, 536, 538, 540, 542, 544, 546, 548, 550, 552, 554, 556, 558, 560, 562, 564, 566, 568, 570, 572, 574, 576, 578, 580, 582, 584, 586, 588, 590, 592, 594, 596, 598, 600]} + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 224875000000 } } + runbook: | + The single-node MI355X TP8 AgentX configuration follows [InferenceX #4393](https://github.com/SemiAnalysisAI/InferenceX/pull/4393): its [master entry](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/configs/amd-master.yaml#L766) and [lane recipe](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4-mtp/agentic.yaml) at `dee7824`. It serves text with tool-call and reasoning parsers, using DCP8/A2A at every point. There is no concurrency-2 point. + + **Per-concurrency tuning.** The Point row renders the lane's ten cells. Every cell requests `FULL_AND_PIECEWISE` with breakable CUDA graphs disabled (`VLLM_USE_BREAKABLE_CUDAGRAPH=0`); because Kimi-K3 has no `torch.compile` submodule, vLLM resolves that mode to `FULL_DECODE_ONLY` at load. The [AgentX TP8](#preset-mi355x_agentx) preset keeps the builder's general values for any concurrency. + + Every cell adds `kda_prefill_backend: fused` in `--additional-config`. Only **resident c1** also sets `decode_backend: latency_specialized` and `VLLM_USE_V2_MODEL_RUNNER=1`. On the pinned image, this [MI355X specialized path](https://github.com/vllm-project/vllm/pull/60683) requires TP8, V2, gfx950, no KV connector and BF16 target projections. It handles pure speculative-verification batches of at most eight rows (c1 with K7); other steps use the normal layers. These lane-specific settings stay out of generic MI355X defaults. + + The `resident` profile keeps KV on GPU; `offload` uses `SimpleCPUOffloadConnector` with 224.875 GB of host DRAM per rank. `--max-num-seqs` is twice the concurrency, except that concurrency 1 admits one sequence. Every cell captures every even CUDA-graph size from 2 through `2 × concurrency × (1 + draft tokens)`. Agentic benchmark runs replace block rejection with the golden acceptance length; synthetic acceptance is not suitable for serving or accuracy evaluation. + + **Image.** The builder keeps the rolling `vllm/vllm-openai-rocm:nightly-rocm100` image. Run the exact commands below inside `vllm/vllm-openai-rocm:nightly-rocm100-df6603f0583375f244d1ac6d36265462b7490392` (vLLM `df6603f0`), the image qualified by #4393 and carrying #60683. The compatibility limits above refer to this pinned runtime. + + ### Reproducing one lane cell + + Run each cell in a fresh shell inside the pinned image. This real-verification form is for serving and accuracy evaluation; `PROFILE` distinguishes the two concurrency-4/DCP8 cells: + + ```bash + PROFILE=resident + CLIENTS=1 + case "$PROFILE:$CLIENTS" in + resident:1) K=7; BATCHED=16384; OFFLOAD=0 ;; + resident:4) K=7; BATCHED=8192; OFFLOAD=0 ;; + offload:4) K=5; BATCHED=8192; OFFLOAD=1 ;; + offload:10|offload:16|offload:24|offload:32|offload:48|offload:52|offload:60) K=4; BATCHED=8192; OFFLOAD=1 ;; + *) echo "PROFILE=$PROFILE CLIENTS=$CLIENTS is not a lane cell" >&2; exit 1 ;; + esac + + export VLLM_ROCM_USE_AITER=1 + export SAFETENSORS_FAST_GPU=1 + export VLLM_ROCM_USE_AITER_MOE_SITUV2=1 + export AITER_BF16_FP8_MOE_BOUND=0 + export AITER_FLYDSL_STAGE2_FP8=1 + export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm + export VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm + export VLLM_USE_DIRECT_DCP_A2A=0 + export VLLM_USE_DIRECT_DCP_Q_GATHER=0 + export VLLM_USE_DIRECT_DCP_KV_GATHER=0 + # Lossy INT4-quantized all-reduce for messages of 16 MiB or more: not in the builder. + export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 + export VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB=16 + # 1.56 TB of weights from shared storage, and long agentic prefills. + export VLLM_ENGINE_READY_TIMEOUT_S=7200 + export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200 + export PYTHONNOUSERSITE=1 + # Required on the qualification nodes' MEC firmware 38; fixed in firmware 177. + export HSA_NO_SCRATCH_RECLAIM=1 + + GRAPH_MODE=FULL_AND_PIECEWISE + export VLLM_USE_BREAKABLE_CUDAGRAPH=0 + MAX_SEQS=$((CLIENTS == 1 ? 1 : 2 * CLIENTS)) + MAX_CAPTURE=$((2 * CLIENTS * (1 + K))) + COMPILATION="{\"mode\":3,\"cudagraph_mode\":\"$GRAPH_MODE\",\"max_cudagraph_capture_size\":$MAX_CAPTURE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$(seq -s, 2 2 "$MAX_CAPTURE")]}" + + EXTRA=(--additional-config '{"kda_prefill_backend":"fused"}') + if [ "$PROFILE:$CLIENTS" = resident:1 ]; then + export VLLM_USE_V2_MODEL_RUNNER=1 + EXTRA=(--additional-config '{"decode_backend":"latency_specialized","kda_prefill_backend":"fused"}') + else + EXTRA+=(--quantization-config '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}') + fi + if [ "$K" -gt 0 ]; then + EXTRA+=(--speculative-config "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$K,\"method\":\"dspark\",\"attention_backend\":\"ROCM_AITER_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"block\"}") + fi + if [ "$OFFLOAD" -eq 1 ]; then + # The lane's node gives the connector 1799 GB of host DRAM (224875000000 bytes + # per rank); size cpu_bytes_to_use_per_rank to your host. The fixed hash seed + # makes identical prefixes map to the same block keys on every rank. + export PYTHONHASHSEED=42 + EXTRA+=(--kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":224875000000,"lazy_offload":false}}') + fi + vllm serve moonshotai/Kimi-K3 \ + --trust-remote-code \ + --moe-backend auto \ + --tensor-parallel-size 8 \ + --decode-context-parallel-size 8 \ + --dcp-comm-backend a2a \ + --gpu-memory-utilization 0.9 \ + --language-model-only \ + --max-num-seqs "$MAX_SEQS" \ + --enable-auto-tool-choice \ + --tool-call-parser kimi_k3 \ + --reasoning-parser kimi_k3 \ + --max-model-len 1048576 \ + --enable-prefix-caching \ + --kv-cache-dtype fp8 \ + --stream-interval 10 \ + --max-num-batched-tokens "$BATCHED" \ + --attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}' \ + --compilation-config "$COMPILATION" \ + "${EXTRA[@]}" + ``` + + The builder's **Simple** option keeps the repository-wide 220 GiB per rank and does not set `PYTHONHASHSEED`. + + ### Throughput measurement only (synthetic acceptance) + + The lane's throughput runs replace block verification with synthetic acceptance, so every concurrency point is measured at InferenceX's golden acceptance length for this draft. Only the `--speculative-config` changes: + + ```bash + --speculative-config "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$K,\"method\":\"dspark\",\"attention_backend\":\"ROCM_AITER_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"synthetic\",\"synthetic_acceptance_length\":$AL}" + ``` + + | Draft tokens | Golden acceptance length | Lane cells | + | ---: | ---: | --- | + | 7 | 3.84 | resident concurrency 1 and 4, DCP8 | + | 5 | 3.62 | offload concurrency 4, DCP8 | + | 4 | 3.36 | dcp8 concurrency 10, 16, 24, 32, 48, 52 and 60 | + + These lengths come from [`kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml): SPEED-Bench coding prompts with thinking on, measured on B300 for `Inferact/Kimi-K3-DSpark` with probabilistic drafting and block verification. They hold only for that draft at those draft lengths, not for `RedHatAI/Kimi-K3-speculator.dspark`. Synthetic acceptance skips real verification: never use it to serve users or to evaluate accuracy. + + ### Evidence and limits + + - InferenceX [run 40037005947, attempt 2](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40037005947/attempts/2) passed all ten performance cells and its accuracy jobs at source `0df7e892b73c53d66a486c2a085e994b67653042`. The Kimi serving recipe, model, image and selected matrix are unchanged at audited head `dee7824`; later changes are unrelated upstream work and changelog text. Accuracy jobs use real verification and are deduplicated by [evaluation identity](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/evals/EVALS.md#L33), so this is not a claim of a separate accuracy job for every matrix cell. + - Performance figures on the InferenceX dashboard are reported by SemiAnalysis and have not been reproduced for this recipe. + - Model policy (vLLM as the native engine, `Inferact/Kimi-K3-DSpark` as the agreed draft): [`MODELS.md`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/docs/MODELS.md). + guide: | ## Overview @@ -728,7 +1418,7 @@ guide: | - **vLLM:** Use the vllm/vllm-openai:kimi-k3 docker - **CUDA:** The kimi-k3 image ships as a CUDA 13 (cu130) build only — there is no -cu129 tag, and the K3-enabled wheels are not on the cu129 nightly index. The host needs an r580+ NVIDIA driver; on a CUDA 12.9 (r575) host, upgrade the driver or build vLLM from the K3 branch against cu129 PyTorch yourself. - **Hardware:** At least 8x GB300. Multi-node for real production traffic. - - **ROCm:** At least 8x MI355X/MI350X. On MI355X the command builder uses the ROCm 10.0 nightly image `vllm/vllm-openai-rocm:nightly-rocm100`; the day-0 `vllm/vllm-openai-rocm:kimi-k3` image predates the DSpark and MoE fixes that path needs (see [AMD MI355X: InferenceX AgentX alignment](#amd-mi355x-inferencex-agentx-alignment)). + - **ROCm:** At least 8x MI355X/MI350X. On MI355X the command builder uses the ROCm 10.0 nightly image `vllm/vllm-openai-rocm:nightly-rocm100`; the day-0 `vllm/vllm-openai-rocm:kimi-k3` image predates the DSpark and MoE fixes that path needs (see the [MI355X AgentX preset](#preset-mi355x_agentx)). ## Client Usage @@ -786,451 +1476,10 @@ guide: | - **mlx5 dmabuf registration failures**: if engine init fails with "NCCL error: unhandled system error" and the log shows `mlx5dv_reg_dmabuf_mr` errno 524, the kernel/driver lacks mlx5 dmabuf support (NCCL 2.28 registers dmabuf by default). Set `NCCL_DMABUF_ENABLE=0` to fall back to `nvidia_peermem` (must be loaded on the nodes) — still GPUDirect RDMA. - **FP8 KV**: If FP8 KV cache is needed, please also add `--attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"TOKENSPEED_MLA"}'` when serving vLLM. `TOKENSPEED_MLA` is what the Blackwell profile above emits (paired with `--attention-backend TOKENSPEED_MLA` for decode), but it is not the only supported choice: vLLM 0.27.1 registers three MLA prefill backends — `FLASHINFER`, `TRTLLM_RAGGED` and `TOKENSPEED_MLA` (see `vllm/v1/attention/backends/mla/prefill/registry.py`) — and any of them may be substituted for `mla_prefill_backend` alongside `use_prefill_query_quantization`. Which one wins depends on the input-length distribution, so benchmark before switching. Those backends do not run on ROCm: on MI355X the builder pairs FP8 KV with `--attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}'` instead. - **Prefix-match unit**: K3 is a hybrid model: its MLA attention layers and KDA (Mamba-like) layers form two KV-cache groups under the hybrid KV-cache manager. With prefix caching on, vLLM pads the attention block size up to match the Mamba state page, so both groups resolve to one (large) block size and the default prefix-cache hit boundary lands on that block — very coarse. The Blackwell profile sets `--prefix-match-unit 128`: the inflation formula always yields a multiple of 128, so 128 divides whatever block size results, and it aligns with the MLA kernel's native block boundary, giving finer prefix-hit granularity. The value is only effective when prefix caching (or a KV connector) is active; it is a no-op otherwise, and it is safe under all tensor-parallel layouts. - - **Decode context parallelism (DCP)**: enable the **Decode context parallelism** opt-in feature in the command builder to shard the decode KV cache across the tensor-parallel ranks for decode-heavy long-context serving. It emits `--decode-context-parallel-size 8 --dcp-comm-backend a2a --attention-backend TOKENSPEED_MLA` with `mla_prefill_backend: TRTLLM_RAGGED` (on MI355X: `--attention-backend ROCM_AITER_MLA` with `mla_prefill_backend: ROCM_AITER_FA`), and is scoped to the TP8 single-node profile because the DCP size must divide the tensor-parallel size. On builds that carry them, the `VLLM_USE_DIRECT_DCP_A2A`, `VLLM_USE_DIRECT_DCP_Q_GATHER` and `VLLM_USE_DIRECT_DCP_KV_GATHER` knobs default to auto and turn the direct DCP ops on where applicable; set them to `1` only to pin that path explicitly, or `0` to disable it. The MI355X DCP option sets all three to `0`, as InferenceX does. On GB200 and GB300 the builder's **Decode context parallel** row (Off, DCP8, DCP16) sets `--decode-context-parallel-size` on every node tab and, under Prefill/Decode, on both roles. It offers only the size equal to the tensor-parallel size of every role, the layouts InferenceX qualified: GB200 TP8 x PP2 aggregated is DCP8 (#4322), TP16 per role disaggregated is DCP16 (#3748, #4322), and GB300 TP8 pools are DCP8 (#4047). Other layouts, such as a TP1 DEP decode pool, show no DCP row and drop the flag. + - **Decode context parallelism (DCP)**: the qualified DCP layouts are part of the presets that ran them: DCP8 on the GB200 aggregated TP8 x PP2 and GB300 disaggregated TP8 pools, DCP16 on the GB200 TP16 disaggregated pools, DCP16/DCP8 at the GB300 low-latency points, and DCP8 in the MI355X AgentX cells. Each pairs `--decode-context-parallel-size` (a divisor of the tensor-parallel size) with `--dcp-comm-backend a2a` and an MLA attention backend that supports it (`TOKENSPEED_MLA` with `mla_prefill_backend: TRTLLM_RAGGED` on Blackwell, `ROCM_AITER_MLA` with `mla_prefill_backend: ROCM_AITER_FA` on MI355X). On builds that carry them, the `VLLM_USE_DIRECT_DCP_A2A`, `VLLM_USE_DIRECT_DCP_Q_GATHER` and `VLLM_USE_DIRECT_DCP_KV_GATHER` knobs default to auto; set them to `1` to pin the direct DCP ops or `0` to disable them, as the MI355X cells do. - Ascend W4A8 checkpoint (ModelScope): https://www.modelscope.cn/models/Eco-Tech/Kimi-K3-w4a8 - vLLM Ascend tutorial: https://docs.vllm.ai/projects/ascend/en/v0.23.0/tutorials/models/Kimi-K3.html - ## Low latency serving - - ### GB300: TP8/TP16 with DSpark - - The `latency_specialized` decode backend provides a CuTe-based implementation for low-concurrency Kimi-K3 serving on Blackwell. The configurations below serve the default [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) MXFP4 checkpoint with [Inferact/Kimi-K3-DSpark](https://huggingface.co/Inferact/Kimi-K3-DSpark), using one aggregated engine across 4-GPU GB300 NVLink compute nodes and the native Python vLLM frontend. - - Use the vLLM nightly image on every node; it contains the implementation and needs no patches: - - ```bash - docker pull vllm/vllm-openai:nightly - ``` - - At client concurrency 1, use TP16/DCP16 with five DSpark draft tokens. At concurrency 2 and 4, use TP8/DCP8 with seven draft tokens. Both presets keep `--max-num-seqs 8` and `--max-num-batched-tokens 16384`. - - The explicit opt-in is `--additional-config '{"decode_backend":"latency_specialized"}'`; `auto` keeps the stock decode path in this image. The specialized path requires the V2 model runner, FULL decode CUDA graphs, and BF16 convolution state with FP32 recurrent state (`--mamba-ssm-cache-dtype float32`). Keep `FULL_AND_PIECEWISE` and the capture sizes together: the lists cover DSpark draft and verification shapes. Prefill and eager/piecewise steps still use the stock implementation. This recipe uses PP1, no expert parallelism, and at most seven draft tokens; do not combine it with `--enforce-eager`, ReplaySSM/RecoverSSM, or adaptive verification. See the pinned [backend requirements](https://github.com/vllm-project/vllm/blob/49cb2bbf71cfd36ee5aa33d1e184b1cddff55294/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py). - - ### Launch on each node - - Make the same target checkpoint available at the same local path on every node, and make `Inferact/Kimi-K3-DSpark` available through the mounted Hugging Face cache. Run the script below **inside the nightly image** on every allocated node, with four visible GPUs, host networking, sufficient shared memory, and the devices and network access required by your NVLink/RDMA deployment. Use the reachable address of rank 0 for `MASTER_ADDR` on every node and a distinct `NODE_RANK` from `0` through `NODES - 1`. - - Save as `kimi-k3-low-latency.sh`: - - ```bash - #!/usr/bin/env bash - set -eo pipefail - - TP=${1:?Usage: bash kimi-k3-low-latency.sh TP MASTER_ADDR NODE_RANK MODEL_PATH} - MASTER_ADDR=${2:?Set the reachable rank-0 address} - NODE_RANK=${3:?Set the node rank} - MODEL_PATH=${4:?Set the mounted Kimi-K3 checkpoint path} - - case "$TP" in - 16) NODES=4; K=5; MAX_CAPTURE=48; CAPTURES='5,6,10,12,20,24,40,48' ;; - 8) NODES=2; K=7; MAX_CAPTURE=64; CAPTURES='7,8,14,16,28,32,56,64' ;; - *) echo 'TP must be 8 or 16' >&2; exit 1 ;; - esac - - export VLLM_USE_V2_MODEL_RUNNER=1 - export VLLM_USE_RUST_FRONTEND=0 - export VLLM_ALLREDUCE_USE_FLASHINFER=1 - export VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION=1 - export VLLM_USE_DIRECT_DCP_A2A=1 - export VLLM_USE_DIRECT_DCP_Q_GATHER=1 - export VLLM_USE_DIRECT_DCP_KV_GATHER=1 - export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0 - export VLLM_RPC_TIMEOUT=600000 - export VLLM_LOG_STATS_INTERVAL=1 - - API_ARGS=(--headless) - if [ "$NODE_RANK" -eq 0 ]; then - API_ARGS=(--host 0.0.0.0 --port 8000) - fi - - exec vllm serve "$MODEL_PATH" \ - --served-model-name moonshotai/Kimi-K3 \ - --nnodes "$NODES" --node-rank "$NODE_RANK" --master-addr "$MASTER_ADDR" \ - --device-ids 0,1,2,3 \ - --tensor-parallel-size "$TP" --decode-context-parallel-size "$TP" \ - --dcp-comm-backend a2a \ - --additional-config '{"decode_backend":"latency_specialized"}' \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype float32 \ - --gpu-memory-utilization 0.92 \ - --max-num-seqs 8 --max-num-batched-tokens 16384 \ - --max-cudagraph-capture-size "$MAX_CAPTURE" \ - --compilation-config "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[$CAPTURES]}" \ - --speculative-config "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"attention_backend\":\"TOKENSPEED_MLA\",\"method\":\"dspark\",\"num_speculative_tokens\":$K,\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"block\"}" \ - --enable-prefix-caching --prefix-match-unit 128 \ - --disable-flashinfer-autotune \ - --moe-backend auto --load-format fastsafetensors \ - --stream-interval 10 --language-model-only --trust-remote-code \ - --enable-auto-tool-choice --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 \ - "${API_ARGS[@]}" - ``` - - ## GB200: InferenceX AgentX alignment - - The GB200 AgentX lanes in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb200-dynamo-vllm-agentic-dspark-mooncake-*`, [SemiAnalysisAI/InferenceX#4322](https://github.com/SemiAnalysisAI/InferenceX/pull/4322), pinned to [`b7eee97`](https://github.com/SemiAnalysisAI/InferenceX/tree/b7eee972451f9297941d4cbe6d42d489d7ec8363/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/gb200-fp4/agentx), which builds on #3748) serve the default `moonshotai/Kimi-K3` checkpoint (MXFP4 weights) with DSpark speculative decoding and a Mooncake KV store on 4-GPU NVL trays. They are benchmark configurations; the settings below are the production-useful parts. The builder's **Mooncake** option generates the `mooncake_master` and store config; the lane-specific connector settings are below. Throughput runs use InferenceX's golden acceptance lengths for DSpark. - - ### Aggregated: TP8 x PP2 + DCP8 (16 GPUs, 4 trays) - - Image `vllm/vllm-openai:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f` (the same nightly as the disaggregated lane), served by vLLM's own frontend (no Dynamo) with an embedded RDMA Mooncake store (`global_segment_size` 200GB, `mlx5_0..mlx5_3`). In the command builder, select GB200, **Multi-Node TP + Pipeline Parallel** and the **Python** frontend (the lane did not enable the Rust frontend): on GB200 it renders this TP8 x PP2 + DCP8 layout across four trays with the K4 draft and `--max-num-seqs 64` (the C=32 row). Leave **Spec decoding** off, because the toggle replaces the K4 draft with the generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft. The builder's **Mooncake** KV Offload option renders a plain `MooncakeStoreConnector` with a 100GB segment; use the lane's connector and 200GB store from the command below. The per-concurrency settings below stay manual. `C` is client concurrency, `K` is `num_speculative_tokens` and `S` is `--max-num-seqs`: - - | C | K | S | `cudagraph_capture_sizes` (before the shared tail) | - |---:|---:|---:|---| - | 1 | 7 | 2 | `8,16,64,128` | - | 4 | 7 | 8 | `8,16,24,32,40,48,56,64,128` | - | 8 | 7 | 16 | `8,16,24,40,56,64,80,104,128` | - | 24 | 4 | 48 | `8,16,24,40,56,64,80,104,128` | - | 32, 48, 56 | 4 | 64, 96, 112 | `1,7,18,34,53,64,75,100,128` | - - Every list ends with the shared tail `256,512,1024,2048,4096,8192`. - - ```bash - export VLLM_USE_DIRECT_DCP_A2A=1 VLLM_USE_DIRECT_DCP_Q_GATHER=1 VLLM_USE_DIRECT_DCP_KV_GATHER=1 - export VLLM_ALLREDUCE_USE_FLASHINFER=1 NCCL_CUMEM_ENABLE=1 VLLM_USE_V2_MODEL_RUNNER=1 - export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0 - export MC_STORE_MEMCPY=1 MC_ENABLE_DEST_DEVICE_AFFINITY=1 MC_WORKERS_PER_CTX=4 - export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 - export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000 - vllm serve moonshotai/Kimi-K3 --trust-remote-code \ - --tensor-parallel-size 8 --pipeline-parallel-size 2 \ - --decode-context-parallel-size 8 --dcp-comm-backend a2a \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 \ - --max-model-len 1048576 --gpu-memory-utilization 0.88 --max-num-seqs --max-num-batched-tokens 8192 \ - --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","method":"dspark","num_speculative_tokens":,"attention_backend":"TOKENSPEED_MLA","draft_sample_method":"probabilistic"}' \ - --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[,256,512,1024,2048,4096,8192]}' \ - --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}' \ - --load-format fastsafetensors --moe-backend auto --no-enable-flashinfer-autotune \ - --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 --enable-auto-tool-choice \ - --language-model-only - ``` - - ### Disaggregated: 1 prefill + 1 decode, TP16 + DCP16 per role (32 GPUs) - - Image `vllm/vllm-openai:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f` with Dynamo `ffbbaaa6c44370e798687881317535ac17d66a76`, etcd and an embedded RDMA Mooncake store. The lane's request plane is TCP (`DYN_REQUEST_PLANE=tcp`) and it starts no NATS service. Each role is one worker across four trays. - - In the command builder, select GB200, **Prefill/Decode Disaggregation** and the **Dynamo** PD router, and keep the default TP prefill and TP decode pools: they render this TP16 + DCP16 layout with the concurrency-16 settings. The profile already includes the DSpark K4 draft and `--language-model-only`, so leave **Spec decoding** off (the toggle swaps in the generic RedHat K8 draft) and untick **Tool calling** (see the parser note below). The DEP decode pool is not supported on GB200: it renders TP1 with `--decode-context-parallel-size 16`, which vLLM rejects. The NIXL + Mooncake store connector, the Mooncake master, prefill `--kv-events-config` and the frontend's router mode and session affinity below are documented here only. - - Run each command on all four trays of the role with the same `--master-addr` (the first tray's address), `--nnodes 4`, `--node-rank 0` to `3`, and `--headless` on ranks 1-3. Prefill (`P_GPU`/`P_CAP` and `S` from the table below): - - ```bash - python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ - --disaggregation-mode prefill --request-plane tcp \ - --nnodes 4 --node-rank <0-3> --master-addr --device-ids 0,1,2,3 \ - --tensor-parallel-size 16 --decode-context-parallel-size 16 --dcp-comm-backend a2a \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ - --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 --language-model-only \ - --gpu-memory-utilization --max-num-seqs --max-cudagraph-capture-size \ - --enable-cumem-allocator --stream-interval 10 --load-format safetensors --safetensors-load-strategy lazy \ - --trust-remote-code \ - --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \ - --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' - ``` - - Decode (`D_GPU`/`D_CAP` and `S` from the table below). `--max-num-batched-tokens 1024` bounds the FlashInfer MoE autotune to the token counts decode can reach (at most 480); at vLLM's 16384-token default the autotune bucket needs a 2.9 GiB workspace that runs out of memory at 0.92, and the other ranks hang: - - ```bash - python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ - --disaggregation-mode decode --request-plane tcp \ - --nnodes 4 --node-rank <0-3> --master-addr --device-ids 0,1,2,3 \ - --tensor-parallel-size 16 --decode-context-parallel-size 16 --dcp-comm-backend a2a \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ - --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic"}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 --enable-prefix-caching --prefix-match-unit 128 --language-model-only \ - --gpu-memory-utilization --max-num-seqs --max-num-batched-tokens 1024 --max-cudagraph-capture-size \ - --enable-cumem-allocator --stream-interval 10 --load-format safetensors --safetensors-load-strategy lazy \ - --trust-remote-code \ - --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' - ``` - - For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there. - - Per-concurrency settings (`--max-num-seqs` is 2x concurrency on both roles; decode `--max-cudagraph-capture-size` is `max-num-seqs x (4 + 1)` tokens): - - | C | `S` (`--max-num-seqs`) | Prefill `P_GPU` / `P_CAP` | Decode `D_GPU` / `D_CAP` | - |---:|---:|---|---| - | 16 | 32 | 0.88 / 512 | 0.92 / 160 | - | 32 | 64 | 0.88 / 512 | 0.92 / 320 | - | 48 | 96 | 0.88 / 512 | 0.90 / 480 | - - Decode memory utilization is 0.92 while the capture size is at most 320 tokens and 0.90 above it. - - - **Environment (both roles):** the three direct-DCP variables, `VLLM_ALLREDUCE_USE_FLASHINFER=1`, `VLLM_USE_V2_MODEL_RUNNER=1`, `VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0`, `VLLM_USE_NCCL_SYMM_MEM=0`, `VLLM_SSM_CONV_STATE_LAYOUT=DS`, `VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0`, `NCCL_MNNVL_ENABLE=1`, `NCCL_NVLS_ENABLE=1`, `NCCL_CUMEM_ENABLE=1`, `NCCL_P2P_LEVEL=NVL`, `NCCL_NET_PLUGIN=none`, `VLLM_MOONCAKE_LOAD_RECV_THREADS=4`, `MC_SLICE_SIZE=1048576`, `MC_ENABLE_DEST_DEVICE_AFFINITY=1`, `UCX_MEMTYPE_CACHE=n`, `UCX_MEMTYPE_REG_WHOLE=n`, `UCX_RCACHE_MAX_UNRELEASED=1024`, `UCX_TCP_AF_PRIO=inet`, `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0`, `VLLM_ENGINE_READY_TIMEOUT_S=3600`, `VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800`, `ETCD_LEASE_TTL=600`, `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512` and `DYN_REQUEST_PLANE=tcp`. - - **Frontend:** the Dynamo frontend runs with `--router-mode least-loaded --router-session-affinity-ttl-secs 900`, so each conversation sticks to one worker and reuses its prefix cache. The Rust frontend is selected with `VLLM_USE_RUST_FRONTEND=1`, with `DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600` and `DYN_TOKENIZER_CACHE_BYTES=8589934592`. AgentX sends the conversation ID in the `X-Dynamo-Session-Id` request header. - - **Mooncake master:** `--default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10`, 128GB global segment, 4GB local buffer, RDMA over `mlx5_0..mlx5_3`. - - ## GB300: InferenceX AgentX disaggregated (xPyD, TP8 + DCP8 prefill and decode) - - The GB300 AgentX disaggregated lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg`, [SemiAnalysisAI/InferenceX#4047](https://github.com/SemiAnalysisAI/InferenceX/pull/4047), points refined in [#4373](https://github.com/SemiAnalysisAI/InferenceX/pull/4373)) runs Dynamo + vLLM with several prefill and decode workers. Prefill and decode workers are both TP8 + DCP8 across two 4-GPU GB300 trays; prefill workers also use EP8 with MegaMoE, and decode workers use no expert parallelism. KV moves over NIXL, and all workers share an RDMA Mooncake store. Image `vllm/vllm-openai:nightly-1b2a666391f2a24228a1436b64d547c34cc6b6ce`, Dynamo `ffbbaaa6c44370e798687881317535ac17d66a76`. - - | Clients | Prefill workers | Decode workers | GPUs | Decode profile | Decode `--max-num-seqs` | - |---:|---:|---:|---:|---|---:| - | 16 | 1 | 3 | 32 | [Specialized](#specialized-decode-low-concurrency) | 12 | - | 24-32 | 1 | 2 | 24 | [Specialized](#specialized-decode-low-concurrency) | 12 | - | 48-64 | 1 | 2 | 24 | Regular | 128 | - | 240 | 2 | 3 | 40 | Regular | 128 | - | 240 | 3 | 4 | 56 | Regular | 128 | - - Prefill workers always use `--max-num-seqs 64`. The regular profile is the Workers commands below; the specialized one changes both roles as described under [Specialized decode](#specialized-decode-low-concurrency). - - The command builder renders one prefill pool and one decode pool of the specialized profile (the 1P2D c24-c32 row): select GB300, **Prefill/Decode Disaggregation** and the **Dynamo** PD router, and keep the default TP prefill and TP decode pools. Its decode uses `--max-num-seqs 12`; for the regular points use the Workers commands below instead. The profile already includes the DSpark K4 draft and `--language-model-only`, so leave **Spec decoding** off: turning it on swaps in the generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft, which this profile was not qualified with and which specialized decode rejects (it allows at most 7 draft tokens). Untick **Tool calling** as well; see the parser note under Workers. Other pool modes are not supported on GB300: the DEP decode pool renders TP1 with `--decode-context-parallel-size 8`, which vLLM rejects. - - Documented here only: worker counts beyond one pool per role, the Mooncake store connector (the builder's Mooncake option renders a plain `MooncakeStoreConnector` without the qualified NIXL MultiConnector, recompute policy and load options), the Mooncake services, prefill `--kv-events-config`, and the frontend's router mode and session affinity (the builder renders round-robin). - - ### Services - - Start etcd, NATS and a Mooncake master once: - - ```bash - etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 - printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf - mooncake_master --enable_http_metadata_server=true \ - --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10 - ``` - - The workers' Mooncake store uses P2P handshake metadata, RDMA over the node's NICs, a 160GB global segment and a 4GB local buffer, with offload disabled. - - ### Environment (every worker) - - ```bash - export VLLM_USE_DIRECT_DCP_A2A=1 VLLM_USE_DIRECT_DCP_Q_GATHER=1 VLLM_USE_DIRECT_DCP_KV_GATHER=1 - export VLLM_ALLREDUCE_USE_FLASHINFER=1 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0 VLLM_USE_NCCL_SYMM_MEM=0 - export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_P2P_LEVEL=NVL - export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=1 VLLM_SSM_CONV_STATE_LAYOUT=DS - export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0 - export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 MC_MAX_MR_SIZE=4294967296 MC_ENABLE_DEST_DEVICE_AFFINITY=1 - export VLLM_CONNECTOR_PREFETCH_DEPTH=8 VLLM_CONNECTOR_PREFETCH_KV_CAP=0.65 - export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n UCX_RCACHE_MAX_UNRELEASED=1024 UCX_TCP_AF_PRIO=inet - export DYN_REQUEST_PLANE=tcp DYN_TCP_REQUEST_TIMEOUT=120 ETCD_LEASE_TTL=600 - export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000 - ``` - - Set `VLLM_ENABLE_GEMM_RS=1` on prefill workers and `VLLM_ENABLE_GEMM_RS=0` on decode workers. Loading the 2.8T checkpoint across trays and autotuning takes a while, which is why the engine-ready, execute-model and RPC timeouts are raised. - - ### Workers - - Each worker spans two trays. Run the command on both trays with the same `--master-addr` (the first tray's address), `--nnodes 2` and `--node-rank 0` or `1`. - - Prefill: - - ```bash - python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ - --disaggregation-mode prefill --request-plane tcp \ - --nnodes 2 --node-rank <0|1> --master-addr --device-ids 0,1,2,3 \ - --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \ - --enable-expert-parallel --moe-backend deep_gemm_mega_moe \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' \ - --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \ - --enable-prefix-caching --prefix-match-unit 128 \ - --gpu-memory-utilization 0.92 --max-num-seqs 64 --max-cudagraph-capture-size 512 \ - --enable-cumem-allocator --enable-flashinfer-autotune \ - --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \ - --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \ - --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' - ``` - - Decode (regular profile, `--max-num-seqs 128`): - - ```bash - python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \ - --disaggregation-mode decode --request-plane tcp \ - --nnodes 2 --node-rank <0|1> --master-addr --device-ids 0,1,2,3 \ - --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \ - --attention-backend TOKENSPEED_MLA \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"tokenspeed_mla_min_split_kv":8}' \ - --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \ - --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \ - --enable-prefix-caching --prefix-match-unit 128 \ - --gpu-memory-utilization 0.92 --max-num-seqs 128 --max-cudagraph-capture-size 512 \ - --enable-cumem-allocator --enable-flashinfer-autotune \ - --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \ - --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"compact_group_io":true,"max_load_batch_keys":2,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}' - ``` - - For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there. - - ### Specialized decode (low concurrency) - - The c16-c32 points use the `latency_specialized` decode backend from [Low latency serving](#low-latency-serving) on the decode workers (1P3D for c16, 1P2D for c24-c32), with the Workers commands above changed as follows. The backend's requirements are listed there: V2 model runner (set in Environment), FULL decode CUDA graphs, BF16 convolution state with FP32 recurrent state, DSpark with at most 7 draft tokens, no pipeline or expert parallelism, and DCP of 1 or TP. The decode workers already meet the last two. - - Both roles: - - ```bash - --mamba-ssm-cache-dtype float32 --cp-kv-cache-interleave-size 1 --no-enable-flashinfer-autotune - # replacing --mamba-ssm-cache-dtype bfloat16 and --enable-flashinfer-autotune - ``` - - Pin `--cp-kv-cache-interleave-size 1` on both roles. Left unset, NIXL P/D raises the DCP KV interleave to the KV block size after the specialized decode has read it, so its spec-verify causal mask uses the stale value. - - Decode only: - - ```bash - --additional-config '{"decode_backend":"latency_specialized"}' \ - --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \ - --max-num-seqs 12 --max-cudagraph-capture-size 64 \ - --compilation-config "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[$(seq -s, 1 64)]}" - # replacing the regular decode's --attention-config (it also sets tokenspeed_mla_min_split_kv), - # --max-num-seqs 128 and --max-cudagraph-capture-size 512 - ``` - - `--max-num-seqs` is `64 // (K + 1)` with the K4 draft, so a full verify batch (12 x 5 = 60 tokens) fits the 64-token graph, and CUDA graphs are captured for every batch size up to 64. The prefill workers do not set `--additional-config`: the backend rejects the expert parallelism they use. The uncapped c48+ points keep the regular profile, so TTFT is not traded for ITL by a standing decode queue. - - ### Frontend - - ```bash - export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TOKENIZER_CACHE_BYTES=8589934592 DYN_TCP_REQUEST_TIMEOUT=120 - python3 -m dynamo.frontend --router-mode least-loaded --router-session-affinity-ttl-secs 900 - ``` - - Session affinity keeps each conversation on one worker so it reuses that worker's prefix cache; send the conversation ID in the `X-Dynamo-Session-Id` request header. - - ## AMD MI355X: InferenceX AgentX alignment - - The single-node MI355X TP8 AgentX configuration follows [InferenceX #4393](https://github.com/SemiAnalysisAI/InferenceX/pull/4393): its [master entry](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/configs/amd-master.yaml#L766) and [lane recipe](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4-mtp/agentic.yaml) at `dee7824`. It serves text with tool-call and reasoning parsers, using DCP8/A2A at every point. There is no concurrency-2 point. - - **In the command builder.** Settings the lane uses at every concurrency are part of the MI355X TP8 command: `--gpu-memory-utilization 0.9`, `--max-model-len 1048576`, `--enable-prefix-caching`, `--kv-cache-dtype fp8`, `--attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}'`, `--stream-interval 10`, `VLLM_ROCM_AITER_MLA_ASM_PADDING=asm`, `VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm` and `AITER_FLYDSL_STAGE2_FP8=1`. Weights load with vLLM's default loader (the AMD profile's `--load-format auto`): with `fastsafetensors`, the lane's servers ran out of GPU memory while loading weights on its image. The lane's `--moe-backend auto` is vLLM's default selection, so the builder leaves that flag unset. Other AMD GPUs keep the general AMD profile. - - - **Spec decoding** serves the lane's draft, the published [`Inferact/Kimi-K3-DSpark`](https://huggingface.co/Inferact/Kimi-K3-DSpark), with its attention on `ROCM_AITER_MLA` ([vllm-project/vllm#55966](https://github.com/vllm-project/vllm/pull/55966)), FP8 KV, probabilistic drafting and block verification. Draft length is tuned per cell in the matrix below. - - **Target quantization** stays in the guide: resident c1 passes no online quantization config, preserving the BF16 target weights required by specialized decode. Every other cell quantizes only the KDA `in_proj_qkvgfab` projections to FP8 per-channel. Routed experts retain their native MXFP4 weights. The DSpark draft remains BF16 with FP8 KV: vLLM builds the named draft from its own unquantized config, so removing c1's target quantization does not change the draft's dtype. - - **Decode context parallelism** emits the lane's DCP8 pairing: `--attention-backend ROCM_AITER_MLA`, `--dcp-comm-backend a2a`, `ROCM_AITER_FA` prefill and the direct DCP ops off. - - The **Simple** KV-offload option emits the connector the lane uses (`SimpleCPUOffloadConnector`, `kv_both`, `lazy_offload: false`). - - Select **Python** in the frontend selector to match the lane, which runs vLLM's default frontend. - - **Per-concurrency tuning (guide only).** The lane runs the following ten cells. The builder keeps general values (`--max-num-seqs 128`, `--max-num-batched-tokens 4096`, and effective `FULL_DECODE_ONLY` CUDA graphs); use the reproduction command below for exact benchmark cells. Every cell requests `FULL_AND_PIECEWISE` with breakable CUDA graphs disabled (`VLLM_USE_BREAKABLE_CUDAGRAPH=0`); because Kimi-K3 has no `torch.compile` submodule, vLLM resolves that mode to `FULL_DECODE_ONLY` at load. - - Every cell adds `kda_prefill_backend: fused` in `--additional-config`. Only **resident c1** also sets `decode_backend: latency_specialized` and `VLLM_USE_V2_MODEL_RUNNER=1`. On the pinned image, this [MI355X specialized path](https://github.com/vllm-project/vllm/pull/60683) requires TP8, V2, gfx950, no KV connector and BF16 target projections. It handles pure speculative-verification batches of at most eight rows (c1 with K7); other steps use the normal layers. These lane-specific settings stay out of generic MI355X defaults. - - | Concurrency | DCP | Profile | Draft tokens | `--max-num-seqs` | `--max-num-batched-tokens` | Largest CUDA graph | Target online quantization | - | ---: | ---: | --- | ---: | ---: | ---: | ---: | --- | - | 1 | 8 | resident | 7 | 1 | 16384 | 16 | none | - | 4 | 8 | resident | 7 | 8 | 8192 | 64 | `in_proj_qkvgfab` | - | 4 | 8 | offload | 5 | 8 | 8192 | 48 | `in_proj_qkvgfab` | - | 10 | 8 | offload | 4 | 20 | 8192 | 100 | `in_proj_qkvgfab` | - | 16 | 8 | offload | 4 | 32 | 8192 | 160 | `in_proj_qkvgfab` | - | 24 | 8 | offload | 4 | 48 | 8192 | 240 | `in_proj_qkvgfab` | - | 32 | 8 | offload | 4 | 64 | 8192 | 320 | `in_proj_qkvgfab` | - | 48 | 8 | offload | 4 | 96 | 8192 | 480 | `in_proj_qkvgfab` | - | 52 | 8 | offload | 4 | 104 | 8192 | 520 | `in_proj_qkvgfab` | - | 60 | 8 | offload | 4 | 120 | 8192 | 600 | `in_proj_qkvgfab` | - - The `resident` profile keeps KV on GPU; `offload` uses `SimpleCPUOffloadConnector` with 224.875 GB of host DRAM per rank. `--max-num-seqs` is twice the concurrency, except that concurrency 1 admits one sequence. Every cell captures every even CUDA-graph size from 2 through `2 × concurrency × (1 + draft tokens)`. Agentic benchmark runs replace block rejection with the golden acceptance length; synthetic acceptance is not suitable for serving or accuracy evaluation. - - **Image.** The builder keeps the rolling `vllm/vllm-openai-rocm:nightly-rocm100` image. Run the exact commands below inside `vllm/vllm-openai-rocm:nightly-rocm100-df6603f0583375f244d1ac6d36265462b7490392` (vLLM `df6603f0`), the image qualified by #4393 and carrying #60683. The compatibility limits above refer to this pinned runtime. - - **MoE environment.** The lane sets `VLLM_ROCM_USE_AITER_MOE_SITUV2=1`, as the AMD profile does, plus `AITER_BF16_FP8_MOE_BOUND=0` and `AITER_FLYDSL_STAGE2_FP8=1`. The SiTUv2 path does not read `AITER_BF16_FP8_MOE_BOUND`. `AITER_FLYDSL_STAGE2_FP8=1` makes the FlyDSL MoE down projection write its per-route outputs as MXFP8 (FP8 values with E8M0 scales) instead of BF16 before they are summed. It affects only the target's MoE layers; the DSpark draft's MLP is dense. - - **QuickReduce.** The lane also enables INT4 QuickReduce (`VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4`) with `VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB=16`, so only tensor-parallel all-reduces of 16 MiB or more use the INT4 codec. INT4 all-reduce is lossy, so it is not in the builder. The DSpark draft's all-reduces carry only its draft-step tokens (at most 600 tokens × 7,168 BF16 values, about 8.2 MiB, at concurrency 60), and its context projection has no all-reduce, so the 16 MiB floor keeps the draft out of the INT4 path in every lane cell. - - ### Reproducing one lane cell - - Run each cell in a fresh shell inside the pinned image. This real-verification form is for serving and accuracy evaluation; `PROFILE` distinguishes the two concurrency-4/DCP8 cells: - - ```bash - PROFILE=resident - CLIENTS=1 - case "$PROFILE:$CLIENTS" in - resident:1) K=7; BATCHED=16384; OFFLOAD=0 ;; - resident:4) K=7; BATCHED=8192; OFFLOAD=0 ;; - offload:4) K=5; BATCHED=8192; OFFLOAD=1 ;; - offload:10|offload:16|offload:24|offload:32|offload:48|offload:52|offload:60) K=4; BATCHED=8192; OFFLOAD=1 ;; - *) echo "PROFILE=$PROFILE CLIENTS=$CLIENTS is not a lane cell" >&2; exit 1 ;; - esac - - export VLLM_ROCM_USE_AITER=1 - export SAFETENSORS_FAST_GPU=1 - export VLLM_ROCM_USE_AITER_MOE_SITUV2=1 - export AITER_BF16_FP8_MOE_BOUND=0 - export AITER_FLYDSL_STAGE2_FP8=1 - export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm - export VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm - export VLLM_USE_DIRECT_DCP_A2A=0 - export VLLM_USE_DIRECT_DCP_Q_GATHER=0 - export VLLM_USE_DIRECT_DCP_KV_GATHER=0 - # Lossy INT4-quantized all-reduce for messages of 16 MiB or more: not in the builder. - export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 - export VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB=16 - # 1.56 TB of weights from shared storage, and long agentic prefills. - export VLLM_ENGINE_READY_TIMEOUT_S=7200 - export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200 - export PYTHONNOUSERSITE=1 - # Required on the qualification nodes' MEC firmware 38; fixed in firmware 177. - export HSA_NO_SCRATCH_RECLAIM=1 - - GRAPH_MODE=FULL_AND_PIECEWISE - export VLLM_USE_BREAKABLE_CUDAGRAPH=0 - MAX_SEQS=$((CLIENTS == 1 ? 1 : 2 * CLIENTS)) - MAX_CAPTURE=$((2 * CLIENTS * (1 + K))) - COMPILATION="{\"mode\":3,\"cudagraph_mode\":\"$GRAPH_MODE\",\"max_cudagraph_capture_size\":$MAX_CAPTURE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$(seq -s, 2 2 "$MAX_CAPTURE")]}" - - EXTRA=(--additional-config '{"kda_prefill_backend":"fused"}') - if [ "$PROFILE:$CLIENTS" = resident:1 ]; then - export VLLM_USE_V2_MODEL_RUNNER=1 - EXTRA=(--additional-config '{"decode_backend":"latency_specialized","kda_prefill_backend":"fused"}') - else - EXTRA+=(--quantization-config '{"targets":{"*in_proj_qkvgfab":"fp8_per_channel"}}') - fi - if [ "$K" -gt 0 ]; then - EXTRA+=(--speculative-config "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$K,\"method\":\"dspark\",\"attention_backend\":\"ROCM_AITER_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"block\"}") - fi - if [ "$OFFLOAD" -eq 1 ]; then - # The lane's node gives the connector 1799 GB of host DRAM (224875000000 bytes - # per rank); size cpu_bytes_to_use_per_rank to your host. The fixed hash seed - # makes identical prefixes map to the same block keys on every rank. - export PYTHONHASHSEED=42 - EXTRA+=(--kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":224875000000,"lazy_offload":false}}') - fi - vllm serve moonshotai/Kimi-K3 \ - --trust-remote-code \ - --moe-backend auto \ - --tensor-parallel-size 8 \ - --decode-context-parallel-size 8 \ - --dcp-comm-backend a2a \ - --gpu-memory-utilization 0.9 \ - --language-model-only \ - --max-num-seqs "$MAX_SEQS" \ - --enable-auto-tool-choice \ - --tool-call-parser kimi_k3 \ - --reasoning-parser kimi_k3 \ - --max-model-len 1048576 \ - --enable-prefix-caching \ - --kv-cache-dtype fp8 \ - --stream-interval 10 \ - --max-num-batched-tokens "$BATCHED" \ - --attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}' \ - --compilation-config "$COMPILATION" \ - "${EXTRA[@]}" - ``` - - The builder's **Simple** option keeps the repository-wide 220 GiB per rank and does not set `PYTHONHASHSEED`. - - ### Throughput measurement only (synthetic acceptance) - - The lane's throughput runs replace block verification with synthetic acceptance, so every concurrency point is measured at InferenceX's golden acceptance length for this draft. Only the `--speculative-config` changes: - - ```bash - --speculative-config "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$K,\"method\":\"dspark\",\"attention_backend\":\"ROCM_AITER_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"synthetic\",\"synthetic_acceptance_length\":$AL}" - ``` - - | Draft tokens | Golden acceptance length | Lane cells | - | ---: | ---: | --- | - | 7 | 3.84 | resident concurrency 1 and 4, DCP8 | - | 5 | 3.62 | offload concurrency 4, DCP8 | - | 4 | 3.36 | dcp8 concurrency 10, 16, 24, 32, 48, 52 and 60 | - - These lengths come from [`kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml): SPEED-Bench coding prompts with thinking on, measured on B300 for `Inferact/Kimi-K3-DSpark` with probabilistic drafting and block verification. They hold only for that draft at those draft lengths, not for `RedHatAI/Kimi-K3-speculator.dspark`. Synthetic acceptance skips real verification: never use it to serve users or to evaluate accuracy. - - ### Evidence and limits - - - InferenceX [run 40037005947, attempt 2](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40037005947/attempts/2) passed all ten performance cells and its accuracy jobs at source `0df7e892b73c53d66a486c2a085e994b67653042`. The Kimi serving recipe, model, image and selected matrix are unchanged at audited head `dee7824`; later changes are unrelated upstream work and changelog text. Accuracy jobs use real verification and are deduplicated by [evaluation identity](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/infx/evals/EVALS.md#L33), so this is not a claim of a separate accuracy job for every matrix cell. - - Performance figures on the InferenceX dashboard are reported by SemiAnalysis and have not been reproduced for this recipe. - - Model policy (vLLM as the native engine, `Inferact/Kimi-K3-DSpark` as the agreed draft): [`MODELS.md`](https://github.com/SemiAnalysisAI/InferenceX/blob/dee7824e8ca3810cc36fe668ed75a927183af41f/inferencex-e2e/docs/MODELS.md). - ## AMD MI355X: disaggregated prefill and decode This example uses one TP8/DCP8 prefill worker and one TP8/DCP8 decode worker on separate 8-GPU MI355X nodes. MoRIIO transfers the KV cache in READ mode; `MultiConnector` combines it with `SimpleCPUOffloadConnector` on prefill only. Both roles use the same target and DSpark draft checkpoints and FP8 KV. From 2b1f75c80f79496b313e28362f0856632dfbf5ea Mon Sep 17 00:00:00 2001 From: Marcel Yeonghyeon Ko Date: Sun, 25 Oct 2026 17:54:10 +0900 Subject: [PATCH 042/631] [Bugfix][KV Offload] Treat failed secondary-tier stores as back-pressure (#61780) Signed-off-by: YeonghyeonKO Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: vllm-agent --- .../kv_offload/tiering/test_backpressure.py | 132 +++++++++++++++++- vllm/v1/kv_offload/tiering/backpressure.py | 74 +++++++++- vllm/v1/kv_offload/tiering/manager.py | 20 +-- 3 files changed, 207 insertions(+), 19 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_backpressure.py b/tests/v1/kv_offload/tiering/test_backpressure.py index 4368f254f72d..a854b0fd0aa2 100644 --- a/tests/v1/kv_offload/tiering/test_backpressure.py +++ b/tests/v1/kv_offload/tiering/test_backpressure.py @@ -45,6 +45,7 @@ # Water marks in s/MiB, scaled for the test's tiny 16-byte blocks. _BP_HIGH_WATER_S = 1000.0 _BP_LOW_WATER_S = 500.0 +_BP_FAILURE_THRESHOLD = EMABackpressureDetector.DEFAULT_FAILURE_THRESHOLD _CTX = ReqContext(req_id="test") _MOCK_OFFLOADING_SPEC = MagicMock() @@ -90,14 +91,22 @@ def __init__( self._held_jobs: list[JobResult] = [] self._released_jobs: list[JobResult] = [] self._request_policy: OffloadPolicy | None = None + # When True, stores complete with success=False, like a tier whose + # writes are rejected (e.g. ENOSPC on a full filesystem). + self.fail_stores = False + self.num_store_submissions = 0 def lookup(self, key, req_context): return LookupResult.HIT if key in self.blocks else LookupResult.MISS def submit_store(self, job_metadata: TransferJob) -> None: - for key in job_metadata.keys: - self.blocks[key] = True - self._held_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + self.num_store_submissions += 1 + if not self.fail_stores: + for key in job_metadata.keys: + self.blocks[key] = True + self._held_jobs.append( + JobResult(job_id=job_metadata.job_id, success=not self.fail_stores) + ) def submit_load(self, job_metadata: TransferJob) -> None: self._released_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) @@ -152,6 +161,18 @@ def test_fs_with_local_locality_gets_local_watermarks(self): assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S + @pytest.mark.parametrize( + "tier_type, expected", + [ + ("fs", EMABackpressureDetector.DEFAULT_FAILURE_THRESHOLD), + ("obj", EMABackpressureDetector.DEFAULT_FAILURE_THRESHOLD), + ("p2p", None), + ], + ) + def test_failure_threshold_only_for_storage_tiers(self, tier_type, expected): + cfg = EMABackpressureDetector.default_config(tier_type) + assert cfg.get("failure_threshold") == expected + class TestIdleDecay: """Tests for EMA idle decay — recovery from pressure without completions.""" @@ -240,6 +261,7 @@ def setup(self): backpressure_detector=EMABackpressureDetector( high_water_s=_BP_HIGH_WATER_S, low_water_s=_BP_LOW_WATER_S, + failure_threshold=_BP_FAILURE_THRESHOLD, ), ) self.manager = TieringOffloadingManager( @@ -556,6 +578,110 @@ def test_loads_continue_under_pressure(self, setup): result = self.manager.lookup(b, _CTX) assert result is LookupResult.HIT_PENDING + def _store_and_complete(self, keys): + self._store_blocks(keys) + self.manager._processed_jobs_this_step = False + self.tier.release_jobs() + self._simulate_on_schedule_end() + + def test_failed_stores_activate_pressure(self, setup): + """A tier whose stores fail fast is put under pressure, so later + stores are dropped instead of being submitted to it.""" + bp = self.tier.bp_detector + assert isinstance(bp, EMABackpressureDetector) + bp._decay_half_life_s = 0.0 + # A healthy tier that starts failing must not keep bypassing the policy. + bp._completions = _BP_WARMUP + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + self.tier.fail_stores = True + + for i in range(_BP_FAILURE_THRESHOLD): + assert bp.is_under_pressure() is False + self._store_and_complete(to_keys([900 + i])) + assert bp.is_under_pressure() is True + assert bp.is_healthy() is False + + keys = to_keys(range(910, 913)) + self._store_blocks(keys) + assert self.tier.num_store_submissions == _BP_FAILURE_THRESHOLD + statuses = [self.primary._policy.get(k) for k in keys] + assert all(s is not None and s.ref_cnt == 0 for s in statuses) + + +class TestStoreFailures: + """Detector behaviour for failed stores.""" + + def _make_detector(self, **kwargs): + defaults = dict( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + cooldown_s=0.0, + failure_threshold=_BP_FAILURE_THRESHOLD, + ) + defaults.update(kwargs) + return EMABackpressureDetector(**defaults) + + def _complete_fast_store(self, bp): + block_bytes = 16 + fast_s_per_mib = _BP_LOW_WATER_S * 0.5 + bp.on_store_completed( + fast_s_per_mib * (block_bytes / EMABackpressureDetector._MIB), + block_bytes, + ) + + def test_failures_ignored_without_threshold(self): + bp = self._make_detector(failure_threshold=None) + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + for _ in range(10): + bp.on_store_failed() + assert bp.is_under_pressure() is False + assert bp.is_healthy() is True + + def test_success_resets_failure_streak(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._last_update = time.monotonic() + for _ in range(_BP_FAILURE_THRESHOLD - 1): + bp.on_store_failed() + self._complete_fast_store(bp) + for _ in range(_BP_FAILURE_THRESHOLD - 1): + bp.on_store_failed() + assert bp.is_under_pressure() is False + + def test_failure_pressure_drops_all_stores(self): + bp = self._make_detector(failure_threshold=1, decay_half_life_s=0.0) + bp._completions = _BP_WARMUP + bp.on_store_failed() + assert isinstance(bp.policy, ThrottledDropPolicy) + assert not any(bp.should_store(1) for _ in range(20)) + + def test_failure_pressure_decays_before_warmup(self): + """A tier that never succeeds never finishes warmup; its pressure + must still decay so it gets probed again.""" + bp = self._make_detector(failure_threshold=1) + bp.on_store_failed() + assert bp.is_under_pressure() is True + with patch("time.monotonic") as mock_time: + mock_time.return_value = bp._last_update + 60.0 + assert bp.is_under_pressure() is False + + @pytest.mark.parametrize("probe_succeeds", [False, True]) + def test_probe_after_decay(self, probe_succeeds): + bp = self._make_detector() + bp._completions = _BP_WARMUP + for _ in range(_BP_FAILURE_THRESHOLD): + bp.on_store_failed() + with patch("time.monotonic") as mock_time: + mock_time.return_value = bp._last_update + 60.0 + assert bp.is_under_pressure() is False + if probe_succeeds: + self._complete_fast_store(bp) + bp.on_store_failed() + assert bp.is_under_pressure() is False + else: + bp.on_store_failed() + assert bp.is_under_pressure() is True + class TestThrottledDropPolicy: """Tests for proportional throttling under pressure.""" diff --git a/vllm/v1/kv_offload/tiering/backpressure.py b/vllm/v1/kv_offload/tiering/backpressure.py index 0da4aa5f3cc1..814b2c88599d 100644 --- a/vllm/v1/kv_offload/tiering/backpressure.py +++ b/vllm/v1/kv_offload/tiering/backpressure.py @@ -132,6 +132,13 @@ def policy(self) -> BackpressurePolicy: @abstractmethod def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: ... + def on_store_failed(self) -> None: + """Observe a store job that completed with ``success=False``. + + Ignored by default; detectors may treat failures as pressure. + """ + return + @abstractmethod def is_under_pressure(self) -> bool: ... @@ -229,6 +236,16 @@ class EMABackpressureDetector(BackpressureDetector): # detector enters a "healthy" fast-path that bypasses policy checks. _HEALTHY_THRESHOLD = 10 + # Consecutive failed stores that put a storage tier under pressure + # (e.g. ENOSPC/EDQUOT on a full filesystem). Not applied to tiers such + # as p2p, whose store failures are per-request rather than tier health. + DEFAULT_FAILURE_THRESHOLD = 3 + + # Failure-induced pressure raises the EMA to this multiple of the high + # water mark, enough for ThrottledDropPolicy (ramp_factor=2.0) to drop + # every store until idle decay lets a probe store through. + _FAILURE_EMA_FACTOR = 4.0 + LOCAL_HIGH_WATER_S = 0.010 LOCAL_LOW_WATER_S = 0.002 @@ -236,19 +253,24 @@ class EMABackpressureDetector(BackpressureDetector): NETWORK_LOW_WATER_S = 0.010 _NETWORK_TIER_TYPES = frozenset({"obj"}) + _STORAGE_TIER_TYPES = frozenset({"fs", "obj"}) @classmethod def default_config(cls, tier_type: str, *, locality: str | None = None) -> dict: is_remote = locality is not None and locality.upper() == "REMOTE" if is_remote or tier_type in cls._NETWORK_TIER_TYPES: - return { + config: dict = { "high_water_s": cls.NETWORK_HIGH_WATER_S, "low_water_s": cls.NETWORK_LOW_WATER_S, } - return { - "high_water_s": cls.LOCAL_HIGH_WATER_S, - "low_water_s": cls.LOCAL_LOW_WATER_S, - } + else: + config = { + "high_water_s": cls.LOCAL_HIGH_WATER_S, + "low_water_s": cls.LOCAL_LOW_WATER_S, + } + if tier_type in cls._STORAGE_TIER_TYPES: + config["failure_threshold"] = cls.DEFAULT_FAILURE_THRESHOLD + return config def __init__( self, @@ -258,6 +280,7 @@ def __init__( warmup_completions: int = DEFAULT_WARMUP_COMPLETIONS, decay_half_life_s: float = DEFAULT_DECAY_HALF_LIFE_S, cooldown_s: float = DEFAULT_COOLDOWN_S, + failure_threshold: int | None = None, policy: BackpressurePolicy | None = None, ): super().__init__(policy=policy) @@ -267,6 +290,8 @@ def __init__( raise ValueError( f"low_water_s ({low_water_s}) must be <= high_water_s ({high_water_s})" ) + if failure_threshold is not None and failure_threshold < 1: + raise ValueError(f"failure_threshold must be >= 1, got {failure_threshold}") self._alpha = alpha self._high = high_water_s @@ -289,6 +314,13 @@ def __init__( # Monotonic timestamp when pressure was last cleared (for cooldown). self._pressure_cleared_at: float = 0.0 self._cooldown_count: int = 0 + # Consecutive failed stores that trigger pressure; None ignores + # failures. + self._failure_threshold = failure_threshold + self._consecutive_failures: int = 0 + # Failure-induced pressure can start before warmup completes, so it + # must be allowed to decay during warmup. + self._failure_tripped: bool = False def _apply_idle_decay(self) -> None: """Decay the EMA toward zero based on elapsed idle time. @@ -309,6 +341,7 @@ def _apply_idle_decay(self) -> None: def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: now = time.monotonic() s_per_mib = elapsed_s / (num_bytes / self._MIB) + self._consecutive_failures = 0 self._completions += 1 if self._completions <= self._warmup_completions: self._warmup_samples.append(s_per_mib) @@ -327,15 +360,37 @@ def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: self._pressure_cleared_at = now self._cooldown_count = 0 self._under_pressure = False + self._failure_tripped = False self._healthy_streak += 1 else: self._healthy_streak = 0 + def on_store_failed(self) -> None: + """Enter pressure after ``failure_threshold`` consecutive failures. + + Idle decay then lets probe stores through: a failed probe re-enters + pressure at once, and a successful one resets the failure streak. + """ + if self._failure_threshold is None: + return + self._consecutive_failures += 1 + self._healthy_streak = 0 + if self._consecutive_failures < self._failure_threshold: + return + self._ema = max(self._ema, self._high * self._FAILURE_EMA_FACTOR) + self._last_update = time.monotonic() + self._under_pressure = True + self._failure_tripped = True + self._pressure_cleared_at = 0.0 + def is_under_pressure(self) -> bool: - if self._under_pressure and self._completions >= self._warmup_completions: + if self._under_pressure and ( + self._completions >= self._warmup_completions or self._failure_tripped + ): self._apply_idle_decay() if self._ema < self._low: self._under_pressure = False + self._failure_tripped = False self._pressure_cleared_at = time.monotonic() self._cooldown_count = 0 return self._under_pressure @@ -376,6 +431,8 @@ def reset(self) -> None: self._healthy_streak = 0 self._pressure_cleared_at = 0.0 self._cooldown_count = 0 + self._consecutive_failures = 0 + self._failure_tripped = False self._policy.reset() @property @@ -388,4 +445,7 @@ def store_latency_ema(self, value: float) -> None: @property def stats(self) -> dict[str, float]: - return {"store_latency_ema": self._ema} + return { + "store_latency_ema": self._ema, + "consecutive_store_failures": self._consecutive_failures, + } diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 89a126d8b106..0231419164c3 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -542,8 +542,7 @@ def _process_finished_jobs( self.primary_tier.complete_read( transfer_job.keys, transfer_job.req_context ) - if completed_job.success: - self._update_backpressure(tier, job_metadata, completed_job) + self._update_backpressure(tier, job_metadata, completed_job) def _should_store_to_tier( self, tier: SecondaryTierManager, num_blocks: int @@ -563,13 +562,16 @@ def _update_backpressure( if detector is None: return was_under_pressure = detector.is_under_pressure() - tj = job_metadata.transfer_job - num_bytes = ( - completed_job.transfer_bytes - if completed_job.transfer_bytes is not None - else len(tj.keys) * tier.block_size_bytes - ) - detector.update(tj.submit_time, num_bytes) + if completed_job.success: + tj = job_metadata.transfer_job + num_bytes = ( + completed_job.transfer_bytes + if completed_job.transfer_bytes is not None + else len(tj.keys) * tier.block_size_bytes + ) + detector.update(tj.submit_time, num_bytes) + else: + detector.on_store_failed() if detector.is_under_pressure() != was_under_pressure: tier_idx = self._tier_index[tier] logger.info( From b70de52965a87125fa103a50834ec6cc1b276cb0 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 25 Oct 2026 04:20:53 -0500 Subject: [PATCH 043/631] [Bugfix][Model] Borrow Gemma4 _get_max_soft_tokens in EmbeddingGemma2 and DiffusionGemma (#63538) Signed-off-by: Andreas Karatzas Co-authored-by: Claude Co-authored-by: Cyrus Leung --- vllm/model_executor/models/diffusion_gemma.py | 1 + vllm/model_executor/models/embedding_gemma2.py | 1 + 2 files changed, 2 insertions(+) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 69866f92ff24..f2f49cf03297 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -286,6 +286,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs ) _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) + _get_max_soft_tokens = Gemma4ForConditionalGeneration._get_max_soft_tokens _process_image_input = Gemma4ForConditionalGeneration._process_image_input _process_video_input = Gemma4ForConditionalGeneration._process_video_input embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py index efde2acd1e34..9841e4f388eb 100644 --- a/vllm/model_executor/models/embedding_gemma2.py +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -1012,6 +1012,7 @@ class EmbeddingGemma2Model(nn.Module, SupportsMultiModal, VllmModelForPooling): _parse_and_validate_multimodal_inputs = ( Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs ) + _get_max_soft_tokens = Gemma4ForConditionalGeneration._get_max_soft_tokens _process_image_input = Gemma4ForConditionalGeneration._process_image_input _process_video_input = Gemma4ForConditionalGeneration._process_video_input _process_audio_input = Gemma4ForConditionalGeneration._process_audio_input From 24324ef0f667c633ee9bbb9412343d829f7e6d4e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 25 Oct 2026 04:22:35 -0500 Subject: [PATCH 044/631] [CI] Fix the Qwen2 embed memory leak and HfRunner no-cache logprobs (#63537) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- tests/conftest.py | 8 ++++---- tests/models/language/pooling/test_embedding.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index ccc40724c4bb..a979968c1621 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -817,10 +817,10 @@ def _hidden_states_to_logprobs( # convert to dict seq_logprobs_lst: list[dict[int, float]] = [] - for tok_idx, tok_logprobs in enumerate(seq_logprobs): - # drop prompt logprobs - if tok_idx == 0: - tok_logprobs = tok_logprobs[-1, :].reshape(1, -1) + for tok_logprobs in seq_logprobs: + # Keep only the last position: the first step (and every step when + # use_cache=False) also returns logprobs for the preceding tokens. + tok_logprobs = tok_logprobs[-1, :].reshape(1, -1) topk = tok_logprobs.topk(num_logprobs) tok_logprobs_dct = {} diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index 47e1a0a7bad4..80978a018c10 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/tests/models/language/pooling/test_embedding.py @@ -88,6 +88,7 @@ def test_models( st_model[1] = Pooling( st_model.get_sentence_embedding_dimension(), pooling_mode="mean" ) + del st_model hf_outputs = hf_model.encode(example_prompts) with vllm_runner( From 75fc05a53a366068f534a61ab3ee6b50485419f2 Mon Sep 17 00:00:00 2001 From: Shantanu Acharya Date: Sun, 25 Oct 2026 07:13:22 -0400 Subject: [PATCH 045/631] [Model] Add Dory support (#60149) Signed-off-by: Shantanu Acharya Co-authored-by: Roger Wang --- .buildkite/test_areas/models_basic.yaml | 3 +- docs/models/supported_models.md | 1 + tests/models/registry.py | 13 + tests/models/test_dory.py | 388 ++++++++++++++++++ tests/transformers_utils/test_dory_config.py | 120 ++++++ vllm/model_executor/models/dory.py | 392 +++++++++++++++++++ vllm/model_executor/models/registry.py | 1 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + vllm/transformers_utils/configs/dory.py | 138 +++++++ 10 files changed, 1058 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_dory.py create mode 100644 tests/transformers_utils/test_dory_config.py create mode 100644 vllm/model_executor/models/dory.py create mode 100644 vllm/transformers_utils/configs/dory.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 59983e0e45db..a9505b7f787d 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -313,6 +313,7 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py + - tests/models/test_dory.py - tests/models/test_alice_ai.py - tests/models/test_alice_ai_config.py - tests/models/test_beam_weight_mapping.py @@ -328,7 +329,7 @@ steps: - tests/models/test_cohere2_moe_parallel_block.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py models/test_cohere2_moe_parallel_block.py + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_dory.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py models/test_cohere2_moe_parallel_block.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 32d417d55dba..6f637415bd32 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -392,6 +392,7 @@ th { | `DeepseekV32ForCausalLM` | DeepSeek-V3.2 | `deepseek-ai/DeepSeek-V3.2`, etc. | ✅︎ | ✅︎ | | `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | ✅︎ | | `DomynEdgeForCausalLM` | Domyn Edge | `domyn/Domyn-Edge-Base`, `domyn/Domyn-Edge-SFT`, etc. | | | +| `DoryForCausalLM` | Dory | `nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16` | | | | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ | | `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. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index dce4d917f580..cfed6f0d5317 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -270,6 +270,19 @@ def check_available_online( trust_remote_code=True, extras={"sft": "domyn/Domyn-Edge-SFT"}, ), + "DoryForCausalLM": _HfExamplesInfo( + "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16", + use_original_num_layers=True, + max_model_len=128, + hf_overrides={ + "num_hidden_layers": 8, + "n_input_layers": 2, + "n_recurrent_layers": 4, + "n_output_layers": 2, + "n_recurrent_loops": 2, + "layer_types": ["sliding_attention", "mlp"] * 3 + ["full_attention", "mlp"], + }, + ), "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT"), "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT"), "ExaoneForCausalLM": _HfExamplesInfo( diff --git a/tests/models/test_dory.py b/tests/models/test_dory.py new file mode 100644 index 000000000000..14f1bd715e53 --- /dev/null +++ b/tests/models/test_dory.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""CPU checks of Dory math and loading, with a reference attention backend.""" + +from types import SimpleNamespace + +import pytest +import torch +from huggingface_hub.errors import StrictDataclassClassValidationError +from torch import nn + +from vllm.config import ( + CompilationConfig, + DeviceConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed import parallel_state +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.models import dory +from vllm.model_executor.models.utils import extract_layer_index +from vllm.transformers_utils.configs.dory import DoryConfig + + +class ReferenceAttention(nn.Module): + """Single-request causal GQA, including sliding-window and decode caching.""" + + def __init__(self, num_heads, head_size, scale, num_kv_heads, **kwargs): + super().__init__() + self.num_heads, self.num_kv_heads = num_heads, num_kv_heads + self.head_size, self.scale = head_size, scale + self.prefix = kwargs["prefix"] + self.window = kwargs["per_layer_sliding_window"] + self.keys = self.values = None + + def forward(self, q, k, v): + q = q.view(-1, self.num_heads, self.head_size).transpose(0, 1) + k = k.view(-1, self.num_kv_heads, self.head_size).transpose(0, 1) + v = v.view(-1, self.num_kv_heads, self.head_size).transpose(0, 1) + offset = self.start_pos + self.keys = ( + k if self.keys is None else torch.cat((self.keys[:, :offset], k), dim=1) + ) + self.values = ( + v if self.values is None else torch.cat((self.values[:, :offset], v), dim=1) + ) + key_pos = torch.arange(self.keys.shape[1]) + query_pos = torch.arange(offset, offset + q.shape[1])[:, None] + mask = key_pos <= query_pos + if self.window is not None: + mask &= key_pos > query_pos - self.window + out = torch.nn.functional.scaled_dot_product_attention( + q, self.keys, self.values, attn_mask=mask, scale=self.scale, enable_gqa=True + ) + return out.transpose(0, 1).reshape(-1, self.num_heads * self.head_size) + + +@pytest.fixture +def model_factory(monkeypatch): + # Keep the actual vLLM linear/embedding modules and their weight loaders. + monkeypatch.setattr( + parallel_state, "_TP", SimpleNamespace(world_size=1, rank_in_group=0) + ) + monkeypatch.setattr(dory, "Attention", ReferenceAttention) + vc = VllmConfig( + device_config=DeviceConfig(device="cpu"), + compilation_config=CompilationConfig(mode=0, custom_ops=["none"]), + ) + + def build(*, hf_overrides=None, quant_config=None, **overrides): + kwargs = dict( + vocab_size=32, + hidden_size=16, + intermediate_size=24, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + num_hidden_layers=8, + n_input_layers=2, + n_recurrent_layers=4, + n_output_layers=2, + n_recurrent_loops=3, + max_position_embeddings=32, + layer_types=["sliding_attention", "mlp"] * 3 + ["full_attention", "mlp"], + sliding_window=3, + rope_parameters={ + kind: { + "rope_type": "default", + "rope_theta": theta, + "partial_rotary_factor": 1.0, + } + for kind, theta in ( + ("sliding_attention", 100), + ("full_attention", 10000), + ) + }, + ) + kwargs.update(overrides) + cfg = DoryConfig(**kwargs) + cfg.update(hf_overrides or {}) + vc.model_config = SimpleNamespace( + hf_config=cfg, head_dtype=None, dtype=torch.float32 + ) + vc.quant_config = quant_config + with set_current_vllm_config(vc): + model = dory.DoryForCausalLM(vllm_config=vc) + + def set_cache_positions(module, args): + # Emulate the runner's slot mapping for contiguous request tokens. + for cache in module.modules(): + if isinstance(cache, ReferenceAttention): + cache.start_pos = int(args[1][0]) + + model.register_forward_pre_hook(set_cache_positions) + return model + + with set_current_vllm_config(vc): + yield build + + +def _prepare_for_inference(model): + # CPU kernels may pack or replace weights, so run after all weight updates. + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, QuantizeMethodBase): + quant_method.process_weights_after_loading(module) + + +def test_invalid_loop_override_is_rejected_before_model_creation(model_factory): + with pytest.raises(StrictDataclassClassValidationError, match="n_recurrent_loops"): + model_factory(hf_overrides={"n_recurrent_loops": 0}) + + +@pytest.mark.parametrize("num_loops", [3, 27]) +@torch.inference_mode() +def test_chunked_prefill_and_decode_keep_loop_caches_separate(model_factory, num_loops): + model = model_factory(n_recurrent_loops=num_loops) + torch.manual_seed(0) + for param in model.parameters(): + param.normal_(std=0.02) + _prepare_for_inference(model) + tokens, positions = torch.arange(7), torch.arange(7) + expected = model(tokens, positions) + caches = [m for m in model.modules() if isinstance(m, ReferenceAttention)] + # Two outer attention layers + two shared recurrent layers per loop. + assert ( + len(caches) + == len({extract_layer_index(m.prefix) for m in caches}) + == (2 + 2 * num_loops) + ) + for cache in caches: + assert cache.keys is not None + assert cache.keys.shape[1] == len(tokens) + cache.keys = cache.values = None + actual = torch.cat( + [model(tokens[:3], positions[:3])] + + [model(tokens[i : i + 1], positions[i : i + 1]) for i in range(3, 7)] + ) + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + + +@pytest.mark.parametrize("cache_mode", ["per_loop", "last_loop"]) +def test_loop_count_does_not_duplicate_projection_parameters(model_factory, cache_mode): + single_loop = model_factory(n_recurrent_loops=1) + many_loops = model_factory(n_recurrent_loops=27, recurrent_kv_cache_mode=cache_mode) + assert {name: p.shape for name, p in single_loop.named_parameters()} == { + name: p.shape for name, p in many_loops.named_parameters() + } + + +@pytest.mark.parametrize("num_loops", [1, 3]) +@torch.inference_mode() +def test_last_loop_kv_matches_final_loop_history(model_factory, num_loops): + """Reuse final-loop history while preserving current-loop KV within a chunk.""" + shared = model_factory( + n_recurrent_loops=num_loops, recurrent_kv_cache_mode="last_loop" + ) + reference = model_factory(n_recurrent_loops=num_loops) + torch.manual_seed(0) + for param in shared.parameters(): + param.normal_(std=0.02) + reference.load_state_dict(shared.state_dict()) + _prepare_for_inference(shared) + _prepare_for_inference(reference) + assert sum(isinstance(m, ReferenceAttention) for m in shared.modules()) == 4 + tokens = positions = torch.arange(7) + # Include prefill chunks, decode, and history beyond the sliding window. + for start, end in ((0, 3), (3, 4), (4, 7)): + expected = reference(tokens[start:end], positions[start:end]) + actual = shared(tokens[start:end], positions[start:end]) + torch.testing.assert_close(actual, expected) + for layer, ref_layer in zip(shared.backbone.layers, reference.backbone.layers): + if not layer.is_attention: + continue + cache, final = layer.mixer.attn[0], ref_layer.mixer.attn[-1] + assert cache.keys.shape[1] == end + torch.testing.assert_close(cache.keys, final.keys) + torch.testing.assert_close(cache.values, final.values) + # Next step, every reference loop starts from final-loop history. + for ref_cache in ref_layer.mixer.attn: + ref_cache.keys = final.keys.clone() + ref_cache.values = final.values.clone() + + +@pytest.mark.parametrize("reverse", [False, True]) +@torch.inference_mode() +def test_split_scales_load_independently_in_either_order(model_factory, reverse): + model = model_factory() + prefix = "backbone.layers.1.mixer" + gate, up = torch.arange(24.0), torch.arange(24.0) + 100 + head = torch.arange(32 * 16.0).reshape(32, 16) / 512 + scale = torch.linspace(0.01, 0.03, 32) + weights = [ + (f"{prefix}.suv_gate", gate), + ("lm_head.weight", head), + (f"{prefix}.suv_up", up), + ("lm_head.logit_scale", scale), + ] + if reverse: + weights.reverse() + loaded = set() + for name, tensor in weights: + loaded_now = model.load_weights([(name, tensor)]) + assert loaded_now == {name} + loaded |= loaded_now + assert loaded == {name for name, _ in weights} + torch.testing.assert_close(model.backbone.layers[1].mixer.suv_gate, gate) + torch.testing.assert_close(model.backbone.layers[1].mixer.suv_up, up) + torch.testing.assert_close(model.lm_head.weight[:32], head) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@torch.inference_mode() +def test_logits_apply_vocab_scale_after_projection(model_factory, monkeypatch, dtype): + model = model_factory().to(dtype) + monkeypatch.setattr(model.logits_processor, "_gather_logits", lambda logits: logits) + torch.manual_seed(0) + head = torch.randn(32, 16, dtype=dtype) + scale = torch.linspace(0.01, 0.03, 32).to(dtype) + hidden = torch.randn(3, 16, dtype=dtype) + model.load_weights([("lm_head.weight", head), ("lm_head.logit_scale", scale)]) + _prepare_for_inference(model) + expected = torch.nn.functional.linear(hidden, head).float() * ( + scale.float() / model.config.initializer_range + ) + torch.testing.assert_close(model.compute_logits(hidden), expected) + + +@torch.inference_mode() +def test_exported_qkv_and_mlp_weights_load_in_correct_order(model_factory): + model = model_factory() + attn = "backbone.layers.0.mixer" + mlp = "backbone.layers.1.mixer" + pieces = [ + (f"{attn}.q_proj.weight", torch.full((16, 16), 1.0)), + (f"{attn}.k_proj.weight", torch.full((8, 16), 2.0)), + (f"{attn}.v_proj.weight", torch.full((8, 16), 3.0)), + (f"{mlp}.gate_proj.weight", torch.full((24, 16), 4.0)), + (f"{mlp}.up_proj.weight", torch.full((24, 16), 5.0)), + ] + loaded = model.load_weights(reversed(pieces)) + assert loaded == {f"{attn}.qkv_proj.weight", f"{mlp}.gate_up_proj.weight"} + torch.testing.assert_close( + model.backbone.layers[0].mixer.qkv_proj.weight, + torch.cat([w for _, w in pieces[:3]]), + ) + torch.testing.assert_close( + model.backbone.layers[1].mixer.gate_up_proj.weight, + torch.cat([w for _, w in pieces[3:]]), + ) + + +@pytest.mark.parametrize( + "name,shape", + [ + ("backbone.layers.0.mixer.q_proj.weight", (32, 16)), + ("backbone.layers.1.mixer.gate_proj.weight", (32, 16)), + ("backbone.layers.1.mixer.down_proj.weight", (16, 32)), + ("backbone.layers.1.mixer.suv_gate", (32,)), + ], +) +def test_weight_loader_rejects_silent_truncation(model_factory, name, shape): + model = model_factory() + with pytest.raises(ValueError, match="expected shape"): + model.load_weights([(name, torch.zeros(shape))]) + + +@pytest.fixture(params=["mxfp8", "nvfp4"]) +def quantized_model(request, monkeypatch, model_factory): + from vllm.model_executor.layers.quantization import modelopt + from vllm.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, + ) + from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( + compressed_tensors_w8a8_mxfp8, + ) + + # Exercise real parameter allocation and loaders without requiring GPU kernels. + kernel = SimpleNamespace(input_quant_key=lambda: None) + monkeypatch.setattr(modelopt, "select_linear_kernel", lambda *a, **kw: kernel) + monkeypatch.setattr( + compressed_tensors_w8a8_mxfp8, + "init_mxfp8_linear_kernel", + lambda **kw: kernel, + ) + CompressedTensorsConfig = compressed_tensors.CompressedTensorsConfig + monkeypatch.setattr( + CompressedTensorsConfig, "_check_scheme_supported", lambda *a: None + ) + if request.param == "nvfp4": + quant = modelopt.ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True + ) + else: + quant = CompressedTensorsConfig.from_config( + { + "format": "mxfp8-quantized", + "config_groups": { + "weights": { + "targets": ["Linear"], + "weights": { + "num_bits": 8, + "type": "float", + "strategy": "group", + "group_size": 32, + "symmetric": True, + "dynamic": False, + "scale_dtype": "torch.uint8", + }, + } + }, + } + ) + quant.packed_modules_mapping = dory.DoryForCausalLM.packed_modules_mapping + return model_factory( + quant_config=quant, hidden_size=64, intermediate_size=96, head_dim=32 + ) + + +@torch.inference_mode() +def test_quantized_projection_weights_and_scales_load(quantized_model): + """Packed columns, block scales and scalar scales retain every shard's values.""" + model = quantized_model + cases = [ + (0, "qkv_proj", ["q_proj", "k_proj", "v_proj"], [64, 32, 32]), + (1, "gate_up_proj", ["gate_proj", "up_proj"], [96, 96]), + (0, "g_proj", ["g_proj"], [64]), + (0, "o_proj", ["o_proj"], [64]), + (1, "down_proj", ["down_proj"], [64]), + ] + for index, target, sources, rows in cases: + prefix = f"backbone.layers.{index}.mixer" + layer = model.get_submodule(f"{prefix}.{target}") + for suffix in ("weight", "weight_scale", "weight_scale_2"): + param = getattr(layer, suffix, None) + if param is None: + continue + scalar = suffix == "weight_scale_2" + pieces = [ + torch.full( + () if scalar else (n, param.shape[1]), i + 1, dtype=param.dtype + ) + for i, n in enumerate(rows) + ] + loaded = model.load_weights( + reversed( + list( + zip( + [f"{prefix}.{source}.{suffix}" for source in sources], + pieces, + ) + ) + ) + ) + assert loaded == {f"{prefix}.{target}.{suffix}"} + expected = torch.stack(pieces) if scalar else torch.cat(pieces) + torch.testing.assert_close(param.float(), expected.float()) + + +@pytest.mark.parametrize("suffix", ["weight", "weight_scale"]) +def test_quantized_loader_rejects_oversized_shards(quantized_model, suffix): + param = getattr(quantized_model.backbone.layers[0].mixer.qkv_proj, suffix) + oversized = torch.zeros((65, param.shape[1]), dtype=param.dtype) + with pytest.raises(ValueError, match="expected shape"): + quantized_model.load_weights( + [(f"backbone.layers.0.mixer.q_proj.{suffix}", oversized)] + ) diff --git a/tests/transformers_utils/test_dory_config.py b/tests/transformers_utils/test_dory_config.py new file mode 100644 index 000000000000..b158f1d20f45 --- /dev/null +++ b/tests/transformers_utils/test_dory_config.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +from huggingface_hub.errors import ( + StrictDataclassClassValidationError, + StrictDataclassFieldValidationError, +) + +from vllm.config import ModelConfig +from vllm.transformers_utils.config import get_config +from vllm.transformers_utils.configs.dory import DoryConfig + + +def test_dory_loads_without_remote_code_and_preserves_head_size(tmp_path): + config = DoryConfig( + n_input_layers=16, + n_recurrent_layers=40, + n_output_layers=16, + architectures=["DoryForCausalLM"], + ) + config.save_pretrained(tmp_path) + restored = get_config(tmp_path, trust_remote_code=False) + assert isinstance(restored, DoryConfig) + assert restored.recurrent_kv_cache_mode == "per_loop" + assert restored.layer_types is not None + assert restored.layer_types[:8] == ["sliding_attention", "mlp"] * 3 + [ + "full_attention", + "mlp", + ] + assert restored.sliding_window == 1024 + assert DoryConfig.from_dict(restored.to_dict()).to_dict() == restored.to_dict() + # 2560 / 16 = 160 is NOT the checkpoint's attention head dimension. + model_config = ModelConfig( + model=str(tmp_path), skip_tokenizer_init=True, max_model_len=128 + ) + assert model_config.get_head_size() == 256 + + +def test_dory_last_loop_cache_can_be_selected_with_hf_overrides(tmp_path): + DoryConfig(architectures=["DoryForCausalLM"]).save_pretrained(tmp_path) + config = ModelConfig( + model=str(tmp_path), + skip_tokenizer_init=True, + max_model_len=128, + hf_overrides={"recurrent_kv_cache_mode": "last_loop"}, + ).hf_config + assert config.recurrent_kv_cache_mode == "last_loop" + assert DoryConfig.from_dict(config.to_dict()).recurrent_kv_cache_mode == "last_loop" + + +@pytest.mark.parametrize("num_layers", [2, 8]) +def test_dory_rope_theta_defaults_and_explicit_values(num_layers): + config = DoryConfig.from_dict({"num_hidden_layers": num_layers}) + assert all(rope["rope_theta"] == 10000 for rope in config.rope_parameters.values()) + + config = DoryConfig.from_dict( + {"num_hidden_layers": num_layers, "rope_theta": 500000} + ) + assert all(rope["rope_theta"] == 500000 for rope in config.rope_parameters.values()) + + # Explicit checkpoint values take precedence over a top-level default. + config.rope_parameters["sliding_attention"]["rope_theta"] = 10000 + restored = DoryConfig.from_dict(config.to_dict() | {"rope_theta": 1000000}) + assert restored.rope_parameters == config.rope_parameters + + +@pytest.mark.parametrize( + "overrides,match", + [ + ({"n_recurrent_layers": 40}, "sum to num_hidden_layers"), + ({"n_recurrent_loops": 0}, "n_recurrent_loops"), + ({"n_recurrent_loops": 1.5}, "n_recurrent_loops"), + ({"recurrent_kv_cache_mode": "loop1"}, "recurrent_kv_cache_mode"), + ({"layer_types": []}, "num_hidden_layers entries"), + ({"layer_types": ["mamba"] * 72}, "layer_types"), + ({"layer_types": [None] * 72}, "layer_types"), + ({"sliding_window": 0}, "sliding_window"), + ({"sliding_window": None}, "sliding_window"), + ({"rope_parameters": {}}, "each attention block type"), + ( + { + "rope_parameters": { + kind: { + "rope_type": "default", + "rope_theta": 10000, + "partial_rotary_factor": 0, + } + for kind in ("sliding_attention", "full_attention") + } + }, + "Rotary dimensions", + ), + ( + { + "rope_parameters": { + kind: {"rope_type": "linear", "rope_theta": 10000, "factor": 2.0} + for kind in ("sliding_attention", "full_attention") + } + }, + "default RoPE", + ), + ({"initializer_range": 0}, "initializer_range"), + ({"num_key_value_heads": 0}, "Attention dimensions"), + ({"head_dim": 0}, "Attention dimensions"), + ({"hidden_act": "relu"}, "SiLU"), + ({"attention_bias": True}, "bias-free"), + ], +) +def test_dory_rejects_incompatible_architecture(overrides, match): + with pytest.raises( + ( + ValueError, + StrictDataclassClassValidationError, + StrictDataclassFieldValidationError, + ), + match=match, + ): + DoryConfig(**overrides) diff --git a/vllm/model_executor/models/dory.py b/vllm/model_executor/models/dory.py new file mode 100644 index 000000000000..c8fb2f25c475 --- /dev/null +++ b/vllm/model_executor/models/dory.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Dory: normalized latent recurrent transformer.""" + +import math +from collections.abc import Iterable + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.attention import Attention +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.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 default_weight_loader +from vllm.model_executor.parameter import PerTensorScaleParameter +from vllm.sequence import IntermediateTensors + +from .utils import maybe_prefix + +EPS = 1e-6 + + +def justnorm(x: torch.Tensor, dim: int = -1) -> torch.Tensor: + """NGPT normalization: x * rsqrt(sum(x^2) + eps), fp32 accumulation.""" + xf = x.float() + y = xf * torch.rsqrt(xf.pow(2).sum(dim=dim, keepdim=True) + EPS) + return y.to(x.dtype) + + +def slerp_norm(a: torch.Tensor, b: torch.Tensor, lr: torch.Tensor) -> torch.Tensor: + ya, yb = justnorm(a).float(), justnorm(b).float() + return justnorm(ya + lr.float() * (yb - ya)).to(a.dtype) + + +def slerp_alpha(a, b, alpha, alpha_scale): + return slerp_norm(a, b, (alpha.float() * alpha_scale).abs()) + + +class DoryAttention(nn.Module): + """GQA attention with nGPT qk-norm, output gate, and fngpt justnorm.""" + + def __init__( + self, config, layer_idx, cache_config, quant_config, prefix, num_loops=1 + ): + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + std = config.initializer_range + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.num_heads, + self.num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.q_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + self.o_proj = RowParallelLinear( + self.q_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.sqk = nn.Parameter(torch.empty(self.q_size)) + self.s_gate = nn.Parameter(torch.empty(self.q_size)) + self._sqk_factor = self.head_dim**0.5 / std + g0 = config.attention_output_gate_logit_scale_init + if g0 is None: + g0 = config.hidden_size**0.5 + self._s_gate_factor = g0**0.5 / std + + layer_type = config.layer_types[layer_idx] + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters[layer_type], + is_neox_style=True, + ) + + self.per_loop_kv = config.recurrent_kv_cache_mode == "per_loop" + # Match LoopCoder's virtual layer indices for independent loop caches. + # Projection prefixes keep the physical layer index to share weights. + self.attn = nn.ModuleList() + for loop_idx in range(num_loops if self.per_loop_kv else 1): + unique_layer_idx = loop_idx * config.num_hidden_layers + layer_idx + unique_prefix = prefix.replace( + f".layers.{layer_idx}.", f".layers.{unique_layer_idx}." + ) + self.attn.append( + Attention( + self.num_heads, + self.head_dim, + scale=1.0 / math.sqrt(self.head_dim), + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=( + config.sliding_window + if layer_type == "sliding_attention" + else None + ), + prefix=f"{unique_prefix}.attn", + ) + ) + + def forward(self, positions, hidden_states, loop_idx=0): + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + gate, _ = self.g_proj(hidden_states) + + q, k = self.rotary_emb(positions, q, k) + + # nGPT qk-norm AFTER rope: q scaled by (sqk*f)^2 post-justnorm, k plain justnorm + t = q.shape[0] + q = q.view(t, self.num_heads, self.head_dim) + sqk = (self.sqk.float() * self._sqk_factor).view( + 1, self.num_heads, self.head_dim + ) + q = (sqk.square() * justnorm(q).float()).to(q.dtype).view(t, self.q_size) + k = k.view(t, self.num_kv_heads, self.head_dim) + k = justnorm(k).view(t, self.kv_size) + + # In last-loop mode every loop writes this step's KV to the same cache, + # overwriting the previous loop's. Loop i thus attends to loop-i KV for + # this step's tokens and last-loop KV for earlier tokens. + out = self.attn[loop_idx if self.per_loop_kv else 0](q, k, v) + + # output gate (sigmoid, pre-o_proj), then fngpt justnorm over np*hn + sg = (self.s_gate.float() * self._s_gate_factor).square() + out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype) + out = justnorm(out) + out, _ = self.o_proj(out) + return out + + +class DoryMLP(nn.Module): + """SwiGLU with nGPT suv scaling and fngpt justnorm.""" + + def __init__(self, config, quant_config, prefix): + super().__init__() + self.ffn = config.intermediate_size + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [self.ffn] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + self.ffn, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + self.suv_gate = nn.Parameter(torch.empty(self.ffn)) + self.suv_up = nn.Parameter(torch.empty(self.ffn)) + self._suv_factor = config.hidden_size**0.5 + + def forward(self, hidden_states): + m, _ = self.gate_up_proj(hidden_states) + g, u = m.chunk(2, dim=-1) + g = (g.float() * (self.suv_gate.float() * self._suv_factor)).to(g.dtype) + u = (u.float() * (self.suv_up.float() * self._suv_factor)).to(u.dtype) + m = (F.silu(g.float()) * u.float()).to(m.dtype) + m = justnorm(m) + out, _ = self.down_proj(m) + return out + + +class DoryDecoderLayer(nn.Module): + """One physical layer: a single sublayer (attention OR mlp) + slerp residual.""" + + def __init__( + self, config, layer_idx, cache_config, quant_config, prefix, num_loops=1 + ): + super().__init__() + self.is_attention = config.layer_types[layer_idx] != "mlp" + self._alpha_scale = config.ngpt_alpha_init / config.initializer_range + if self.is_attention: + self.mixer = DoryAttention( + config, + layer_idx, + cache_config, + quant_config, + f"{prefix}.mixer", + num_loops=num_loops, + ) + self.attn_alpha = nn.Parameter(torch.empty(config.hidden_size)) + else: + self.mixer = DoryMLP(config, quant_config, f"{prefix}.mixer") + self.mlp_alpha = nn.Parameter(torch.empty(config.hidden_size)) + + def forward(self, positions, hidden_states, loop_idx=0): + residual = hidden_states + if self.is_attention: + out = self.mixer(positions, hidden_states, loop_idx=loop_idx) + return slerp_alpha(residual, out, self.attn_alpha, self._alpha_scale) + out = self.mixer(hidden_states) + return slerp_alpha(residual, out, self.mlp_alpha, self._alpha_scale) + + +@support_torch_compile +class DoryBackbone(nn.Module): + def __init__(self, config, cache_config, quant_config, prefix): + super().__init__() + self.config = config + self.num_loops = config.n_recurrent_loops + n_in = config.n_input_layers + n_rec = config.n_recurrent_layers + self.embeddings = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + DoryDecoderLayer( + config, + i, + cache_config, + quant_config, + f"{prefix}.layers.{i}", + num_loops=self.num_loops if n_in <= i < n_in + n_rec else 1, + ) + for i in range(config.num_hidden_layers) + ] + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + h = inputs_embeds if inputs_embeds is not None else self.embeddings(input_ids) + c = self.config + n_in, n_rec = c.n_input_layers, c.n_recurrent_layers + + for layer in self.layers[:n_in]: + h = layer(positions, h) + input_h = h + + state = input_h + for i in range(self.num_loops): + combined = justnorm(state + input_h) + h = combined + for layer in self.layers[n_in : n_in + n_rec]: + h = layer(positions, h, loop_idx=i) + state = justnorm(h) + + h = state + for layer in self.layers[n_in + n_rec :]: + h = layer(positions, h) + return h + + +class DoryForCausalLM(nn.Module): + 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 + self.config = config + # hf_overrides can change architecture fields after config initialization. + config.validate() + if get_tensor_model_parallel_world_size() != 1: + raise ValueError("Dory currently supports tensor_parallel_size=1 only") + self.backbone = DoryBackbone( + config, + vllm_config.cache_config, + vllm_config.quant_config, + prefix=maybe_prefix(prefix, "backbone"), + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + bias=False, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.lm_head.logit_scale = nn.Parameter(torch.empty(config.vocab_size)) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids): + return self.backbone.embeddings(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ): + return self.backbone(input_ids, positions, inputs_embeds=inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is not None: + logits = logits.float() * ( + self.lm_head.logit_scale.float() / self.config.initializer_range + ) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked = [ + # (target, artifact shard name, shard id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params = dict(self.named_parameters()) + loaded: set[str] = set() + for name, w in weights: + mapped = False + for target, shard_name, shard_id in stacked: + if f".{shard_name}." not in name: + continue + pname = name.replace(f".{shard_name}.", f".{target}.") + param = params[pname] + if target == "qkv_proj": + heads = ( + self.config.num_attention_heads + if shard_id == "q" + else self.config.num_key_value_heads + ) + rows = heads * self.config.head_dim + else: + rows = self.config.intermediate_size + if isinstance(param, PerTensorScaleParameter): + shape = () if w.ndim == 0 else (1,) + else: + # Preserve packed input columns and per-block scale columns. + shape = (rows, *param.shape[1:]) + self._check_weight_shape(name, w, shape) + param.weight_loader(param, w, shard_id) + loaded.add(pname) # report the MODEL param name, not the artifact name + mapped = True + break + if mapped: + continue + + param = params[name] + shape = ( + (self.config.vocab_size, self.config.hidden_size) + if name in ("backbone.embeddings.weight", "lm_head.weight") + else tuple(param.shape) + ) + if isinstance(param, PerTensorScaleParameter) and w.ndim == 0: + shape = () + self._check_weight_shape(name, w, shape) + getattr(param, "weight_loader", default_weight_loader)(param, w) + loaded.add(name) + return loaded + + @staticmethod + def _check_weight_shape(name, weight, expected): + # TP loaders may narrow oversized tensors; reject incompatible exports. + if tuple(weight.shape) != expected: + raise ValueError( + f"Dory weight {name}: expected shape {expected}, " + f"got {tuple(weight.shape)}" + ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 55a18ceba6ff..22bae6d37d69 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -98,6 +98,7 @@ ), "DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"), "DomynEdgeForCausalLM": ("domyn_edge", "DomynEdgeForCausalLM"), + "DoryForCausalLM": ("dory", "DoryForCausalLM"), "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"), "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c5c9438ac56a..e18fc2753cc0 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -131,6 +131,7 @@ def __getitem__(self, key): qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig", deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", + dory="DoryConfig", dots3_note="Dots3NoteConfig", k3_dspark="K3DSparkConfig", funaudiochat="FunAudioChatConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 188e39c115c6..3042e1f5aa10 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -27,6 +27,7 @@ "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV41Config": "vllm.transformers_utils.configs.deepseek_v41", + "DoryConfig": "vllm.transformers_utils.configs.dory", "Dots3NoteConfig": "vllm.transformers_utils.configs.dots3_note", "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark", "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr", @@ -100,6 +101,7 @@ "Qwen3VLNemotronEmbedConfig", "DeepseekVLV2Config", "DeepseekV3Config", + "DoryConfig", "Dots3NoteConfig", "K3DSparkConfig", "DotsOCRConfig", diff --git a/vllm/transformers_utils/configs/dory.py b/vllm/transformers_utils/configs/dory.py new file mode 100644 index 000000000000..ed47dded1f5a --- /dev/null +++ b/vllm/transformers_utils/configs/dory.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Configuration for the normalized latent recurrent Dory architecture.""" + +from huggingface_hub.dataclasses import strict +from transformers.configuration_utils import PreTrainedConfig + + +@strict +class DoryConfig(PreTrainedConfig): + model_type = "dory" + default_theta = 10_000.0 + + vocab_size: int = 131072 + hidden_size: int = 2560 + num_hidden_layers: int = 72 + intermediate_size: int = 7168 + num_attention_heads: int = 16 + num_key_value_heads: int = 4 + head_dim: int = 256 + max_position_embeddings: int = 8192 + tie_word_embeddings: bool = False + hidden_act: str = "silu" + # Also normalizes the learned nGPT scales at inference time. + initializer_range: float = 0.02 + sliding_window: int | None = 1024 + rope_parameters: dict | None = None + layer_types: list[str] | None = None + ngpt_alpha_init: float = 0.05 + attention_output_gate_logit_scale_init: float | None = None + mlp_bias: bool = False + attention_bias: bool = False + n_input_layers: int = 0 + n_recurrent_layers: int | None = None + n_output_layers: int = 0 + n_recurrent_loops: int = 1 + # "last_loop" reuses one cache, overwritten by each recurrent loop. + recurrent_kv_cache_mode: str = "per_loop" + + def __post_init__(self, **kwargs): + if self.layer_types is None: + self.layer_types = [ + "mlp" + if i % 2 + else "full_attention" + if i % 8 == 6 + else "sliding_attention" + for i in range(self.num_hidden_layers) + ] + if self.rope_parameters is None: + self.rope_parameters = { + kind: { + "rope_type": "default", + "partial_rotary_factor": 1.0, + } + for kind in ("sliding_attention", "full_attention") + if kind in self.layer_types + } + if self.n_recurrent_layers is None: + self.n_recurrent_layers = ( + self.num_hidden_layers - self.n_input_layers - self.n_output_layers + ) + super().__post_init__(**kwargs) + + def validate_layer_type(self): + # Each physical block is attention-only or MLP-only. + valid_types = {"sliding_attention", "full_attention", "mlp"} + if self.layer_types is None or any( + kind not in valid_types for kind in self.layer_types + ): + raise ValueError( + "layer_types must be a list of sliding_attention, " + "full_attention, or mlp" + ) + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError("layer_types must have num_hidden_layers entries") + + def validate_architecture(self): + if self.layer_types is None or self.n_recurrent_layers is None: + raise ValueError("Dory layer layout must be initialized") + partitions = ( + self.n_input_layers, + self.n_recurrent_layers, + self.n_output_layers, + ) + if any(n < 0 for n in partitions) or sum(partitions) != self.num_hidden_layers: + raise ValueError( + "Dory input, recurrent, and output layers must sum to num_hidden_layers" + ) + if self.n_recurrent_loops < 1: + raise ValueError("n_recurrent_loops must be a positive integer") + if self.recurrent_kv_cache_mode not in ("per_loop", "last_loop"): + raise ValueError("recurrent_kv_cache_mode must be per_loop or last_loop") + if any( + n <= 0 + for n in (self.head_dim, self.num_attention_heads, self.num_key_value_heads) + ): + raise ValueError("Attention dimensions must be positive integers") + if self.num_attention_heads % self.num_key_value_heads: + raise ValueError( + "num_attention_heads must be divisible by num_key_value_heads" + ) + if self.initializer_range <= 0: + raise ValueError("initializer_range must be positive") + if "sliding_attention" in self.layer_types and ( + self.sliding_window is None or self.sliding_window <= 0 + ): + raise ValueError("sliding_window must be a positive integer") + if self.attention_bias or self.mlp_bias: + raise ValueError("Dory requires bias-free projections") + if self.hidden_act != "silu" or self.tie_word_embeddings: + raise ValueError("Dory requires SiLU and untied embeddings") + + def validate_rope(self): + if self.layer_types is None: + raise ValueError("Dory layer layout must be initialized") + attention_types = set(self.layer_types) - {"mlp"} + if self.rope_parameters is None or ( + set(self.rope_parameters) - {"sliding_attention", "full_attention"} + or attention_types - self.rope_parameters.keys() + ): + raise ValueError("rope_parameters must define each attention block type") + for rope in self.rope_parameters.values(): + if ( + not isinstance(rope, dict) + or rope.get("rope_type") != "default" + or rope.get("rope_theta", 0) <= 0 + ): + raise ValueError( + "Dory requires default RoPE with a positive rope_theta" + ) + fraction = rope.get("partial_rotary_factor", 1.0) + rotary_dim = int(self.head_dim * fraction) + if not 0 < fraction <= 1 or rotary_dim <= 0 or rotary_dim % 2: + raise ValueError( + "Rotary dimensions must be positive, even, and <= head_dim" + ) From 4a81a5097631ad9851a7045ebb4f2f8c4873eec0 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Sun, 25 Oct 2026 20:28:31 +0800 Subject: [PATCH 046/631] [Bugfix][KV Offload] Enable NIXL telemetry for the OBJ tier (#55576) Signed-off-by: Wang Zupeng Co-authored-by: Codex Co-authored-by: vllm-agent --- tests/v1/kv_offload/tiering/test_obj_tier.py | 25 +++++++++++++------- vllm/v1/kv_offload/tiering/obj/manager.py | 2 +- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index fd500c0d5000..2dba5baf3d6f 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -139,7 +139,8 @@ class MockNixlAgent: check_xfer_state: Callable query_memory: Callable - def __init__(self): + def __init__(self, agent_name, config): + self._capture_telemetry = config.capture_telemetry self._stored_obj_keys: set[str] = set() # handle_id -> (op, [obj_keys]) self._pending: dict[int, tuple[str, list[str]]] = {} @@ -203,6 +204,8 @@ def release_dlist_handle(self, handle): pass def get_xfer_telemetry(self, handle): + if not self._capture_telemetry: + raise RuntimeError("Telemetry capture is disabled") return SimpleNamespace(xferDuration=1000) def _query_memory(self, queries, mem_type, agent_name): @@ -232,15 +235,19 @@ def _make_tier( **tier_kwargs, ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: """Create a tier backed by a fresh MockNixlAgent.""" - mock_agent = MockNixlAgent() 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( + "vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config", + side_effect=lambda backends, capture_telemetry=False: SimpleNamespace( + capture_telemetry=capture_telemetry + ), + ), patch( "vllm.v1.kv_offload.tiering.obj.manager.nixl_agent", - return_value=mock_agent, + side_effect=MockNixlAgent, ), ): tier = ObjectStoreSecondaryTierManager( @@ -251,7 +258,7 @@ def _make_tier( prefix=_RUN_PREFIX, **tier_kwargs, ) - return tier, mock_agent + return tier, tier._agent def drain( @@ -315,14 +322,16 @@ def test_lookup_unrelated_key_returns_false(self): assert lookup_and_wait(self.tier, [key(999)]) == [LookupResult.MISS] def test_store_then_load_roundtrip(self): + """Completed stores and loads report transfer timing without an env override.""" self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) results = drain(self.tier) - assert results[0].success + assert results == [JobResult(job_id=1, success=True, transfer_time=0.001)] self.tier.submit_load(make_job(2, [key(1), key(2)], [0, 1])) results = drain(self.tier) - assert len(results) == 1 - assert results[0].success + assert results == [JobResult(job_id=2, success=True, transfer_time=0.001)] + assert not self.tier._transfers + assert list(self.tier.get_finished_jobs()) == [] def test_multiple_jobs_tracked_independently(self): self.tier.submit_store(make_job(1, [key(1)], [0])) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 1886e3b81168..26e78894742a 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -155,7 +155,7 @@ def __init__( # mark its own cached lookup verdicts False (see get_finished_jobs). self._load_job_keys: dict[JobId, list[OffloadKey]] = {} - agent_config = nixl_agent_config(backends=[]) + agent_config = nixl_agent_config(backends=[], capture_telemetry=True) self._agent = nixl_agent("ObjAgent", agent_config) obj_config = ObjStoreConfig(**store_config) params = {**obj_config.to_nixl_params(), "num_threads": str(io_threads)} From 503e4e16ad8d4c4b7b1290c9f2792af0ca622c47 Mon Sep 17 00:00:00 2001 From: Pepsi-Cola-qqq <168319770+Pepsi-Cola-qqq@users.noreply.github.com> Date: Sun, 25 Oct 2026 20:28:46 +0800 Subject: [PATCH 047/631] [Bugfix][KV Offload] Catch async lookup iterator errors (#63126) Signed-off-by: w_wanghaiqiao Co-authored-by: w_wanghaiqiao Co-authored-by: vllm-agent --- .../kv_offload/tiering/test_async_lookup.py | 32 +++++++++++++++++++ vllm/v1/kv_offload/tiering/async_lookup.py | 2 +- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index 97f9e1539348..d23d837371bd 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -49,6 +49,38 @@ def batch_lookup( class TestAsyncLookupManager: + @pytest.mark.parametrize("fail_after", [0, 1]) + def test_iterator_failure_returns_misses_and_preserves_progress(self, fail_after): + class FailingIteratorLookupManager(AsyncLookupManager): + def batch_lookup( + self, keys: list[OffloadKey], req_context: ReqContext + ) -> Iterable[bool]: + for index, _ in enumerate(keys): + if req_context.req_id == "failing" and index == fail_after: + raise OSError("storage lookup failed during iteration") + yield True + + mgr = FailingIteratorLookupManager(tier_type="test") + failing_ctx = _ctx("failing") + healthy_ctx = _ctx("healthy") + later_ctx = _ctx("later") + try: + assert mgr.lookup(_key(1), failing_ctx) is None + assert mgr.lookup(_key(2), failing_ctx) is None + assert mgr.lookup(_key(3), healthy_ctx) is None + mgr.flush() + + assert mgr.lookup(_key(4), later_ctx) is None + mgr.flush() + finally: + # The FIFO shutdown sentinel waits for both batches and drains results. + mgr.shutdown() + + assert mgr.lookup(_key(1), failing_ctx) is False + assert mgr.lookup(_key(2), failing_ctx) is False + assert mgr.lookup(_key(3), healthy_ctx) is True + assert mgr.lookup(_key(4), later_ctx) is True + def test_new_key_returns_none(self): mgr = InMemoryLookupManager() assert mgr.lookup(_key(1), _ctx()) is None diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index 168a194cf7c8..69a513f55de8 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -272,7 +272,7 @@ def _worker(self) -> None: for req_context, entries in batches.values(): keys = [key for key, _ in entries] try: - hits = self.batch_lookup(keys, req_context) + hits: Iterable[bool] = list(self.batch_lookup(keys, req_context)) except Exception as exc: logger.warning( "batch_lookup failed on tier %s for %d keys: %s", From 94a29d316e8c1d6dcf6ff75645c2e61c7e3c83fd Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 25 Oct 2026 21:03:40 +0800 Subject: [PATCH 048/631] [Fast Start] Add daemon-to-daemon warmup cache seeding for Preload (#63088) Signed-off-by: Isotr0py Co-authored-by: Xun Sun Co-authored-by: liusy58 Co-authored-by: Kimi Code --- docs/features/preload.md | 55 ++++ .../model_loader/test_weight_cache.py | 279 +++++++++++++++++- .../test_flashinfer_autotune_warmup.py | 31 +- vllm/entrypoints/cli/preload.py | 140 +++++++-- .../model_loader/weight_cache/daemon.py | 53 ++-- .../model_loader/weight_cache/protocol.py | 42 ++- .../model_loader/weight_cache/seed.py | 232 +++++++++++++++ .../model_loader/weight_cache/warmup_cache.py | 183 ++++++++++++ .../warmup/flashinfer_autotune_cache.py | 44 ++- vllm/model_executor/warmup/kernel_warmup.py | 45 +-- 10 files changed, 990 insertions(+), 114 deletions(-) create mode 100644 vllm/model_executor/model_loader/weight_cache/seed.py create mode 100644 vllm/model_executor/model_loader/weight_cache/warmup_cache.py diff --git a/docs/features/preload.md b/docs/features/preload.md index 31aa9e2db3ec..5318cc447daa 100644 --- a/docs/features/preload.md +++ b/docs/features/preload.md @@ -70,6 +70,50 @@ loading and before binding its socket, then writes the cache: vllm preload --model /path/to/model --tensor-parallel-size 4 ``` +### Seeding warmup caches between daemons (D2D) + +A fresh daemon replica — a new node joining a fleet, a restarted preload +set — would normally pay the autotune pass itself. With a seed source +configured, each daemon rank fetches its rank-matched warmup caches from the +source replica over an authenticated TCP control plane before tuning, +validates them, and installs them into its own on-disk cache, so its warmup +skips the profiling sweep. Any failure (unreachable, missing, incompatible) +falls back to local tuning. + +Source replica: + +```bash +vllm preload --model /path/to/model --tensor-parallel-size 4 \ + --seed-listen-host 0.0.0.0 --seed-listen-port 29600 \ + --seed-token "$TOKEN" +``` + +Mirror replica (possibly on another node): + +```bash +vllm preload --model /path/to/model --tensor-parallel-size 4 \ + --seed-source-host --seed-source-port 29600 \ + --seed-token "$TOKEN" +``` + +Each rank binds and dials `base_port + group offset + its global rank`, so one +`host:base_port` pair covers a whole replica, including the draft daemon group +when speculative decoding is configured. + +The exchange is keyed by the on-disk cache layout: a request names the cache +by kind, configuration hash, DP rank and world rank, and the source resolves +it to files under its cache root. Before installing, the target rejects a +payload whose FlashInfer environment metadata (GPU, CUDA/cuDNN stack) differs +from its own — such a table would never load anyway, and a hard conflict is +never overwritten, which would poison the cache file into re-tuning on every +start. Tables are shareable across vLLM versions as long as the configuration +hash matches. + +The channel is a generic seeding point: warmup caches are the first kind, and +further daemon state (e.g. Triton warmup caches or weight manifests) plugs in +as additional kinds — small JSON payloads travel inline, while large binary +state rides the data plane. + ## How it works 1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its @@ -193,6 +237,10 @@ Socket paths are derived from the GPU UUID, so they are stable regardless of passes, and a table tuned on the draft model would key on the draft's config, which no engine looks up), so those deployments' first engine tunes and fills the cache instead. +- **Warmup cache seeding**: only the v1 per-rank FlashInfer autotune cache is + seeded; single-rank daemons use FlashInfer's managed v2 store and always + tune in place. Seeding requires a trusted network — the TCP channel is + authenticated by a shared token but unencrypted. - **Quantization**: every quantization method in the model must declare support for pre-processed weights (the daemon transfers weights *after* quantization post-processing). Unsupported methods raise @@ -214,6 +262,13 @@ processes owned by the same user: directory is additionally rejected if it is group/world accessible. - On Linux the daemon verifies the connecting peer's UID via `SO_PEERCRED`. +The optional daemon-to-daemon listener (`--seed-listen-*`) is a +separate TCP channel: requests are JSON-only and size-capped, authenticate +with the shared `--seed-token` (constant-time comparison), and +an unauthenticated peer that stops sending is cut off by a whole-message +deadline instead of blocking the accept loop. The channel is unencrypted — +bind it to a trusted network. + See the [security documentation](../usage/security.md) for vLLM's general threat model. diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 5b8112477422..fc9af29c1f5e 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -7,6 +7,7 @@ identical outputs. """ +import os import shutil import subprocess import sys @@ -15,7 +16,7 @@ import time import urllib.error import urllib.request -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any import pytest @@ -36,6 +37,8 @@ def __init__( model: str, tp_size: int, extra_args: list[str] | None = None, + env_extra: dict[str, str] | None = None, + log_file: str | None = None, ): # Short base path: Unix socket paths are limited to ~107 characters. self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_") @@ -58,21 +61,41 @@ def __init__( str(self.health_port), *(extra_args or []), ] + self._env_extra = env_extra + self._log_file = log_file + self._log_chunks: list[str] = [] self._proc: subprocess.Popen | None = None def __enter__(self) -> "WeightCacheDaemon": + env = {**os.environ, **self._env_extra} if self._env_extra is not None else None self._proc = subprocess.Popen( - self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True + self._cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, ) - threading.Thread(target=self._drain_stderr, daemon=True).start() + threading.Thread(target=self._drain_output, daemon=True).start() return self def __exit__(self, *exc_info) -> None: self._stop() - def _drain_stderr(self) -> None: - assert self._proc is not None and self._proc.stderr is not None - self._proc.stderr.read() + @property + def logs(self) -> str: + return "".join(self._log_chunks) + + def _drain_output(self) -> None: + assert self._proc is not None and self._proc.stdout is not None + if self._log_file is None: + for line in self._proc.stdout: + self._log_chunks.append(line) + return + with open(self._log_file, "w") as out: + for line in self._proc.stdout: + self._log_chunks.append(line) + out.write(line) + out.flush() def check_health(self) -> None: url = f"http://127.0.0.1:{self.health_port}/health" @@ -270,6 +293,85 @@ def test_ipc_cache_cold_start_and_warm_restart_tp2_ep(vllm_runner, tmp_path): run_cold_start_and_warm_restart(vllm_runner, DSV2_EP_CASE, tmp_path) +@create_new_process_for_each_test() +@pytest.mark.distributed(num_gpus=2) +def test_d2d_warmup_cache_seeding(vllm_runner, tmp_path): + """A mirror daemon adopts the source's autotune cache over D2D, skips + re-tuning, and serves outputs identical to the default-loader baseline.""" + if not current_platform.is_cuda_alike(): + pytest.skip("Weight cache IPC sharing requires CUDA or ROCm") + if not current_platform.is_device_capability_family(100): + pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend") + from vllm.utils.flashinfer import has_flashinfer + + if not (has_flashinfer() and current_platform.has_device_capability(90)): + pytest.skip("FlashInfer autotune requires FlashInfer and SM90+") + + case = replace( + K3_CASE, + llm_kwargs={**K3_CASE.llm_kwargs, "tensor_parallel_size": 2}, + ) + baseline_outputs = generate(vllm_runner, case, None, fallback=True) + assert all(text for _, texts in baseline_outputs for text in texts) + + token = "d2d-test-token" + base_port = get_open_port() + log_dir = os.environ.get("D2D_E2E_LOG_DIR") + log_file = ( + (lambda name: os.path.join(log_dir, name)) if log_dir else lambda name: None + ) + # The v1 per-rank cache path needs a multi-rank tuning group (v2 is gated + # to single-rank), so both daemons run at tp=2. + with WeightCacheDaemon( + case.model, + tp_size=2, + extra_args=[ + *case.daemon_args, + "--seed-listen-host", + "127.0.0.1", + "--seed-listen-port", + str(base_port), + "--seed-token", + token, + ], + env_extra={"VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR": str(tmp_path / "source")}, + log_file=log_file("source.log"), + ) as source: + source.wait_until_ready() + # The source had an empty cache dir, so it must have tuned in place. + assert "Running FlashInfer autotune" in source.logs + + with WeightCacheDaemon( + case.model, + tp_size=2, + extra_args=[ + *case.daemon_args, + "--seed-source-host", + "127.0.0.1", + "--seed-source-port", + str(base_port), + "--seed-token", + token, + ], + env_extra={"VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR": str(tmp_path / "mirror")}, + log_file=log_file("mirror.log"), + ) as mirror: + mirror.wait_until_ready() + # The mirror adopted both ranks' tables and skipped the sweep. + assert mirror.logs.count("Adopted remote flashinfer_autotune seed") == 2 + assert "Running FlashInfer autotune" not in mirror.logs + mirror_tables = sorted( + (tmp_path / "mirror").rglob("autotune_configs_dp0_rank*.json") + ) + assert len(mirror_tables) == 2 + + mirror_outputs = generate( + vllm_runner, case, mirror.socket_dir, fallback=False + ) + + assert mirror_outputs == baseline_outputs + + def run_cold_start_and_warm_restart(vllm_runner, case: ModelCase, tmp_path, capfd=None): # Baseline: plain disk loading with the default loader. baseline_outputs = generate( @@ -802,3 +904,168 @@ def request_state(key, *, is_reload=False): assert len(requested_keys) == 2 for key in requested_keys: assert server.cache_key.mismatched_fields(key) == mismatched_fields + + +def _warmup_cache_request(config_hash: str = "a" * 64, **overrides) -> dict: + request = { + "kind": "flashinfer_autotune", + "config_hash": config_hash, + "dp_rank": 0, + "rank": 0, + } + request.update(overrides) + return request + + +def test_warmup_cache_request_validation_and_miss(tmp_path, monkeypatch): + """The source validates the wire key before touching the filesystem.""" + monkeypatch.setenv("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)) + from vllm.model_executor.model_loader.weight_cache.seed import SeedTransport + from vllm.model_executor.model_loader.weight_cache.warmup_cache import ( + FlashInferWarmupCache, + ) + + def handle(request): + return SeedTransport([FlashInferWarmupCache()]).handle_request(request) + + assert handle("not a dict")["status"] == "error" + assert handle(_warmup_cache_request(kind="bogus"))["status"] == "error" + # A request missing the kind entirely is rejected, not resolved. + assert handle({"token": "t"})["status"] == "error" + # A path traversal or malformed hash must never reach the filesystem. + assert handle(_warmup_cache_request(config_hash="../x"))["status"] == "error" + assert handle(_warmup_cache_request(config_hash="A" * 64))["status"] == "error" + # A garbage rank just names a file that does not exist: a clean miss. + assert handle(_warmup_cache_request(rank=-1))["status"] == "miss" + # A well-formed request for a cache that does not exist is a clean miss. + assert handle(_warmup_cache_request())["status"] == "miss" + + +def test_warmup_cache_metadata_gate(): + """Tables flashinfer would refuse to load are rejected before install.""" + pytest.importorskip("flashinfer") + from flashinfer.autotuner.autotuner import _collect_metadata + + from vllm.model_executor.model_loader.weight_cache.warmup_cache import ( + FlashInferWarmupCache, + ) + + kind = FlashInferWarmupCache() + filename = "autotune_configs_dp0_rank0.json" + + matching = {filename: {"_metadata": _collect_metadata(), "records": {}}} + assert kind.validate_payload(matching) == matching + + # Pre-metadata flashinfer files are accepted, like load_configs does. + legacy: dict[str, dict] = {filename: {"records": {}}} + assert kind.validate_payload(legacy) == legacy + + foreign = {"_metadata": {**_collect_metadata(), "gpu": "Not A Real GPU"}} + with pytest.raises(ValueError, match="environment mismatch"): + kind.validate_payload({filename: foreign}) + with pytest.raises(ValueError, match="JSON object"): + kind.validate_payload({filename: [1, 2]}) + with pytest.raises(ValueError, match="exactly one"): + kind.validate_payload({filename: {}, "extra.json": {}}) + + +def test_remote_warmup_cache_listener(tmp_path, monkeypatch): + """Peers need the token, and only get warmup caches for their rank.""" + import json + import socket + + monkeypatch.setenv("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)) + from vllm.model_executor.model_loader.weight_cache.protocol import ( + recv_json, + send_json, + ) + from vllm.model_executor.model_loader.weight_cache.seed import ( + SeedConfig, + SeedTransport, + ) + from vllm.model_executor.model_loader.weight_cache.warmup_cache import ( + FlashInferWarmupCache, + ) + + config_hash = "b" * 64 + table = {"records": {"op": {"shape": 1}}} + cache_dir = tmp_path / config_hash + cache_dir.mkdir() + (cache_dir / "autotune_configs_dp0_rank0.json").write_text(json.dumps(table)) + + transport = SeedTransport( + [FlashInferWarmupCache()], + SeedConfig(listen=("127.0.0.1", 0), token="test-secret"), + ) + address = transport.start_listener() + assert address is not None + try: + cases = [ + (_warmup_cache_request(config_hash, token="wrong"), "error"), + ( + _warmup_cache_request(config_hash, kind="weights", token="test-secret"), + "error", + ), + (_warmup_cache_request(config_hash, token="test-secret"), "ok"), + (_warmup_cache_request("f" * 64, token="test-secret"), "miss"), + ] + for request, expected in cases: + with socket.create_connection(address, timeout=2.0) as peer: + send_json(peer, request) + response = recv_json(peer) + assert response["status"] == expected + if expected == "ok": + assert response["payload"] == {"autotune_configs_dp0_rank0.json": table} + finally: + transport.close() + + +def test_transport_cuts_off_trickling_peer(monkeypatch): + """A peer that never finishes its message cannot block the accept loop.""" + import socket + import time + + import vllm.model_executor.model_loader.weight_cache.seed as seed_mod + from vllm.model_executor.model_loader.weight_cache.protocol import ( + recv_json, + send_json, + ) + from vllm.model_executor.model_loader.weight_cache.seed import ( + SeedConfig, + SeedTransport, + ) + + monkeypatch.setattr(seed_mod, "_REMOTE_REQUEST_TIMEOUT_S", 0.5) + transport = SeedTransport([], SeedConfig(listen=("127.0.0.1", 0), token="t")) + address = transport.start_listener() + assert address is not None + try: + with socket.create_connection(address) as idle: + idle.sendall(b"\x00") # a partial frame, then silence + time.sleep(1.5) # past the whole-message deadline + with socket.create_connection(address, timeout=2.0) as peer: + send_json(peer, {"cmd": "anything", "token": "t"}) + response = recv_json(peer) + # The trickler was cut off and this peer got served (unknown + # command, but the accept loop answered). + assert response["status"] == "error" + finally: + transport.close() + + +def test_warmup_cache_install_writes_table(tmp_path, monkeypatch): + """Install writes the fetched table to the rank's on-disk cache path.""" + import json + + from vllm.model_executor.model_loader.weight_cache.warmup_cache import ( + FlashInferWarmupCache, + ) + from vllm.model_executor.warmup import flashinfer_autotune_cache as fi_cache + + target = tmp_path / "autotune_configs_dp0_rank0.json" + monkeypatch.setattr( + fi_cache, "resolve_rank_flashinfer_autotune_file", lambda _: target + ) + table = {"records": {"op": {"shape": 1}}} + FlashInferWarmupCache().install(None, {"autotune_configs_dp0_rank0.json": table}) + assert json.loads(target.read_text()) == table diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index f8c86cd1b5c1..74ae494ed8f7 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -176,7 +176,11 @@ def test_flashinfer_autotune_lifecycle( model = runner.get_model() cache_path = tmp_path / "autotune.json" rank_cache_path = cache_path.with_name(f"autotune_dp0_rank{rank}.json") - rank_cache_path.write_text(json.dumps({"_metadata": {}, "op": 1})) + # A matching cache lets every rank skip the tuning sweep, so the dummy + # failure case must run on a cache miss to reach the sweep at all. + cache_hit = failure != "dummy" + if cache_hit: + rank_cache_path.write_text(json.dumps({"_metadata": {}, "op": 1})) events = Mock() events.tuner.load_configs.return_value = True events.tuner._dirty = True @@ -226,9 +230,9 @@ def autotune_context(**kwargs): monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: has_v2) monkeypatch.setattr(fi_utils, "autotune", events.autotune) monkeypatch.setattr(warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS", skip_ops) - resolve_file = Mock(return_value=cache_path) + resolve_file = Mock(return_value=rank_cache_path) resolve_root = Mock(return_value=str(tmp_path)) - monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_file", resolve_file) + monkeypatch.setattr(warmup, "resolve_rank_flashinfer_autotune_file", resolve_file) monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_v2_root", resolve_root) def gather(out, value, *, group): @@ -257,22 +261,21 @@ def gather(out, value, *, group): else: resolve_root.assert_not_called() resolve_file.assert_called_once_with(runner) - expected.extend( - [ - call.tuner.load_configs(str(rank_cache_path)), - call.autotune(tune_mode=True, **kwargs), - ] - ) + if cache_hit: + expected.append(call.tuner.load_configs(str(rank_cache_path))) + expected.append(call.autotune(tune_mode=True, **kwargs)) expected.extend( [call.set_group(world.cpu_group if world_size > 1 else None), call.enter()] ) if hisparse: expected.append(call.hisparse(runner)) - for name, arg, kwargs in [ - ("dummy", runner, {"skip_attn": hisparse}), + warmup_calls: list[tuple[str, object, dict]] = [ ("replayssm", runner, {}), ("kimi", model, {}), - ]: + ] + if use_v2 or not cache_hit: + warmup_calls.insert(0, ("dummy", runner, {"skip_attn": hisparse})) + for name, arg, kwargs in warmup_calls: expected.append(getattr(call, name)(arg, **kwargs)) if failure == name: break @@ -567,8 +570,8 @@ def make_run(*, pp=2, tp=4): ) monkeypatch.setattr( warmup, - "resolve_flashinfer_autotune_file", - lambda runner: tmp_path / "autotune_configs.json", + "resolve_rank_flashinfer_autotune_file", + lambda runner: run.cache_path(run.rank), ) monkeypatch.setattr( warmup, "_flashinfer_autotune_skip_ops", lambda runner: None diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index bf5a9fc88c65..ef54163ec5c3 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/vllm/entrypoints/cli/preload.py @@ -100,6 +100,37 @@ def _start_health_server( return server, thread +def _validate_seed_args( + args: argparse.Namespace, *, world_size: int, num_groups: int +) -> None: + """Validate the seed CLI options before any daemon spawns.""" + for flag, host, port in ( + ("--seed-listen", args.seed_listen_host, args.seed_listen_port), + ("--seed-source", args.seed_source_host, args.seed_source_port), + ): + if (host is None) != (port is None): + raise ValueError(f"{flag}-host and {flag}-port must be provided together") + has_endpoint = any( + value is not None + for value in ( + args.seed_listen_host, + args.seed_listen_port, + args.seed_source_host, + args.seed_source_port, + ) + ) + if has_endpoint != bool(args.seed_token): + raise ValueError("Preload daemons must set --seed-token to enable D2D transfer") + for base_port in (args.seed_listen_port, args.seed_source_port): + if base_port is None: + continue + max_port = base_port + num_groups * world_size - 1 + if not 1 <= base_port <= 65535 or max_port > 65535: + raise ValueError( + f"Daemon-to-daemon port range {base_port}-{max_port} exceeds 65535" + ) + + class PreloadSubcommand(CLISubcommand): """The `preload` subcommand for the vLLM CLI.""" @@ -143,6 +174,39 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None: default="0.0.0.0", help="Host for the parent /health endpoint (default: 0.0.0.0).", ) + parser.add_argument( + "--seed-listen-host", + type=str, + default=None, + help="Host to bind for authenticated daemon-to-daemon requests.", + ) + parser.add_argument( + "--seed-listen-port", + type=int, + default=None, + help="Base port for daemon-to-daemon requests; each daemon binds " + "base_port + group offset + its global rank.", + ) + parser.add_argument( + "--seed-source-host", + type=str, + default=None, + help="Host of a source daemon replica to seed warmup caches from.", + ) + parser.add_argument( + "--seed-source-port", + type=int, + default=None, + help="Base port of the source daemon replica; each daemon dials " + "base_port + group offset + its global rank.", + ) + parser.add_argument( + "--seed-token", + type=str, + default=None, + help="Shared secret for daemon-to-daemon requests. The JSON " + "connection is unencrypted; use a trusted network.", + ) @staticmethod def cmd(args: argparse.Namespace) -> None: @@ -153,6 +217,7 @@ def cmd(args: argparse.Namespace) -> None: get_draft_daemon_config, plan_local_ranks, ) + from vllm.model_executor.model_loader.weight_cache.seed import SeedConfig from vllm.model_executor.model_loader.weight_cache.utils import ( format_daemon_role, ) @@ -179,6 +244,7 @@ def cmd(args: argparse.Namespace) -> None: tp_size = parallel_config.tensor_parallel_size pp_size = parallel_config.pipeline_parallel_size dp_size = parallel_config.data_parallel_size + world_size = dp_size * pp_size * tp_size placements = plan_local_ranks(parallel_config) local_world_size = len(placements) @@ -248,32 +314,54 @@ def cmd(args: argparse.Namespace) -> None: groups, placements ) } - procs = [ - ctx.Process( - target=_run_daemon, - args=( - dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank, - local_rank, - config, - init_method, - args.weight_cache_socket_dir, - ready_queue, - is_draft, - dp_rank, - pp_rank, - ), - name=( - f"vllm-weight-cache-{format_daemon_role(is_draft)}-" - f"{(dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank)}" - ), - ) - for (is_draft, config, init_method), ( - local_rank, - dp_rank, - pp_rank, - tp_rank, - ) in product(groups, placements) - ] + _validate_seed_args(args, world_size=world_size, num_groups=len(groups)) + + def _seed_endpoint( + host: str | None, base_port: int | None, offset: int + ) -> tuple[str, int] | None: + # Validated above: host and port are provided together. + if host is None or base_port is None: + return None + return host, base_port + offset + + procs = [] + for group_index, (is_draft, config, init_method) in enumerate(groups): + # Draft daemons live past the target group's port block so both + # roles can be reached over a single seed host:base_port. + port_offset = group_index * world_size + for local_rank, dp_rank, pp_rank, tp_rank in placements: + global_rank = dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank + offset = port_offset + global_rank + seed_config = SeedConfig( + listen=_seed_endpoint( + args.seed_listen_host, args.seed_listen_port, offset + ), + source=_seed_endpoint( + args.seed_source_host, args.seed_source_port, offset + ), + token=args.seed_token, + ) + procs.append( + ctx.Process( + target=_run_daemon, + args=( + global_rank, + local_rank, + config, + init_method, + args.weight_cache_socket_dir, + ready_queue, + is_draft, + dp_rank, + pp_rank, + seed_config, + ), + name=( + f"vllm-weight-cache-{format_daemon_role(is_draft)}-" + f"{global_rank}" + ), + ) + ) health_state = _HealthState(procs) health_server = None health_thread = None diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index b7bf8195da3a..77f6540c4c5e 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -119,15 +119,19 @@ send_msg, verify_peer_is_owner, ) +from vllm.model_executor.model_loader.weight_cache.seed import ( + SeedConfig, + SeedTransport, +) from vllm.model_executor.model_loader.weight_cache.utils import ( - build_warmup_runner, export_model_attrs, format_daemon_role, is_draft_model_cacheable, ) -from vllm.model_executor.warmup.kernel_warmup import flashinfer_autotune +from vllm.model_executor.model_loader.weight_cache.warmup_cache import ( + FlashInferWarmupCache, +) from vllm.platforms import current_platform -from vllm.utils.flashinfer import has_flashinfer from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.worker.workspace import init_workspace_manager @@ -212,10 +216,13 @@ def __init__( is_draft: bool = False, dp_rank: int = 0, pp_rank: int = 0, + seed_config: SeedConfig | None = None, ): parallel_config = vllm_config.parallel_config if parallel_config.enable_eplb: raise ValueError("The weight cache daemon does not support EPLB.") + self.warmup_caches = [FlashInferWarmupCache()] + self.d2d = SeedTransport(self.warmup_caches, seed_config) self.tp_size = parallel_config.tensor_parallel_size self.pp_size = parallel_config.pipeline_parallel_size self.dp_size = parallel_config.data_parallel_size @@ -277,39 +284,9 @@ def __init__( ) def warmup(self) -> None: - """Run the FlashInfer autotune pass against the cached model.""" - vllm_config = self.vllm_config - # Mirrors the gate in the engine's kernel warmup. - if ( - vllm_config.kernel_config.enable_flashinfer_autotune is False - or not has_flashinfer() - or not current_platform.has_device_capability(90) - or is_draft_model_cacheable(vllm_config.speculative_config) - ): - return - try: - current_platform.update_block_size_for_backend(vllm_config) - runner = build_warmup_runner( - vllm_config, - self.local_rank, - is_draft=self.is_draft, - model_config=self.model_config, - ) - assert self.model is not None, "warmup ran before load_model" - runner.load_model(model=self.model) - # The V2 runner requires skip_attn when no KV cache exists - flashinfer_autotune(runner, skip_attn=vllm_config.use_v2_model_runner) - logger.info( - "Weight cache %s daemon rank %d tuned FlashInfer; the tuned " - "configs are in the on-disk autotune cache", - self.role, - self.global_rank, - ) - finally: - # Free the runner and the dummy-run activations; the daemon's - # weights are untouched. - gc.collect() - torch.accelerator.empty_cache() + """Run each warmup cache's pass (adopting remote caches via D2D).""" + for cache in self.warmup_caches: + cache.warmup(self, self.d2d) def load_model(self) -> None: torch.accelerator.set_device_index(self.local_rank) @@ -418,6 +395,7 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non server.bind(socket_path) os.chmod(socket_path, 0o600) server.listen() + self.d2d.start_listener() logger.info( "Weight cache %s daemon rank %d serving on %s", self.role, @@ -445,6 +423,7 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non send_msg(conn, {"status": "error", "message": str(e)}) finally: server.close() + self.d2d.close() if os.path.exists(socket_path): os.unlink(socket_path) os.close(lock_fd) @@ -553,6 +532,7 @@ def _run_daemon( is_draft: bool = False, dp_rank: int = 0, pp_rank: int = 0, + seed_config: SeedConfig | None = None, ) -> None: # Spawned daemons start with logging unconfigured, like engine workers. if logging_config := getattr(vllm_config, "logging_config", None): @@ -566,6 +546,7 @@ def _run_daemon( is_draft, dp_rank, pp_rank, + seed_config=seed_config, ) daemon.load_model() daemon.warmup() diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index b2bca701ea57..fc923e9f2b25 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -19,6 +19,7 @@ import stat import struct import tempfile +import time from dataclasses import dataclass, fields from typing import Any, NamedTuple @@ -46,6 +47,9 @@ # Sanity bound for a single message. IPC handles are tiny; only small # non-CUDA tensors are ever shipped by value. MAX_MSG_SIZE = 1 << 34 +# Bound for one JSON message on the remote warmup-cache control plane. +# Warmup cache tables are small JSON documents. +MAX_JSON_MSG_SIZE = 1 << 26 def _current_uid() -> int: @@ -452,9 +456,45 @@ def recv_msg(sock: socket.socket) -> Any: return pickle.loads(_recv_exact(sock, length)) -def _recv_exact(sock: socket.socket, num_bytes: int) -> bytes: +def send_json(sock: socket.socket, obj: Any) -> None: + payload = json.dumps(obj).encode() + if len(payload) > MAX_JSON_MSG_SIZE: + raise ValueError("JSON message exceeds configured limit") + sock.sendall(_LEN_STRUCT.pack(len(payload))) + sock.sendall(payload) + + +def recv_json( + sock: socket.socket, + max_size: int = MAX_JSON_MSG_SIZE, + deadline: float | None = None, +) -> Any: + """Receive one JSON message. + + Args: + sock: Connected socket. + max_size: Largest payload accepted. + deadline: ``time.monotonic()`` by which the whole message must have + arrived; a per-read timeout alone lets a peer trickle bytes + forever. + + """ + (length,) = _LEN_STRUCT.unpack(_recv_exact(sock, _LEN_STRUCT.size, deadline)) + if length > max_size: + raise ValueError(f"Message size {length} exceeds limit {max_size}") + return json.loads(_recv_exact(sock, length, deadline)) + + +def _recv_exact( + sock: socket.socket, num_bytes: int, deadline: float | None = None +) -> bytes: buf = bytearray() while len(buf) < num_bytes: + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Peer did not send its message in time") + sock.settimeout(remaining) chunk = sock.recv(num_bytes - len(buf)) if not chunk: raise ConnectionError("Socket closed while receiving message") diff --git a/vllm/model_executor/model_loader/weight_cache/seed.py b/vllm/model_executor/model_loader/weight_cache/seed.py new file mode 100644 index 000000000000..9ca542f4405b --- /dev/null +++ b/vllm/model_executor/model_loader/weight_cache/seed.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Daemon-to-daemon seeding over an authenticated JSON channel. + +A :class:`Seedable` is one kind of state a daemon can seed to or from a peer +daemon: warmup caches today, weight shards later. The target names what it +wants by kind and key fields; the source resolves the key to a payload; the +target validates and installs it: + + request: {"kind": "flashinfer_autotune", ..., "token": "..."} + response: {"status": "ok", "payload": } + {"status": "miss"} | {"status": "error", "message": "..."} + +Payloads are JSON values. Large binary state (weights) must not be inlined: +the payload then describes the transfer (e.g. a storage manifest) and the +kind's ``install`` drives the data plane itself. + +The connection is unencrypted beyond the shared token; keep the listener on a +trusted network. +""" + +import contextlib +import hmac +import socket +import threading +import time +from collections.abc import Iterable +from typing import Any, NamedTuple, Protocol + +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_cache.protocol import ( + recv_json, + send_json, +) + +logger = init_logger(__name__) + +# Requests carry a kind and small key fields; payloads flow in responses. +_MAX_REMOTE_REQUEST_SIZE = 64 << 10 +_REQUEST_TIMEOUT_S = 30.0 +# An unauthenticated peer gets a whole-message deadline, not a per-read +# timeout: per-read timeouts let a peer trickle bytes forever and stall the +# serial accept loop. Authenticated handlers (e.g. weight manifests) may be +# slow, so the budget relaxes once the token checks out. +_REMOTE_REQUEST_TIMEOUT_S = 10.0 +_AUTHENTICATED_TIMEOUT_S = 300.0 + + +class Seedable(Protocol): + """One kind of state a daemon can seed to or from a peer daemon.""" + + kind: str + + def request_params(self, ctx: Any) -> dict[str, Any]: + """Target side: the key fields identifying what this rank needs.""" + ... + + def resolve(self, request: dict[str, Any]) -> Any | None: + """Source side: produce the payload named by ``request``. + + Validate the wire key first — raise an error on any malformed field + (``ValueError`` for bad values, ``KeyError`` for missing ones), before + touching local state. The transport turns it into an error response. + Return None on a miss. + """ + ... + + def validate_payload(self, payload: Any) -> Any: + """Target side: validate the fetched payload, or raise ValueError.""" + ... + + def install(self, ctx: Any, payload: Any) -> None: + """Target side: install a validated payload.""" + ... + + +class SeedConfig(NamedTuple): + """One daemon's D2D seeding configuration: its listener, its source + replica, and the shared token authenticating both.""" + + listen: tuple[str, int] | None = None + source: tuple[str, int] | None = None + token: str | None = None + + +class SeedTransport: + """The daemon's D2D seeding endpoint: serves and fetches seedables.""" + + def __init__( + self, + seedables: Iterable[Seedable], + seed_config: SeedConfig | None = None, + ): + seed_config = seed_config or SeedConfig() + listen, source, token = seed_config + if (listen is not None or source is not None) and not token: + raise ValueError("A daemon-to-daemon listener/source requires a token") + self._seedables = list(seedables) + self._seed_config = seed_config + self._server: socket.socket | None = None + + def adopt(self, seedable: Seedable, ctx: Any) -> None: + """Best-effort: fetch and install ``seedable`` from the source daemon. + + Any failure degrades to the local path with a log line. + """ + response = self.request( + {"kind": seedable.kind, **seedable.request_params(ctx)}, + ) + if response is None: + return + status = response.get("status") + if status != "ok": + log = logger.info if status == "miss" else logger.warning + log( + "No remote %s seed adopted (status %s); using the local path", + seedable.kind, + status, + ) + return + try: + seedable.install(ctx, seedable.validate_payload(response.get("payload"))) + except Exception as e: + logger.warning( + "Rejected the remote %s seed; using the local path: %s", + seedable.kind, + e, + ) + return + logger.info("Adopted remote %s seed", seedable.kind) + + def handle_request(self, request: Any) -> dict[str, Any]: + """Source side: answer one authenticated seed request.""" + kind_name = request.get("kind") if isinstance(request, dict) else None + seedable = next((s for s in self._seedables if s.kind == kind_name), None) + if seedable is None: + return {"status": "error", "message": "unsupported seed kind"} + try: + payload = seedable.resolve(request) + except (OSError, KeyError, TypeError, ValueError) as e: + return {"status": "error", "message": str(e)} + if payload is None: + return {"status": "miss"} + return {"status": "ok", "payload": payload} + + def request(self, payload: dict[str, Any]) -> dict[str, Any] | None: + """Send one authenticated request to the source daemon. + + Returns the response, or None with a log line on any transport + failure; the caller falls back to its local path then. + """ + source = self._seed_config.source + if source is None: + return None + label = payload.get("kind", "payload") + try: + with socket.create_connection(source, timeout=_REQUEST_TIMEOUT_S) as sock: + send_json(sock, {**payload, "token": self._seed_config.token}) + response = recv_json(sock) + except (OSError, ValueError) as e: + logger.warning( + "Cannot fetch %s from %s over D2D; falling back locally: %s", + label, + source, + e, + ) + return None + if not isinstance(response, dict): + logger.warning( + "Source %s sent a malformed response for %s; falling back locally", + source, + label, + ) + return None + return response + + def start_listener(self) -> tuple[str, int] | None: + """Start serving this rank's seedables to peer daemons. + + Returns the bound address, or None when no listener is configured. + """ + listen = self._seed_config.listen + if listen is None: + return None + family = socket.AF_INET6 if ":" in listen[0] else socket.AF_INET + self._server = socket.create_server(listen, family=family) + threading.Thread( + target=self._serve, + name="vllm-weight-cache-d2d", + daemon=True, + ).start() + host, port = self._server.getsockname()[:2] + return host, port + + def close(self) -> None: + if self._server is not None: + self._server.close() + self._server = None + + def _serve(self) -> None: + assert self._server is not None + while True: + try: + conn, _ = self._server.accept() + except OSError: + return + with conn: + try: + request = recv_json( + conn, + max_size=_MAX_REMOTE_REQUEST_SIZE, + deadline=time.monotonic() + _REMOTE_REQUEST_TIMEOUT_S, + ) + presented = ( + request.get("token") if isinstance(request, dict) else None + ) + if ( + not self._seed_config.token + or not isinstance(presented, str) + or not hmac.compare_digest(presented, self._seed_config.token) + ): + send_json( + conn, + {"status": "error", "message": "Unauthorized request"}, + ) + continue + # Authenticated handlers (e.g. weight manifests) may be slow. + conn.settimeout(_AUTHENTICATED_TIMEOUT_S) + send_json(conn, self.handle_request(request)) + except (OSError, ValueError) as e: + with contextlib.suppress(OSError): + send_json(conn, {"status": "error", "message": str(e)}) diff --git a/vllm/model_executor/model_loader/weight_cache/warmup_cache.py b/vllm/model_executor/model_loader/weight_cache/warmup_cache.py new file mode 100644 index 000000000000..640108200302 --- /dev/null +++ b/vllm/model_executor/model_loader/weight_cache/warmup_cache.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Daemon warmup caches, seeded daemon-to-daemon over the shared transport. + +Each warmup-cache object (currently only :class:`FlashInferWarmupCache`) +drives its own warmup pass on the daemon; the seeding flow itself +(:class:`seed.SeedTransport`) is shared with future kinds such as weight +shards. A target daemon installs its rank-matched cache before tuning, so +both the daemon's own warmup and the later engine start skip re-tuning. The +disk layout is the index: a request names the cache by the same key the +layout uses (kind, config hash, dp rank, world rank) and the source resolves +it to files under the cache root. There is no separate artifact store or +registry. + +The FlashInfer payload is a ``{filename: parsed JSON}`` map rather than a +single document because upcoming kinds (Triton tuned configs) consist of +several JSON files per cache. + +TODO: install valid files individually instead of falling back to local +warmup when any file of a multi-file kind fails validation. +""" + +import gc +import json +from typing import TYPE_CHECKING, Any, Protocol + +import regex as re +import torch + +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_cache.seed import Seedable +from vllm.model_executor.model_loader.weight_cache.utils import ( + build_warmup_runner, + is_draft_model_cacheable, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer + +if TYPE_CHECKING: + from vllm.model_executor.model_loader.weight_cache.daemon import ( + WeightCacheDaemon, + ) + from vllm.model_executor.model_loader.weight_cache.seed import ( + SeedTransport, + ) + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +logger = init_logger(__name__) + +_CONFIG_HASH_RE = re.compile(r"[0-9a-f]{64}") + + +class WarmupCache(Seedable, Protocol): + """One daemon warmup cache: its warmup pass and its cache's semantics. + + The warmup layer; the D2D exchange lives in seed.py (SeedTransport). + """ + + def warmup(self, daemon: "WeightCacheDaemon", transport: "SeedTransport") -> None: + """Run this rank's warmup pass against the cached model.""" + ... + + +class FlashInferWarmupCache(WarmupCache): + """FlashInfer's per-rank autotune table (a single JSON file).""" + + kind = "flashinfer_autotune" + + def warmup(self, daemon: "WeightCacheDaemon", transport: "SeedTransport") -> None: + vllm_config = daemon.vllm_config + # Mirrors the gate in the engine's kernel warmup. + if ( + vllm_config.kernel_config.enable_flashinfer_autotune is False + or not has_flashinfer() + or not current_platform.has_device_capability(90) + or is_draft_model_cacheable(vllm_config.speculative_config) + ): + return + from vllm.model_executor.warmup.kernel_warmup import flashinfer_autotune + + try: + current_platform.update_block_size_for_backend(vllm_config) + runner = build_warmup_runner( + vllm_config, + daemon.local_rank, + is_draft=daemon.is_draft, + model_config=daemon.model_config, + ) + assert daemon.model is not None, "warmup ran before load_model" + runner.load_model(model=daemon.model) + transport.adopt(self, runner) + # The V2 runner requires skip_attn when no KV cache exists + flashinfer_autotune(runner, skip_attn=vllm_config.use_v2_model_runner) + logger.info( + "Weight cache %s daemon rank %d tuned FlashInfer; the tuned " + "configs are in the on-disk autotune cache", + daemon.role, + daemon.global_rank, + ) + finally: + # Free the runner and the dummy-run activations; the daemon's + # weights are untouched. + gc.collect() + torch.accelerator.empty_cache() + + def request_params(self, runner: "GPUModelRunner") -> dict[str, Any]: + from vllm.distributed.parallel_state import get_world_group + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + flashinfer_autotune_cache_hash, + ) + + return { + "config_hash": flashinfer_autotune_cache_hash(runner), + "dp_rank": runner.vllm_config.parallel_config.data_parallel_rank, + "rank": get_world_group().rank_in_group, + } + + def resolve(self, request: dict[str, Any]) -> dict[str, Any] | None: + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + flashinfer_autotune_filename, + flashinfer_autotune_root, + ) + + config_hash = request.get("config_hash") + if not isinstance(config_hash, str) or not _CONFIG_HASH_RE.fullmatch( + config_hash + ): + raise ValueError("invalid config_hash") + + # Garbage ranks simply name a file that does not exist (a miss); the + # hex check above is what keeps the request inside the cache root. + filename = flashinfer_autotune_filename(request["dp_rank"], request["rank"]) + path = flashinfer_autotune_root() / config_hash / filename + if not path.is_file(): + return None + return {filename: json.loads(path.read_bytes())} + + def validate_payload(self, payload: Any) -> dict[str, Any]: + if not isinstance(payload, dict) or len(payload) != 1: + raise ValueError("expected exactly one autotune table") + ((filename, table),) = payload.items() + if not isinstance(table, dict): + raise ValueError("autotune table must be a JSON object") + if not filename.startswith("autotune_configs_"): + raise ValueError(f"unexpected file {filename!r}") + # Reject tables flashinfer would refuse to load: a hard _metadata + # conflict would also never be overwritten, poisoning this cache + # path into re-tuning on every start. + saved_meta = table.get("_metadata") + if saved_meta is None: + # Pre-metadata flashinfer file; load_configs accepts these. + return payload + try: + from flashinfer.autotuner.autotuner import ( + _classify_metadata_mismatches, + _collect_metadata, + ) + except ImportError: + logger.warning( + "flashinfer metadata classifier unavailable; " + "accepting the remote autotune table unchecked" + ) + return payload + hard, soft = _classify_metadata_mismatches(saved_meta, _collect_metadata()) + mismatches = {**hard, **soft} + if mismatches: + details = ", ".join( + f"{key}: remote={old} vs local={new}" + for key, (old, new) in mismatches.items() + ) + raise ValueError(f"autotune table environment mismatch: {details}") + return payload + + def install(self, runner: "GPUModelRunner", payload: dict[str, Any]) -> None: + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + resolve_rank_flashinfer_autotune_file, + write_flashinfer_autotune_cache, + ) + + (data,) = payload.values() + cache_path = resolve_rank_flashinfer_autotune_file(runner) + write_flashinfer_autotune_cache(cache_path, json.dumps(data).encode()) + logger.info("Installed remote FlashInfer autotune table: %s", cache_path) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index 3c4174303379..61b2f58ded7f 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -38,26 +38,44 @@ def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: return hashlib.sha256(config_hash.encode()).hexdigest() -def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: +def flashinfer_autotune_root() -> Path: + """Root directory of the FlashInfer autotune cache (no runner needed).""" override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR if override_dir: - root = Path(override_dir).expanduser() - else: - from flashinfer.jit import env as flashinfer_jit_env - - flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR - root = ( - Path(envs.VLLM_CACHE_ROOT) - / "flashinfer_autotune_cache" - / flashinfer_workspace.parent.name - / flashinfer_workspace.name - ) + return Path(override_dir).expanduser() + from flashinfer.jit import env as flashinfer_jit_env + + flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR + return ( + Path(envs.VLLM_CACHE_ROOT) + / "flashinfer_autotune_cache" + / flashinfer_workspace.parent.name + / flashinfer_workspace.name + ) + + +def flashinfer_autotune_filename(dp_rank: int, rank: int) -> str: + """Per-rank cache filename, shared by the engine reader and D2D exchange.""" + return f"autotune_configs_dp{dp_rank}_rank{rank}.json" + - output_dir = root / flashinfer_autotune_cache_hash(runner) +def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: + output_dir = flashinfer_autotune_root() / flashinfer_autotune_cache_hash(runner) output_dir.mkdir(parents=True, exist_ok=True) return output_dir / "autotune_configs.json" +def resolve_rank_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: + """The per-rank variant of :func:`resolve_flashinfer_autotune_file`.""" + from vllm.distributed.parallel_state import get_world_group + + path = resolve_flashinfer_autotune_file(runner) + dp_rank = runner.vllm_config.parallel_config.data_parallel_rank + return path.with_name( + flashinfer_autotune_filename(dp_rank, get_world_group().rank_in_group) + ) + + def use_flashinfer_autotune_v2(runner: "GPUModelRunner") -> bool: from vllm.distributed.parallel_state import get_node_count, get_world_group from vllm.utils.flashinfer import has_flashinfer_autotune_v2 diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 84c7321557a0..0f1058a98661 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -22,8 +22,8 @@ from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( - resolve_flashinfer_autotune_file, resolve_flashinfer_autotune_v2_root, + resolve_rank_flashinfer_autotune_file, use_flashinfer_autotune_v2, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( @@ -405,6 +405,27 @@ def _all_ranks_have_matching_cache(path: Path, group) -> bool: return fingerprint is not None and all(f == fingerprint for f in gathered) +def _load_shared_autotune_cache(cache_path: Path, tuner, tune_group) -> bool: + """Load the per-rank autotune cache into the tuner, group-wide. + + Returns True only when every rank of the tune group had a matching file + and loaded it: the dummy-run sweep contains collectives, so skipping it + is safe only if the decision is unanimous. + """ + if not _all_ranks_have_matching_cache(cache_path, tune_group): + return False + loaded = tuner.load_configs(str(cache_path)) + if tune_group.world_size > 1: + loaded_by_rank: list[bool | None] = [None] * tune_group.world_size + torch.distributed.all_gather_object( + loaded_by_rank, loaded, group=tune_group.cpu_group + ) + loaded = all(loaded_by_rank) + if not loaded: + tuner.clear_cache() + return loaded + + def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> None: """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, @@ -455,6 +476,7 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> autotune_kwargs["skip_ops"] = skip_ops use_autotune_v2 = use_flashinfer_autotune_v2(runner) + cache_hit = False if use_autotune_v2: from flashinfer import autotune_v2, autotune_v2_reload @@ -469,12 +491,7 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> mode="tune", cache_root=cache_root, **autotune_kwargs ) else: - cache_path = resolve_flashinfer_autotune_file(runner) - # The world group only spans DP ranks when vLLM folds DP into it. - dp_rank = runner.vllm_config.parallel_config.data_parallel_rank - cache_path = cache_path.with_name( - f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" - ) + cache_path = resolve_rank_flashinfer_autotune_file(runner) if is_leader: logger.info_once( "Using FlashInfer autotune cache file: %s. Pass " @@ -482,16 +499,7 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> cache_path, ) - if _all_ranks_have_matching_cache(cache_path, tune_group): - loaded = tuner.load_configs(str(cache_path)) - if tune_group.world_size > 1: - loaded_by_rank: list[bool | None] = [None] * tune_group.world_size - torch.distributed.all_gather_object( - loaded_by_rank, loaded, group=tune_group.cpu_group - ) - loaded = all(loaded_by_rank) - if not loaded: - tuner.clear_cache() + cache_hit = _load_shared_autotune_cache(cache_path, tuner, tune_group) autotune_context = fi_utils.autotune(tune_mode=True, **autotune_kwargs) group = tune_group.cpu_group if tune_group.world_size > 1 else None @@ -506,7 +514,8 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> # HiSparse hot-buffer attention is bounded by decode batch # size, not the prefill-sized batch used for the full model. autotune_hisparse_flashinfer_attention(runner) - _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) + if not cache_hit: + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: From 6f51433838b858a9cd0478b2ffee4d2c610a1712 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sun, 25 Oct 2026 09:40:14 -0400 Subject: [PATCH 049/631] [Bugfix][KV Offload] Reset the store cursor when a request restarts from token 0 (#63010) Signed-off-by: Advay Monga Co-authored-by: Claude Fable 5.1 Co-authored-by: vllm-agent --- .../offloading_connector/test_scheduler.py | 25 +++++++++++++------ .../kv_connector/v1/offloading/scheduler.py | 2 ++ 2 files changed, 19 insertions(+), 8 deletions(-) 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 f1b483af408d..bf612f6aa088 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1846,7 +1846,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_request_preemption(request_runner, async_scheduling: bool): +@pytest.mark.parametrize("cpu_hit", [True, False]) +def test_request_preemption(request_runner, async_scheduling: bool, cpu_hit: bool): + """Resume reloads the chunks, or re-stores them if the CPU tier lost them.""" block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -1897,16 +1899,23 @@ def test_request_preemption(request_runner, async_scheduling: bool): free_block_queue.num_free_blocks = num_free_blocks_empty runner.scheduler.reset_prefix_cache() - # 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 keys, ctx, *_: 3 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) - runner.run( - decoded_tokens=[0] * block_size, - expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8), - ) + if cpu_hit: + # 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 keys, ctx, *_: 3 + runner.run( + decoded_tokens=[0] * block_size, + expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8), + ) + else: + # lookup misses: recompute and re-store [0, ..., 8], then store [9, 10, 11] + runner.run( + decoded_tokens=[0] * block_size, + expected_stored=(0, 1, 2, 3, 4, 5, 6, 7, 8), + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 6bed634430a8..d7a24341a89f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1246,6 +1246,8 @@ def get_num_new_matched_tokens( req_status = self._req_status[request.request_id] for group_state in req_status.group_states: group_state.block_ids.clear() + # Restarting at token 0: offer every chunk again. + group_state.next_stored_chunk_idx = 0 if req_status.transfer_jobs: logger.debug( From b3466364b69db84ddeca30435f56c822551f8b14 Mon Sep 17 00:00:00 2001 From: Chongmin Li Date: Sun, 25 Oct 2026 23:12:42 +0800 Subject: [PATCH 050/631] [Bugfix] Release cached cuBLAS workspaces after daemon warmup (#63451) Signed-off-by: lichongmin Co-authored-by: vllm-agent --- docs/features/preload.md | 5 +++ .../model_loader/test_weight_cache.py | 36 +++++++++++++++++++ .../model_loader/weight_cache/daemon.py | 15 +++++--- 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/docs/features/preload.md b/docs/features/preload.md index 5318cc447daa..5787fdb51c1c 100644 --- a/docs/features/preload.md +++ b/docs/features/preload.md @@ -70,6 +70,11 @@ loading and before binding its socket, then writes the cache: vllm preload --model /path/to/model --tensor-parallel-size 4 ``` +After warmup, the daemon releases its communication resources and, on CUDA, +cached cuBLAS workspaces before accepting engine connections. The model weights +remain resident for CUDA IPC sharing. This reduces the daemon's GPU memory +overhead; the amount reclaimed depends on the model and allocator state. + ### Seeding warmup caches between daemons (D2D) A fresh daemon replica — a new node joining a fleet, a restarted preload diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index fc9af29c1f5e..7d25b85f157f 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -788,6 +788,42 @@ def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): assert loader.get_external_weight_memory(None) == 0 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@create_new_process_for_each_test("spawn") +def test_daemon_releases_cublas_workspaces(): + """Startup cleanup releases cached GEMM scratch without freeing live tensors.""" + import os + + import torch + + from vllm.config import ModelConfig, VllmConfig + from vllm.distributed import init_distributed_environment + from vllm.model_executor.model_loader.weight_cache.daemon import WeightCacheDaemon + + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:2" + os.environ["TORCH_CUBLAS_WORKSPACE_CACHE"] = "1" + config = VllmConfig(model_config=ModelConfig(model=QWEN_CASE.model)) + address = f"tcp://127.0.0.1:{get_open_port()}" + server = WeightCacheDaemon(config, 0, 0, address) + torch.accelerator.set_device_index(0) + init_distributed_environment(1, 0, address, 0, backend="nccl") + torch._C._cuda_clearCublasWorkspaces() + matrix = torch.ones(128, 128, device="cuda") + result = torch.empty_like(matrix) + before = torch.accelerator.memory_allocated() + torch.mm(matrix, matrix, out=result) + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() > before + + server.release_distributed() + after = torch.accelerator.memory_allocated() + torch._C._cuda_clearCublasWorkspaces() + assert torch.accelerator.memory_allocated() == after, ( + "daemon retained BLAS workspaces" + ) + torch.testing.assert_close(result, torch.full_like(result, 128), rtol=0, atol=0) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_daemon_reports_exported_weight_bytes(): """The engine's budget excludes exactly the storages it maps from the diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 77f6540c4c5e..1c41213d20c8 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -319,11 +319,11 @@ def load_model(self) -> None: ) def release_distributed(self) -> None: - """Tear down the daemon's process groups after loading and warmup. + """Release process groups and BLAS workspaces after loading and warmup. Serving only exports CUDA IPC handles, so the communicators (NCCL, - custom all-reduce and FlashInfer workspaces) would otherwise pin GPU - memory outside the engine's budget for the daemon's lifetime. + custom all-reduce and FlashInfer workspaces) and cached BLAS scratch + would otherwise pin GPU memory for the daemon's lifetime. """ free_before, _ = torch.accelerator.get_memory_info() # No rank may tear down a communicator while a peer is still in a @@ -334,10 +334,17 @@ def release_distributed(self) -> None: destroy_model_parallel() destroy_distributed_environment() gc.collect() + if current_platform.is_cuda(): + # ATen caches these separately from the model's exported storages. + clear_workspaces = getattr(torch._C, "_cuda_clearCublasWorkspaces", None) + if clear_workspaces is not None: + clear_workspaces() + else: + logger.debug("PyTorch does not expose BLAS workspace cleanup") torch.accelerator.empty_cache() free_after, _ = torch.accelerator.get_memory_info() logger.info( - "Weight cache %s daemon rank %d released its process groups (%s GiB freed)", + "Weight cache %s daemon rank %d released startup resources (%s GiB freed)", self.role, self.global_rank, format_gib(free_after - free_before), From fa2881ed19bef2957096788ce84d3b12cbd48948 Mon Sep 17 00:00:00 2001 From: John Calderon <81483067+johncalesp@users.noreply.github.com> Date: Sun, 25 Oct 2026 11:31:31 -0400 Subject: [PATCH 051/631] [CI][Benchmarks] Fix attention smoke test backend names and failure exit status (#63177) Signed-off-by: John Calderon Co-authored-by: Artem Perevedentsev --- .buildkite/test_areas/benchmarks.yaml | 2 +- benchmarks/attention_benchmarks/benchmark.py | 5 +++++ docs/benchmarking/cli.md | 14 ++++++++++++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 7c95577e0373..bd381080094f 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -34,7 +34,7 @@ steps: - benchmarks/attention_benchmarks/ - vllm/v1/attention/ commands: - - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" + - python3 benchmarks/attention_benchmarks/benchmark.py --backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k" mirror: amd: label: ":amd: (MI355) Attention Benchmark Smoke" diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 943133252be0..2242e231bb1f 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -1036,6 +1036,8 @@ def main(): pbar.update(1) if args.ncu_profile: + if any(not result.success for result in all_results): + raise SystemExit(1) return # Display decode vs prefill results @@ -1435,6 +1437,9 @@ def main(): if args.output_json: formatter.save_json(all_results, args.output_json) + if any(not result.success for result in all_results): + raise SystemExit(1) + if __name__ == "__main__": main() diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index a38ef0ab66af..4f7d8a4192d2 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -9,6 +9,20 @@ It's a living document, updated as new features and datasets become available. For benchmarking production vLLM servers, we recommend [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns. +## Attention Kernel Benchmark + +Run the attention kernel smoke benchmark from the repository root using the +exact backend names: + +```bash +python3 benchmarks/attention_benchmarks/benchmark.py \ + --backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k" +``` + +If any configuration fails, the benchmark exits with status `1` after reporting +results and writing any requested CSV or JSON output. Successful runs exit with +status `0`. Profiling runs also propagate failures, but do not save timing metrics. + ## Dataset Overview